Spaces:
Runtime error
Runtime error
Add files
Browse files- PyPatchMatch/.gitignore +4 -0
- PyPatchMatch/LICENSE +21 -0
- PyPatchMatch/Makefile +54 -0
- PyPatchMatch/README.md +64 -0
- PyPatchMatch/csrc/inpaint.cpp +234 -0
- PyPatchMatch/csrc/inpaint.h +27 -0
- PyPatchMatch/csrc/masked_image.cpp +138 -0
- PyPatchMatch/csrc/masked_image.h +112 -0
- PyPatchMatch/csrc/nnf.cpp +268 -0
- PyPatchMatch/csrc/nnf.h +133 -0
- PyPatchMatch/csrc/pyinterface.cpp +107 -0
- PyPatchMatch/csrc/pyinterface.h +38 -0
- PyPatchMatch/examples/.gitignore +2 -0
- PyPatchMatch/examples/cpp_example.cpp +31 -0
- PyPatchMatch/examples/cpp_example_run.sh +18 -0
- PyPatchMatch/examples/images/forest.bmp +0 -0
- PyPatchMatch/examples/images/forest_pruned.bmp +0 -0
- PyPatchMatch/examples/py_example.py +21 -0
- PyPatchMatch/examples/py_example_global_mask.py +27 -0
- PyPatchMatch/patch_match.py +201 -0
- PyPatchMatch/travis.sh +9 -0
- app.py +390 -0
- canvas.py +547 -0
- js/mode.js +6 -0
- js/outpaint.js +24 -0
- js/proceed.js +22 -0
- js/setup.js +22 -0
- js/upload.js +19 -0
- packages.txt +2 -0
- perlin2d.py +45 -0
- requirements.txt +12 -0
- utils.py +154 -0
PyPatchMatch/.gitignore
ADDED
@@ -0,0 +1,4 @@
|
|
|
|
|
|
|
|
|
|
|
1 |
+
/build/
|
2 |
+
/*.so
|
3 |
+
__pycache__
|
4 |
+
*.py[cod]
|
PyPatchMatch/LICENSE
ADDED
@@ -0,0 +1,21 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
MIT License
|
2 |
+
|
3 |
+
Copyright (c) 2020 Jiayuan Mao
|
4 |
+
|
5 |
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
6 |
+
of this software and associated documentation files (the "Software"), to deal
|
7 |
+
in the Software without restriction, including without limitation the rights
|
8 |
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
9 |
+
copies of the Software, and to permit persons to whom the Software is
|
10 |
+
furnished to do so, subject to the following conditions:
|
11 |
+
|
12 |
+
The above copyright notice and this permission notice shall be included in all
|
13 |
+
copies or substantial portions of the Software.
|
14 |
+
|
15 |
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
16 |
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
17 |
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
18 |
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
19 |
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
20 |
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
21 |
+
SOFTWARE.
|
PyPatchMatch/Makefile
ADDED
@@ -0,0 +1,54 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
#
|
2 |
+
# Makefile
|
3 |
+
# Jiayuan Mao, 2019-01-09 13:59
|
4 |
+
#
|
5 |
+
|
6 |
+
SRC_DIR = csrc
|
7 |
+
INC_DIR = csrc
|
8 |
+
OBJ_DIR = build/obj
|
9 |
+
TARGET = libpatchmatch.so
|
10 |
+
|
11 |
+
LIB_TARGET = $(TARGET)
|
12 |
+
INCLUDE_DIR = -I $(SRC_DIR) -I $(INC_DIR)
|
13 |
+
|
14 |
+
CXX = $(ENVIRONMENT_OPTIONS) g++
|
15 |
+
CXXFLAGS = -std=c++14
|
16 |
+
CXXFLAGS += -Ofast -ffast-math -w
|
17 |
+
# CXXFLAGS += -g
|
18 |
+
CXXFLAGS += $(shell pkg-config --cflags opencv) -fPIC
|
19 |
+
CXXFLAGS += $(INCLUDE_DIR)
|
20 |
+
LDFLAGS = $(shell pkg-config --cflags --libs opencv) -shared -fPIC
|
21 |
+
|
22 |
+
|
23 |
+
CXXSOURCES = $(shell find $(SRC_DIR)/ -name "*.cpp")
|
24 |
+
OBJS = $(addprefix $(OBJ_DIR)/,$(CXXSOURCES:.cpp=.o))
|
25 |
+
DEPFILES = $(OBJS:.o=.d)
|
26 |
+
|
27 |
+
.PHONY: all clean rebuild test
|
28 |
+
|
29 |
+
all: $(LIB_TARGET)
|
30 |
+
|
31 |
+
$(OBJ_DIR)/%.o: %.cpp
|
32 |
+
@echo "[CC] $< ..."
|
33 |
+
@$(CXX) -c $< $(CXXFLAGS) -o $@
|
34 |
+
|
35 |
+
$(OBJ_DIR)/%.d: %.cpp
|
36 |
+
@mkdir -pv $(dir $@)
|
37 |
+
@echo "[dep] $< ..."
|
38 |
+
@$(CXX) $(INCLUDE_DIR) $(CXXFLAGS) -MM -MT "$(OBJ_DIR)/$(<:.cpp=.o) $(OBJ_DIR)/$(<:.cpp=.d)" "$<" > "$@"
|
39 |
+
|
40 |
+
sinclude $(DEPFILES)
|
41 |
+
|
42 |
+
$(LIB_TARGET): $(OBJS)
|
43 |
+
@echo "[link] $(LIB_TARGET) ..."
|
44 |
+
@$(CXX) $(OBJS) -o $@ $(CXXFLAGS) $(LDFLAGS)
|
45 |
+
|
46 |
+
clean:
|
47 |
+
rm -rf $(OBJ_DIR) $(LIB_TARGET)
|
48 |
+
|
49 |
+
rebuild:
|
50 |
+
+@make clean
|
51 |
+
+@make
|
52 |
+
|
53 |
+
# vim:ft=make
|
54 |
+
#
|
PyPatchMatch/README.md
ADDED
@@ -0,0 +1,64 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
PatchMatch based Inpainting
|
2 |
+
=====================================
|
3 |
+
This library implements the PatchMatch based inpainting algorithm. It provides both C++ and Python interfaces.
|
4 |
+
This implementation is heavily based on the implementation by Younesse ANDAM:
|
5 |
+
(younesse-cv/PatchMatch)[https://github.com/younesse-cv/PatchMatch], with some bugs fix.
|
6 |
+
|
7 |
+
Usage
|
8 |
+
-------------------------------------
|
9 |
+
|
10 |
+
You need to first install OpenCV to compile the C++ libraries. Then, run `make` to compile the
|
11 |
+
shared library `libpatchmatch.so`.
|
12 |
+
|
13 |
+
For Python users (example available at `examples/py_example.py`)
|
14 |
+
|
15 |
+
```python
|
16 |
+
import patch_match
|
17 |
+
|
18 |
+
image = ... # either a numpy ndarray or a PIL Image object.
|
19 |
+
mask = ... # either a numpy ndarray or a PIL Image object.
|
20 |
+
result = patch_match.inpaint(image, mask, patch_size=5)
|
21 |
+
```
|
22 |
+
|
23 |
+
For C++ users (examples available at `examples/cpp_example.cpp`)
|
24 |
+
|
25 |
+
```cpp
|
26 |
+
#include "inpaint.h"
|
27 |
+
|
28 |
+
int main() {
|
29 |
+
cv::Mat image = ...
|
30 |
+
cv::Mat mask = ...
|
31 |
+
|
32 |
+
cv::Mat result = Inpainting(image, mask, 5).run();
|
33 |
+
|
34 |
+
return 0;
|
35 |
+
}
|
36 |
+
```
|
37 |
+
|
38 |
+
|
39 |
+
README and COPYRIGHT by Younesse ANDAM
|
40 |
+
-------------------------------------
|
41 |
+
@Author: Younesse ANDAM
|
42 |
+
|
43 |
+
@Contact: [email protected]
|
44 |
+
|
45 |
+
Description: This project is a personal implementation of an algorithm called PATCHMATCH that restores missing areas in an image.
|
46 |
+
The algorithm is presented in the following paper
|
47 |
+
PatchMatch A Randomized Correspondence Algorithm
|
48 |
+
for Structural Image Editing
|
49 |
+
by C.Barnes,E.Shechtman,A.Finkelstein and Dan B.Goldman
|
50 |
+
ACM Transactions on Graphics (Proc. SIGGRAPH), vol.28, aug-2009
|
51 |
+
|
52 |
+
For more information please refer to
|
53 |
+
http://www.cs.princeton.edu/gfx/pubs/Barnes_2009_PAR/index.php
|
54 |
+
|
55 |
+
Copyright (c) 2010-2011
|
56 |
+
|
57 |
+
|
58 |
+
Requirements
|
59 |
+
-------------------------------------
|
60 |
+
|
61 |
+
To run the project you need to install Opencv library and link it to your project.
|
62 |
+
Opencv can be download it here
|
63 |
+
http://opencv.org/downloads.html
|
64 |
+
|
PyPatchMatch/csrc/inpaint.cpp
ADDED
@@ -0,0 +1,234 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
#include <algorithm>
|
2 |
+
#include <iostream>
|
3 |
+
#include <opencv2/imgcodecs.hpp>
|
4 |
+
#include <opencv2/imgproc.hpp>
|
5 |
+
#include <opencv2/highgui.hpp>
|
6 |
+
|
7 |
+
#include "inpaint.h"
|
8 |
+
|
9 |
+
namespace {
|
10 |
+
static std::vector<double> kDistance2Similarity;
|
11 |
+
|
12 |
+
void init_kDistance2Similarity() {
|
13 |
+
double base[11] = {1.0, 0.99, 0.96, 0.83, 0.38, 0.11, 0.02, 0.005, 0.0006, 0.0001, 0};
|
14 |
+
int length = (PatchDistanceMetric::kDistanceScale + 1);
|
15 |
+
kDistance2Similarity.resize(length);
|
16 |
+
for (int i = 0; i < length; ++i) {
|
17 |
+
double t = (double) i / length;
|
18 |
+
int j = (int) (100 * t);
|
19 |
+
int k = j + 1;
|
20 |
+
double vj = (j < 11) ? base[j] : 0;
|
21 |
+
double vk = (k < 11) ? base[k] : 0;
|
22 |
+
kDistance2Similarity[i] = vj + (100 * t - j) * (vk - vj);
|
23 |
+
}
|
24 |
+
}
|
25 |
+
|
26 |
+
|
27 |
+
inline void _weighted_copy(const MaskedImage &source, int ys, int xs, cv::Mat &target, int yt, int xt, double weight) {
|
28 |
+
if (source.is_masked(ys, xs)) return;
|
29 |
+
if (source.is_globally_masked(ys, xs)) return;
|
30 |
+
|
31 |
+
auto source_ptr = source.get_image(ys, xs);
|
32 |
+
auto target_ptr = target.ptr<double>(yt, xt);
|
33 |
+
|
34 |
+
#pragma unroll
|
35 |
+
for (int c = 0; c < 3; ++c)
|
36 |
+
target_ptr[c] += static_cast<double>(source_ptr[c]) * weight;
|
37 |
+
target_ptr[3] += weight;
|
38 |
+
}
|
39 |
+
}
|
40 |
+
|
41 |
+
/**
|
42 |
+
* This algorithme uses a version proposed by Xavier Philippeau.
|
43 |
+
*/
|
44 |
+
|
45 |
+
Inpainting::Inpainting(cv::Mat image, cv::Mat mask, const PatchDistanceMetric *metric)
|
46 |
+
: m_initial(image, mask), m_distance_metric(metric), m_pyramid(), m_source2target(), m_target2source() {
|
47 |
+
_initialize_pyramid();
|
48 |
+
}
|
49 |
+
|
50 |
+
Inpainting::Inpainting(cv::Mat image, cv::Mat mask, cv::Mat global_mask, const PatchDistanceMetric *metric)
|
51 |
+
: m_initial(image, mask, global_mask), m_distance_metric(metric), m_pyramid(), m_source2target(), m_target2source() {
|
52 |
+
_initialize_pyramid();
|
53 |
+
}
|
54 |
+
|
55 |
+
void Inpainting::_initialize_pyramid() {
|
56 |
+
auto source = m_initial;
|
57 |
+
m_pyramid.push_back(source);
|
58 |
+
while (source.size().height > m_distance_metric->patch_size() && source.size().width > m_distance_metric->patch_size()) {
|
59 |
+
source = source.downsample();
|
60 |
+
m_pyramid.push_back(source);
|
61 |
+
}
|
62 |
+
|
63 |
+
if (kDistance2Similarity.size() == 0) {
|
64 |
+
init_kDistance2Similarity();
|
65 |
+
}
|
66 |
+
}
|
67 |
+
|
68 |
+
cv::Mat Inpainting::run(bool verbose, bool verbose_visualize, unsigned int random_seed) {
|
69 |
+
srand(random_seed);
|
70 |
+
const int nr_levels = m_pyramid.size();
|
71 |
+
|
72 |
+
MaskedImage source, target;
|
73 |
+
for (int level = nr_levels - 1; level >= 0; --level) {
|
74 |
+
if (verbose) std::cerr << "Inpainting level: " << level << std::endl;
|
75 |
+
|
76 |
+
source = m_pyramid[level];
|
77 |
+
|
78 |
+
if (level == nr_levels - 1) {
|
79 |
+
target = source.clone();
|
80 |
+
target.clear_mask();
|
81 |
+
m_source2target = NearestNeighborField(source, target, m_distance_metric);
|
82 |
+
m_target2source = NearestNeighborField(target, source, m_distance_metric);
|
83 |
+
} else {
|
84 |
+
m_source2target = NearestNeighborField(source, target, m_distance_metric, m_source2target);
|
85 |
+
m_target2source = NearestNeighborField(target, source, m_distance_metric, m_target2source);
|
86 |
+
}
|
87 |
+
|
88 |
+
if (verbose) std::cerr << "Initialization done." << std::endl;
|
89 |
+
|
90 |
+
if (verbose_visualize) {
|
91 |
+
auto visualize_size = m_initial.size();
|
92 |
+
cv::Mat source_visualize(visualize_size, m_initial.image().type());
|
93 |
+
cv::resize(source.image(), source_visualize, visualize_size);
|
94 |
+
cv::imshow("Source", source_visualize);
|
95 |
+
cv::Mat target_visualize(visualize_size, m_initial.image().type());
|
96 |
+
cv::resize(target.image(), target_visualize, visualize_size);
|
97 |
+
cv::imshow("Target", target_visualize);
|
98 |
+
cv::waitKey(0);
|
99 |
+
}
|
100 |
+
|
101 |
+
target = _expectation_maximization(source, target, level, verbose);
|
102 |
+
}
|
103 |
+
|
104 |
+
return target.image();
|
105 |
+
}
|
106 |
+
|
107 |
+
// EM-Like algorithm (see "PatchMatch" - page 6).
|
108 |
+
// Returns a double sized target image (unless level = 0).
|
109 |
+
MaskedImage Inpainting::_expectation_maximization(MaskedImage source, MaskedImage target, int level, bool verbose) {
|
110 |
+
const int nr_iters_em = 1 + 2 * level;
|
111 |
+
const int nr_iters_nnf = static_cast<int>(std::min(7, 1 + level));
|
112 |
+
const int patch_size = m_distance_metric->patch_size();
|
113 |
+
|
114 |
+
MaskedImage new_source, new_target;
|
115 |
+
|
116 |
+
for (int iter_em = 0; iter_em < nr_iters_em; ++iter_em) {
|
117 |
+
if (iter_em != 0) {
|
118 |
+
m_source2target.set_target(new_target);
|
119 |
+
m_target2source.set_source(new_target);
|
120 |
+
target = new_target;
|
121 |
+
}
|
122 |
+
|
123 |
+
if (verbose) std::cerr << "EM Iteration: " << iter_em << std::endl;
|
124 |
+
|
125 |
+
auto size = source.size();
|
126 |
+
for (int i = 0; i < size.height; ++i) {
|
127 |
+
for (int j = 0; j < size.width; ++j) {
|
128 |
+
if (!source.contains_mask(i, j, patch_size)) {
|
129 |
+
m_source2target.set_identity(i, j);
|
130 |
+
m_target2source.set_identity(i, j);
|
131 |
+
}
|
132 |
+
}
|
133 |
+
}
|
134 |
+
if (verbose) std::cerr << " NNF minimization started." << std::endl;
|
135 |
+
m_source2target.minimize(nr_iters_nnf);
|
136 |
+
m_target2source.minimize(nr_iters_nnf);
|
137 |
+
if (verbose) std::cerr << " NNF minimization finished." << std::endl;
|
138 |
+
|
139 |
+
// Instead of upsizing the final target, we build the last target from the next level source image.
|
140 |
+
// Thus, the final target is less blurry (see "Space-Time Video Completion" - page 5).
|
141 |
+
bool upscaled = false;
|
142 |
+
if (level >= 1 && iter_em == nr_iters_em - 1) {
|
143 |
+
new_source = m_pyramid[level - 1];
|
144 |
+
new_target = target.upsample(new_source.size().width, new_source.size().height, m_pyramid[level - 1].global_mask());
|
145 |
+
upscaled = true;
|
146 |
+
} else {
|
147 |
+
new_source = m_pyramid[level];
|
148 |
+
new_target = target.clone();
|
149 |
+
}
|
150 |
+
|
151 |
+
auto vote = cv::Mat(new_target.size(), CV_64FC4);
|
152 |
+
vote.setTo(cv::Scalar::all(0));
|
153 |
+
|
154 |
+
// Votes for best patch from NNF Source->Target (completeness) and Target->Source (coherence).
|
155 |
+
_expectation_step(m_source2target, 1, vote, new_source, upscaled);
|
156 |
+
if (verbose) std::cerr << " Expectation source to target finished." << std::endl;
|
157 |
+
_expectation_step(m_target2source, 0, vote, new_source, upscaled);
|
158 |
+
if (verbose) std::cerr << " Expectation target to source finished." << std::endl;
|
159 |
+
|
160 |
+
// Compile votes and update pixel values.
|
161 |
+
_maximization_step(new_target, vote);
|
162 |
+
if (verbose) std::cerr << " Minimization step finished." << std::endl;
|
163 |
+
}
|
164 |
+
|
165 |
+
return new_target;
|
166 |
+
}
|
167 |
+
|
168 |
+
// Expectation step: vote for best estimations of each pixel.
|
169 |
+
void Inpainting::_expectation_step(
|
170 |
+
const NearestNeighborField &nnf, bool source2target,
|
171 |
+
cv::Mat &vote, const MaskedImage &source, bool upscaled
|
172 |
+
) {
|
173 |
+
auto source_size = nnf.source_size();
|
174 |
+
auto target_size = nnf.target_size();
|
175 |
+
const int patch_size = m_distance_metric->patch_size();
|
176 |
+
|
177 |
+
for (int i = 0; i < source_size.height; ++i) {
|
178 |
+
for (int j = 0; j < source_size.width; ++j) {
|
179 |
+
if (nnf.source().is_globally_masked(i, j)) continue;
|
180 |
+
int yp = nnf.at(i, j, 0), xp = nnf.at(i, j, 1), dp = nnf.at(i, j, 2);
|
181 |
+
double w = kDistance2Similarity[dp];
|
182 |
+
|
183 |
+
for (int di = -patch_size; di <= patch_size; ++di) {
|
184 |
+
for (int dj = -patch_size; dj <= patch_size; ++dj) {
|
185 |
+
int ys = i + di, xs = j + dj, yt = yp + di, xt = xp + dj;
|
186 |
+
if (!(ys >= 0 && ys < source_size.height && xs >= 0 && xs < source_size.width)) continue;
|
187 |
+
if (nnf.source().is_globally_masked(ys, xs)) continue;
|
188 |
+
if (!(yt >= 0 && yt < target_size.height && xt >= 0 && xt < target_size.width)) continue;
|
189 |
+
if (nnf.target().is_globally_masked(yt, xt)) continue;
|
190 |
+
|
191 |
+
if (!source2target) {
|
192 |
+
std::swap(ys, yt);
|
193 |
+
std::swap(xs, xt);
|
194 |
+
}
|
195 |
+
|
196 |
+
if (upscaled) {
|
197 |
+
for (int uy = 0; uy < 2; ++uy) {
|
198 |
+
for (int ux = 0; ux < 2; ++ux) {
|
199 |
+
_weighted_copy(source, 2 * ys + uy, 2 * xs + ux, vote, 2 * yt + uy, 2 * xt + ux, w);
|
200 |
+
}
|
201 |
+
}
|
202 |
+
} else {
|
203 |
+
_weighted_copy(source, ys, xs, vote, yt, xt, w);
|
204 |
+
}
|
205 |
+
}
|
206 |
+
}
|
207 |
+
}
|
208 |
+
}
|
209 |
+
}
|
210 |
+
|
211 |
+
// Maximization Step: maximum likelihood of target pixel.
|
212 |
+
void Inpainting::_maximization_step(MaskedImage &target, const cv::Mat &vote) {
|
213 |
+
auto target_size = target.size();
|
214 |
+
for (int i = 0; i < target_size.height; ++i) {
|
215 |
+
for (int j = 0; j < target_size.width; ++j) {
|
216 |
+
const double *source_ptr = vote.ptr<double>(i, j);
|
217 |
+
unsigned char *target_ptr = target.get_mutable_image(i, j);
|
218 |
+
|
219 |
+
if (target.is_globally_masked(i, j)) {
|
220 |
+
continue;
|
221 |
+
}
|
222 |
+
|
223 |
+
if (source_ptr[3] > 0) {
|
224 |
+
unsigned char r = cv::saturate_cast<unsigned char>(source_ptr[0] / source_ptr[3]);
|
225 |
+
unsigned char g = cv::saturate_cast<unsigned char>(source_ptr[1] / source_ptr[3]);
|
226 |
+
unsigned char b = cv::saturate_cast<unsigned char>(source_ptr[2] / source_ptr[3]);
|
227 |
+
target_ptr[0] = r, target_ptr[1] = g, target_ptr[2] = b;
|
228 |
+
} else {
|
229 |
+
target.set_mask(i, j, 0);
|
230 |
+
}
|
231 |
+
}
|
232 |
+
}
|
233 |
+
}
|
234 |
+
|
PyPatchMatch/csrc/inpaint.h
ADDED
@@ -0,0 +1,27 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
#pragma once
|
2 |
+
|
3 |
+
#include <vector>
|
4 |
+
|
5 |
+
#include "masked_image.h"
|
6 |
+
#include "nnf.h"
|
7 |
+
|
8 |
+
class Inpainting {
|
9 |
+
public:
|
10 |
+
Inpainting(cv::Mat image, cv::Mat mask, const PatchDistanceMetric *metric);
|
11 |
+
Inpainting(cv::Mat image, cv::Mat mask, cv::Mat global_mask, const PatchDistanceMetric *metric);
|
12 |
+
cv::Mat run(bool verbose = false, bool verbose_visualize = false, unsigned int random_seed = 1212);
|
13 |
+
|
14 |
+
private:
|
15 |
+
void _initialize_pyramid(void);
|
16 |
+
MaskedImage _expectation_maximization(MaskedImage source, MaskedImage target, int level, bool verbose);
|
17 |
+
void _expectation_step(const NearestNeighborField &nnf, bool source2target, cv::Mat &vote, const MaskedImage &source, bool upscaled);
|
18 |
+
void _maximization_step(MaskedImage &target, const cv::Mat &vote);
|
19 |
+
|
20 |
+
MaskedImage m_initial;
|
21 |
+
std::vector<MaskedImage> m_pyramid;
|
22 |
+
|
23 |
+
NearestNeighborField m_source2target;
|
24 |
+
NearestNeighborField m_target2source;
|
25 |
+
const PatchDistanceMetric *m_distance_metric;
|
26 |
+
};
|
27 |
+
|
PyPatchMatch/csrc/masked_image.cpp
ADDED
@@ -0,0 +1,138 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
#include "masked_image.h"
|
2 |
+
#include <algorithm>
|
3 |
+
#include <iostream>
|
4 |
+
|
5 |
+
const cv::Size MaskedImage::kDownsampleKernelSize = cv::Size(6, 6);
|
6 |
+
const int MaskedImage::kDownsampleKernel[6] = {1, 5, 10, 10, 5, 1};
|
7 |
+
|
8 |
+
bool MaskedImage::contains_mask(int y, int x, int patch_size) const {
|
9 |
+
auto mask_size = size();
|
10 |
+
for (int dy = -patch_size; dy <= patch_size; ++dy) {
|
11 |
+
for (int dx = -patch_size; dx <= patch_size; ++dx) {
|
12 |
+
int yy = y + dy, xx = x + dx;
|
13 |
+
if (yy >= 0 && yy < mask_size.height && xx >= 0 && xx < mask_size.width) {
|
14 |
+
if (is_masked(yy, xx) && !is_globally_masked(yy, xx)) return true;
|
15 |
+
}
|
16 |
+
}
|
17 |
+
}
|
18 |
+
return false;
|
19 |
+
}
|
20 |
+
|
21 |
+
MaskedImage MaskedImage::downsample() const {
|
22 |
+
const auto &kernel_size = MaskedImage::kDownsampleKernelSize;
|
23 |
+
const auto &kernel = MaskedImage::kDownsampleKernel;
|
24 |
+
|
25 |
+
const auto size = this->size();
|
26 |
+
const auto new_size = cv::Size(size.width / 2, size.height / 2);
|
27 |
+
|
28 |
+
auto ret = MaskedImage(new_size.width, new_size.height);
|
29 |
+
if (!m_global_mask.empty()) ret.init_global_mask_mat();
|
30 |
+
for (int y = 0; y < size.height - 1; y += 2) {
|
31 |
+
for (int x = 0; x < size.width - 1; x += 2) {
|
32 |
+
int r = 0, g = 0, b = 0, ksum = 0;
|
33 |
+
bool is_gmasked = true;
|
34 |
+
|
35 |
+
for (int dy = -kernel_size.height / 2 + 1; dy <= kernel_size.height / 2; ++dy) {
|
36 |
+
for (int dx = -kernel_size.width / 2 + 1; dx <= kernel_size.width / 2; ++dx) {
|
37 |
+
int yy = y + dy, xx = x + dx;
|
38 |
+
if (yy >= 0 && yy < size.height && xx >= 0 && xx < size.width) {
|
39 |
+
if (!is_globally_masked(yy, xx)) {
|
40 |
+
is_gmasked = false;
|
41 |
+
}
|
42 |
+
if (!is_masked(yy, xx)) {
|
43 |
+
auto source_ptr = get_image(yy, xx);
|
44 |
+
int k = kernel[kernel_size.height / 2 - 1 + dy] * kernel[kernel_size.width / 2 - 1 + dx];
|
45 |
+
r += source_ptr[0] * k, g += source_ptr[1] * k, b += source_ptr[2] * k;
|
46 |
+
ksum += k;
|
47 |
+
}
|
48 |
+
}
|
49 |
+
}
|
50 |
+
}
|
51 |
+
|
52 |
+
if (ksum > 0) r /= ksum, g /= ksum, b /= ksum;
|
53 |
+
|
54 |
+
if (!m_global_mask.empty()) {
|
55 |
+
ret.set_global_mask(y / 2, x / 2, is_gmasked);
|
56 |
+
}
|
57 |
+
if (ksum > 0) {
|
58 |
+
auto target_ptr = ret.get_mutable_image(y / 2, x / 2);
|
59 |
+
target_ptr[0] = r, target_ptr[1] = g, target_ptr[2] = b;
|
60 |
+
ret.set_mask(y / 2, x / 2, 0);
|
61 |
+
} else {
|
62 |
+
ret.set_mask(y / 2, x / 2, 1);
|
63 |
+
}
|
64 |
+
}
|
65 |
+
}
|
66 |
+
|
67 |
+
return ret;
|
68 |
+
}
|
69 |
+
|
70 |
+
MaskedImage MaskedImage::upsample(int new_w, int new_h) const {
|
71 |
+
const auto size = this->size();
|
72 |
+
auto ret = MaskedImage(new_w, new_h);
|
73 |
+
if (!m_global_mask.empty()) ret.init_global_mask_mat();
|
74 |
+
for (int y = 0; y < new_h; ++y) {
|
75 |
+
for (int x = 0; x < new_w; ++x) {
|
76 |
+
int yy = y * size.height / new_h;
|
77 |
+
int xx = x * size.width / new_w;
|
78 |
+
|
79 |
+
if (is_globally_masked(yy, xx)) {
|
80 |
+
ret.set_global_mask(y, x, 1);
|
81 |
+
ret.set_mask(y, x, 1);
|
82 |
+
} else {
|
83 |
+
if (!m_global_mask.empty()) ret.set_global_mask(y, x, 0);
|
84 |
+
|
85 |
+
if (is_masked(yy, xx)) {
|
86 |
+
ret.set_mask(y, x, 1);
|
87 |
+
} else {
|
88 |
+
auto source_ptr = get_image(yy, xx);
|
89 |
+
auto target_ptr = ret.get_mutable_image(y, x);
|
90 |
+
for (int c = 0; c < 3; ++c)
|
91 |
+
target_ptr[c] = source_ptr[c];
|
92 |
+
ret.set_mask(y, x, 0);
|
93 |
+
}
|
94 |
+
}
|
95 |
+
}
|
96 |
+
}
|
97 |
+
|
98 |
+
return ret;
|
99 |
+
}
|
100 |
+
|
101 |
+
MaskedImage MaskedImage::upsample(int new_w, int new_h, const cv::Mat &new_global_mask) const {
|
102 |
+
auto ret = upsample(new_w, new_h);
|
103 |
+
ret.set_global_mask_mat(new_global_mask);
|
104 |
+
return ret;
|
105 |
+
}
|
106 |
+
|
107 |
+
void MaskedImage::compute_image_gradients() {
|
108 |
+
if (m_image_grad_computed) {
|
109 |
+
return;
|
110 |
+
}
|
111 |
+
|
112 |
+
const auto size = m_image.size();
|
113 |
+
m_image_grady = cv::Mat(size, CV_8UC3);
|
114 |
+
m_image_gradx = cv::Mat(size, CV_8UC3);
|
115 |
+
m_image_grady = cv::Scalar::all(0);
|
116 |
+
m_image_gradx = cv::Scalar::all(0);
|
117 |
+
|
118 |
+
for (int i = 1; i < size.height - 1; ++i) {
|
119 |
+
const auto *ptr = m_image.ptr<unsigned char>(i, 0);
|
120 |
+
const auto *ptry1 = m_image.ptr<unsigned char>(i + 1, 0);
|
121 |
+
const auto *ptry2 = m_image.ptr<unsigned char>(i - 1, 0);
|
122 |
+
const auto *ptrx1 = m_image.ptr<unsigned char>(i, 0) + 3;
|
123 |
+
const auto *ptrx2 = m_image.ptr<unsigned char>(i, 0) - 3;
|
124 |
+
auto *mptry = m_image_grady.ptr<unsigned char>(i, 0);
|
125 |
+
auto *mptrx = m_image_gradx.ptr<unsigned char>(i, 0);
|
126 |
+
for (int j = 3; j < size.width * 3 - 3; ++j) {
|
127 |
+
mptry[j] = (ptry1[j] / 2 - ptry2[j] / 2) + 128;
|
128 |
+
mptrx[j] = (ptrx1[j] / 2 - ptrx2[j] / 2) + 128;
|
129 |
+
}
|
130 |
+
}
|
131 |
+
|
132 |
+
m_image_grad_computed = true;
|
133 |
+
}
|
134 |
+
|
135 |
+
void MaskedImage::compute_image_gradients() const {
|
136 |
+
const_cast<MaskedImage *>(this)->compute_image_gradients();
|
137 |
+
}
|
138 |
+
|
PyPatchMatch/csrc/masked_image.h
ADDED
@@ -0,0 +1,112 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
#pragma once
|
2 |
+
|
3 |
+
#include <opencv2/core.hpp>
|
4 |
+
|
5 |
+
class MaskedImage {
|
6 |
+
public:
|
7 |
+
MaskedImage() : m_image(), m_mask(), m_global_mask(), m_image_grady(), m_image_gradx(), m_image_grad_computed(false) {
|
8 |
+
// pass
|
9 |
+
}
|
10 |
+
MaskedImage(cv::Mat image, cv::Mat mask) : m_image(image), m_mask(mask), m_image_grad_computed(false) {
|
11 |
+
// pass
|
12 |
+
}
|
13 |
+
MaskedImage(cv::Mat image, cv::Mat mask, cv::Mat global_mask) : m_image(image), m_mask(mask), m_global_mask(global_mask), m_image_grad_computed(false) {
|
14 |
+
// pass
|
15 |
+
}
|
16 |
+
MaskedImage(cv::Mat image, cv::Mat mask, cv::Mat global_mask, cv::Mat grady, cv::Mat gradx, bool grad_computed) :
|
17 |
+
m_image(image), m_mask(mask), m_global_mask(global_mask),
|
18 |
+
m_image_grady(grady), m_image_gradx(gradx), m_image_grad_computed(grad_computed) {
|
19 |
+
// pass
|
20 |
+
}
|
21 |
+
MaskedImage(int width, int height) : m_global_mask(), m_image_grady(), m_image_gradx() {
|
22 |
+
m_image = cv::Mat(cv::Size(width, height), CV_8UC3);
|
23 |
+
m_image = cv::Scalar::all(0);
|
24 |
+
|
25 |
+
m_mask = cv::Mat(cv::Size(width, height), CV_8U);
|
26 |
+
m_mask = cv::Scalar::all(0);
|
27 |
+
}
|
28 |
+
inline MaskedImage clone() {
|
29 |
+
return MaskedImage(
|
30 |
+
m_image.clone(), m_mask.clone(), m_global_mask.clone(),
|
31 |
+
m_image_grady.clone(), m_image_gradx.clone(), m_image_grad_computed
|
32 |
+
);
|
33 |
+
}
|
34 |
+
|
35 |
+
inline cv::Size size() const {
|
36 |
+
return m_image.size();
|
37 |
+
}
|
38 |
+
inline const cv::Mat &image() const {
|
39 |
+
return m_image;
|
40 |
+
}
|
41 |
+
inline const cv::Mat &mask() const {
|
42 |
+
return m_mask;
|
43 |
+
}
|
44 |
+
inline const cv::Mat &global_mask() const {
|
45 |
+
return m_global_mask;
|
46 |
+
}
|
47 |
+
inline const cv::Mat &grady() const {
|
48 |
+
assert(m_image_grad_computed);
|
49 |
+
return m_image_grady;
|
50 |
+
}
|
51 |
+
inline const cv::Mat &gradx() const {
|
52 |
+
assert(m_image_grad_computed);
|
53 |
+
return m_image_gradx;
|
54 |
+
}
|
55 |
+
|
56 |
+
inline void init_global_mask_mat() {
|
57 |
+
m_global_mask = cv::Mat(m_mask.size(), CV_8U);
|
58 |
+
m_global_mask.setTo(cv::Scalar(0));
|
59 |
+
}
|
60 |
+
inline void set_global_mask_mat(const cv::Mat &other) {
|
61 |
+
m_global_mask = other;
|
62 |
+
}
|
63 |
+
|
64 |
+
inline bool is_masked(int y, int x) const {
|
65 |
+
return static_cast<bool>(m_mask.at<unsigned char>(y, x));
|
66 |
+
}
|
67 |
+
inline bool is_globally_masked(int y, int x) const {
|
68 |
+
return !m_global_mask.empty() && static_cast<bool>(m_global_mask.at<unsigned char>(y, x));
|
69 |
+
}
|
70 |
+
inline void set_mask(int y, int x, bool value) {
|
71 |
+
m_mask.at<unsigned char>(y, x) = static_cast<unsigned char>(value);
|
72 |
+
}
|
73 |
+
inline void set_global_mask(int y, int x, bool value) {
|
74 |
+
m_global_mask.at<unsigned char>(y, x) = static_cast<unsigned char>(value);
|
75 |
+
}
|
76 |
+
inline void clear_mask() {
|
77 |
+
m_mask.setTo(cv::Scalar(0));
|
78 |
+
}
|
79 |
+
|
80 |
+
inline const unsigned char *get_image(int y, int x) const {
|
81 |
+
return m_image.ptr<unsigned char>(y, x);
|
82 |
+
}
|
83 |
+
inline unsigned char *get_mutable_image(int y, int x) {
|
84 |
+
return m_image.ptr<unsigned char>(y, x);
|
85 |
+
}
|
86 |
+
|
87 |
+
inline unsigned char get_image(int y, int x, int c) const {
|
88 |
+
return m_image.ptr<unsigned char>(y, x)[c];
|
89 |
+
}
|
90 |
+
inline int get_image_int(int y, int x, int c) const {
|
91 |
+
return static_cast<int>(m_image.ptr<unsigned char>(y, x)[c]);
|
92 |
+
}
|
93 |
+
|
94 |
+
bool contains_mask(int y, int x, int patch_size) const;
|
95 |
+
MaskedImage downsample() const;
|
96 |
+
MaskedImage upsample(int new_w, int new_h) const;
|
97 |
+
MaskedImage upsample(int new_w, int new_h, const cv::Mat &new_global_mask) const;
|
98 |
+
void compute_image_gradients();
|
99 |
+
void compute_image_gradients() const;
|
100 |
+
|
101 |
+
static const cv::Size kDownsampleKernelSize;
|
102 |
+
static const int kDownsampleKernel[6];
|
103 |
+
|
104 |
+
private:
|
105 |
+
cv::Mat m_image;
|
106 |
+
cv::Mat m_mask;
|
107 |
+
cv::Mat m_global_mask;
|
108 |
+
cv::Mat m_image_grady;
|
109 |
+
cv::Mat m_image_gradx;
|
110 |
+
bool m_image_grad_computed = false;
|
111 |
+
};
|
112 |
+
|
PyPatchMatch/csrc/nnf.cpp
ADDED
@@ -0,0 +1,268 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
#include <algorithm>
|
2 |
+
#include <iostream>
|
3 |
+
#include <cmath>
|
4 |
+
|
5 |
+
#include "masked_image.h"
|
6 |
+
#include "nnf.h"
|
7 |
+
|
8 |
+
/**
|
9 |
+
* Nearest-Neighbor Field (see PatchMatch algorithm).
|
10 |
+
* This algorithme uses a version proposed by Xavier Philippeau.
|
11 |
+
*
|
12 |
+
*/
|
13 |
+
|
14 |
+
template <typename T>
|
15 |
+
T clamp(T value, T min_value, T max_value) {
|
16 |
+
return std::min(std::max(value, min_value), max_value);
|
17 |
+
}
|
18 |
+
|
19 |
+
void NearestNeighborField::_randomize_field(int max_retry, bool reset) {
|
20 |
+
auto this_size = source_size();
|
21 |
+
for (int i = 0; i < this_size.height; ++i) {
|
22 |
+
for (int j = 0; j < this_size.width; ++j) {
|
23 |
+
if (m_source.is_globally_masked(i, j)) continue;
|
24 |
+
|
25 |
+
auto this_ptr = mutable_ptr(i, j);
|
26 |
+
int distance = reset ? PatchDistanceMetric::kDistanceScale : this_ptr[2];
|
27 |
+
if (distance < PatchDistanceMetric::kDistanceScale) {
|
28 |
+
continue;
|
29 |
+
}
|
30 |
+
|
31 |
+
int i_target = 0, j_target = 0;
|
32 |
+
for (int t = 0; t < max_retry; ++t) {
|
33 |
+
i_target = rand() % this_size.height;
|
34 |
+
j_target = rand() % this_size.width;
|
35 |
+
if (m_target.is_globally_masked(i_target, j_target)) continue;
|
36 |
+
|
37 |
+
distance = _distance(i, j, i_target, j_target);
|
38 |
+
if (distance < PatchDistanceMetric::kDistanceScale)
|
39 |
+
break;
|
40 |
+
}
|
41 |
+
|
42 |
+
this_ptr[0] = i_target, this_ptr[1] = j_target, this_ptr[2] = distance;
|
43 |
+
}
|
44 |
+
}
|
45 |
+
}
|
46 |
+
|
47 |
+
void NearestNeighborField::_initialize_field_from(const NearestNeighborField &other, int max_retry) {
|
48 |
+
const auto &this_size = source_size();
|
49 |
+
const auto &other_size = other.source_size();
|
50 |
+
double fi = static_cast<double>(this_size.height) / other_size.height;
|
51 |
+
double fj = static_cast<double>(this_size.width) / other_size.width;
|
52 |
+
|
53 |
+
for (int i = 0; i < this_size.height; ++i) {
|
54 |
+
for (int j = 0; j < this_size.width; ++j) {
|
55 |
+
if (m_source.is_globally_masked(i, j)) continue;
|
56 |
+
|
57 |
+
int ilow = static_cast<int>(std::min(i / fi, static_cast<double>(other_size.height - 1)));
|
58 |
+
int jlow = static_cast<int>(std::min(j / fj, static_cast<double>(other_size.width - 1)));
|
59 |
+
auto this_value = mutable_ptr(i, j);
|
60 |
+
auto other_value = other.ptr(ilow, jlow);
|
61 |
+
|
62 |
+
this_value[0] = static_cast<int>(other_value[0] * fi);
|
63 |
+
this_value[1] = static_cast<int>(other_value[1] * fj);
|
64 |
+
this_value[2] = _distance(i, j, this_value[0], this_value[1]);
|
65 |
+
}
|
66 |
+
}
|
67 |
+
|
68 |
+
_randomize_field(max_retry, false);
|
69 |
+
}
|
70 |
+
|
71 |
+
void NearestNeighborField::minimize(int nr_pass) {
|
72 |
+
const auto &this_size = source_size();
|
73 |
+
while (nr_pass--) {
|
74 |
+
for (int i = 0; i < this_size.height; ++i)
|
75 |
+
for (int j = 0; j < this_size.width; ++j) {
|
76 |
+
if (m_source.is_globally_masked(i, j)) continue;
|
77 |
+
if (at(i, j, 2) > 0) _minimize_link(i, j, +1);
|
78 |
+
}
|
79 |
+
for (int i = this_size.height - 1; i >= 0; --i)
|
80 |
+
for (int j = this_size.width - 1; j >= 0; --j) {
|
81 |
+
if (m_source.is_globally_masked(i, j)) continue;
|
82 |
+
if (at(i, j, 2) > 0) _minimize_link(i, j, -1);
|
83 |
+
}
|
84 |
+
}
|
85 |
+
}
|
86 |
+
|
87 |
+
void NearestNeighborField::_minimize_link(int y, int x, int direction) {
|
88 |
+
const auto &this_size = source_size();
|
89 |
+
const auto &this_target_size = target_size();
|
90 |
+
auto this_ptr = mutable_ptr(y, x);
|
91 |
+
|
92 |
+
// propagation along the y direction.
|
93 |
+
if (y - direction >= 0 && y - direction < this_size.height && !m_source.is_globally_masked(y - direction, x)) {
|
94 |
+
int yp = at(y - direction, x, 0) + direction;
|
95 |
+
int xp = at(y - direction, x, 1);
|
96 |
+
int dp = _distance(y, x, yp, xp);
|
97 |
+
if (dp < at(y, x, 2)) {
|
98 |
+
this_ptr[0] = yp, this_ptr[1] = xp, this_ptr[2] = dp;
|
99 |
+
}
|
100 |
+
}
|
101 |
+
|
102 |
+
// propagation along the x direction.
|
103 |
+
if (x - direction >= 0 && x - direction < this_size.width && !m_source.is_globally_masked(y, x - direction)) {
|
104 |
+
int yp = at(y, x - direction, 0);
|
105 |
+
int xp = at(y, x - direction, 1) + direction;
|
106 |
+
int dp = _distance(y, x, yp, xp);
|
107 |
+
if (dp < at(y, x, 2)) {
|
108 |
+
this_ptr[0] = yp, this_ptr[1] = xp, this_ptr[2] = dp;
|
109 |
+
}
|
110 |
+
}
|
111 |
+
|
112 |
+
// random search with a progressive step size.
|
113 |
+
int random_scale = (std::min(this_target_size.height, this_target_size.width) - 1) / 2;
|
114 |
+
while (random_scale > 0) {
|
115 |
+
int yp = this_ptr[0] + (rand() % (2 * random_scale + 1) - random_scale);
|
116 |
+
int xp = this_ptr[1] + (rand() % (2 * random_scale + 1) - random_scale);
|
117 |
+
yp = clamp(yp, 0, target_size().height - 1);
|
118 |
+
xp = clamp(xp, 0, target_size().width - 1);
|
119 |
+
|
120 |
+
if (m_target.is_globally_masked(yp, xp)) {
|
121 |
+
random_scale /= 2;
|
122 |
+
}
|
123 |
+
|
124 |
+
int dp = _distance(y, x, yp, xp);
|
125 |
+
if (dp < at(y, x, 2)) {
|
126 |
+
this_ptr[0] = yp, this_ptr[1] = xp, this_ptr[2] = dp;
|
127 |
+
}
|
128 |
+
random_scale /= 2;
|
129 |
+
}
|
130 |
+
}
|
131 |
+
|
132 |
+
const int PatchDistanceMetric::kDistanceScale = 65535;
|
133 |
+
const int PatchSSDDistanceMetric::kSSDScale = 9 * 255 * 255;
|
134 |
+
|
135 |
+
namespace {
|
136 |
+
|
137 |
+
inline int pow2(int i) {
|
138 |
+
return i * i;
|
139 |
+
}
|
140 |
+
|
141 |
+
int distance_masked_images(
|
142 |
+
const MaskedImage &source, int ys, int xs,
|
143 |
+
const MaskedImage &target, int yt, int xt,
|
144 |
+
int patch_size
|
145 |
+
) {
|
146 |
+
long double distance = 0;
|
147 |
+
long double wsum = 0;
|
148 |
+
|
149 |
+
source.compute_image_gradients();
|
150 |
+
target.compute_image_gradients();
|
151 |
+
|
152 |
+
auto source_size = source.size();
|
153 |
+
auto target_size = target.size();
|
154 |
+
|
155 |
+
for (int dy = -patch_size; dy <= patch_size; ++dy) {
|
156 |
+
const int yys = ys + dy, yyt = yt + dy;
|
157 |
+
|
158 |
+
if (yys <= 0 || yys >= source_size.height - 1 || yyt <= 0 || yyt >= target_size.height - 1) {
|
159 |
+
distance += (long double)(PatchSSDDistanceMetric::kSSDScale) * (2 * patch_size + 1);
|
160 |
+
wsum += 2 * patch_size + 1;
|
161 |
+
continue;
|
162 |
+
}
|
163 |
+
|
164 |
+
const auto *p_si = source.image().ptr<unsigned char>(yys, 0);
|
165 |
+
const auto *p_ti = target.image().ptr<unsigned char>(yyt, 0);
|
166 |
+
const auto *p_sm = source.mask().ptr<unsigned char>(yys, 0);
|
167 |
+
const auto *p_tm = target.mask().ptr<unsigned char>(yyt, 0);
|
168 |
+
|
169 |
+
const unsigned char *p_sgm = nullptr;
|
170 |
+
const unsigned char *p_tgm = nullptr;
|
171 |
+
if (!source.global_mask().empty()) {
|
172 |
+
p_sgm = source.global_mask().ptr<unsigned char>(yys, 0);
|
173 |
+
p_tgm = target.global_mask().ptr<unsigned char>(yyt, 0);
|
174 |
+
}
|
175 |
+
|
176 |
+
const auto *p_sgy = source.grady().ptr<unsigned char>(yys, 0);
|
177 |
+
const auto *p_tgy = target.grady().ptr<unsigned char>(yyt, 0);
|
178 |
+
const auto *p_sgx = source.gradx().ptr<unsigned char>(yys, 0);
|
179 |
+
const auto *p_tgx = target.gradx().ptr<unsigned char>(yyt, 0);
|
180 |
+
|
181 |
+
for (int dx = -patch_size; dx <= patch_size; ++dx) {
|
182 |
+
int xxs = xs + dx, xxt = xt + dx;
|
183 |
+
wsum += 1;
|
184 |
+
|
185 |
+
if (xxs <= 0 || xxs >= source_size.width - 1 || xxt <= 0 || xxt >= source_size.width - 1) {
|
186 |
+
distance += PatchSSDDistanceMetric::kSSDScale;
|
187 |
+
continue;
|
188 |
+
}
|
189 |
+
|
190 |
+
if (p_sm[xxs] || p_tm[xxt] || (p_sgm && p_sgm[xxs]) || (p_tgm && p_tgm[xxt]) ) {
|
191 |
+
distance += PatchSSDDistanceMetric::kSSDScale;
|
192 |
+
continue;
|
193 |
+
}
|
194 |
+
|
195 |
+
int ssd = 0;
|
196 |
+
for (int c = 0; c < 3; ++c) {
|
197 |
+
int s_value = p_si[xxs * 3 + c];
|
198 |
+
int t_value = p_ti[xxt * 3 + c];
|
199 |
+
int s_gy = p_sgy[xxs * 3 + c];
|
200 |
+
int t_gy = p_tgy[xxt * 3 + c];
|
201 |
+
int s_gx = p_sgx[xxs * 3 + c];
|
202 |
+
int t_gx = p_tgx[xxt * 3 + c];
|
203 |
+
|
204 |
+
ssd += pow2(static_cast<int>(s_value) - t_value);
|
205 |
+
ssd += pow2(static_cast<int>(s_gx) - t_gx);
|
206 |
+
ssd += pow2(static_cast<int>(s_gy) - t_gy);
|
207 |
+
}
|
208 |
+
distance += ssd;
|
209 |
+
}
|
210 |
+
}
|
211 |
+
|
212 |
+
distance /= (long double)(PatchSSDDistanceMetric::kSSDScale);
|
213 |
+
|
214 |
+
int res = int(PatchDistanceMetric::kDistanceScale * distance / wsum);
|
215 |
+
if (res < 0 || res > PatchDistanceMetric::kDistanceScale) return PatchDistanceMetric::kDistanceScale;
|
216 |
+
return res;
|
217 |
+
}
|
218 |
+
|
219 |
+
}
|
220 |
+
|
221 |
+
int PatchSSDDistanceMetric::operator ()(const MaskedImage &source, int source_y, int source_x, const MaskedImage &target, int target_y, int target_x) const {
|
222 |
+
return distance_masked_images(source, source_y, source_x, target, target_y, target_x, m_patch_size);
|
223 |
+
}
|
224 |
+
|
225 |
+
int DebugPatchSSDDistanceMetric::operator ()(const MaskedImage &source, int source_y, int source_x, const MaskedImage &target, int target_y, int target_x) const {
|
226 |
+
fprintf(stderr, "DebugPatchSSDDistanceMetric: %d %d %d %d\n", source.size().width, source.size().height, m_width, m_height);
|
227 |
+
return distance_masked_images(source, source_y, source_x, target, target_y, target_x, m_patch_size);
|
228 |
+
}
|
229 |
+
|
230 |
+
int RegularityGuidedPatchDistanceMetricV1::operator ()(const MaskedImage &source, int source_y, int source_x, const MaskedImage &target, int target_y, int target_x) const {
|
231 |
+
double dx = remainder(double(source_x - target_x) / source.size().width, m_dx1);
|
232 |
+
double dy = remainder(double(source_y - target_y) / source.size().height, m_dy2);
|
233 |
+
|
234 |
+
double score1 = sqrt(dx * dx + dy *dy) / m_scale;
|
235 |
+
if (score1 < 0 || score1 > 1) score1 = 1;
|
236 |
+
score1 *= PatchDistanceMetric::kDistanceScale;
|
237 |
+
|
238 |
+
double score2 = distance_masked_images(source, source_y, source_x, target, target_y, target_x, m_patch_size);
|
239 |
+
double score = score1 * m_weight + score2 / (1 + m_weight);
|
240 |
+
return static_cast<int>(score / (1 + m_weight));
|
241 |
+
}
|
242 |
+
|
243 |
+
int RegularityGuidedPatchDistanceMetricV2::operator ()(const MaskedImage &source, int source_y, int source_x, const MaskedImage &target, int target_y, int target_x) const {
|
244 |
+
if (target_y < 0 || target_y >= target.size().height || target_x < 0 || target_x >= target.size().width)
|
245 |
+
return PatchDistanceMetric::kDistanceScale;
|
246 |
+
|
247 |
+
int source_scale = m_ijmap.size().height / source.size().height;
|
248 |
+
int target_scale = m_ijmap.size().height / target.size().height;
|
249 |
+
|
250 |
+
// fprintf(stderr, "RegularityGuidedPatchDistanceMetricV2 %d %d %d %d\n", source_y * source_scale, m_ijmap.size().height, source_x * source_scale, m_ijmap.size().width);
|
251 |
+
|
252 |
+
double score1 = PatchDistanceMetric::kDistanceScale;
|
253 |
+
if (!source.is_globally_masked(source_y, source_x) && !target.is_globally_masked(target_y, target_x)) {
|
254 |
+
auto source_ij = m_ijmap.ptr<float>(source_y * source_scale, source_x * source_scale);
|
255 |
+
auto target_ij = m_ijmap.ptr<float>(target_y * target_scale, target_x * target_scale);
|
256 |
+
|
257 |
+
float di = fabs(source_ij[0] - target_ij[0]); if (di > 0.5) di = 1 - di;
|
258 |
+
float dj = fabs(source_ij[1] - target_ij[1]); if (dj > 0.5) dj = 1 - dj;
|
259 |
+
score1 = sqrt(di * di + dj *dj) / 0.707;
|
260 |
+
if (score1 < 0 || score1 > 1) score1 = 1;
|
261 |
+
score1 *= PatchDistanceMetric::kDistanceScale;
|
262 |
+
}
|
263 |
+
|
264 |
+
double score2 = distance_masked_images(source, source_y, source_x, target, target_y, target_x, m_patch_size);
|
265 |
+
double score = score1 * m_weight + score2;
|
266 |
+
return int(score / (1 + m_weight));
|
267 |
+
}
|
268 |
+
|
PyPatchMatch/csrc/nnf.h
ADDED
@@ -0,0 +1,133 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
#pragma once
|
2 |
+
|
3 |
+
#include <opencv2/core.hpp>
|
4 |
+
#include "masked_image.h"
|
5 |
+
|
6 |
+
class PatchDistanceMetric {
|
7 |
+
public:
|
8 |
+
PatchDistanceMetric(int patch_size) : m_patch_size(patch_size) {}
|
9 |
+
virtual ~PatchDistanceMetric() = default;
|
10 |
+
|
11 |
+
inline int patch_size() const { return m_patch_size; }
|
12 |
+
virtual int operator()(const MaskedImage &source, int source_y, int source_x, const MaskedImage &target, int target_y, int target_x) const = 0;
|
13 |
+
static const int kDistanceScale;
|
14 |
+
|
15 |
+
protected:
|
16 |
+
int m_patch_size;
|
17 |
+
};
|
18 |
+
|
19 |
+
class NearestNeighborField {
|
20 |
+
public:
|
21 |
+
NearestNeighborField() : m_source(), m_target(), m_field(), m_distance_metric(nullptr) {
|
22 |
+
// pass
|
23 |
+
}
|
24 |
+
NearestNeighborField(const MaskedImage &source, const MaskedImage &target, const PatchDistanceMetric *metric, int max_retry = 20)
|
25 |
+
: m_source(source), m_target(target), m_distance_metric(metric) {
|
26 |
+
m_field = cv::Mat(m_source.size(), CV_32SC3);
|
27 |
+
_randomize_field(max_retry);
|
28 |
+
}
|
29 |
+
NearestNeighborField(const MaskedImage &source, const MaskedImage &target, const PatchDistanceMetric *metric, const NearestNeighborField &other, int max_retry = 20)
|
30 |
+
: m_source(source), m_target(target), m_distance_metric(metric) {
|
31 |
+
m_field = cv::Mat(m_source.size(), CV_32SC3);
|
32 |
+
_initialize_field_from(other, max_retry);
|
33 |
+
}
|
34 |
+
|
35 |
+
const MaskedImage &source() const {
|
36 |
+
return m_source;
|
37 |
+
}
|
38 |
+
const MaskedImage &target() const {
|
39 |
+
return m_target;
|
40 |
+
}
|
41 |
+
inline cv::Size source_size() const {
|
42 |
+
return m_source.size();
|
43 |
+
}
|
44 |
+
inline cv::Size target_size() const {
|
45 |
+
return m_target.size();
|
46 |
+
}
|
47 |
+
inline void set_source(const MaskedImage &source) {
|
48 |
+
m_source = source;
|
49 |
+
}
|
50 |
+
inline void set_target(const MaskedImage &target) {
|
51 |
+
m_target = target;
|
52 |
+
}
|
53 |
+
|
54 |
+
inline int *mutable_ptr(int y, int x) {
|
55 |
+
return m_field.ptr<int>(y, x);
|
56 |
+
}
|
57 |
+
inline const int *ptr(int y, int x) const {
|
58 |
+
return m_field.ptr<int>(y, x);
|
59 |
+
}
|
60 |
+
|
61 |
+
inline int at(int y, int x, int c) const {
|
62 |
+
return m_field.ptr<int>(y, x)[c];
|
63 |
+
}
|
64 |
+
inline int &at(int y, int x, int c) {
|
65 |
+
return m_field.ptr<int>(y, x)[c];
|
66 |
+
}
|
67 |
+
inline void set_identity(int y, int x) {
|
68 |
+
auto ptr = mutable_ptr(y, x);
|
69 |
+
ptr[0] = y, ptr[1] = x, ptr[2] = 0;
|
70 |
+
}
|
71 |
+
|
72 |
+
void minimize(int nr_pass);
|
73 |
+
|
74 |
+
private:
|
75 |
+
inline int _distance(int source_y, int source_x, int target_y, int target_x) {
|
76 |
+
return (*m_distance_metric)(m_source, source_y, source_x, m_target, target_y, target_x);
|
77 |
+
}
|
78 |
+
|
79 |
+
void _randomize_field(int max_retry = 20, bool reset = true);
|
80 |
+
void _initialize_field_from(const NearestNeighborField &other, int max_retry);
|
81 |
+
void _minimize_link(int y, int x, int direction);
|
82 |
+
|
83 |
+
MaskedImage m_source;
|
84 |
+
MaskedImage m_target;
|
85 |
+
cv::Mat m_field; // { y_target, x_target, distance_scaled }
|
86 |
+
const PatchDistanceMetric *m_distance_metric;
|
87 |
+
};
|
88 |
+
|
89 |
+
|
90 |
+
class PatchSSDDistanceMetric : public PatchDistanceMetric {
|
91 |
+
public:
|
92 |
+
using PatchDistanceMetric::PatchDistanceMetric;
|
93 |
+
virtual int operator ()(const MaskedImage &source, int source_y, int source_x, const MaskedImage &target, int target_y, int target_x) const;
|
94 |
+
static const int kSSDScale;
|
95 |
+
};
|
96 |
+
|
97 |
+
class DebugPatchSSDDistanceMetric : public PatchDistanceMetric {
|
98 |
+
public:
|
99 |
+
DebugPatchSSDDistanceMetric(int patch_size, int width, int height) : PatchDistanceMetric(patch_size), m_width(width), m_height(height) {}
|
100 |
+
virtual int operator ()(const MaskedImage &source, int source_y, int source_x, const MaskedImage &target, int target_y, int target_x) const;
|
101 |
+
protected:
|
102 |
+
int m_width, m_height;
|
103 |
+
};
|
104 |
+
|
105 |
+
class RegularityGuidedPatchDistanceMetricV1 : public PatchDistanceMetric {
|
106 |
+
public:
|
107 |
+
RegularityGuidedPatchDistanceMetricV1(int patch_size, double dx1, double dy1, double dx2, double dy2, double weight)
|
108 |
+
: PatchDistanceMetric(patch_size), m_dx1(dx1), m_dy1(dy1), m_dx2(dx2), m_dy2(dy2), m_weight(weight) {
|
109 |
+
|
110 |
+
assert(m_dy1 == 0);
|
111 |
+
assert(m_dx2 == 0);
|
112 |
+
m_scale = sqrt(m_dx1 * m_dx1 + m_dy2 * m_dy2) / 4;
|
113 |
+
}
|
114 |
+
virtual int operator ()(const MaskedImage &source, int source_y, int source_x, const MaskedImage &target, int target_y, int target_x) const;
|
115 |
+
|
116 |
+
protected:
|
117 |
+
double m_dx1, m_dy1, m_dx2, m_dy2;
|
118 |
+
double m_scale, m_weight;
|
119 |
+
};
|
120 |
+
|
121 |
+
class RegularityGuidedPatchDistanceMetricV2 : public PatchDistanceMetric {
|
122 |
+
public:
|
123 |
+
RegularityGuidedPatchDistanceMetricV2(int patch_size, cv::Mat ijmap, double weight)
|
124 |
+
: PatchDistanceMetric(patch_size), m_ijmap(ijmap), m_weight(weight) {
|
125 |
+
|
126 |
+
}
|
127 |
+
virtual int operator ()(const MaskedImage &source, int source_y, int source_x, const MaskedImage &target, int target_y, int target_x) const;
|
128 |
+
|
129 |
+
protected:
|
130 |
+
cv::Mat m_ijmap;
|
131 |
+
double m_width, m_height, m_weight;
|
132 |
+
};
|
133 |
+
|
PyPatchMatch/csrc/pyinterface.cpp
ADDED
@@ -0,0 +1,107 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
#include "pyinterface.h"
|
2 |
+
#include "inpaint.h"
|
3 |
+
|
4 |
+
static unsigned int PM_seed = 1212;
|
5 |
+
static bool PM_verbose = false;
|
6 |
+
|
7 |
+
int _dtype_py_to_cv(int dtype_py);
|
8 |
+
int _dtype_cv_to_py(int dtype_cv);
|
9 |
+
cv::Mat _py_to_cv2(PM_mat_t pymat);
|
10 |
+
PM_mat_t _cv2_to_py(cv::Mat cvmat);
|
11 |
+
|
12 |
+
void PM_set_random_seed(unsigned int seed) {
|
13 |
+
PM_seed = seed;
|
14 |
+
}
|
15 |
+
|
16 |
+
void PM_set_verbose(int value) {
|
17 |
+
PM_verbose = static_cast<bool>(value);
|
18 |
+
}
|
19 |
+
|
20 |
+
void PM_free_pymat(PM_mat_t pymat) {
|
21 |
+
free(pymat.data_ptr);
|
22 |
+
}
|
23 |
+
|
24 |
+
PM_mat_t PM_inpaint(PM_mat_t source_py, PM_mat_t mask_py, int patch_size) {
|
25 |
+
cv::Mat source = _py_to_cv2(source_py);
|
26 |
+
cv::Mat mask = _py_to_cv2(mask_py);
|
27 |
+
auto metric = PatchSSDDistanceMetric(patch_size);
|
28 |
+
cv::Mat result = Inpainting(source, mask, &metric).run(PM_verbose, false, PM_seed);
|
29 |
+
return _cv2_to_py(result);
|
30 |
+
}
|
31 |
+
|
32 |
+
PM_mat_t PM_inpaint_regularity(PM_mat_t source_py, PM_mat_t mask_py, PM_mat_t ijmap_py, int patch_size, float guide_weight) {
|
33 |
+
cv::Mat source = _py_to_cv2(source_py);
|
34 |
+
cv::Mat mask = _py_to_cv2(mask_py);
|
35 |
+
cv::Mat ijmap = _py_to_cv2(ijmap_py);
|
36 |
+
|
37 |
+
auto metric = RegularityGuidedPatchDistanceMetricV2(patch_size, ijmap, guide_weight);
|
38 |
+
cv::Mat result = Inpainting(source, mask, &metric).run(PM_verbose, false, PM_seed);
|
39 |
+
return _cv2_to_py(result);
|
40 |
+
}
|
41 |
+
|
42 |
+
PM_mat_t PM_inpaint2(PM_mat_t source_py, PM_mat_t mask_py, PM_mat_t global_mask_py, int patch_size) {
|
43 |
+
cv::Mat source = _py_to_cv2(source_py);
|
44 |
+
cv::Mat mask = _py_to_cv2(mask_py);
|
45 |
+
cv::Mat global_mask = _py_to_cv2(global_mask_py);
|
46 |
+
|
47 |
+
auto metric = PatchSSDDistanceMetric(patch_size);
|
48 |
+
cv::Mat result = Inpainting(source, mask, global_mask, &metric).run(PM_verbose, false, PM_seed);
|
49 |
+
return _cv2_to_py(result);
|
50 |
+
}
|
51 |
+
|
52 |
+
PM_mat_t PM_inpaint2_regularity(PM_mat_t source_py, PM_mat_t mask_py, PM_mat_t global_mask_py, PM_mat_t ijmap_py, int patch_size, float guide_weight) {
|
53 |
+
cv::Mat source = _py_to_cv2(source_py);
|
54 |
+
cv::Mat mask = _py_to_cv2(mask_py);
|
55 |
+
cv::Mat global_mask = _py_to_cv2(global_mask_py);
|
56 |
+
cv::Mat ijmap = _py_to_cv2(ijmap_py);
|
57 |
+
|
58 |
+
auto metric = RegularityGuidedPatchDistanceMetricV2(patch_size, ijmap, guide_weight);
|
59 |
+
cv::Mat result = Inpainting(source, mask, global_mask, &metric).run(PM_verbose, false, PM_seed);
|
60 |
+
return _cv2_to_py(result);
|
61 |
+
}
|
62 |
+
|
63 |
+
int _dtype_py_to_cv(int dtype_py) {
|
64 |
+
switch (dtype_py) {
|
65 |
+
case PM_UINT8: return CV_8U;
|
66 |
+
case PM_INT8: return CV_8S;
|
67 |
+
case PM_UINT16: return CV_16U;
|
68 |
+
case PM_INT16: return CV_16S;
|
69 |
+
case PM_INT32: return CV_32S;
|
70 |
+
case PM_FLOAT32: return CV_32F;
|
71 |
+
case PM_FLOAT64: return CV_64F;
|
72 |
+
}
|
73 |
+
|
74 |
+
return CV_8U;
|
75 |
+
}
|
76 |
+
|
77 |
+
int _dtype_cv_to_py(int dtype_cv) {
|
78 |
+
switch (dtype_cv) {
|
79 |
+
case CV_8U: return PM_UINT8;
|
80 |
+
case CV_8S: return PM_INT8;
|
81 |
+
case CV_16U: return PM_UINT16;
|
82 |
+
case CV_16S: return PM_INT16;
|
83 |
+
case CV_32S: return PM_INT32;
|
84 |
+
case CV_32F: return PM_FLOAT32;
|
85 |
+
case CV_64F: return PM_FLOAT64;
|
86 |
+
}
|
87 |
+
|
88 |
+
return PM_UINT8;
|
89 |
+
}
|
90 |
+
|
91 |
+
cv::Mat _py_to_cv2(PM_mat_t pymat) {
|
92 |
+
int dtype = _dtype_py_to_cv(pymat.dtype);
|
93 |
+
dtype = CV_MAKETYPE(pymat.dtype, pymat.shape.channels);
|
94 |
+
return cv::Mat(cv::Size(pymat.shape.width, pymat.shape.height), dtype, pymat.data_ptr).clone();
|
95 |
+
}
|
96 |
+
|
97 |
+
PM_mat_t _cv2_to_py(cv::Mat cvmat) {
|
98 |
+
PM_shape_t shape = {cvmat.size().width, cvmat.size().height, cvmat.channels()};
|
99 |
+
int dtype = _dtype_cv_to_py(cvmat.depth());
|
100 |
+
size_t dsize = cvmat.total() * cvmat.elemSize();
|
101 |
+
|
102 |
+
void *data_ptr = reinterpret_cast<void *>(malloc(dsize));
|
103 |
+
memcpy(data_ptr, reinterpret_cast<void *>(cvmat.data), dsize);
|
104 |
+
|
105 |
+
return PM_mat_t {data_ptr, shape, dtype};
|
106 |
+
}
|
107 |
+
|
PyPatchMatch/csrc/pyinterface.h
ADDED
@@ -0,0 +1,38 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
#include <opencv2/core.hpp>
|
2 |
+
#include <cstdlib>
|
3 |
+
#include <cstdio>
|
4 |
+
#include <cstring>
|
5 |
+
|
6 |
+
extern "C" {
|
7 |
+
|
8 |
+
struct PM_shape_t {
|
9 |
+
int width, height, channels;
|
10 |
+
};
|
11 |
+
|
12 |
+
enum PM_dtype_e {
|
13 |
+
PM_UINT8,
|
14 |
+
PM_INT8,
|
15 |
+
PM_UINT16,
|
16 |
+
PM_INT16,
|
17 |
+
PM_INT32,
|
18 |
+
PM_FLOAT32,
|
19 |
+
PM_FLOAT64,
|
20 |
+
};
|
21 |
+
|
22 |
+
struct PM_mat_t {
|
23 |
+
void *data_ptr;
|
24 |
+
PM_shape_t shape;
|
25 |
+
int dtype;
|
26 |
+
};
|
27 |
+
|
28 |
+
void PM_set_random_seed(unsigned int seed);
|
29 |
+
void PM_set_verbose(int value);
|
30 |
+
|
31 |
+
void PM_free_pymat(PM_mat_t pymat);
|
32 |
+
PM_mat_t PM_inpaint(PM_mat_t image, PM_mat_t mask, int patch_size);
|
33 |
+
PM_mat_t PM_inpaint_regularity(PM_mat_t image, PM_mat_t mask, PM_mat_t ijmap, int patch_size, float guide_weight);
|
34 |
+
PM_mat_t PM_inpaint2(PM_mat_t image, PM_mat_t mask, PM_mat_t global_mask, int patch_size);
|
35 |
+
PM_mat_t PM_inpaint2_regularity(PM_mat_t image, PM_mat_t mask, PM_mat_t global_mask, PM_mat_t ijmap, int patch_size, float guide_weight);
|
36 |
+
|
37 |
+
} /* extern "C" */
|
38 |
+
|
PyPatchMatch/examples/.gitignore
ADDED
@@ -0,0 +1,2 @@
|
|
|
|
|
|
|
1 |
+
/cpp_example.exe
|
2 |
+
/images/*recovered.bmp
|
PyPatchMatch/examples/cpp_example.cpp
ADDED
@@ -0,0 +1,31 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
#include <iostream>
|
2 |
+
#include <opencv2/imgcodecs.hpp>
|
3 |
+
#include <opencv2/highgui.hpp>
|
4 |
+
|
5 |
+
#include "masked_image.h"
|
6 |
+
#include "nnf.h"
|
7 |
+
#include "inpaint.h"
|
8 |
+
|
9 |
+
int main() {
|
10 |
+
auto source = cv::imread("./images/forest_pruned.bmp", cv::IMREAD_COLOR);
|
11 |
+
|
12 |
+
auto mask = cv::Mat(source.size(), CV_8UC1);
|
13 |
+
mask = cv::Scalar::all(0);
|
14 |
+
for (int i = 0; i < source.size().height; ++i) {
|
15 |
+
for (int j = 0; j < source.size().width; ++j) {
|
16 |
+
auto source_ptr = source.ptr<unsigned char>(i, j);
|
17 |
+
if (source_ptr[0] == 255 && source_ptr[1] == 255 && source_ptr[2] == 255) {
|
18 |
+
mask.at<unsigned char>(i, j) = 1;
|
19 |
+
}
|
20 |
+
}
|
21 |
+
}
|
22 |
+
|
23 |
+
auto metric = PatchSSDDistanceMetric(3);
|
24 |
+
auto result = Inpainting(source, mask, &metric).run(true, true);
|
25 |
+
// cv::imwrite("./images/forest_recovered.bmp", result);
|
26 |
+
// cv::imshow("Result", result);
|
27 |
+
// cv::waitKey();
|
28 |
+
|
29 |
+
return 0;
|
30 |
+
}
|
31 |
+
|
PyPatchMatch/examples/cpp_example_run.sh
ADDED
@@ -0,0 +1,18 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
#! /bin/bash
|
2 |
+
#
|
3 |
+
# cpp_example_run.sh
|
4 |
+
# Copyright (C) 2020 Jiayuan Mao <[email protected]>
|
5 |
+
#
|
6 |
+
# Distributed under terms of the MIT license.
|
7 |
+
#
|
8 |
+
|
9 |
+
set -x
|
10 |
+
|
11 |
+
CFLAGS="-std=c++14 -O2 $(pkg-config --cflags opencv)"
|
12 |
+
LDFLAGS="$(pkg-config --libs opencv)"
|
13 |
+
g++ $CFLAGS cpp_example.cpp -I../csrc/ -L../ -lpatchmatch $LDFLAGS -o cpp_example.exe
|
14 |
+
|
15 |
+
export DYLD_LIBRARY_PATH=../:$DYLD_LIBRARY_PATH # For macOS
|
16 |
+
export LD_LIBRARY_PATH=../:$LD_LIBRARY_PATH # For Linux
|
17 |
+
time ./cpp_example.exe
|
18 |
+
|
PyPatchMatch/examples/images/forest.bmp
ADDED
PyPatchMatch/examples/images/forest_pruned.bmp
ADDED
PyPatchMatch/examples/py_example.py
ADDED
@@ -0,0 +1,21 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
#! /usr/bin/env python3
|
2 |
+
# -*- coding: utf-8 -*-
|
3 |
+
# File : test.py
|
4 |
+
# Author : Jiayuan Mao
|
5 |
+
# Email : [email protected]
|
6 |
+
# Date : 01/09/2020
|
7 |
+
#
|
8 |
+
# Distributed under terms of the MIT license.
|
9 |
+
|
10 |
+
from PIL import Image
|
11 |
+
|
12 |
+
import sys
|
13 |
+
sys.path.insert(0, '../')
|
14 |
+
import patch_match
|
15 |
+
|
16 |
+
|
17 |
+
if __name__ == '__main__':
|
18 |
+
source = Image.open('./images/forest_pruned.bmp')
|
19 |
+
result = patch_match.inpaint(source, patch_size=3)
|
20 |
+
Image.fromarray(result).save('./images/forest_recovered.bmp')
|
21 |
+
|
PyPatchMatch/examples/py_example_global_mask.py
ADDED
@@ -0,0 +1,27 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
#! /usr/bin/env python3
|
2 |
+
# -*- coding: utf-8 -*-
|
3 |
+
# File : test.py
|
4 |
+
# Author : Jiayuan Mao
|
5 |
+
# Email : [email protected]
|
6 |
+
# Date : 01/09/2020
|
7 |
+
#
|
8 |
+
# Distributed under terms of the MIT license.
|
9 |
+
|
10 |
+
import numpy as np
|
11 |
+
from PIL import Image
|
12 |
+
|
13 |
+
import sys
|
14 |
+
sys.path.insert(0, '../')
|
15 |
+
import patch_match
|
16 |
+
|
17 |
+
|
18 |
+
if __name__ == '__main__':
|
19 |
+
patch_match.set_verbose(True)
|
20 |
+
source = Image.open('./images/forest_pruned.bmp')
|
21 |
+
source = np.array(source)
|
22 |
+
source[:100, :100] = 255
|
23 |
+
global_mask = np.zeros_like(source[..., 0])
|
24 |
+
global_mask[:100, :100] = 1
|
25 |
+
result = patch_match.inpaint(source, global_mask=global_mask, patch_size=3)
|
26 |
+
Image.fromarray(result).save('./images/forest_recovered.bmp')
|
27 |
+
|
PyPatchMatch/patch_match.py
ADDED
@@ -0,0 +1,201 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
#! /usr/bin/env python3
|
2 |
+
# -*- coding: utf-8 -*-
|
3 |
+
# File : patch_match.py
|
4 |
+
# Author : Jiayuan Mao
|
5 |
+
# Email : [email protected]
|
6 |
+
# Date : 01/09/2020
|
7 |
+
#
|
8 |
+
# Distributed under terms of the MIT license.
|
9 |
+
|
10 |
+
import ctypes
|
11 |
+
import os.path as osp
|
12 |
+
from typing import Optional, Union
|
13 |
+
|
14 |
+
import numpy as np
|
15 |
+
from PIL import Image
|
16 |
+
|
17 |
+
try:
|
18 |
+
# If the Jacinle library (https://github.com/vacancy/Jacinle) is present, use its auto_travis feature.
|
19 |
+
from jacinle.jit.cext import auto_travis
|
20 |
+
auto_travis(__file__, required_files=['*.so'])
|
21 |
+
except ImportError as e:
|
22 |
+
# Otherwise, fall back to the subprocess.
|
23 |
+
import subprocess
|
24 |
+
print('Compiling and loading c extensions from "{}".'.format(osp.realpath(osp.dirname(__file__))))
|
25 |
+
subprocess.check_call(['./travis.sh'], cwd=osp.dirname(__file__))
|
26 |
+
|
27 |
+
|
28 |
+
__all__ = ['set_random_seed', 'set_verbose', 'inpaint', 'inpaint_regularity']
|
29 |
+
|
30 |
+
|
31 |
+
class CShapeT(ctypes.Structure):
|
32 |
+
_fields_ = [
|
33 |
+
('width', ctypes.c_int),
|
34 |
+
('height', ctypes.c_int),
|
35 |
+
('channels', ctypes.c_int),
|
36 |
+
]
|
37 |
+
|
38 |
+
|
39 |
+
class CMatT(ctypes.Structure):
|
40 |
+
_fields_ = [
|
41 |
+
('data_ptr', ctypes.c_void_p),
|
42 |
+
('shape', CShapeT),
|
43 |
+
('dtype', ctypes.c_int)
|
44 |
+
]
|
45 |
+
|
46 |
+
|
47 |
+
PMLIB = ctypes.CDLL(osp.join(osp.dirname(__file__), 'libpatchmatch.so'))
|
48 |
+
|
49 |
+
PMLIB.PM_set_random_seed.argtypes = [ctypes.c_uint]
|
50 |
+
PMLIB.PM_set_verbose.argtypes = [ctypes.c_int]
|
51 |
+
PMLIB.PM_free_pymat.argtypes = [CMatT]
|
52 |
+
PMLIB.PM_inpaint.argtypes = [CMatT, CMatT, ctypes.c_int]
|
53 |
+
PMLIB.PM_inpaint.restype = CMatT
|
54 |
+
PMLIB.PM_inpaint_regularity.argtypes = [CMatT, CMatT, CMatT, ctypes.c_int, ctypes.c_float]
|
55 |
+
PMLIB.PM_inpaint_regularity.restype = CMatT
|
56 |
+
PMLIB.PM_inpaint2.argtypes = [CMatT, CMatT, CMatT, ctypes.c_int]
|
57 |
+
PMLIB.PM_inpaint2.restype = CMatT
|
58 |
+
PMLIB.PM_inpaint2_regularity.argtypes = [CMatT, CMatT, CMatT, CMatT, ctypes.c_int, ctypes.c_float]
|
59 |
+
PMLIB.PM_inpaint2_regularity.restype = CMatT
|
60 |
+
|
61 |
+
|
62 |
+
def set_random_seed(seed: int):
|
63 |
+
PMLIB.PM_set_random_seed(ctypes.c_uint(seed))
|
64 |
+
|
65 |
+
|
66 |
+
def set_verbose(verbose: bool):
|
67 |
+
PMLIB.PM_set_verbose(ctypes.c_int(verbose))
|
68 |
+
|
69 |
+
|
70 |
+
def inpaint(
|
71 |
+
image: Union[np.ndarray, Image.Image],
|
72 |
+
mask: Optional[Union[np.ndarray, Image.Image]] = None,
|
73 |
+
*,
|
74 |
+
global_mask: Optional[Union[np.ndarray, Image.Image]] = None,
|
75 |
+
patch_size: int = 15
|
76 |
+
) -> np.ndarray:
|
77 |
+
"""
|
78 |
+
PatchMatch based inpainting proposed in:
|
79 |
+
|
80 |
+
PatchMatch : A Randomized Correspondence Algorithm for Structural Image Editing
|
81 |
+
C.Barnes, E.Shechtman, A.Finkelstein and Dan B.Goldman
|
82 |
+
SIGGRAPH 2009
|
83 |
+
|
84 |
+
Args:
|
85 |
+
image (Union[np.ndarray, Image.Image]): the input image, should be 3-channel RGB/BGR.
|
86 |
+
mask (Union[np.array, Image.Image], optional): the mask of the hole(s) to be filled, should be 1-channel.
|
87 |
+
If not provided (None), the algorithm will treat all purely white pixels as the holes (255, 255, 255).
|
88 |
+
global_mask (Union[np.array, Image.Image], optional): the target mask of the output image.
|
89 |
+
patch_size (int): the patch size for the inpainting algorithm.
|
90 |
+
|
91 |
+
Return:
|
92 |
+
result (np.ndarray): the repaired image, of the same size as the input image.
|
93 |
+
"""
|
94 |
+
|
95 |
+
if isinstance(image, Image.Image):
|
96 |
+
image = np.array(image)
|
97 |
+
image = np.ascontiguousarray(image)
|
98 |
+
assert image.ndim == 3 and image.shape[2] == 3 and image.dtype == 'uint8'
|
99 |
+
|
100 |
+
if mask is None:
|
101 |
+
mask = (image == (255, 255, 255)).all(axis=2, keepdims=True).astype('uint8')
|
102 |
+
mask = np.ascontiguousarray(mask)
|
103 |
+
else:
|
104 |
+
mask = _canonize_mask_array(mask)
|
105 |
+
|
106 |
+
if global_mask is None:
|
107 |
+
ret_pymat = PMLIB.PM_inpaint(np_to_pymat(image), np_to_pymat(mask), ctypes.c_int(patch_size))
|
108 |
+
else:
|
109 |
+
global_mask = _canonize_mask_array(global_mask)
|
110 |
+
ret_pymat = PMLIB.PM_inpaint2(np_to_pymat(image), np_to_pymat(mask), np_to_pymat(global_mask), ctypes.c_int(patch_size))
|
111 |
+
|
112 |
+
ret_npmat = pymat_to_np(ret_pymat)
|
113 |
+
PMLIB.PM_free_pymat(ret_pymat)
|
114 |
+
|
115 |
+
return ret_npmat
|
116 |
+
|
117 |
+
|
118 |
+
def inpaint_regularity(
|
119 |
+
image: Union[np.ndarray, Image.Image],
|
120 |
+
mask: Optional[Union[np.ndarray, Image.Image]],
|
121 |
+
ijmap: np.ndarray,
|
122 |
+
*,
|
123 |
+
global_mask: Optional[Union[np.ndarray, Image.Image]] = None,
|
124 |
+
patch_size: int = 15, guide_weight: float = 0.25
|
125 |
+
) -> np.ndarray:
|
126 |
+
if isinstance(image, Image.Image):
|
127 |
+
image = np.array(image)
|
128 |
+
image = np.ascontiguousarray(image)
|
129 |
+
|
130 |
+
assert isinstance(ijmap, np.ndarray) and ijmap.ndim == 3 and ijmap.shape[2] == 3 and ijmap.dtype == 'float32'
|
131 |
+
ijmap = np.ascontiguousarray(ijmap)
|
132 |
+
|
133 |
+
assert image.ndim == 3 and image.shape[2] == 3 and image.dtype == 'uint8'
|
134 |
+
if mask is None:
|
135 |
+
mask = (image == (255, 255, 255)).all(axis=2, keepdims=True).astype('uint8')
|
136 |
+
mask = np.ascontiguousarray(mask)
|
137 |
+
else:
|
138 |
+
mask = _canonize_mask_array(mask)
|
139 |
+
|
140 |
+
|
141 |
+
if global_mask is None:
|
142 |
+
ret_pymat = PMLIB.PM_inpaint_regularity(np_to_pymat(image), np_to_pymat(mask), np_to_pymat(ijmap), ctypes.c_int(patch_size), ctypes.c_float(guide_weight))
|
143 |
+
else:
|
144 |
+
global_mask = _canonize_mask_array(global_mask)
|
145 |
+
ret_pymat = PMLIB.PM_inpaint2_regularity(np_to_pymat(image), np_to_pymat(mask), np_to_pymat(global_mask), np_to_pymat(ijmap), ctypes.c_int(patch_size), ctypes.c_float(guide_weight))
|
146 |
+
|
147 |
+
ret_npmat = pymat_to_np(ret_pymat)
|
148 |
+
PMLIB.PM_free_pymat(ret_pymat)
|
149 |
+
|
150 |
+
return ret_npmat
|
151 |
+
|
152 |
+
|
153 |
+
def _canonize_mask_array(mask):
|
154 |
+
if isinstance(mask, Image.Image):
|
155 |
+
mask = np.array(mask)
|
156 |
+
if mask.ndim == 2 and mask.dtype == 'uint8':
|
157 |
+
mask = mask[..., np.newaxis]
|
158 |
+
assert mask.ndim == 3 and mask.shape[2] == 1 and mask.dtype == 'uint8'
|
159 |
+
return np.ascontiguousarray(mask)
|
160 |
+
|
161 |
+
|
162 |
+
dtype_pymat_to_ctypes = [
|
163 |
+
ctypes.c_uint8,
|
164 |
+
ctypes.c_int8,
|
165 |
+
ctypes.c_uint16,
|
166 |
+
ctypes.c_int16,
|
167 |
+
ctypes.c_int32,
|
168 |
+
ctypes.c_float,
|
169 |
+
ctypes.c_double,
|
170 |
+
]
|
171 |
+
|
172 |
+
|
173 |
+
dtype_np_to_pymat = {
|
174 |
+
'uint8': 0,
|
175 |
+
'int8': 1,
|
176 |
+
'uint16': 2,
|
177 |
+
'int16': 3,
|
178 |
+
'int32': 4,
|
179 |
+
'float32': 5,
|
180 |
+
'float64': 6,
|
181 |
+
}
|
182 |
+
|
183 |
+
|
184 |
+
def np_to_pymat(npmat):
|
185 |
+
assert npmat.ndim == 3
|
186 |
+
return CMatT(
|
187 |
+
ctypes.cast(npmat.ctypes.data, ctypes.c_void_p),
|
188 |
+
CShapeT(npmat.shape[1], npmat.shape[0], npmat.shape[2]),
|
189 |
+
dtype_np_to_pymat[str(npmat.dtype)]
|
190 |
+
)
|
191 |
+
|
192 |
+
|
193 |
+
def pymat_to_np(pymat):
|
194 |
+
npmat = np.ctypeslib.as_array(
|
195 |
+
ctypes.cast(pymat.data_ptr, ctypes.POINTER(dtype_pymat_to_ctypes[pymat.dtype])),
|
196 |
+
(pymat.shape.height, pymat.shape.width, pymat.shape.channels)
|
197 |
+
)
|
198 |
+
ret = np.empty(npmat.shape, npmat.dtype)
|
199 |
+
ret[:] = npmat
|
200 |
+
return ret
|
201 |
+
|
PyPatchMatch/travis.sh
ADDED
@@ -0,0 +1,9 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
#! /bin/bash
|
2 |
+
#
|
3 |
+
# travis.sh
|
4 |
+
# Copyright (C) 2020 Jiayuan Mao <[email protected]>
|
5 |
+
#
|
6 |
+
# Distributed under terms of the MIT license.
|
7 |
+
#
|
8 |
+
|
9 |
+
make clean && make
|
app.py
ADDED
@@ -0,0 +1,390 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import io
|
2 |
+
import base64
|
3 |
+
import os
|
4 |
+
|
5 |
+
import numpy as np
|
6 |
+
import torch
|
7 |
+
from torch import autocast
|
8 |
+
from diffusers import StableDiffusionPipeline, StableDiffusionInpaintPipeline
|
9 |
+
from PIL import Image
|
10 |
+
from PIL import ImageOps
|
11 |
+
import gradio as gr
|
12 |
+
import base64
|
13 |
+
import skimage
|
14 |
+
import skimage.measure
|
15 |
+
from utils import *
|
16 |
+
|
17 |
+
|
18 |
+
def load_html():
|
19 |
+
body, canvaspy = "", ""
|
20 |
+
with open("index.html", encoding="utf8") as f:
|
21 |
+
body = f.read()
|
22 |
+
with open("canvas.py", encoding="utf8") as f:
|
23 |
+
canvaspy = f.read()
|
24 |
+
body = body.replace("- paths:\n", "")
|
25 |
+
body = body.replace(" - ./canvas.py\n", "")
|
26 |
+
body = body.replace("from canvas import InfCanvas", canvaspy)
|
27 |
+
return body
|
28 |
+
|
29 |
+
|
30 |
+
def test(x):
|
31 |
+
x = load_html()
|
32 |
+
return f"""<iframe id="sdinfframe" style="width: 100%; height: 700px" name="result" allow="midi; geolocation; microphone; camera;
|
33 |
+
display-capture; encrypted-media;" sandbox="allow-modals allow-forms
|
34 |
+
allow-scripts allow-same-origin allow-popups
|
35 |
+
allow-top-navigation-by-user-activation allow-downloads" allowfullscreen=""
|
36 |
+
allowpaymentrequest="" frameborder="0" srcdoc='{x}'></iframe>"""
|
37 |
+
|
38 |
+
|
39 |
+
DEBUG_MODE = False
|
40 |
+
|
41 |
+
try:
|
42 |
+
SAMPLING_MODE = Image.Resampling.LANCZOS
|
43 |
+
except Exception as e:
|
44 |
+
SAMPLING_MODE = Image.LANCZOS
|
45 |
+
|
46 |
+
try:
|
47 |
+
contain_func = ImageOps.contain
|
48 |
+
except Exception as e:
|
49 |
+
|
50 |
+
def contain_func(image, size, method=SAMPLING_MODE):
|
51 |
+
# from PIL: https://pillow.readthedocs.io/en/stable/reference/ImageOps.html#PIL.ImageOps.contain
|
52 |
+
im_ratio = image.width / image.height
|
53 |
+
dest_ratio = size[0] / size[1]
|
54 |
+
if im_ratio != dest_ratio:
|
55 |
+
if im_ratio > dest_ratio:
|
56 |
+
new_height = int(image.height / image.width * size[0])
|
57 |
+
if new_height != size[1]:
|
58 |
+
size = (size[0], new_height)
|
59 |
+
else:
|
60 |
+
new_width = int(image.width / image.height * size[1])
|
61 |
+
if new_width != size[0]:
|
62 |
+
size = (new_width, size[1])
|
63 |
+
return image.resize(size, resample=method)
|
64 |
+
|
65 |
+
|
66 |
+
PAINT_SELECTION = "✥"
|
67 |
+
IMAGE_SELECTION = "🖼️"
|
68 |
+
BRUSH_SELECTION = "🖌️"
|
69 |
+
blocks = gr.Blocks()
|
70 |
+
model = {}
|
71 |
+
model["width"] = 1500
|
72 |
+
model["height"] = 600
|
73 |
+
model["sel_size"] = 256
|
74 |
+
|
75 |
+
def get_token():
|
76 |
+
token = ""
|
77 |
+
token = os.environ.get("hftoken", token)
|
78 |
+
return token
|
79 |
+
|
80 |
+
|
81 |
+
def save_token(token):
|
82 |
+
return
|
83 |
+
|
84 |
+
|
85 |
+
def get_model(token=""):
|
86 |
+
if "text2img" not in model:
|
87 |
+
text2img = StableDiffusionPipeline.from_pretrained(
|
88 |
+
"CompVis/stable-diffusion-v1-4",
|
89 |
+
revision="fp16",
|
90 |
+
torch_dtype=torch.float16,
|
91 |
+
use_auth_token=token,
|
92 |
+
).to("cuda")
|
93 |
+
model["safety_checker"] = text2img.safety_checker
|
94 |
+
inpaint = StableDiffusionInpaintPipeline(
|
95 |
+
vae=text2img.vae,
|
96 |
+
text_encoder=text2img.text_encoder,
|
97 |
+
tokenizer=text2img.tokenizer,
|
98 |
+
unet=text2img.unet,
|
99 |
+
scheduler=text2img.scheduler,
|
100 |
+
safety_checker=text2img.safety_checker,
|
101 |
+
feature_extractor=text2img.feature_extractor,
|
102 |
+
).to("cuda")
|
103 |
+
save_token(token)
|
104 |
+
try:
|
105 |
+
total_memory = torch.cuda.get_device_properties(0).total_memory // (
|
106 |
+
1024 ** 3
|
107 |
+
)
|
108 |
+
if total_memory <= 5:
|
109 |
+
inpaint.enable_attention_slicing()
|
110 |
+
except:
|
111 |
+
pass
|
112 |
+
model["text2img"] = text2img
|
113 |
+
model["inpaint"] = inpaint
|
114 |
+
return model["text2img"], model["inpaint"]
|
115 |
+
|
116 |
+
|
117 |
+
def run_outpaint(
|
118 |
+
sel_buffer_str,
|
119 |
+
prompt_text,
|
120 |
+
strength,
|
121 |
+
guidance,
|
122 |
+
step,
|
123 |
+
resize_check,
|
124 |
+
fill_mode,
|
125 |
+
enable_safety,
|
126 |
+
state,
|
127 |
+
):
|
128 |
+
base64_str = "base64"
|
129 |
+
if True:
|
130 |
+
text2img, inpaint = get_model()
|
131 |
+
if enable_safety:
|
132 |
+
text2img.safety_checker = model["safety_checker"]
|
133 |
+
inpaint.safety_checker = model["safety_checker"]
|
134 |
+
else:
|
135 |
+
text2img.safety_checker = lambda images, **kwargs: (images, False)
|
136 |
+
inpaint.safety_checker = lambda images, **kwargs: (images, False)
|
137 |
+
data = base64.b64decode(str(sel_buffer_str))
|
138 |
+
pil = Image.open(io.BytesIO(data))
|
139 |
+
# base.output.clear_output()
|
140 |
+
# base.read_selection_from_buffer()
|
141 |
+
sel_buffer = np.array(pil)
|
142 |
+
img = sel_buffer[:, :, 0:3]
|
143 |
+
mask = sel_buffer[:, :, -1]
|
144 |
+
process_size = 512 if resize_check else model["sel_size"]
|
145 |
+
if mask.sum() > 0:
|
146 |
+
img, mask = functbl[fill_mode](img, mask)
|
147 |
+
init_image = Image.fromarray(img)
|
148 |
+
mask = 255 - mask
|
149 |
+
mask = skimage.measure.block_reduce(mask, (8, 8), np.max)
|
150 |
+
mask = mask.repeat(8, axis=0).repeat(8, axis=1)
|
151 |
+
mask_image = Image.fromarray(mask)
|
152 |
+
# mask_image=mask_image.filter(ImageFilter.GaussianBlur(radius = 8))
|
153 |
+
with autocast("cuda"):
|
154 |
+
images = inpaint(
|
155 |
+
prompt=prompt_text,
|
156 |
+
init_image=init_image.resize(
|
157 |
+
(process_size, process_size), resample=SAMPLING_MODE
|
158 |
+
),
|
159 |
+
mask_image=mask_image.resize((process_size, process_size)),
|
160 |
+
strength=strength,
|
161 |
+
num_inference_steps=step,
|
162 |
+
guidance_scale=guidance,
|
163 |
+
)["sample"]
|
164 |
+
else:
|
165 |
+
with autocast("cuda"):
|
166 |
+
images = text2img(
|
167 |
+
prompt=prompt_text, height=process_size, width=process_size,
|
168 |
+
)["sample"]
|
169 |
+
out = sel_buffer.copy()
|
170 |
+
out[:, :, 0:3] = np.array(
|
171 |
+
images[0].resize(
|
172 |
+
(model["sel_size"], model["sel_size"]), resample=SAMPLING_MODE,
|
173 |
+
)
|
174 |
+
)
|
175 |
+
out[:, :, -1] = 255
|
176 |
+
out_pil = Image.fromarray(out)
|
177 |
+
out_buffer = io.BytesIO()
|
178 |
+
out_pil.save(out_buffer, format="PNG")
|
179 |
+
out_buffer.seek(0)
|
180 |
+
base64_bytes = base64.b64encode(out_buffer.read())
|
181 |
+
base64_str = base64_bytes.decode("ascii")
|
182 |
+
return (
|
183 |
+
gr.update(label=str(state + 1), value=base64_str,),
|
184 |
+
gr.update(label="Prompt"),
|
185 |
+
state + 1,
|
186 |
+
)
|
187 |
+
|
188 |
+
|
189 |
+
def load_js(name):
|
190 |
+
if name in ["export", "commit", "undo"]:
|
191 |
+
return f"""
|
192 |
+
function (x)
|
193 |
+
{{
|
194 |
+
let frame=document.querySelector("gradio-app").shadowRoot.querySelector("#sdinfframe").contentWindow.document;
|
195 |
+
let button=frame.querySelector("#{name}");
|
196 |
+
button.click();
|
197 |
+
return x;
|
198 |
+
}}
|
199 |
+
"""
|
200 |
+
ret = ""
|
201 |
+
with open(f"./js/{name}.js", "r") as f:
|
202 |
+
ret = f.read()
|
203 |
+
return ret
|
204 |
+
|
205 |
+
|
206 |
+
upload_button_js = load_js("upload")
|
207 |
+
outpaint_button_js = load_js("outpaint")
|
208 |
+
proceed_button_js = load_js("proceed")
|
209 |
+
mode_js = load_js("mode")
|
210 |
+
setup_button_js = load_js("setup")
|
211 |
+
|
212 |
+
get_model(get_token())
|
213 |
+
|
214 |
+
with blocks as demo:
|
215 |
+
# title
|
216 |
+
title = gr.Markdown(
|
217 |
+
"""
|
218 |
+
**stablediffusion-infinity**: Outpainting with Stable Diffusion on an infinite canvas: [https://github.com/lkwq007/stablediffusion-infinity](https://github.com/lkwq007/stablediffusion-infinity)
|
219 |
+
"""
|
220 |
+
)
|
221 |
+
# frame
|
222 |
+
frame = gr.HTML(test(2), visible=True)
|
223 |
+
# setup
|
224 |
+
# with gr.Row():
|
225 |
+
# token = gr.Textbox(
|
226 |
+
# label="Huggingface token",
|
227 |
+
# value="",
|
228 |
+
# placeholder="Input your token here",
|
229 |
+
# )
|
230 |
+
# canvas_width = gr.Number(
|
231 |
+
# label="Canvas width", value=1024, precision=0, elem_id="canvas_width"
|
232 |
+
# )
|
233 |
+
# canvas_height = gr.Number(
|
234 |
+
# label="Canvas height", value=600, precision=0, elem_id="canvas_height"
|
235 |
+
# )
|
236 |
+
# selection_size = gr.Number(
|
237 |
+
# label="Selection box size", value=256, precision=0, elem_id="selection_size"
|
238 |
+
# )
|
239 |
+
# setup_button = gr.Button("Start (may take a while)", variant="primary")
|
240 |
+
with gr.Row():
|
241 |
+
with gr.Column(scale=3, min_width=270):
|
242 |
+
# canvas control
|
243 |
+
canvas_control = gr.Radio(
|
244 |
+
label="Control",
|
245 |
+
choices=[PAINT_SELECTION, IMAGE_SELECTION, BRUSH_SELECTION],
|
246 |
+
value=PAINT_SELECTION,
|
247 |
+
elem_id="control",
|
248 |
+
)
|
249 |
+
with gr.Box():
|
250 |
+
with gr.Group():
|
251 |
+
run_button = gr.Button(value="Outpaint")
|
252 |
+
export_button = gr.Button(value="Export")
|
253 |
+
commit_button = gr.Button(value="✓")
|
254 |
+
retry_button = gr.Button(value="⟳")
|
255 |
+
undo_button = gr.Button(value="↶")
|
256 |
+
with gr.Column(scale=3, min_width=270):
|
257 |
+
sd_prompt = gr.Textbox(
|
258 |
+
label="Prompt", placeholder="input your prompt here", lines=4
|
259 |
+
)
|
260 |
+
with gr.Column(scale=2, min_width=150):
|
261 |
+
with gr.Box():
|
262 |
+
sd_resize = gr.Checkbox(label="Resize input to 515x512", value=True)
|
263 |
+
safety_check = gr.Checkbox(label="Enable Safety Checker", value=True)
|
264 |
+
sd_strength = gr.Slider(
|
265 |
+
label="Strength", minimum=0.0, maximum=1.0, value=0.75, step=0.01
|
266 |
+
)
|
267 |
+
with gr.Column(scale=1, min_width=150):
|
268 |
+
sd_step = gr.Number(label="Step", value=50, precision=0)
|
269 |
+
sd_guidance = gr.Number(label="Guidance", value=7.5)
|
270 |
+
with gr.Row():
|
271 |
+
with gr.Column(scale=4, min_width=600):
|
272 |
+
init_mode = gr.Radio(
|
273 |
+
label="Init mode",
|
274 |
+
choices=[
|
275 |
+
"patchmatch",
|
276 |
+
"edge_pad",
|
277 |
+
"cv2_ns",
|
278 |
+
"cv2_telea",
|
279 |
+
"gaussian",
|
280 |
+
"perlin",
|
281 |
+
],
|
282 |
+
value="patchmatch",
|
283 |
+
type="value",
|
284 |
+
)
|
285 |
+
|
286 |
+
proceed_button = gr.Button("Proceed", elem_id="proceed", visible=DEBUG_MODE)
|
287 |
+
# sd pipeline parameters
|
288 |
+
with gr.Accordion("Upload image", open=False):
|
289 |
+
image_box = gr.Image(image_mode="RGBA", source="upload", type="pil")
|
290 |
+
upload_button = gr.Button(
|
291 |
+
"Upload"
|
292 |
+
)
|
293 |
+
model_output = gr.Textbox(visible=DEBUG_MODE, elem_id="output", label="0")
|
294 |
+
model_input = gr.Textbox(visible=DEBUG_MODE, elem_id="input", label="Input")
|
295 |
+
upload_output = gr.Textbox(visible=DEBUG_MODE, elem_id="upload", label="0")
|
296 |
+
model_output_state = gr.State(value=0)
|
297 |
+
upload_output_state = gr.State(value=0)
|
298 |
+
# canvas_state = gr.State({"width":1024,"height":600,"selection_size":384})
|
299 |
+
|
300 |
+
def upload_func(image, state):
|
301 |
+
pil = image.convert("RGBA")
|
302 |
+
w, h = pil.size
|
303 |
+
if w > model["width"] - 100 or h > model["height"] - 100:
|
304 |
+
pil = contain_func(pil, (model["width"] - 100, model["height"] - 100))
|
305 |
+
out_buffer = io.BytesIO()
|
306 |
+
pil.save(out_buffer, format="PNG")
|
307 |
+
out_buffer.seek(0)
|
308 |
+
base64_bytes = base64.b64encode(out_buffer.read())
|
309 |
+
base64_str = base64_bytes.decode("ascii")
|
310 |
+
return (
|
311 |
+
gr.update(label=str(state + 1), value=base64_str),
|
312 |
+
state + 1,
|
313 |
+
)
|
314 |
+
|
315 |
+
upload_button.click(
|
316 |
+
fn=upload_func,
|
317 |
+
inputs=[image_box, upload_output_state],
|
318 |
+
outputs=[upload_output, upload_output_state],
|
319 |
+
_js=upload_button_js,
|
320 |
+
)
|
321 |
+
|
322 |
+
def setup_func(token_val, width, height, size):
|
323 |
+
model["width"] = width
|
324 |
+
model["height"] = height
|
325 |
+
model["sel_size"] = size
|
326 |
+
try:
|
327 |
+
get_model(token_val)
|
328 |
+
except Exception as e:
|
329 |
+
return {token: gr.update(value="Invalid token!")}
|
330 |
+
return {
|
331 |
+
token: gr.update(visible=False),
|
332 |
+
canvas_width: gr.update(visible=False),
|
333 |
+
canvas_height: gr.update(visible=False),
|
334 |
+
selection_size: gr.update(visible=False),
|
335 |
+
setup_button: gr.update(visible=False),
|
336 |
+
frame: gr.update(visible=True),
|
337 |
+
upload_button: gr.update(value="Upload"),
|
338 |
+
}
|
339 |
+
|
340 |
+
# setup_button.click(
|
341 |
+
# fn=setup_func,
|
342 |
+
# inputs=[token, canvas_width, canvas_height, selection_size],
|
343 |
+
# outputs=[
|
344 |
+
# token,
|
345 |
+
# canvas_width,
|
346 |
+
# canvas_height,
|
347 |
+
# selection_size,
|
348 |
+
# setup_button,
|
349 |
+
# frame,
|
350 |
+
# upload_button,
|
351 |
+
# ],
|
352 |
+
# _js=setup_button_js,
|
353 |
+
# )
|
354 |
+
run_button.click(
|
355 |
+
fn=None, inputs=[run_button], outputs=[run_button], _js=outpaint_button_js,
|
356 |
+
)
|
357 |
+
retry_button.click(
|
358 |
+
fn=None, inputs=[run_button], outputs=[run_button], _js=outpaint_button_js,
|
359 |
+
)
|
360 |
+
proceed_button.click(
|
361 |
+
fn=run_outpaint,
|
362 |
+
inputs=[
|
363 |
+
model_input,
|
364 |
+
sd_prompt,
|
365 |
+
sd_strength,
|
366 |
+
sd_guidance,
|
367 |
+
sd_step,
|
368 |
+
sd_resize,
|
369 |
+
init_mode,
|
370 |
+
safety_check,
|
371 |
+
model_output_state,
|
372 |
+
],
|
373 |
+
outputs=[model_output, sd_prompt, model_output_state],
|
374 |
+
_js=proceed_button_js,
|
375 |
+
)
|
376 |
+
export_button.click(
|
377 |
+
fn=None, inputs=[export_button], outputs=[export_button], _js=load_js("export")
|
378 |
+
)
|
379 |
+
commit_button.click(
|
380 |
+
fn=None, inputs=[export_button], outputs=[export_button], _js=load_js("commit")
|
381 |
+
)
|
382 |
+
undo_button.click(
|
383 |
+
fn=None, inputs=[export_button], outputs=[export_button], _js=load_js("undo")
|
384 |
+
)
|
385 |
+
canvas_control.change(
|
386 |
+
fn=None, inputs=[canvas_control], outputs=[canvas_control], _js=mode_js,
|
387 |
+
)
|
388 |
+
|
389 |
+
demo.launch()
|
390 |
+
|
canvas.py
ADDED
@@ -0,0 +1,547 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import base64
|
2 |
+
import io
|
3 |
+
import numpy as np
|
4 |
+
from PIL import Image
|
5 |
+
from pyodide import to_js, create_proxy
|
6 |
+
from js import (
|
7 |
+
console,
|
8 |
+
document,
|
9 |
+
devicePixelRatio,
|
10 |
+
ImageData,
|
11 |
+
Uint8ClampedArray,
|
12 |
+
CanvasRenderingContext2D as Context2d,
|
13 |
+
requestAnimationFrame,
|
14 |
+
)
|
15 |
+
|
16 |
+
PAINT_SELECTION = "✥"
|
17 |
+
IMAGE_SELECTION = "🖼️"
|
18 |
+
BRUSH_SELECTION = "🖌️"
|
19 |
+
NOP_MODE = 0
|
20 |
+
PAINT_MODE = 1
|
21 |
+
IMAGE_MODE = 2
|
22 |
+
BRUSH_MODE = 3
|
23 |
+
|
24 |
+
|
25 |
+
def hold_canvas():
|
26 |
+
pass
|
27 |
+
|
28 |
+
|
29 |
+
def prepare_canvas(width, height, canvas) -> Context2d:
|
30 |
+
ctx = canvas.getContext("2d")
|
31 |
+
|
32 |
+
canvas.style.width = f"{width}px"
|
33 |
+
canvas.style.height = f"{height}px"
|
34 |
+
|
35 |
+
canvas.width = width
|
36 |
+
canvas.height = height
|
37 |
+
|
38 |
+
ctx.clearRect(0, 0, width, height)
|
39 |
+
|
40 |
+
return ctx
|
41 |
+
|
42 |
+
|
43 |
+
# class MultiCanvas:
|
44 |
+
# def __init__(self,layer,width=800, height=600) -> None:
|
45 |
+
# pass
|
46 |
+
def multi_canvas(layer, width=800, height=600):
|
47 |
+
lst = [
|
48 |
+
CanvasProxy(document.querySelector(f"#canvas{i}"), width, height)
|
49 |
+
for i in range(layer)
|
50 |
+
]
|
51 |
+
return lst
|
52 |
+
|
53 |
+
|
54 |
+
class CanvasProxy:
|
55 |
+
def __init__(self, canvas, width=800, height=600) -> None:
|
56 |
+
self.canvas = canvas
|
57 |
+
self.ctx = prepare_canvas(width, height, canvas)
|
58 |
+
self.width = width
|
59 |
+
self.height = height
|
60 |
+
|
61 |
+
def clear_rect(self, x, y, w, h):
|
62 |
+
self.ctx.clearRect(x, y, w, h)
|
63 |
+
|
64 |
+
def clear(self,):
|
65 |
+
self.clear_rect(0, 0, self.width, self.height)
|
66 |
+
|
67 |
+
def stroke_rect(self, x, y, w, h):
|
68 |
+
self.ctx.strokeRect(x, y, w, h)
|
69 |
+
|
70 |
+
def fill_rect(self, x, y, w, h):
|
71 |
+
self.ctx.fillRect(x, y, w, h)
|
72 |
+
|
73 |
+
def put_image_data(self, image, x, y):
|
74 |
+
data = Uint8ClampedArray.new(to_js(image.tobytes()))
|
75 |
+
height, width, _ = image.shape
|
76 |
+
image_data = ImageData.new(data, width, height)
|
77 |
+
self.ctx.putImageData(image_data, x, y)
|
78 |
+
|
79 |
+
@property
|
80 |
+
def stroke_style(self):
|
81 |
+
return self.ctx.strokeStyle
|
82 |
+
|
83 |
+
@stroke_style.setter
|
84 |
+
def stroke_style(self, value):
|
85 |
+
self.ctx.strokeStyle = value
|
86 |
+
|
87 |
+
@property
|
88 |
+
def fill_style(self):
|
89 |
+
return self.ctx.strokeStyle
|
90 |
+
|
91 |
+
@fill_style.setter
|
92 |
+
def fill_style(self, value):
|
93 |
+
self.ctx.fillStyle = value
|
94 |
+
|
95 |
+
|
96 |
+
# RGBA for masking
|
97 |
+
class InfCanvas:
|
98 |
+
def __init__(
|
99 |
+
self,
|
100 |
+
width,
|
101 |
+
height,
|
102 |
+
selection_size=256,
|
103 |
+
grid_size=32,
|
104 |
+
patch_size=4096,
|
105 |
+
test_mode=False,
|
106 |
+
) -> None:
|
107 |
+
assert selection_size < min(height, width)
|
108 |
+
self.width = width
|
109 |
+
self.height = height
|
110 |
+
self.canvas = multi_canvas(5, width=width, height=height)
|
111 |
+
# self.canvas = Canvas(width=width, height=height)
|
112 |
+
self.view_pos = [0, 0]
|
113 |
+
self.cursor = [
|
114 |
+
width // 2 - selection_size // 2,
|
115 |
+
height // 2 - selection_size // 2,
|
116 |
+
]
|
117 |
+
self.data = {}
|
118 |
+
self.grid_size = grid_size
|
119 |
+
self.selection_size = selection_size
|
120 |
+
self.patch_size = patch_size
|
121 |
+
# note that for image data, the height comes before width
|
122 |
+
self.buffer = np.zeros((height, width, 4), dtype=np.uint8)
|
123 |
+
self.sel_buffer = np.zeros((selection_size, selection_size, 4), dtype=np.uint8)
|
124 |
+
self.sel_buffer_bak = np.zeros(
|
125 |
+
(selection_size, selection_size, 4), dtype=np.uint8
|
126 |
+
)
|
127 |
+
self.sel_dirty = False
|
128 |
+
self.buffer_dirty = False
|
129 |
+
self.mouse_pos = [-1, -1]
|
130 |
+
self.mouse_state = 0
|
131 |
+
# self.output = widgets.Output()
|
132 |
+
self.test_mode = test_mode
|
133 |
+
self.buffer_updated = False
|
134 |
+
self.image_move_freq = 1
|
135 |
+
self.show_brush = False
|
136 |
+
# inpaint pipeline from diffuser
|
137 |
+
|
138 |
+
def setup_mouse(self):
|
139 |
+
self.image_move_cnt = 0
|
140 |
+
|
141 |
+
def get_mouse_mode():
|
142 |
+
mode = document.querySelector("#mode").value
|
143 |
+
if mode == PAINT_SELECTION:
|
144 |
+
return PAINT_MODE
|
145 |
+
elif mode == IMAGE_SELECTION:
|
146 |
+
return IMAGE_MODE
|
147 |
+
return BRUSH_MODE
|
148 |
+
|
149 |
+
def get_event_pos(event):
|
150 |
+
canvas = self.canvas[-1].canvas
|
151 |
+
rect = canvas.getBoundingClientRect()
|
152 |
+
x = (canvas.width * (event.clientX - rect.left)) / rect.width
|
153 |
+
y = (canvas.height * (event.clientY - rect.top)) / rect.height
|
154 |
+
return x, y
|
155 |
+
|
156 |
+
def handle_mouse_down(event):
|
157 |
+
self.mouse_state = get_mouse_mode()
|
158 |
+
|
159 |
+
def handle_mouse_out(event):
|
160 |
+
last_state = self.mouse_state
|
161 |
+
self.mouse_state = NOP_MODE
|
162 |
+
self.image_move_cnt = 0
|
163 |
+
if last_state == IMAGE_MODE:
|
164 |
+
if True:
|
165 |
+
self.clear_background()
|
166 |
+
self.draw_buffer()
|
167 |
+
self.canvas[2].clear()
|
168 |
+
self.draw_selection_box()
|
169 |
+
if self.show_brush:
|
170 |
+
self.canvas[-2].clear()
|
171 |
+
self.show_brush = False
|
172 |
+
|
173 |
+
def handle_mouse_up(event):
|
174 |
+
last_state = self.mouse_state
|
175 |
+
self.mouse_state = NOP_MODE
|
176 |
+
self.image_move_cnt = 0
|
177 |
+
if last_state == IMAGE_MODE:
|
178 |
+
if True:
|
179 |
+
self.clear_background()
|
180 |
+
self.draw_buffer()
|
181 |
+
self.canvas[2].clear()
|
182 |
+
self.draw_selection_box()
|
183 |
+
|
184 |
+
async def handle_mouse_move(event):
|
185 |
+
x, y = get_event_pos(event)
|
186 |
+
x0, y0 = self.mouse_pos
|
187 |
+
xo = x - x0
|
188 |
+
yo = y - y0
|
189 |
+
if self.mouse_state == PAINT_MODE:
|
190 |
+
self.update_cursor(int(xo), int(yo))
|
191 |
+
if True:
|
192 |
+
# self.clear_background()
|
193 |
+
# console.log(self.buffer_updated)
|
194 |
+
if self.buffer_updated:
|
195 |
+
self.draw_buffer()
|
196 |
+
self.buffer_updated = False
|
197 |
+
self.draw_selection_box()
|
198 |
+
elif self.mouse_state == IMAGE_MODE:
|
199 |
+
self.image_move_cnt += 1
|
200 |
+
self.update_view_pos(int(xo), int(yo))
|
201 |
+
if self.image_move_cnt == self.image_move_freq:
|
202 |
+
if True:
|
203 |
+
self.clear_background()
|
204 |
+
self.draw_buffer()
|
205 |
+
self.canvas[2].clear()
|
206 |
+
self.draw_selection_box()
|
207 |
+
self.image_move_cnt = 0
|
208 |
+
elif self.mouse_state == BRUSH_MODE:
|
209 |
+
if self.sel_dirty:
|
210 |
+
self.write_selection_to_buffer()
|
211 |
+
self.canvas[2].clear()
|
212 |
+
self.buffer_dirty=True
|
213 |
+
bx0,by0=int(x)-self.grid_size//2,int(y)-self.grid_size//2
|
214 |
+
bx1,by1=bx0+self.grid_size,by0+self.grid_size
|
215 |
+
bx0,by0=max(0,bx0),max(0,by0)
|
216 |
+
bx1,by1=min(self.width,bx1),min(self.height,by1)
|
217 |
+
self.buffer[by0:by1,bx0:bx1,:]*=0
|
218 |
+
self.draw_buffer()
|
219 |
+
self.draw_selection_box()
|
220 |
+
|
221 |
+
mode = document.querySelector("#mode").value
|
222 |
+
if mode == BRUSH_SELECTION:
|
223 |
+
self.canvas[-2].clear()
|
224 |
+
self.canvas[-2].fill_style = "#ffffff"
|
225 |
+
self.canvas[-2].fill_rect(x-self.grid_size//2,y-self.grid_size//2,self.grid_size,self.grid_size)
|
226 |
+
self.canvas[-2].stroke_rect(x-self.grid_size//2,y-self.grid_size//2,self.grid_size,self.grid_size)
|
227 |
+
self.show_brush = True
|
228 |
+
elif self.show_brush:
|
229 |
+
self.canvas[-2].clear()
|
230 |
+
self.show_brush = False
|
231 |
+
self.mouse_pos[0] = x
|
232 |
+
self.mouse_pos[1] = y
|
233 |
+
|
234 |
+
self.canvas[-1].canvas.addEventListener(
|
235 |
+
"mousedown", create_proxy(handle_mouse_down)
|
236 |
+
)
|
237 |
+
self.canvas[-1].canvas.addEventListener(
|
238 |
+
"mousemove", create_proxy(handle_mouse_move)
|
239 |
+
)
|
240 |
+
self.canvas[-1].canvas.addEventListener(
|
241 |
+
"mouseup", create_proxy(handle_mouse_up)
|
242 |
+
)
|
243 |
+
self.canvas[-1].canvas.addEventListener(
|
244 |
+
"mouseout", create_proxy(handle_mouse_out)
|
245 |
+
)
|
246 |
+
|
247 |
+
def setup_widgets(self):
|
248 |
+
self.mode_button = widgets.ToggleButtons(
|
249 |
+
options=[PAINT_SELECTION, IMAGE_SELECTION],
|
250 |
+
disabled=False,
|
251 |
+
button_style="",
|
252 |
+
style={"button_width": "50px", "font_weight": "bold"},
|
253 |
+
tooltips=["Outpaint region", "Image"],
|
254 |
+
)
|
255 |
+
self.test_button = widgets.ToggleButtons(
|
256 |
+
options=["r", "g", "b"],
|
257 |
+
disabled=False,
|
258 |
+
style={"button_width": "50px", "font_weight": "bold", "font_size": "36px"},
|
259 |
+
)
|
260 |
+
self.text_input = widgets.Textarea(
|
261 |
+
value="",
|
262 |
+
placeholder="input your prompt here",
|
263 |
+
description="Prompt:",
|
264 |
+
disabled=False,
|
265 |
+
)
|
266 |
+
self.run_button = widgets.Button(
|
267 |
+
description="Outpaint",
|
268 |
+
tooltip="Run outpainting",
|
269 |
+
icon="pen",
|
270 |
+
button_style="primary",
|
271 |
+
)
|
272 |
+
self.export_button = widgets.Button(
|
273 |
+
description="Export",
|
274 |
+
tooltip="Export the image",
|
275 |
+
icon="save",
|
276 |
+
button_style="success",
|
277 |
+
)
|
278 |
+
self.fill_button = widgets.ToggleButtons(
|
279 |
+
description="Init mode:",
|
280 |
+
options=[
|
281 |
+
"patchmatch",
|
282 |
+
"edge_pad",
|
283 |
+
"cv2_ns",
|
284 |
+
"cv2_telea",
|
285 |
+
"gaussian",
|
286 |
+
"perlin",
|
287 |
+
],
|
288 |
+
disabled=False,
|
289 |
+
button_style="",
|
290 |
+
style={"button_width": "80px", "font_weight": "bold"},
|
291 |
+
)
|
292 |
+
|
293 |
+
if self.test_mode:
|
294 |
+
|
295 |
+
def test_button_clicked(btn):
|
296 |
+
# lst.append(tuple(base.cursor))
|
297 |
+
with self.output:
|
298 |
+
val = self.test_button.value
|
299 |
+
if val == "r":
|
300 |
+
self.fill_selection(
|
301 |
+
np.tile(
|
302 |
+
np.array([255, 0, 0, 255], dtype=np.uint8),
|
303 |
+
(self.selection_size, self.selection_size, 1),
|
304 |
+
)
|
305 |
+
)
|
306 |
+
if val == "g":
|
307 |
+
self.fill_selection(
|
308 |
+
np.tile(
|
309 |
+
np.array([0, 255, 0, 255], dtype=np.uint8),
|
310 |
+
(self.selection_size, self.selection_size, 1),
|
311 |
+
)
|
312 |
+
)
|
313 |
+
if val == "b":
|
314 |
+
self.fill_selection(
|
315 |
+
np.tile(
|
316 |
+
np.array([0, 0, 255, 255], dtype=np.uint8),
|
317 |
+
(self.selection_size, self.selection_size, 1),
|
318 |
+
)
|
319 |
+
)
|
320 |
+
if True:
|
321 |
+
self.clear_background()
|
322 |
+
self.draw_buffer()
|
323 |
+
self.draw_selection_box()
|
324 |
+
|
325 |
+
self.run_button.on_click(test_button_clicked)
|
326 |
+
|
327 |
+
def display(self):
|
328 |
+
if True:
|
329 |
+
self.clear_background()
|
330 |
+
self.draw_buffer()
|
331 |
+
self.draw_selection_box()
|
332 |
+
if self.test_mode:
|
333 |
+
return [
|
334 |
+
self.test_button,
|
335 |
+
self.mode_button,
|
336 |
+
self.canvas,
|
337 |
+
widgets.HBox([self.run_button, self.text_input]),
|
338 |
+
self.output,
|
339 |
+
]
|
340 |
+
return [
|
341 |
+
self.fill_button,
|
342 |
+
self.canvas,
|
343 |
+
widgets.HBox(
|
344 |
+
[self.mode_button, self.run_button, self.export_button, self.text_input]
|
345 |
+
),
|
346 |
+
self.output,
|
347 |
+
]
|
348 |
+
|
349 |
+
def clear_background(self):
|
350 |
+
# fake transparent background
|
351 |
+
h, w, step = self.height, self.width, self.grid_size
|
352 |
+
stride = step * 2
|
353 |
+
x0, y0 = self.view_pos
|
354 |
+
x0 = (-x0) % stride
|
355 |
+
y0 = (-y0) % stride
|
356 |
+
# self.canvas.clear()
|
357 |
+
self.canvas[0].fill_style = "#ffffff"
|
358 |
+
self.canvas[0].fill_rect(0, 0, w, h)
|
359 |
+
self.canvas[0].fill_style = "#aaaaaa"
|
360 |
+
for y in range(y0 - stride, h + step, step):
|
361 |
+
start = (x0 - stride) if y // step % 2 == 0 else (x0 - step)
|
362 |
+
for x in range(start, w + step, stride):
|
363 |
+
self.canvas[0].fill_rect(x, y, step, step)
|
364 |
+
self.canvas[0].stroke_rect(0, 0, w, h)
|
365 |
+
|
366 |
+
def update_view_pos(self, xo, yo):
|
367 |
+
if abs(xo) + abs(yo) == 0:
|
368 |
+
return
|
369 |
+
if self.sel_dirty:
|
370 |
+
self.write_selection_to_buffer()
|
371 |
+
if self.buffer_dirty:
|
372 |
+
self.buffer2data()
|
373 |
+
self.view_pos[0] -= xo
|
374 |
+
self.view_pos[1] -= yo
|
375 |
+
self.data2buffer()
|
376 |
+
# self.read_selection_from_buffer()
|
377 |
+
|
378 |
+
def update_cursor(self, xo, yo):
|
379 |
+
if abs(xo) + abs(yo) == 0:
|
380 |
+
return
|
381 |
+
if self.sel_dirty:
|
382 |
+
self.write_selection_to_buffer()
|
383 |
+
self.cursor[0] += xo
|
384 |
+
self.cursor[1] += yo
|
385 |
+
self.cursor[0] = max(min(self.width - self.selection_size, self.cursor[0]), 0)
|
386 |
+
self.cursor[1] = max(min(self.height - self.selection_size, self.cursor[1]), 0)
|
387 |
+
# self.read_selection_from_buffer()
|
388 |
+
|
389 |
+
def data2buffer(self):
|
390 |
+
x, y = self.view_pos
|
391 |
+
h, w = self.height, self.width
|
392 |
+
# fill four parts
|
393 |
+
for i in range(4):
|
394 |
+
pos_src, pos_dst, data = self.select(x, y, i)
|
395 |
+
xs0, xs1 = pos_src[0]
|
396 |
+
ys0, ys1 = pos_src[1]
|
397 |
+
xd0, xd1 = pos_dst[0]
|
398 |
+
yd0, yd1 = pos_dst[1]
|
399 |
+
self.buffer[yd0:yd1, xd0:xd1, :] = data[ys0:ys1, xs0:xs1, :]
|
400 |
+
|
401 |
+
def buffer2data(self):
|
402 |
+
x, y = self.view_pos
|
403 |
+
h, w = self.height, self.width
|
404 |
+
# fill four parts
|
405 |
+
for i in range(4):
|
406 |
+
pos_src, pos_dst, data = self.select(x, y, i)
|
407 |
+
xs0, xs1 = pos_src[0]
|
408 |
+
ys0, ys1 = pos_src[1]
|
409 |
+
xd0, xd1 = pos_dst[0]
|
410 |
+
yd0, yd1 = pos_dst[1]
|
411 |
+
data[ys0:ys1, xs0:xs1, :] = self.buffer[yd0:yd1, xd0:xd1, :]
|
412 |
+
self.buffer_dirty = False
|
413 |
+
|
414 |
+
def select(self, x, y, idx):
|
415 |
+
w, h = self.width, self.height
|
416 |
+
lst = [(0, 0), (0, h), (w, 0), (w, h)]
|
417 |
+
if idx == 0:
|
418 |
+
x0, y0 = x % self.patch_size, y % self.patch_size
|
419 |
+
x1 = min(x0 + w, self.patch_size)
|
420 |
+
y1 = min(y0 + h, self.patch_size)
|
421 |
+
elif idx == 1:
|
422 |
+
y += h
|
423 |
+
x0, y0 = x % self.patch_size, y % self.patch_size
|
424 |
+
x1 = min(x0 + w, self.patch_size)
|
425 |
+
y1 = max(y0 - h, 0)
|
426 |
+
elif idx == 2:
|
427 |
+
x += w
|
428 |
+
x0, y0 = x % self.patch_size, y % self.patch_size
|
429 |
+
x1 = max(x0 - w, 0)
|
430 |
+
y1 = min(y0 + h, self.patch_size)
|
431 |
+
else:
|
432 |
+
x += w
|
433 |
+
y += h
|
434 |
+
x0, y0 = x % self.patch_size, y % self.patch_size
|
435 |
+
x1 = max(x0 - w, 0)
|
436 |
+
y1 = max(y0 - h, 0)
|
437 |
+
xi, yi = x // self.patch_size, y // self.patch_size
|
438 |
+
cur = self.data.setdefault(
|
439 |
+
(xi, yi), np.zeros((self.patch_size, self.patch_size, 4), dtype=np.uint8)
|
440 |
+
)
|
441 |
+
x0_img, y0_img = lst[idx]
|
442 |
+
x1_img = x0_img + x1 - x0
|
443 |
+
y1_img = y0_img + y1 - y0
|
444 |
+
sort = lambda a, b: ((a, b) if a < b else (b, a))
|
445 |
+
return (
|
446 |
+
(sort(x0, x1), sort(y0, y1)),
|
447 |
+
(sort(x0_img, x1_img), sort(y0_img, y1_img)),
|
448 |
+
cur,
|
449 |
+
)
|
450 |
+
|
451 |
+
def draw_buffer(self):
|
452 |
+
self.canvas[1].clear()
|
453 |
+
self.canvas[1].put_image_data(self.buffer, 0, 0)
|
454 |
+
|
455 |
+
def fill_selection(self, img):
|
456 |
+
self.sel_buffer = img
|
457 |
+
self.sel_dirty = True
|
458 |
+
|
459 |
+
def draw_selection_box(self):
|
460 |
+
x0, y0 = self.cursor
|
461 |
+
size = self.selection_size
|
462 |
+
if self.sel_dirty:
|
463 |
+
self.canvas[2].clear()
|
464 |
+
self.canvas[2].put_image_data(self.sel_buffer, x0, y0)
|
465 |
+
self.canvas[-1].clear()
|
466 |
+
self.canvas[-1].stroke_style = "#0a0a0a"
|
467 |
+
self.canvas[-1].stroke_rect(x0, y0, size, size)
|
468 |
+
self.canvas[-1].stroke_style = "#ffffff"
|
469 |
+
self.canvas[-1].stroke_rect(x0 - 1, y0 - 1, size + 2, size + 2)
|
470 |
+
self.canvas[-1].stroke_style = "#000000"
|
471 |
+
self.canvas[-1].stroke_rect(x0 - 2, y0 - 2, size + 4, size + 4)
|
472 |
+
|
473 |
+
def write_selection_to_buffer(self):
|
474 |
+
x0, y0 = self.cursor
|
475 |
+
x1, y1 = x0 + self.selection_size, y0 + self.selection_size
|
476 |
+
self.buffer[y0:y1, x0:x1] = self.sel_buffer
|
477 |
+
self.sel_dirty = False
|
478 |
+
self.sel_buffer = self.sel_buffer_bak.copy()
|
479 |
+
self.buffer_dirty = True
|
480 |
+
self.buffer_updated = True
|
481 |
+
|
482 |
+
def read_selection_from_buffer(self):
|
483 |
+
x0, y0 = self.cursor
|
484 |
+
x1, y1 = x0 + self.selection_size, y0 + self.selection_size
|
485 |
+
self.sel_buffer = self.buffer[y0:y1, x0:x1]
|
486 |
+
self.sel_dirty = False
|
487 |
+
|
488 |
+
def base64_to_numpy(self, base64_str):
|
489 |
+
try:
|
490 |
+
data = base64.b64decode(str(base64_str))
|
491 |
+
pil = Image.open(io.BytesIO(data))
|
492 |
+
arr = np.array(pil)
|
493 |
+
ret = arr
|
494 |
+
except:
|
495 |
+
ret = np.tile(
|
496 |
+
np.array([255, 0, 0, 255], dtype=np.uint8),
|
497 |
+
(self.selection_size, self.selection_size, 1),
|
498 |
+
)
|
499 |
+
return ret
|
500 |
+
|
501 |
+
def numpy_to_base64(self, arr):
|
502 |
+
out_pil = Image.fromarray(arr)
|
503 |
+
out_buffer = io.BytesIO()
|
504 |
+
out_pil.save(out_buffer, format="PNG")
|
505 |
+
out_buffer.seek(0)
|
506 |
+
base64_bytes = base64.b64encode(out_buffer.read())
|
507 |
+
base64_str = base64_bytes.decode("ascii")
|
508 |
+
return base64_str
|
509 |
+
|
510 |
+
def export(self):
|
511 |
+
if self.sel_dirty:
|
512 |
+
self.write_selection_to_buffer()
|
513 |
+
if self.buffer_dirty:
|
514 |
+
self.buffer2data()
|
515 |
+
xmin, xmax, ymin, ymax = 0, 0, 0, 0
|
516 |
+
if len(self.data.keys()) == 0:
|
517 |
+
return np.zeros(
|
518 |
+
(self.selection_size, self.selection_size, 4), dtype=np.uint8
|
519 |
+
)
|
520 |
+
for xi, yi in self.data.keys():
|
521 |
+
buf = self.data[(xi, yi)]
|
522 |
+
if buf.sum() > 0:
|
523 |
+
xmin = min(xi, xmin)
|
524 |
+
xmax = max(xi, xmax)
|
525 |
+
ymin = min(yi, ymin)
|
526 |
+
ymax = max(yi, ymax)
|
527 |
+
yn = ymax - ymin + 1
|
528 |
+
xn = xmax - xmin + 1
|
529 |
+
image = np.zeros(
|
530 |
+
(yn * self.patch_size, xn * self.patch_size, 4), dtype=np.uint8
|
531 |
+
)
|
532 |
+
for xi, yi in self.data.keys():
|
533 |
+
buf = self.data[(xi, yi)]
|
534 |
+
if buf.sum() > 0:
|
535 |
+
y0 = (yi - ymin) * self.patch_size
|
536 |
+
x0 = (xi - xmin) * self.patch_size
|
537 |
+
image[y0 : y0 + self.patch_size, x0 : x0 + self.patch_size] = buf
|
538 |
+
ylst, xlst = image[:, :, -1].nonzero()
|
539 |
+
if len(ylst) > 0:
|
540 |
+
yt, xt = ylst.min(), xlst.min()
|
541 |
+
yb, xb = ylst.max(), xlst.max()
|
542 |
+
image = image[yt : yb + 1, xt : xb + 1]
|
543 |
+
return image
|
544 |
+
else:
|
545 |
+
return np.zeros(
|
546 |
+
(self.selection_size, self.selection_size, 4), dtype=np.uint8
|
547 |
+
)
|
js/mode.js
ADDED
@@ -0,0 +1,6 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
function(mode){
|
2 |
+
let app=document.querySelector("gradio-app").shadowRoot;
|
3 |
+
let frame=app.querySelector("#sdinfframe").contentWindow.document;
|
4 |
+
frame.querySelector("#mode").value=mode;
|
5 |
+
return mode;
|
6 |
+
}
|
js/outpaint.js
ADDED
@@ -0,0 +1,24 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
function(a){
|
2 |
+
if(!window.my_observe_outpaint)
|
3 |
+
{
|
4 |
+
console.log("setup outpaint here");
|
5 |
+
window.my_observe_outpaint = new MutationObserver(function (event) {
|
6 |
+
console.log(event);
|
7 |
+
let app=document.querySelector("gradio-app").shadowRoot;
|
8 |
+
let frame=app.querySelector("#sdinfframe").contentWindow.document;
|
9 |
+
frame.querySelector("#outpaint").click();
|
10 |
+
});
|
11 |
+
window.my_observe_outpaint_target=document.querySelector("gradio-app").shadowRoot.querySelector("#output span")
|
12 |
+
window.my_observe_outpaint.observe(window.my_observe_outpaint_target, {
|
13 |
+
attributes: false,
|
14 |
+
subtree: true,
|
15 |
+
childList: true,
|
16 |
+
characterData: true
|
17 |
+
});
|
18 |
+
}
|
19 |
+
let app=document.querySelector("gradio-app").shadowRoot;
|
20 |
+
let frame=app.querySelector("#sdinfframe").contentWindow.document;
|
21 |
+
let button=frame.querySelector("#transfer");
|
22 |
+
button.click();
|
23 |
+
return a;
|
24 |
+
}
|
js/proceed.js
ADDED
@@ -0,0 +1,22 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
function(sel_buffer_str,
|
2 |
+
prompt_text,
|
3 |
+
strength,
|
4 |
+
guidance,
|
5 |
+
step,
|
6 |
+
resize_check,
|
7 |
+
fill_mode,
|
8 |
+
enable_safety,
|
9 |
+
state){
|
10 |
+
sel_buffer = document.querySelector("gradio-app").shadowRoot.querySelector("#input textarea").value;
|
11 |
+
return [
|
12 |
+
sel_buffer,
|
13 |
+
prompt_text,
|
14 |
+
strength,
|
15 |
+
guidance,
|
16 |
+
step,
|
17 |
+
resize_check,
|
18 |
+
fill_mode,
|
19 |
+
enable_safety,
|
20 |
+
state
|
21 |
+
]
|
22 |
+
}
|
js/setup.js
ADDED
@@ -0,0 +1,22 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
function(token_val, width, height, size){
|
2 |
+
let app=document.querySelector("gradio-app").shadowRoot;
|
3 |
+
app.querySelector("#sdinfframe").style.height=height+"px";
|
4 |
+
let frame=app.querySelector("#sdinfframe").contentWindow.document;
|
5 |
+
if(frame.querySelector("#setup").value=="0")
|
6 |
+
{
|
7 |
+
window.my_setup=setInterval(function(){
|
8 |
+
let frame=document.querySelector("gradio-app").shadowRoot.querySelector("#sdinfframe").contentWindow.document;
|
9 |
+
console.log("Check PyScript...")
|
10 |
+
if(frame.querySelector("#setup").value=="1")
|
11 |
+
{
|
12 |
+
frame.querySelector("#draw").click();
|
13 |
+
clearInterval(window.my_setup);
|
14 |
+
}
|
15 |
+
},100)
|
16 |
+
}
|
17 |
+
else
|
18 |
+
{
|
19 |
+
frame.querySelector("#draw").click();
|
20 |
+
}
|
21 |
+
return [token_val, width, height, size];
|
22 |
+
}
|
js/upload.js
ADDED
@@ -0,0 +1,19 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
function(a,b){
|
2 |
+
if(!window.my_observe_upload)
|
3 |
+
{
|
4 |
+
console.log("setup upload here");
|
5 |
+
window.my_observe_upload = new MutationObserver(function (event) {
|
6 |
+
console.log(event);
|
7 |
+
var frame=document.querySelector("gradio-app").shadowRoot.querySelector("#sdinfframe").contentWindow.document;
|
8 |
+
frame.querySelector("#upload").click();
|
9 |
+
});
|
10 |
+
window.my_observe_upload_target = document.querySelector("gradio-app").shadowRoot.querySelector("#upload span");
|
11 |
+
window.my_observe_upload.observe(window.my_observe_upload_target, {
|
12 |
+
attributes: false,
|
13 |
+
subtree: true,
|
14 |
+
childList: true,
|
15 |
+
characterData: true
|
16 |
+
});
|
17 |
+
}
|
18 |
+
return [a,b];
|
19 |
+
}
|
packages.txt
ADDED
@@ -0,0 +1,2 @@
|
|
|
|
|
|
|
1 |
+
build-essential
|
2 |
+
libopencv-dev
|
perlin2d.py
ADDED
@@ -0,0 +1,45 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import numpy as np
|
2 |
+
|
3 |
+
##########
|
4 |
+
# https://stackoverflow.com/questions/42147776/producing-2d-perlin-noise-with-numpy/42154921#42154921
|
5 |
+
def perlin(x, y, seed=0):
|
6 |
+
# permutation table
|
7 |
+
np.random.seed(seed)
|
8 |
+
p = np.arange(256, dtype=int)
|
9 |
+
np.random.shuffle(p)
|
10 |
+
p = np.stack([p, p]).flatten()
|
11 |
+
# coordinates of the top-left
|
12 |
+
xi, yi = x.astype(int), y.astype(int)
|
13 |
+
# internal coordinates
|
14 |
+
xf, yf = x - xi, y - yi
|
15 |
+
# fade factors
|
16 |
+
u, v = fade(xf), fade(yf)
|
17 |
+
# noise components
|
18 |
+
n00 = gradient(p[p[xi] + yi], xf, yf)
|
19 |
+
n01 = gradient(p[p[xi] + yi + 1], xf, yf - 1)
|
20 |
+
n11 = gradient(p[p[xi + 1] + yi + 1], xf - 1, yf - 1)
|
21 |
+
n10 = gradient(p[p[xi + 1] + yi], xf - 1, yf)
|
22 |
+
# combine noises
|
23 |
+
x1 = lerp(n00, n10, u)
|
24 |
+
x2 = lerp(n01, n11, u) # FIX1: I was using n10 instead of n01
|
25 |
+
return lerp(x1, x2, v) # FIX2: I also had to reverse x1 and x2 here
|
26 |
+
|
27 |
+
|
28 |
+
def lerp(a, b, x):
|
29 |
+
"linear interpolation"
|
30 |
+
return a + x * (b - a)
|
31 |
+
|
32 |
+
|
33 |
+
def fade(t):
|
34 |
+
"6t^5 - 15t^4 + 10t^3"
|
35 |
+
return 6 * t ** 5 - 15 * t ** 4 + 10 * t ** 3
|
36 |
+
|
37 |
+
|
38 |
+
def gradient(h, x, y):
|
39 |
+
"grad converts h to the right gradient vector and return the dot product with (x,y)"
|
40 |
+
vectors = np.array([[0, 1], [0, -1], [1, 0], [-1, 0]])
|
41 |
+
g = vectors[h % 4]
|
42 |
+
return g[:, :, 0] * x + g[:, :, 1] * y
|
43 |
+
|
44 |
+
|
45 |
+
##########
|
requirements.txt
ADDED
@@ -0,0 +1,12 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
--extra-index-url https://download.pytorch.org/whl/cu113
|
2 |
+
imageio==2.19.5
|
3 |
+
imageio-ffmpeg==0.4.7
|
4 |
+
numpy==1.22.4
|
5 |
+
opencv-python-headless==4.6.0.66
|
6 |
+
torch==1.12.0+cu113
|
7 |
+
torchvision==0.13.0+cu113
|
8 |
+
scipy
|
9 |
+
scikit-image
|
10 |
+
diffusers==0.3.0
|
11 |
+
transformers
|
12 |
+
ftfy
|
utils.py
ADDED
@@ -0,0 +1,154 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from PIL import Image
|
2 |
+
from PIL import ImageFilter
|
3 |
+
import cv2
|
4 |
+
import numpy as np
|
5 |
+
import scipy
|
6 |
+
import scipy.signal
|
7 |
+
from scipy.spatial import cKDTree
|
8 |
+
|
9 |
+
import os
|
10 |
+
from perlin2d import *
|
11 |
+
|
12 |
+
patch_match_compiled = True
|
13 |
+
if os.name != "nt":
|
14 |
+
try:
|
15 |
+
from PyPatchMatch import patch_match
|
16 |
+
except Exception as e:
|
17 |
+
try:
|
18 |
+
import patch_match
|
19 |
+
except Exception as e:
|
20 |
+
patch_match_compiled = False
|
21 |
+
|
22 |
+
try:
|
23 |
+
patch_match
|
24 |
+
except NameError:
|
25 |
+
print("patch_match compiling failed")
|
26 |
+
patch_match_compiled = False
|
27 |
+
|
28 |
+
|
29 |
+
|
30 |
+
|
31 |
+
def edge_pad(img, mask, mode=1):
|
32 |
+
if mode == 0:
|
33 |
+
nmask = mask.copy()
|
34 |
+
nmask[nmask > 0] = 1
|
35 |
+
res0 = 1 - nmask
|
36 |
+
res1 = nmask
|
37 |
+
p0 = np.stack(res0.nonzero(), axis=0).transpose()
|
38 |
+
p1 = np.stack(res1.nonzero(), axis=0).transpose()
|
39 |
+
min_dists, min_dist_idx = cKDTree(p1).query(p0, 1)
|
40 |
+
loc = p1[min_dist_idx]
|
41 |
+
for (a, b), (c, d) in zip(p0, loc):
|
42 |
+
img[a, b] = img[c, d]
|
43 |
+
elif mode == 1:
|
44 |
+
record = {}
|
45 |
+
kernel = [[1] * 3 for _ in range(3)]
|
46 |
+
nmask = mask.copy()
|
47 |
+
nmask[nmask > 0] = 1
|
48 |
+
res = scipy.signal.convolve2d(
|
49 |
+
nmask, kernel, mode="same", boundary="fill", fillvalue=1
|
50 |
+
)
|
51 |
+
res[nmask < 1] = 0
|
52 |
+
res[res == 9] = 0
|
53 |
+
res[res > 0] = 1
|
54 |
+
ylst, xlst = res.nonzero()
|
55 |
+
queue = [(y, x) for y, x in zip(ylst, xlst)]
|
56 |
+
# bfs here
|
57 |
+
cnt = res.astype(np.float32)
|
58 |
+
acc = img.astype(np.float32)
|
59 |
+
step = 1
|
60 |
+
h = acc.shape[0]
|
61 |
+
w = acc.shape[1]
|
62 |
+
offset = [(1, 0), (-1, 0), (0, 1), (0, -1)]
|
63 |
+
while queue:
|
64 |
+
target = []
|
65 |
+
for y, x in queue:
|
66 |
+
val = acc[y][x]
|
67 |
+
for yo, xo in offset:
|
68 |
+
yn = y + yo
|
69 |
+
xn = x + xo
|
70 |
+
if 0 <= yn < h and 0 <= xn < w and nmask[yn][xn] < 1:
|
71 |
+
if record.get((yn, xn), step) == step:
|
72 |
+
acc[yn][xn] = acc[yn][xn] * cnt[yn][xn] + val
|
73 |
+
cnt[yn][xn] += 1
|
74 |
+
acc[yn][xn] /= cnt[yn][xn]
|
75 |
+
if (yn, xn) not in record:
|
76 |
+
record[(yn, xn)] = step
|
77 |
+
target.append((yn, xn))
|
78 |
+
step += 1
|
79 |
+
queue = target
|
80 |
+
img = acc.astype(np.uint8)
|
81 |
+
else:
|
82 |
+
nmask = mask.copy()
|
83 |
+
ylst, xlst = nmask.nonzero()
|
84 |
+
yt, xt = ylst.min(), xlst.min()
|
85 |
+
yb, xb = ylst.max(), xlst.max()
|
86 |
+
content = img[yt : yb + 1, xt : xb + 1]
|
87 |
+
img = np.pad(
|
88 |
+
content,
|
89 |
+
((yt, mask.shape[0] - yb - 1), (xt, mask.shape[1] - xb - 1), (0, 0)),
|
90 |
+
mode="edge",
|
91 |
+
)
|
92 |
+
return img, mask
|
93 |
+
|
94 |
+
|
95 |
+
def perlin_noise(img, mask):
|
96 |
+
lin = np.linspace(0, 5, mask.shape[0], endpoint=False)
|
97 |
+
x, y = np.meshgrid(lin, lin)
|
98 |
+
avg = img.mean(axis=0).mean(axis=0)
|
99 |
+
# noise=[((perlin(x, y)+1)*128+avg[i]).astype(np.uint8) for i in range(3)]
|
100 |
+
noise = [((perlin(x, y) + 1) * 0.5 * 255).astype(np.uint8) for i in range(3)]
|
101 |
+
noise = np.stack(noise, axis=-1)
|
102 |
+
# mask=skimage.measure.block_reduce(mask,(8,8),np.min)
|
103 |
+
# mask=mask.repeat(8, axis=0).repeat(8, axis=1)
|
104 |
+
# mask_image=Image.fromarray(mask)
|
105 |
+
# mask_image=mask_image.filter(ImageFilter.GaussianBlur(radius = 4))
|
106 |
+
# mask=np.array(mask_image)
|
107 |
+
nmask = mask.copy()
|
108 |
+
# nmask=nmask/255.0
|
109 |
+
nmask[mask > 0] = 1
|
110 |
+
img = nmask[:, :, np.newaxis] * img + (1 - nmask[:, :, np.newaxis]) * noise
|
111 |
+
# img=img.astype(np.uint8)
|
112 |
+
return img, mask
|
113 |
+
|
114 |
+
|
115 |
+
def gaussian_noise(img, mask):
|
116 |
+
noise = np.random.randn(mask.shape[0], mask.shape[1], 3)
|
117 |
+
noise = (noise + 1) / 2 * 255
|
118 |
+
noise = noise.astype(np.uint8)
|
119 |
+
nmask = mask.copy()
|
120 |
+
nmask[mask > 0] = 1
|
121 |
+
img = nmask[:, :, np.newaxis] * img + (1 - nmask[:, :, np.newaxis]) * noise
|
122 |
+
return img, mask
|
123 |
+
|
124 |
+
|
125 |
+
def cv2_telea(img, mask):
|
126 |
+
ret = cv2.inpaint(img, 255 - mask, 5, cv2.INPAINT_TELEA)
|
127 |
+
return ret, mask
|
128 |
+
|
129 |
+
|
130 |
+
def cv2_ns(img, mask):
|
131 |
+
ret = cv2.inpaint(img, 255 - mask, 5, cv2.INPAINT_NS)
|
132 |
+
return ret, mask
|
133 |
+
|
134 |
+
|
135 |
+
def patch_match_func(img, mask):
|
136 |
+
ret = patch_match.inpaint(img, mask=255 - mask, patch_size=3)
|
137 |
+
return ret, mask
|
138 |
+
|
139 |
+
|
140 |
+
def mean_fill(img, mask):
|
141 |
+
avg = img.mean(axis=0).mean(axis=0)
|
142 |
+
img[mask < 1] = avg
|
143 |
+
return img, mask
|
144 |
+
|
145 |
+
|
146 |
+
functbl = {
|
147 |
+
"gaussian": gaussian_noise,
|
148 |
+
"perlin": perlin_noise,
|
149 |
+
"edge_pad": edge_pad,
|
150 |
+
"patchmatch": patch_match_func if (os.name != "nt" and patch_match_compiled) else edge_pad,
|
151 |
+
"cv2_ns": cv2_ns,
|
152 |
+
"cv2_telea": cv2_telea,
|
153 |
+
"mean_fill": mean_fill,
|
154 |
+
}
|