lnyan commited on
Commit
5704551
1 Parent(s): 7f0ff5a
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: younesse.andam@gmail.com
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 <maojiayuan@gmail.com>
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 : maojiayuan@gmail.com
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 : maojiayuan@gmail.com
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 : maojiayuan@gmail.com
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 <maojiayuan@gmail.com>
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
+ }