Update files

PyPatchMatch/.gitignore

@@ -1,4 +1,4 @@
-/build/
-/*.so
-__pycache__
-*.py[cod]
+/build/
+/*.so
+__pycache__
+*.py[cod]

PyPatchMatch/LICENSE

@@ -1,21 +1,21 @@
-MIT License
-
-Copyright (c) 2020 Jiayuan Mao
-
-Permission is hereby granted, free of charge, to any person obtaining a copy
-of this software and associated documentation files (the "Software"), to deal
-in the Software without restriction, including without limitation the rights
-to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-copies of the Software, and to permit persons to whom the Software is
-furnished to do so, subject to the following conditions:
-
-The above copyright notice and this permission notice shall be included in all
-copies or substantial portions of the Software.
-
-THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
-SOFTWARE.
+MIT License
+
+Copyright (c) 2020 Jiayuan Mao
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

PyPatchMatch/Makefile

@@ -1,54 +1,54 @@
-#
-# Makefile
-# Jiayuan Mao, 2019-01-09 13:59
-#
-
-SRC_DIR = csrc
-INC_DIR = csrc
-OBJ_DIR = build/obj
-TARGET = libpatchmatch.so
-
-LIB_TARGET = $(TARGET)
-INCLUDE_DIR = -I $(SRC_DIR) -I $(INC_DIR)
-
-CXX = $(ENVIRONMENT_OPTIONS) g++
-CXXFLAGS = -std=c++14
-CXXFLAGS += -Ofast -ffast-math -w
-# CXXFLAGS += -g
-CXXFLAGS += $(shell pkg-config --cflags opencv) -fPIC
-CXXFLAGS += $(INCLUDE_DIR)
-LDFLAGS = $(shell pkg-config --cflags --libs opencv) -shared -fPIC
-
-
-CXXSOURCES = $(shell find $(SRC_DIR)/ -name "*.cpp")
-OBJS = $(addprefix $(OBJ_DIR)/,$(CXXSOURCES:.cpp=.o))
-DEPFILES = $(OBJS:.o=.d)
-
-.PHONY: all clean rebuild test
-
-all: $(LIB_TARGET)
-
-$(OBJ_DIR)/%.o: %.cpp
-	@echo "[CC] $< ..."
-	@$(CXX) -c $< $(CXXFLAGS) -o $@
-
-$(OBJ_DIR)/%.d: %.cpp
-	@mkdir -pv $(dir $@)
-	@echo "[dep] $< ..."
-	@$(CXX) $(INCLUDE_DIR) $(CXXFLAGS) -MM -MT "$(OBJ_DIR)/$(<:.cpp=.o) $(OBJ_DIR)/$(<:.cpp=.d)" "$<" > "$@"
-
-sinclude $(DEPFILES)
-
-$(LIB_TARGET): $(OBJS)
-	@echo "[link] $(LIB_TARGET) ..."
-	@$(CXX) $(OBJS) -o $@ $(CXXFLAGS) $(LDFLAGS)
-
-clean:
-	rm -rf $(OBJ_DIR) $(LIB_TARGET)
-
-rebuild:
-	+@make clean
-	+@make
-
-# vim:ft=make
-#
+#
+# Makefile
+# Jiayuan Mao, 2019-01-09 13:59
+#
+
+SRC_DIR = csrc
+INC_DIR = csrc
+OBJ_DIR = build/obj
+TARGET = libpatchmatch.so
+
+LIB_TARGET = $(TARGET)
+INCLUDE_DIR = -I $(SRC_DIR) -I $(INC_DIR)
+
+CXX = $(ENVIRONMENT_OPTIONS) g++
+CXXFLAGS = -std=c++14
+CXXFLAGS += -Ofast -ffast-math -w
+# CXXFLAGS += -g
+CXXFLAGS += $(shell pkg-config --cflags opencv) -fPIC
+CXXFLAGS += $(INCLUDE_DIR)
+LDFLAGS = $(shell pkg-config --cflags --libs opencv) -shared -fPIC
+
+
+CXXSOURCES = $(shell find $(SRC_DIR)/ -name "*.cpp")
+OBJS = $(addprefix $(OBJ_DIR)/,$(CXXSOURCES:.cpp=.o))
+DEPFILES = $(OBJS:.o=.d)
+
+.PHONY: all clean rebuild test
+
+all: $(LIB_TARGET)
+
+$(OBJ_DIR)/%.o: %.cpp
+	@echo "[CC] $< ..."
+	@$(CXX) -c $< $(CXXFLAGS) -o $@
+
+$(OBJ_DIR)/%.d: %.cpp
+	@mkdir -pv $(dir $@)
+	@echo "[dep] $< ..."
+	@$(CXX) $(INCLUDE_DIR) $(CXXFLAGS) -MM -MT "$(OBJ_DIR)/$(<:.cpp=.o) $(OBJ_DIR)/$(<:.cpp=.d)" "$<" > "$@"
+
+sinclude $(DEPFILES)
+
+$(LIB_TARGET): $(OBJS)
+	@echo "[link] $(LIB_TARGET) ..."
+	@$(CXX) $(OBJS) -o $@ $(CXXFLAGS) $(LDFLAGS)
+
+clean:
+	rm -rf $(OBJ_DIR) $(LIB_TARGET)
+
+rebuild:
+	+@make clean
+	+@make
+
+# vim:ft=make
+#

PyPatchMatch/README.md

@@ -1,64 +1,64 @@
-PatchMatch based Inpainting
-=====================================
-This library implements the PatchMatch based inpainting algorithm. It provides both C++ and Python interfaces.
-This implementation is heavily based on the implementation by Younesse ANDAM:
-(younesse-cv/PatchMatch)[https://github.com/younesse-cv/PatchMatch], with some bugs fix.
-
-Usage
--------------------------------------
-
-You need to first install OpenCV to compile the C++ libraries. Then, run `make` to compile the
-shared library `libpatchmatch.so`.
-
-For Python users (example available at `examples/py_example.py`)
-
-```python
-import patch_match
-
-image = ...  # either a numpy ndarray or a PIL Image object.
-mask = ...   # either a numpy ndarray or a PIL Image object.
-result = patch_match.inpaint(image, mask, patch_size=5)
-```
-
-For C++ users (examples available at `examples/cpp_example.cpp`)
-
-```cpp
-#include "inpaint.h"
-
-int main() {
-    cv::Mat image = ...
-    cv::Mat mask = ...
-
-    cv::Mat result = Inpainting(image, mask, 5).run();
-
-    return 0;
-}
-```
-
-
-README and COPYRIGHT by Younesse ANDAM
--------------------------------------
-@Author: Younesse ANDAM
-
-@Contact: younesse.andam@gmail.com
-
-Description: This project is a personal implementation of an algorithm called PATCHMATCH that restores missing areas in an image.
-The algorithm is presented in the following paper
- PatchMatch  A Randomized Correspondence Algorithm
-               for Structural Image Editing
-   by C.Barnes,E.Shechtman,A.Finkelstein and Dan B.Goldman
-   ACM Transactions on Graphics (Proc. SIGGRAPH), vol.28, aug-2009
-
- For more information please refer to
- http://www.cs.princeton.edu/gfx/pubs/Barnes_2009_PAR/index.php
-
-Copyright (c) 2010-2011
-
-
-Requirements
--------------------------------------
-
-To run the project you need to install Opencv library and link it to your project.
-Opencv can be download it here
-http://opencv.org/downloads.html
-
+PatchMatch based Inpainting
+=====================================
+This library implements the PatchMatch based inpainting algorithm. It provides both C++ and Python interfaces.
+This implementation is heavily based on the implementation by Younesse ANDAM:
+(younesse-cv/PatchMatch)[https://github.com/younesse-cv/PatchMatch], with some bugs fix.
+
+Usage
+-------------------------------------
+
+You need to first install OpenCV to compile the C++ libraries. Then, run `make` to compile the
+shared library `libpatchmatch.so`.
+
+For Python users (example available at `examples/py_example.py`)
+
+```python
+import patch_match
+
+image = ...  # either a numpy ndarray or a PIL Image object.
+mask = ...   # either a numpy ndarray or a PIL Image object.
+result = patch_match.inpaint(image, mask, patch_size=5)
+```
+
+For C++ users (examples available at `examples/cpp_example.cpp`)
+
+```cpp
+#include "inpaint.h"
+
+int main() {
+    cv::Mat image = ...
+    cv::Mat mask = ...
+
+    cv::Mat result = Inpainting(image, mask, 5).run();
+
+    return 0;
+}
+```
+
+
+README and COPYRIGHT by Younesse ANDAM
+-------------------------------------
+@Author: Younesse ANDAM
+
+@Contact: younesse.andam@gmail.com
+
+Description: This project is a personal implementation of an algorithm called PATCHMATCH that restores missing areas in an image.
+The algorithm is presented in the following paper
+ PatchMatch  A Randomized Correspondence Algorithm
+               for Structural Image Editing
+   by C.Barnes,E.Shechtman,A.Finkelstein and Dan B.Goldman
+   ACM Transactions on Graphics (Proc. SIGGRAPH), vol.28, aug-2009
+
+ For more information please refer to
+ http://www.cs.princeton.edu/gfx/pubs/Barnes_2009_PAR/index.php
+
+Copyright (c) 2010-2011
+
+
+Requirements
+-------------------------------------
+
+To run the project you need to install Opencv library and link it to your project.
+Opencv can be download it here
+http://opencv.org/downloads.html
+

PyPatchMatch/csrc/inpaint.cpp

@@ -1,234 +1,234 @@
-#include <algorithm>
-#include <iostream>
-#include <opencv2/imgcodecs.hpp>
-#include <opencv2/imgproc.hpp>
-#include <opencv2/highgui.hpp>
-
-#include "inpaint.h"
-
-namespace {
-    static std::vector<double> kDistance2Similarity;
-
-    void init_kDistance2Similarity() {
-        double base[11] = {1.0, 0.99, 0.96, 0.83, 0.38, 0.11, 0.02, 0.005, 0.0006, 0.0001, 0};
-        int length = (PatchDistanceMetric::kDistanceScale + 1);
-        kDistance2Similarity.resize(length);
-        for (int i = 0; i < length; ++i) {
-            double t = (double) i / length;
-            int j = (int) (100 * t);
-            int k = j + 1;
-            double vj = (j < 11) ? base[j] : 0;
-            double vk = (k < 11) ? base[k] : 0;
-            kDistance2Similarity[i] = vj + (100 * t - j) * (vk - vj);
-        }
-    }
-
-
-    inline void _weighted_copy(const MaskedImage &source, int ys, int xs, cv::Mat &target, int yt, int xt, double weight) {
-        if (source.is_masked(ys, xs)) return;
-        if (source.is_globally_masked(ys, xs)) return;
-
-        auto source_ptr = source.get_image(ys, xs);
-        auto target_ptr = target.ptr<double>(yt, xt);
-
-#pragma unroll
-        for (int c = 0; c < 3; ++c)
-            target_ptr[c] += static_cast<double>(source_ptr[c]) * weight;
-        target_ptr[3] += weight;
-    }
-}
-
-/**
- * This algorithme uses a version proposed by Xavier Philippeau.
- */
-
-Inpainting::Inpainting(cv::Mat image, cv::Mat mask, const PatchDistanceMetric *metric)
-    : m_initial(image, mask), m_distance_metric(metric), m_pyramid(), m_source2target(), m_target2source() {
-    _initialize_pyramid();
-}
-
-Inpainting::Inpainting(cv::Mat image, cv::Mat mask, cv::Mat global_mask, const PatchDistanceMetric *metric)
-    : m_initial(image, mask, global_mask), m_distance_metric(metric), m_pyramid(), m_source2target(), m_target2source() {
-    _initialize_pyramid();
-}
-
-void Inpainting::_initialize_pyramid() {
-    auto source = m_initial;
-    m_pyramid.push_back(source);
-    while (source.size().height > m_distance_metric->patch_size() && source.size().width > m_distance_metric->patch_size()) {
-        source = source.downsample();
-        m_pyramid.push_back(source);
-    }
-
-    if (kDistance2Similarity.size() == 0) {
-        init_kDistance2Similarity();
-    }
-}
-
-cv::Mat Inpainting::run(bool verbose, bool verbose_visualize, unsigned int random_seed) {
-    srand(random_seed);
-    const int nr_levels = m_pyramid.size();
-
-    MaskedImage source, target;
-    for (int level = nr_levels - 1; level >= 0; --level) {
-        if (verbose) std::cerr << "Inpainting level: " << level << std::endl;
-
-        source = m_pyramid[level];
-
-        if (level == nr_levels - 1) {
-            target = source.clone();
-            target.clear_mask();
-            m_source2target = NearestNeighborField(source, target, m_distance_metric);
-            m_target2source = NearestNeighborField(target, source, m_distance_metric);
-        } else {
-            m_source2target = NearestNeighborField(source, target, m_distance_metric, m_source2target);
-            m_target2source = NearestNeighborField(target, source, m_distance_metric, m_target2source);
-        }
-
-        if (verbose) std::cerr << "Initialization done." << std::endl;
-
-        if (verbose_visualize) {
-            auto visualize_size = m_initial.size();
-            cv::Mat source_visualize(visualize_size, m_initial.image().type());
-            cv::resize(source.image(), source_visualize, visualize_size);
-            cv::imshow("Source", source_visualize);
-            cv::Mat target_visualize(visualize_size, m_initial.image().type());
-            cv::resize(target.image(), target_visualize, visualize_size);
-            cv::imshow("Target", target_visualize);
-            cv::waitKey(0);
-        }
-
-        target = _expectation_maximization(source, target, level, verbose);
-    }
-
-    return target.image();
-}
-
-// EM-Like algorithm (see "PatchMatch" - page 6).
-// Returns a double sized target image (unless level = 0).
-MaskedImage Inpainting::_expectation_maximization(MaskedImage source, MaskedImage target, int level, bool verbose) {
-    const int nr_iters_em = 1 + 2 * level;
-    const int nr_iters_nnf = static_cast<int>(std::min(7, 1 + level));
-    const int patch_size = m_distance_metric->patch_size();
-
-    MaskedImage new_source, new_target;
-
-    for (int iter_em = 0; iter_em < nr_iters_em; ++iter_em) {
-        if (iter_em != 0) {
-            m_source2target.set_target(new_target);
-            m_target2source.set_source(new_target);
-            target = new_target;
-        }
-
-        if (verbose) std::cerr << "EM Iteration: " << iter_em << std::endl;
-
-        auto size = source.size();
-        for (int i = 0; i < size.height; ++i) {
-            for (int j = 0; j < size.width; ++j) {
-                if (!source.contains_mask(i, j, patch_size)) {
-                    m_source2target.set_identity(i, j);
-                    m_target2source.set_identity(i, j);
-                }
-            }
-        }
-        if (verbose) std::cerr << "  NNF minimization started." << std::endl;
-        m_source2target.minimize(nr_iters_nnf);
-        m_target2source.minimize(nr_iters_nnf);
-        if (verbose) std::cerr << "  NNF minimization finished." << std::endl;
-
-        // Instead of upsizing the final target, we build the last target from the next level source image.
-        // Thus, the final target is less blurry (see "Space-Time Video Completion" - page 5).
-        bool upscaled = false;
-        if (level >= 1 && iter_em == nr_iters_em - 1) {
-            new_source = m_pyramid[level - 1];
-            new_target = target.upsample(new_source.size().width, new_source.size().height, m_pyramid[level - 1].global_mask());
-            upscaled = true;
-        } else {
-            new_source = m_pyramid[level];
-            new_target = target.clone();
-        }
-
-        auto vote = cv::Mat(new_target.size(), CV_64FC4);
-        vote.setTo(cv::Scalar::all(0));
-
-        // Votes for best patch from NNF Source->Target (completeness) and Target->Source (coherence).
-        _expectation_step(m_source2target, 1, vote, new_source, upscaled);
-        if (verbose) std::cerr << "  Expectation source to target finished." << std::endl;
-        _expectation_step(m_target2source, 0, vote, new_source, upscaled);
-        if (verbose) std::cerr << "  Expectation target to source finished." << std::endl;
-
-        // Compile votes and update pixel values.
-        _maximization_step(new_target, vote);
-        if (verbose) std::cerr << "  Minimization step finished." << std::endl;
-    }
-
-    return new_target;
-}
-
-// Expectation step: vote for best estimations of each pixel.
-void Inpainting::_expectation_step(
-    const NearestNeighborField &nnf, bool source2target,
-    cv::Mat &vote, const MaskedImage &source, bool upscaled
-) {
-    auto source_size = nnf.source_size();
-    auto target_size = nnf.target_size();
-    const int patch_size = m_distance_metric->patch_size();
-
-    for (int i = 0; i < source_size.height; ++i) {
-        for (int j = 0; j < source_size.width; ++j) {
-            if (nnf.source().is_globally_masked(i, j)) continue;
-            int yp = nnf.at(i, j, 0), xp = nnf.at(i, j, 1), dp = nnf.at(i, j, 2);
-            double w = kDistance2Similarity[dp];
-
-            for (int di = -patch_size; di <= patch_size; ++di) {
-                for (int dj = -patch_size; dj <= patch_size; ++dj) {
-                    int ys = i + di, xs = j + dj, yt = yp + di, xt = xp + dj;
-                    if (!(ys >= 0 && ys < source_size.height && xs >= 0 && xs < source_size.width)) continue;
-                    if (nnf.source().is_globally_masked(ys, xs)) continue;
-                    if (!(yt >= 0 && yt < target_size.height && xt >= 0 && xt < target_size.width)) continue;
-                    if (nnf.target().is_globally_masked(yt, xt)) continue;
-
-                    if (!source2target) {
-                        std::swap(ys, yt);
-                        std::swap(xs, xt);
-                    }
-
-                    if (upscaled) {
-                        for (int uy = 0; uy < 2; ++uy) {
-                            for (int ux = 0; ux < 2; ++ux) {
-                                _weighted_copy(source, 2 * ys + uy, 2 * xs + ux, vote, 2 * yt + uy, 2 * xt + ux, w);
-                            }
-                        }
-                    } else {
-                        _weighted_copy(source, ys, xs, vote, yt, xt, w);
-                    }
-                }
-            }
-        }
-    }
-}
-
-// Maximization Step: maximum likelihood of target pixel.
-void Inpainting::_maximization_step(MaskedImage &target, const cv::Mat &vote) {
-    auto target_size = target.size();
-    for (int i = 0; i < target_size.height; ++i) {
-        for (int j = 0; j < target_size.width; ++j) {
-            const double *source_ptr = vote.ptr<double>(i, j);
-            unsigned char *target_ptr = target.get_mutable_image(i, j);
-
-            if (target.is_globally_masked(i, j)) {
-                continue;
-            }
-
-            if (source_ptr[3] > 0) {
-                unsigned char r = cv::saturate_cast<unsigned char>(source_ptr[0] / source_ptr[3]);
-                unsigned char g = cv::saturate_cast<unsigned char>(source_ptr[1] / source_ptr[3]);
-                unsigned char b = cv::saturate_cast<unsigned char>(source_ptr[2] / source_ptr[3]);
-                target_ptr[0] = r, target_ptr[1] = g, target_ptr[2] = b;
-            } else {
-                target.set_mask(i, j, 0);
-            }
-        }
-    }
-}
-
+#include <algorithm>
+#include <iostream>
+#include <opencv2/imgcodecs.hpp>
+#include <opencv2/imgproc.hpp>
+#include <opencv2/highgui.hpp>
+
+#include "inpaint.h"
+
+namespace {
+    static std::vector<double> kDistance2Similarity;
+
+    void init_kDistance2Similarity() {
+        double base[11] = {1.0, 0.99, 0.96, 0.83, 0.38, 0.11, 0.02, 0.005, 0.0006, 0.0001, 0};
+        int length = (PatchDistanceMetric::kDistanceScale + 1);
+        kDistance2Similarity.resize(length);
+        for (int i = 0; i < length; ++i) {
+            double t = (double) i / length;
+            int j = (int) (100 * t);
+            int k = j + 1;
+            double vj = (j < 11) ? base[j] : 0;
+            double vk = (k < 11) ? base[k] : 0;
+            kDistance2Similarity[i] = vj + (100 * t - j) * (vk - vj);
+        }
+    }
+
+
+    inline void _weighted_copy(const MaskedImage &source, int ys, int xs, cv::Mat &target, int yt, int xt, double weight) {
+        if (source.is_masked(ys, xs)) return;
+        if (source.is_globally_masked(ys, xs)) return;
+
+        auto source_ptr = source.get_image(ys, xs);
+        auto target_ptr = target.ptr<double>(yt, xt);
+
+#pragma unroll
+        for (int c = 0; c < 3; ++c)
+            target_ptr[c] += static_cast<double>(source_ptr[c]) * weight;
+        target_ptr[3] += weight;
+    }
+}
+
+/**
+ * This algorithme uses a version proposed by Xavier Philippeau.
+ */
+
+Inpainting::Inpainting(cv::Mat image, cv::Mat mask, const PatchDistanceMetric *metric)
+    : m_initial(image, mask), m_distance_metric(metric), m_pyramid(), m_source2target(), m_target2source() {
+    _initialize_pyramid();
+}
+
+Inpainting::Inpainting(cv::Mat image, cv::Mat mask, cv::Mat global_mask, const PatchDistanceMetric *metric)
+    : m_initial(image, mask, global_mask), m_distance_metric(metric), m_pyramid(), m_source2target(), m_target2source() {
+    _initialize_pyramid();
+}
+
+void Inpainting::_initialize_pyramid() {
+    auto source = m_initial;
+    m_pyramid.push_back(source);
+    while (source.size().height > m_distance_metric->patch_size() && source.size().width > m_distance_metric->patch_size()) {
+        source = source.downsample();
+        m_pyramid.push_back(source);
+    }
+
+    if (kDistance2Similarity.size() == 0) {
+        init_kDistance2Similarity();
+    }
+}
+
+cv::Mat Inpainting::run(bool verbose, bool verbose_visualize, unsigned int random_seed) {
+    srand(random_seed);
+    const int nr_levels = m_pyramid.size();
+
+    MaskedImage source, target;
+    for (int level = nr_levels - 1; level >= 0; --level) {
+        if (verbose) std::cerr << "Inpainting level: " << level << std::endl;
+
+        source = m_pyramid[level];
+
+        if (level == nr_levels - 1) {
+            target = source.clone();
+            target.clear_mask();
+            m_source2target = NearestNeighborField(source, target, m_distance_metric);
+            m_target2source = NearestNeighborField(target, source, m_distance_metric);
+        } else {
+            m_source2target = NearestNeighborField(source, target, m_distance_metric, m_source2target);
+            m_target2source = NearestNeighborField(target, source, m_distance_metric, m_target2source);
+        }
+
+        if (verbose) std::cerr << "Initialization done." << std::endl;
+
+        if (verbose_visualize) {
+            auto visualize_size = m_initial.size();
+            cv::Mat source_visualize(visualize_size, m_initial.image().type());
+            cv::resize(source.image(), source_visualize, visualize_size);
+            cv::imshow("Source", source_visualize);
+            cv::Mat target_visualize(visualize_size, m_initial.image().type());
+            cv::resize(target.image(), target_visualize, visualize_size);
+            cv::imshow("Target", target_visualize);
+            cv::waitKey(0);
+        }
+
+        target = _expectation_maximization(source, target, level, verbose);
+    }
+
+    return target.image();
+}
+
+// EM-Like algorithm (see "PatchMatch" - page 6).
+// Returns a double sized target image (unless level = 0).
+MaskedImage Inpainting::_expectation_maximization(MaskedImage source, MaskedImage target, int level, bool verbose) {
+    const int nr_iters_em = 1 + 2 * level;
+    const int nr_iters_nnf = static_cast<int>(std::min(7, 1 + level));
+    const int patch_size = m_distance_metric->patch_size();
+
+    MaskedImage new_source, new_target;
+
+    for (int iter_em = 0; iter_em < nr_iters_em; ++iter_em) {
+        if (iter_em != 0) {
+            m_source2target.set_target(new_target);
+            m_target2source.set_source(new_target);
+            target = new_target;
+        }
+
+        if (verbose) std::cerr << "EM Iteration: " << iter_em << std::endl;
+
+        auto size = source.size();
+        for (int i = 0; i < size.height; ++i) {
+            for (int j = 0; j < size.width; ++j) {
+                if (!source.contains_mask(i, j, patch_size)) {
+                    m_source2target.set_identity(i, j);
+                    m_target2source.set_identity(i, j);
+                }
+            }
+        }
+        if (verbose) std::cerr << "  NNF minimization started." << std::endl;
+        m_source2target.minimize(nr_iters_nnf);
+        m_target2source.minimize(nr_iters_nnf);
+        if (verbose) std::cerr << "  NNF minimization finished." << std::endl;
+
+        // Instead of upsizing the final target, we build the last target from the next level source image.
+        // Thus, the final target is less blurry (see "Space-Time Video Completion" - page 5).
+        bool upscaled = false;
+        if (level >= 1 && iter_em == nr_iters_em - 1) {
+            new_source = m_pyramid[level - 1];
+            new_target = target.upsample(new_source.size().width, new_source.size().height, m_pyramid[level - 1].global_mask());
+            upscaled = true;
+        } else {
+            new_source = m_pyramid[level];
+            new_target = target.clone();
+        }
+
+        auto vote = cv::Mat(new_target.size(), CV_64FC4);
+        vote.setTo(cv::Scalar::all(0));
+
+        // Votes for best patch from NNF Source->Target (completeness) and Target->Source (coherence).
+        _expectation_step(m_source2target, 1, vote, new_source, upscaled);
+        if (verbose) std::cerr << "  Expectation source to target finished." << std::endl;
+        _expectation_step(m_target2source, 0, vote, new_source, upscaled);
+        if (verbose) std::cerr << "  Expectation target to source finished." << std::endl;
+
+        // Compile votes and update pixel values.
+        _maximization_step(new_target, vote);
+        if (verbose) std::cerr << "  Minimization step finished." << std::endl;
+    }
+
+    return new_target;
+}
+
+// Expectation step: vote for best estimations of each pixel.
+void Inpainting::_expectation_step(
+    const NearestNeighborField &nnf, bool source2target,
+    cv::Mat &vote, const MaskedImage &source, bool upscaled
+) {
+    auto source_size = nnf.source_size();
+    auto target_size = nnf.target_size();
+    const int patch_size = m_distance_metric->patch_size();
+
+    for (int i = 0; i < source_size.height; ++i) {
+        for (int j = 0; j < source_size.width; ++j) {
+            if (nnf.source().is_globally_masked(i, j)) continue;
+            int yp = nnf.at(i, j, 0), xp = nnf.at(i, j, 1), dp = nnf.at(i, j, 2);
+            double w = kDistance2Similarity[dp];
+
+            for (int di = -patch_size; di <= patch_size; ++di) {
+                for (int dj = -patch_size; dj <= patch_size; ++dj) {
+                    int ys = i + di, xs = j + dj, yt = yp + di, xt = xp + dj;
+                    if (!(ys >= 0 && ys < source_size.height && xs >= 0 && xs < source_size.width)) continue;
+                    if (nnf.source().is_globally_masked(ys, xs)) continue;
+                    if (!(yt >= 0 && yt < target_size.height && xt >= 0 && xt < target_size.width)) continue;
+                    if (nnf.target().is_globally_masked(yt, xt)) continue;
+
+                    if (!source2target) {
+                        std::swap(ys, yt);
+                        std::swap(xs, xt);
+                    }
+
+                    if (upscaled) {
+                        for (int uy = 0; uy < 2; ++uy) {
+                            for (int ux = 0; ux < 2; ++ux) {
+                                _weighted_copy(source, 2 * ys + uy, 2 * xs + ux, vote, 2 * yt + uy, 2 * xt + ux, w);
+                            }
+                        }
+                    } else {
+                        _weighted_copy(source, ys, xs, vote, yt, xt, w);
+                    }
+                }
+            }
+        }
+    }
+}
+
+// Maximization Step: maximum likelihood of target pixel.
+void Inpainting::_maximization_step(MaskedImage &target, const cv::Mat &vote) {
+    auto target_size = target.size();
+    for (int i = 0; i < target_size.height; ++i) {
+        for (int j = 0; j < target_size.width; ++j) {
+            const double *source_ptr = vote.ptr<double>(i, j);
+            unsigned char *target_ptr = target.get_mutable_image(i, j);
+
+            if (target.is_globally_masked(i, j)) {
+                continue;
+            }
+
+            if (source_ptr[3] > 0) {
+                unsigned char r = cv::saturate_cast<unsigned char>(source_ptr[0] / source_ptr[3]);
+                unsigned char g = cv::saturate_cast<unsigned char>(source_ptr[1] / source_ptr[3]);
+                unsigned char b = cv::saturate_cast<unsigned char>(source_ptr[2] / source_ptr[3]);
+                target_ptr[0] = r, target_ptr[1] = g, target_ptr[2] = b;
+            } else {
+                target.set_mask(i, j, 0);
+            }
+        }
+    }
+}
+

PyPatchMatch/csrc/inpaint.h

@@ -1,27 +1,27 @@
-#pragma once
-
-#include <vector>
-
-#include "masked_image.h"
-#include "nnf.h"
-
-class Inpainting {
-public:
-    Inpainting(cv::Mat image, cv::Mat mask, const PatchDistanceMetric *metric);
-    Inpainting(cv::Mat image, cv::Mat mask, cv::Mat global_mask, const PatchDistanceMetric *metric);
-    cv::Mat run(bool verbose = false, bool verbose_visualize = false, unsigned int random_seed = 1212);
-
-private:
-    void _initialize_pyramid(void);
-    MaskedImage _expectation_maximization(MaskedImage source, MaskedImage target, int level, bool verbose);
-    void _expectation_step(const NearestNeighborField &nnf, bool source2target, cv::Mat &vote, const MaskedImage &source, bool upscaled);
-    void _maximization_step(MaskedImage &target, const cv::Mat &vote);
-
-    MaskedImage m_initial;
-    std::vector<MaskedImage> m_pyramid;
-
-    NearestNeighborField m_source2target;
-    NearestNeighborField m_target2source;
-    const PatchDistanceMetric *m_distance_metric;
-};
-
+#pragma once
+
+#include <vector>
+
+#include "masked_image.h"
+#include "nnf.h"
+
+class Inpainting {
+public:
+    Inpainting(cv::Mat image, cv::Mat mask, const PatchDistanceMetric *metric);
+    Inpainting(cv::Mat image, cv::Mat mask, cv::Mat global_mask, const PatchDistanceMetric *metric);
+    cv::Mat run(bool verbose = false, bool verbose_visualize = false, unsigned int random_seed = 1212);
+
+private:
+    void _initialize_pyramid(void);
+    MaskedImage _expectation_maximization(MaskedImage source, MaskedImage target, int level, bool verbose);
+    void _expectation_step(const NearestNeighborField &nnf, bool source2target, cv::Mat &vote, const MaskedImage &source, bool upscaled);
+    void _maximization_step(MaskedImage &target, const cv::Mat &vote);
+
+    MaskedImage m_initial;
+    std::vector<MaskedImage> m_pyramid;
+
+    NearestNeighborField m_source2target;
+    NearestNeighborField m_target2source;
+    const PatchDistanceMetric *m_distance_metric;
+};
+

PyPatchMatch/csrc/masked_image.cpp

@@ -1,138 +1,138 @@
-#include "masked_image.h"
-#include <algorithm>
-#include <iostream>
-
-const cv::Size MaskedImage::kDownsampleKernelSize = cv::Size(6, 6);
-const int MaskedImage::kDownsampleKernel[6] = {1, 5, 10, 10, 5, 1};
-
-bool MaskedImage::contains_mask(int y, int x, int patch_size) const {
-    auto mask_size = size();
-    for (int dy = -patch_size; dy <= patch_size; ++dy) {
-        for (int dx = -patch_size; dx <= patch_size; ++dx) {
-            int yy = y + dy, xx = x + dx;
-            if (yy >= 0 && yy < mask_size.height && xx >= 0 && xx < mask_size.width) {
-                if (is_masked(yy, xx) && !is_globally_masked(yy, xx)) return true;
-            }
-        }
-    }
-    return false;
-}
-
-MaskedImage MaskedImage::downsample() const {
-    const auto &kernel_size = MaskedImage::kDownsampleKernelSize;
-    const auto &kernel = MaskedImage::kDownsampleKernel;
-
-    const auto size = this->size();
-    const auto new_size = cv::Size(size.width / 2, size.height / 2);
-
-    auto ret = MaskedImage(new_size.width, new_size.height);
-    if (!m_global_mask.empty()) ret.init_global_mask_mat();
-    for (int y = 0; y < size.height - 1; y += 2) {
-        for (int x = 0; x < size.width - 1; x += 2) {
-            int r = 0, g = 0, b = 0, ksum = 0;
-            bool is_gmasked = true;
-
-            for (int dy = -kernel_size.height / 2 + 1; dy <= kernel_size.height / 2; ++dy) {
-                for (int dx = -kernel_size.width / 2 + 1; dx <= kernel_size.width / 2; ++dx) {
-                    int yy = y + dy, xx = x + dx;
-                    if (yy >= 0 && yy < size.height && xx >= 0 && xx < size.width) {
-                        if (!is_globally_masked(yy, xx)) {
-                            is_gmasked = false;
-                        }
-                        if (!is_masked(yy, xx)) {
-                            auto source_ptr = get_image(yy, xx);
-                            int k = kernel[kernel_size.height / 2 - 1 + dy] * kernel[kernel_size.width / 2 - 1 + dx];
-                            r += source_ptr[0] * k, g += source_ptr[1] * k, b += source_ptr[2] * k;
-                            ksum += k;
-                        }
-                    }
-                }
-            }
-
-            if (ksum > 0) r /= ksum, g /= ksum, b /= ksum;
-
-            if (!m_global_mask.empty()) {
-                ret.set_global_mask(y / 2, x / 2, is_gmasked);
-            }
-            if (ksum > 0) {
-                auto target_ptr = ret.get_mutable_image(y / 2, x / 2);
-                target_ptr[0] = r, target_ptr[1] = g, target_ptr[2] = b;
-                ret.set_mask(y / 2, x / 2, 0);
-            } else {
-                ret.set_mask(y / 2, x / 2, 1);
-            }
-        }
-    }
-
-    return ret;
-}
-
-MaskedImage MaskedImage::upsample(int new_w, int new_h) const {
-    const auto size = this->size();
-    auto ret = MaskedImage(new_w, new_h);
-    if (!m_global_mask.empty()) ret.init_global_mask_mat();
-    for (int y = 0; y < new_h; ++y) {
-        for (int x = 0; x < new_w; ++x) {
-            int yy = y * size.height / new_h;
-            int xx = x * size.width / new_w;
-
-            if (is_globally_masked(yy, xx)) {
-                ret.set_global_mask(y, x, 1);
-                ret.set_mask(y, x, 1);
-            } else {
-                if (!m_global_mask.empty()) ret.set_global_mask(y, x, 0);
-
-                if (is_masked(yy, xx)) {
-                    ret.set_mask(y, x, 1);
-                } else {
-                    auto source_ptr = get_image(yy, xx);
-                    auto target_ptr = ret.get_mutable_image(y, x);
-                    for (int c = 0; c < 3; ++c)
-                        target_ptr[c] = source_ptr[c];
-                    ret.set_mask(y, x, 0);
-                }
-            }
-        }
-    }
-
-    return ret;
-}
-
-MaskedImage MaskedImage::upsample(int new_w, int new_h, const cv::Mat &new_global_mask) const {
-    auto ret = upsample(new_w, new_h);
-    ret.set_global_mask_mat(new_global_mask);
-    return ret;
-}
-
-void MaskedImage::compute_image_gradients() {
-    if (m_image_grad_computed) {
-        return;
-    }
-
-    const auto size = m_image.size();
-    m_image_grady = cv::Mat(size, CV_8UC3);
-    m_image_gradx = cv::Mat(size, CV_8UC3);
-    m_image_grady = cv::Scalar::all(0);
-    m_image_gradx = cv::Scalar::all(0);
-
-    for (int i = 1; i < size.height - 1; ++i) {
-        const auto *ptr = m_image.ptr<unsigned char>(i, 0);
-        const auto *ptry1 = m_image.ptr<unsigned char>(i + 1, 0);
-        const auto *ptry2 = m_image.ptr<unsigned char>(i - 1, 0);
-        const auto *ptrx1 = m_image.ptr<unsigned char>(i, 0) + 3;
-        const auto *ptrx2 = m_image.ptr<unsigned char>(i, 0) - 3;
-        auto *mptry = m_image_grady.ptr<unsigned char>(i, 0);
-        auto *mptrx = m_image_gradx.ptr<unsigned char>(i, 0);
-        for (int j = 3; j < size.width * 3 - 3; ++j) {
-            mptry[j] = (ptry1[j] / 2 - ptry2[j] / 2) + 128;
-            mptrx[j] = (ptrx1[j] / 2 - ptrx2[j] / 2) + 128;
-        }
-    }
-
-    m_image_grad_computed = true;
-}
-
-void MaskedImage::compute_image_gradients() const {
-    const_cast<MaskedImage *>(this)->compute_image_gradients();
-}
-
+#include "masked_image.h"
+#include <algorithm>
+#include <iostream>
+
+const cv::Size MaskedImage::kDownsampleKernelSize = cv::Size(6, 6);
+const int MaskedImage::kDownsampleKernel[6] = {1, 5, 10, 10, 5, 1};
+
+bool MaskedImage::contains_mask(int y, int x, int patch_size) const {
+    auto mask_size = size();
+    for (int dy = -patch_size; dy <= patch_size; ++dy) {
+        for (int dx = -patch_size; dx <= patch_size; ++dx) {
+            int yy = y + dy, xx = x + dx;
+            if (yy >= 0 && yy < mask_size.height && xx >= 0 && xx < mask_size.width) {
+                if (is_masked(yy, xx) && !is_globally_masked(yy, xx)) return true;
+            }
+        }
+    }
+    return false;
+}
+
+MaskedImage MaskedImage::downsample() const {
+    const auto &kernel_size = MaskedImage::kDownsampleKernelSize;
+    const auto &kernel = MaskedImage::kDownsampleKernel;
+
+    const auto size = this->size();
+    const auto new_size = cv::Size(size.width / 2, size.height / 2);
+
+    auto ret = MaskedImage(new_size.width, new_size.height);
+    if (!m_global_mask.empty()) ret.init_global_mask_mat();
+    for (int y = 0; y < size.height - 1; y += 2) {
+        for (int x = 0; x < size.width - 1; x += 2) {
+            int r = 0, g = 0, b = 0, ksum = 0;
+            bool is_gmasked = true;
+
+            for (int dy = -kernel_size.height / 2 + 1; dy <= kernel_size.height / 2; ++dy) {
+                for (int dx = -kernel_size.width / 2 + 1; dx <= kernel_size.width / 2; ++dx) {
+                    int yy = y + dy, xx = x + dx;
+                    if (yy >= 0 && yy < size.height && xx >= 0 && xx < size.width) {
+                        if (!is_globally_masked(yy, xx)) {
+                            is_gmasked = false;
+                        }
+                        if (!is_masked(yy, xx)) {
+                            auto source_ptr = get_image(yy, xx);
+                            int k = kernel[kernel_size.height / 2 - 1 + dy] * kernel[kernel_size.width / 2 - 1 + dx];
+                            r += source_ptr[0] * k, g += source_ptr[1] * k, b += source_ptr[2] * k;
+                            ksum += k;
+                        }
+                    }
+                }
+            }
+
+            if (ksum > 0) r /= ksum, g /= ksum, b /= ksum;
+
+            if (!m_global_mask.empty()) {
+                ret.set_global_mask(y / 2, x / 2, is_gmasked);
+            }
+            if (ksum > 0) {
+                auto target_ptr = ret.get_mutable_image(y / 2, x / 2);
+                target_ptr[0] = r, target_ptr[1] = g, target_ptr[2] = b;
+                ret.set_mask(y / 2, x / 2, 0);
+            } else {
+                ret.set_mask(y / 2, x / 2, 1);
+            }
+        }
+    }
+
+    return ret;
+}
+
+MaskedImage MaskedImage::upsample(int new_w, int new_h) const {
+    const auto size = this->size();
+    auto ret = MaskedImage(new_w, new_h);
+    if (!m_global_mask.empty()) ret.init_global_mask_mat();
+    for (int y = 0; y < new_h; ++y) {
+        for (int x = 0; x < new_w; ++x) {
+            int yy = y * size.height / new_h;
+            int xx = x * size.width / new_w;
+
+            if (is_globally_masked(yy, xx)) {
+                ret.set_global_mask(y, x, 1);
+                ret.set_mask(y, x, 1);
+            } else {
+                if (!m_global_mask.empty()) ret.set_global_mask(y, x, 0);
+
+                if (is_masked(yy, xx)) {
+                    ret.set_mask(y, x, 1);
+                } else {
+                    auto source_ptr = get_image(yy, xx);
+                    auto target_ptr = ret.get_mutable_image(y, x);
+                    for (int c = 0; c < 3; ++c)
+                        target_ptr[c] = source_ptr[c];
+                    ret.set_mask(y, x, 0);
+                }
+            }
+        }
+    }
+
+    return ret;
+}
+
+MaskedImage MaskedImage::upsample(int new_w, int new_h, const cv::Mat &new_global_mask) const {
+    auto ret = upsample(new_w, new_h);
+    ret.set_global_mask_mat(new_global_mask);
+    return ret;
+}
+
+void MaskedImage::compute_image_gradients() {
+    if (m_image_grad_computed) {
+        return;
+    }
+
+    const auto size = m_image.size();
+    m_image_grady = cv::Mat(size, CV_8UC3);
+    m_image_gradx = cv::Mat(size, CV_8UC3);
+    m_image_grady = cv::Scalar::all(0);
+    m_image_gradx = cv::Scalar::all(0);
+
+    for (int i = 1; i < size.height - 1; ++i) {
+        const auto *ptr = m_image.ptr<unsigned char>(i, 0);
+        const auto *ptry1 = m_image.ptr<unsigned char>(i + 1, 0);
+        const auto *ptry2 = m_image.ptr<unsigned char>(i - 1, 0);
+        const auto *ptrx1 = m_image.ptr<unsigned char>(i, 0) + 3;
+        const auto *ptrx2 = m_image.ptr<unsigned char>(i, 0) - 3;
+        auto *mptry = m_image_grady.ptr<unsigned char>(i, 0);
+        auto *mptrx = m_image_gradx.ptr<unsigned char>(i, 0);
+        for (int j = 3; j < size.width * 3 - 3; ++j) {
+            mptry[j] = (ptry1[j] / 2 - ptry2[j] / 2) + 128;
+            mptrx[j] = (ptrx1[j] / 2 - ptrx2[j] / 2) + 128;
+        }
+    }
+
+    m_image_grad_computed = true;
+}
+
+void MaskedImage::compute_image_gradients() const {
+    const_cast<MaskedImage *>(this)->compute_image_gradients();
+}
+

PyPatchMatch/csrc/masked_image.h

@@ -1,112 +1,112 @@
-#pragma once
-
-#include <opencv2/core.hpp>
-
-class MaskedImage {
-public:
-    MaskedImage() : m_image(), m_mask(), m_global_mask(), m_image_grady(), m_image_gradx(), m_image_grad_computed(false) {
-        // pass
-    }
-    MaskedImage(cv::Mat image, cv::Mat mask) : m_image(image), m_mask(mask), m_image_grad_computed(false) {
-        // pass
-    }
-    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) {
-        // pass
-    }
-    MaskedImage(cv::Mat image, cv::Mat mask, cv::Mat global_mask, cv::Mat grady, cv::Mat gradx, bool grad_computed) :
-        m_image(image), m_mask(mask), m_global_mask(global_mask),
-        m_image_grady(grady), m_image_gradx(gradx), m_image_grad_computed(grad_computed) {
-        // pass
-    }
-    MaskedImage(int width, int height) : m_global_mask(), m_image_grady(), m_image_gradx() {
-        m_image = cv::Mat(cv::Size(width, height), CV_8UC3);
-        m_image = cv::Scalar::all(0);
-
-        m_mask = cv::Mat(cv::Size(width, height), CV_8U);
-        m_mask = cv::Scalar::all(0);
-    }
-    inline MaskedImage clone() {
-        return MaskedImage(
-            m_image.clone(), m_mask.clone(), m_global_mask.clone(),
-            m_image_grady.clone(), m_image_gradx.clone(), m_image_grad_computed
-        );
-    }
-
-    inline cv::Size size() const {
-        return m_image.size();
-    }
-    inline const cv::Mat &image() const {
-        return m_image;
-    }
-    inline const cv::Mat &mask() const {
-        return m_mask;
-    }
-    inline const cv::Mat &global_mask() const {
-        return m_global_mask;
-    }
-    inline const cv::Mat &grady() const {
-        assert(m_image_grad_computed);
-        return m_image_grady;
-    }
-    inline const cv::Mat &gradx() const {
-        assert(m_image_grad_computed);
-        return m_image_gradx;
-    }
-
-    inline void init_global_mask_mat() {
-        m_global_mask = cv::Mat(m_mask.size(), CV_8U);
-        m_global_mask.setTo(cv::Scalar(0));
-    }
-    inline void set_global_mask_mat(const cv::Mat &other) {
-        m_global_mask = other;
-    }
-
-    inline bool is_masked(int y, int x) const {
-        return static_cast<bool>(m_mask.at<unsigned char>(y, x));
-    }
-    inline bool is_globally_masked(int y, int x) const {
-        return !m_global_mask.empty() && static_cast<bool>(m_global_mask.at<unsigned char>(y, x));
-    }
-    inline void set_mask(int y, int x, bool value) {
-        m_mask.at<unsigned char>(y, x) = static_cast<unsigned char>(value);
-    }
-    inline void set_global_mask(int y, int x, bool value) {
-        m_global_mask.at<unsigned char>(y, x) = static_cast<unsigned char>(value);
-    }
-    inline void clear_mask() {
-        m_mask.setTo(cv::Scalar(0));
-    }
-
-    inline const unsigned char *get_image(int y, int x) const {
-        return m_image.ptr<unsigned char>(y, x);
-    }
-    inline unsigned char *get_mutable_image(int y, int x) {
-        return m_image.ptr<unsigned char>(y, x);
-    }
-
-    inline unsigned char get_image(int y, int x, int c) const {
-        return m_image.ptr<unsigned char>(y, x)[c];
-    }
-    inline int get_image_int(int y, int x, int c) const {
-        return static_cast<int>(m_image.ptr<unsigned char>(y, x)[c]);
-    }
-
-    bool contains_mask(int y, int x, int patch_size) const;
-    MaskedImage downsample() const;
-    MaskedImage upsample(int new_w, int new_h) const;
-    MaskedImage upsample(int new_w, int new_h, const cv::Mat &new_global_mask) const;
-    void compute_image_gradients();
-    void compute_image_gradients() const;
-
-    static const cv::Size kDownsampleKernelSize;
-    static const int kDownsampleKernel[6];
-
-private:
-	cv::Mat m_image;
-	cv::Mat m_mask;
-    cv::Mat m_global_mask;
-    cv::Mat m_image_grady;
-    cv::Mat m_image_gradx;
-    bool m_image_grad_computed = false;
-};
-
+#pragma once
+
+#include <opencv2/core.hpp>
+
+class MaskedImage {
+public:
+    MaskedImage() : m_image(), m_mask(), m_global_mask(), m_image_grady(), m_image_gradx(), m_image_grad_computed(false) {
+        // pass
+    }
+    MaskedImage(cv::Mat image, cv::Mat mask) : m_image(image), m_mask(mask), m_image_grad_computed(false) {
+        // pass
+    }
+    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) {
+        // pass
+    }
+    MaskedImage(cv::Mat image, cv::Mat mask, cv::Mat global_mask, cv::Mat grady, cv::Mat gradx, bool grad_computed) :
+        m_image(image), m_mask(mask), m_global_mask(global_mask),
+        m_image_grady(grady), m_image_gradx(gradx), m_image_grad_computed(grad_computed) {
+        // pass
+    }
+    MaskedImage(int width, int height) : m_global_mask(), m_image_grady(), m_image_gradx() {
+        m_image = cv::Mat(cv::Size(width, height), CV_8UC3);
+        m_image = cv::Scalar::all(0);
+
+        m_mask = cv::Mat(cv::Size(width, height), CV_8U);
+        m_mask = cv::Scalar::all(0);
+    }
+    inline MaskedImage clone() {
+        return MaskedImage(
+            m_image.clone(), m_mask.clone(), m_global_mask.clone(),
+            m_image_grady.clone(), m_image_gradx.clone(), m_image_grad_computed
+        );
+    }
+
+    inline cv::Size size() const {
+        return m_image.size();
+    }
+    inline const cv::Mat &image() const {
+        return m_image;
+    }
+    inline const cv::Mat &mask() const {
+        return m_mask;
+    }
+    inline const cv::Mat &global_mask() const {
+        return m_global_mask;
+    }
+    inline const cv::Mat &grady() const {
+        assert(m_image_grad_computed);
+        return m_image_grady;
+    }
+    inline const cv::Mat &gradx() const {
+        assert(m_image_grad_computed);
+        return m_image_gradx;
+    }
+
+    inline void init_global_mask_mat() {
+        m_global_mask = cv::Mat(m_mask.size(), CV_8U);
+        m_global_mask.setTo(cv::Scalar(0));
+    }
+    inline void set_global_mask_mat(const cv::Mat &other) {
+        m_global_mask = other;
+    }
+
+    inline bool is_masked(int y, int x) const {
+        return static_cast<bool>(m_mask.at<unsigned char>(y, x));
+    }
+    inline bool is_globally_masked(int y, int x) const {
+        return !m_global_mask.empty() && static_cast<bool>(m_global_mask.at<unsigned char>(y, x));
+    }
+    inline void set_mask(int y, int x, bool value) {
+        m_mask.at<unsigned char>(y, x) = static_cast<unsigned char>(value);
+    }
+    inline void set_global_mask(int y, int x, bool value) {
+        m_global_mask.at<unsigned char>(y, x) = static_cast<unsigned char>(value);
+    }
+    inline void clear_mask() {
+        m_mask.setTo(cv::Scalar(0));
+    }
+
+    inline const unsigned char *get_image(int y, int x) const {
+        return m_image.ptr<unsigned char>(y, x);
+    }
+    inline unsigned char *get_mutable_image(int y, int x) {
+        return m_image.ptr<unsigned char>(y, x);
+    }
+
+    inline unsigned char get_image(int y, int x, int c) const {
+        return m_image.ptr<unsigned char>(y, x)[c];
+    }
+    inline int get_image_int(int y, int x, int c) const {
+        return static_cast<int>(m_image.ptr<unsigned char>(y, x)[c]);
+    }
+
+    bool contains_mask(int y, int x, int patch_size) const;
+    MaskedImage downsample() const;
+    MaskedImage upsample(int new_w, int new_h) const;
+    MaskedImage upsample(int new_w, int new_h, const cv::Mat &new_global_mask) const;
+    void compute_image_gradients();
+    void compute_image_gradients() const;
+
+    static const cv::Size kDownsampleKernelSize;
+    static const int kDownsampleKernel[6];
+
+private:
+	cv::Mat m_image;
+	cv::Mat m_mask;
+    cv::Mat m_global_mask;
+    cv::Mat m_image_grady;
+    cv::Mat m_image_gradx;
+    bool m_image_grad_computed = false;
+};
+

PyPatchMatch/csrc/nnf.cpp

@@ -1,268 +1,268 @@
-#include <algorithm>
-#include <iostream>
-#include <cmath>
-
-#include "masked_image.h"
-#include "nnf.h"
-
-/**
-* Nearest-Neighbor Field (see PatchMatch algorithm).
-* This algorithme uses a version proposed by Xavier Philippeau.
-*
-*/
-
-template <typename T>
-T clamp(T value, T min_value, T max_value) {
-    return std::min(std::max(value, min_value), max_value);
-}
-
-void NearestNeighborField::_randomize_field(int max_retry, bool reset) {
-    auto this_size = source_size();
-    for (int i = 0; i < this_size.height; ++i) {
-        for (int j = 0; j < this_size.width; ++j) {
-            if (m_source.is_globally_masked(i, j)) continue;
-
-            auto this_ptr = mutable_ptr(i, j);
-            int distance = reset ? PatchDistanceMetric::kDistanceScale : this_ptr[2];
-            if (distance < PatchDistanceMetric::kDistanceScale) {
-                continue;
-            }
-
-            int i_target = 0, j_target = 0;
-            for (int t = 0; t < max_retry; ++t) {
-                i_target = rand() % this_size.height;
-                j_target = rand() % this_size.width;
-                if (m_target.is_globally_masked(i_target, j_target)) continue;
-
-                distance = _distance(i, j, i_target, j_target);
-                if (distance < PatchDistanceMetric::kDistanceScale)
-                    break;
-            }
-
-            this_ptr[0] = i_target, this_ptr[1] = j_target, this_ptr[2] = distance;
-        }
-    }
-}
-
-void NearestNeighborField::_initialize_field_from(const NearestNeighborField &other, int max_retry) {
-    const auto &this_size = source_size();
-    const auto &other_size = other.source_size();
-    double fi = static_cast<double>(this_size.height) / other_size.height;
-    double fj = static_cast<double>(this_size.width) / other_size.width;
-
-    for (int i = 0; i < this_size.height; ++i) {
-        for (int j = 0; j < this_size.width; ++j) {
-            if (m_source.is_globally_masked(i, j)) continue;
-
-            int ilow = static_cast<int>(std::min(i / fi, static_cast<double>(other_size.height - 1)));
-            int jlow = static_cast<int>(std::min(j / fj, static_cast<double>(other_size.width - 1)));
-            auto this_value = mutable_ptr(i, j);
-            auto other_value = other.ptr(ilow, jlow);
-
-            this_value[0] = static_cast<int>(other_value[0] * fi);
-            this_value[1] = static_cast<int>(other_value[1] * fj);
-            this_value[2] = _distance(i, j, this_value[0], this_value[1]);
-        }
-    }
-
-    _randomize_field(max_retry, false);
-}
-
-void NearestNeighborField::minimize(int nr_pass) {
-    const auto &this_size = source_size();
-    while (nr_pass--) {
-        for (int i = 0; i < this_size.height; ++i)
-            for (int j = 0; j < this_size.width; ++j) {
-                if (m_source.is_globally_masked(i, j)) continue;
-                if (at(i, j, 2) > 0) _minimize_link(i, j, +1);
-            }
-        for (int i = this_size.height - 1; i >= 0; --i)
-            for (int j = this_size.width - 1; j >= 0; --j) {
-                if (m_source.is_globally_masked(i, j)) continue;
-                if (at(i, j, 2) > 0) _minimize_link(i, j, -1);
-            }
-    }
-}
-
-void NearestNeighborField::_minimize_link(int y, int x, int direction) {
-    const auto &this_size = source_size();
-    const auto &this_target_size = target_size();
-    auto this_ptr = mutable_ptr(y, x);
-
-    // propagation along the y direction.
-    if (y - direction >= 0 && y - direction < this_size.height && !m_source.is_globally_masked(y - direction, x)) {
-        int yp = at(y - direction, x, 0) + direction;
-        int xp = at(y - direction, x, 1);
-        int dp = _distance(y, x, yp, xp);
-        if (dp < at(y, x, 2)) {
-            this_ptr[0] = yp, this_ptr[1] = xp, this_ptr[2] = dp;
-        }
-    }
-
-    // propagation along the x direction.
-    if (x - direction >= 0 && x - direction < this_size.width && !m_source.is_globally_masked(y, x - direction)) {
-        int yp = at(y, x - direction, 0);
-        int xp = at(y, x - direction, 1) + direction;
-        int dp = _distance(y, x, yp, xp);
-        if (dp < at(y, x, 2)) {
-            this_ptr[0] = yp, this_ptr[1] = xp, this_ptr[2] = dp;
-        }
-    }
-
-    // random search with a progressive step size.
-    int random_scale = (std::min(this_target_size.height, this_target_size.width) - 1) / 2;
-    while (random_scale > 0) {
-        int yp = this_ptr[0] + (rand() % (2 * random_scale + 1) - random_scale);
-        int xp = this_ptr[1] + (rand() % (2 * random_scale + 1) - random_scale);
-        yp = clamp(yp, 0, target_size().height - 1);
-        xp = clamp(xp, 0, target_size().width - 1);
-
-        if (m_target.is_globally_masked(yp, xp)) {
-            random_scale /= 2;
-        }
-
-        int dp = _distance(y, x, yp, xp);
-        if (dp < at(y, x, 2)) {
-            this_ptr[0] = yp, this_ptr[1] = xp, this_ptr[2] = dp;
-        }
-        random_scale /= 2;
-    }
-}
-
-const int PatchDistanceMetric::kDistanceScale = 65535;
-const int PatchSSDDistanceMetric::kSSDScale = 9 * 255 * 255;
-
-namespace {
-
-inline int pow2(int i) {
-    return i * i;
-}
-
-int distance_masked_images(
-    const MaskedImage &source, int ys, int xs,
-    const MaskedImage &target, int yt, int xt,
-    int patch_size
-) {
-    long double distance = 0;
-    long double wsum = 0;
-
-    source.compute_image_gradients();
-    target.compute_image_gradients();
-
-    auto source_size = source.size();
-    auto target_size = target.size();
-
-    for (int dy = -patch_size; dy <= patch_size; ++dy) {
-        const int yys = ys + dy, yyt = yt + dy;
-
-        if (yys <= 0 || yys >= source_size.height - 1 || yyt <= 0 || yyt >= target_size.height - 1) {
-            distance += (long double)(PatchSSDDistanceMetric::kSSDScale) * (2 * patch_size + 1);
-            wsum += 2 * patch_size + 1;
-            continue;
-        }
-
-        const auto *p_si = source.image().ptr<unsigned char>(yys, 0);
-        const auto *p_ti = target.image().ptr<unsigned char>(yyt, 0);
-        const auto *p_sm = source.mask().ptr<unsigned char>(yys, 0);
-        const auto *p_tm = target.mask().ptr<unsigned char>(yyt, 0);
-
-        const unsigned char *p_sgm = nullptr;
-        const unsigned char *p_tgm = nullptr;
-        if (!source.global_mask().empty()) {
-            p_sgm = source.global_mask().ptr<unsigned char>(yys, 0);
-            p_tgm = target.global_mask().ptr<unsigned char>(yyt, 0);
-        }
-
-        const auto *p_sgy = source.grady().ptr<unsigned char>(yys, 0);
-        const auto *p_tgy = target.grady().ptr<unsigned char>(yyt, 0);
-        const auto *p_sgx = source.gradx().ptr<unsigned char>(yys, 0);
-        const auto *p_tgx = target.gradx().ptr<unsigned char>(yyt, 0);
-
-        for (int dx = -patch_size; dx <= patch_size; ++dx) {
-            int xxs = xs + dx, xxt = xt + dx;
-            wsum += 1;
-
-            if (xxs <= 0 || xxs >= source_size.width - 1 || xxt <= 0 || xxt >= source_size.width - 1) {
-                distance += PatchSSDDistanceMetric::kSSDScale;
-                continue;
-            }
-
-            if (p_sm[xxs] || p_tm[xxt] || (p_sgm && p_sgm[xxs]) || (p_tgm && p_tgm[xxt]) ) {
-                distance += PatchSSDDistanceMetric::kSSDScale;
-                continue;
-            }
-
-            int ssd = 0;
-            for (int c = 0; c < 3; ++c) {
-                int s_value = p_si[xxs * 3 + c];
-                int t_value = p_ti[xxt * 3 + c];
-                int s_gy = p_sgy[xxs * 3 + c];
-                int t_gy = p_tgy[xxt * 3 + c];
-                int s_gx = p_sgx[xxs * 3 + c];
-                int t_gx = p_tgx[xxt * 3 + c];
-
-                ssd += pow2(static_cast<int>(s_value) - t_value);
-                ssd += pow2(static_cast<int>(s_gx) - t_gx);
-                ssd += pow2(static_cast<int>(s_gy) - t_gy);
-            }
-            distance += ssd;
-        }
-    }
-
-    distance /= (long double)(PatchSSDDistanceMetric::kSSDScale);
-
-    int res = int(PatchDistanceMetric::kDistanceScale * distance / wsum);
-    if (res < 0 || res > PatchDistanceMetric::kDistanceScale) return PatchDistanceMetric::kDistanceScale;
-    return res;
-}
-
-}
-
-int PatchSSDDistanceMetric::operator ()(const MaskedImage &source, int source_y, int source_x, const MaskedImage &target, int target_y, int target_x) const {
-    return distance_masked_images(source, source_y, source_x, target, target_y, target_x, m_patch_size);
-}
-
-int DebugPatchSSDDistanceMetric::operator ()(const MaskedImage &source, int source_y, int source_x, const MaskedImage &target, int target_y, int target_x) const {
-    fprintf(stderr, "DebugPatchSSDDistanceMetric: %d %d %d %d\n", source.size().width, source.size().height, m_width, m_height);
-    return distance_masked_images(source, source_y, source_x, target, target_y, target_x, m_patch_size);
-}
-
-int RegularityGuidedPatchDistanceMetricV1::operator ()(const MaskedImage &source, int source_y, int source_x, const MaskedImage &target, int target_y, int target_x) const {
-    double dx = remainder(double(source_x - target_x) / source.size().width, m_dx1);
-    double dy = remainder(double(source_y - target_y) / source.size().height, m_dy2);
-
-    double score1 = sqrt(dx * dx + dy *dy) / m_scale;
-    if (score1 < 0 || score1 > 1) score1 = 1;
-    score1 *= PatchDistanceMetric::kDistanceScale;
-
-    double score2 = distance_masked_images(source, source_y, source_x, target, target_y, target_x, m_patch_size);
-    double score = score1 * m_weight + score2 / (1 + m_weight);
-    return static_cast<int>(score / (1 + m_weight));
-}
-
-int RegularityGuidedPatchDistanceMetricV2::operator ()(const MaskedImage &source, int source_y, int source_x, const MaskedImage &target, int target_y, int target_x) const {
-    if (target_y < 0 || target_y >= target.size().height || target_x < 0 || target_x >= target.size().width)
-        return PatchDistanceMetric::kDistanceScale;
-
-    int source_scale = m_ijmap.size().height / source.size().height;
-    int target_scale = m_ijmap.size().height / target.size().height;
-
-    // fprintf(stderr, "RegularityGuidedPatchDistanceMetricV2 %d %d %d %d\n", source_y * source_scale, m_ijmap.size().height, source_x * source_scale, m_ijmap.size().width);
-
-    double score1 = PatchDistanceMetric::kDistanceScale;
-    if (!source.is_globally_masked(source_y, source_x) && !target.is_globally_masked(target_y, target_x)) {
-        auto source_ij = m_ijmap.ptr<float>(source_y * source_scale, source_x * source_scale);
-        auto target_ij = m_ijmap.ptr<float>(target_y * target_scale, target_x * target_scale);
-
-        float di = fabs(source_ij[0] - target_ij[0]); if (di > 0.5) di = 1 - di;
-        float dj = fabs(source_ij[1] - target_ij[1]); if (dj > 0.5) dj = 1 - dj;
-        score1 = sqrt(di * di + dj *dj) / 0.707;
-        if (score1 < 0 || score1 > 1) score1 = 1;
-        score1 *= PatchDistanceMetric::kDistanceScale;
-    }
-
-    double score2 = distance_masked_images(source, source_y, source_x, target, target_y, target_x, m_patch_size);
-    double score = score1 * m_weight + score2;
-    return int(score / (1 + m_weight));
-}
-
+#include <algorithm>
+#include <iostream>
+#include <cmath>
+
+#include "masked_image.h"
+#include "nnf.h"
+
+/**
+* Nearest-Neighbor Field (see PatchMatch algorithm).
+* This algorithme uses a version proposed by Xavier Philippeau.
+*
+*/
+
+template <typename T>
+T clamp(T value, T min_value, T max_value) {
+    return std::min(std::max(value, min_value), max_value);
+}
+
+void NearestNeighborField::_randomize_field(int max_retry, bool reset) {
+    auto this_size = source_size();
+    for (int i = 0; i < this_size.height; ++i) {
+        for (int j = 0; j < this_size.width; ++j) {
+            if (m_source.is_globally_masked(i, j)) continue;
+
+            auto this_ptr = mutable_ptr(i, j);
+            int distance = reset ? PatchDistanceMetric::kDistanceScale : this_ptr[2];
+            if (distance < PatchDistanceMetric::kDistanceScale) {
+                continue;
+            }
+
+            int i_target = 0, j_target = 0;
+            for (int t = 0; t < max_retry; ++t) {
+                i_target = rand() % this_size.height;
+                j_target = rand() % this_size.width;
+                if (m_target.is_globally_masked(i_target, j_target)) continue;
+
+                distance = _distance(i, j, i_target, j_target);
+                if (distance < PatchDistanceMetric::kDistanceScale)
+                    break;
+            }
+
+            this_ptr[0] = i_target, this_ptr[1] = j_target, this_ptr[2] = distance;
+        }
+    }
+}
+
+void NearestNeighborField::_initialize_field_from(const NearestNeighborField &other, int max_retry) {
+    const auto &this_size = source_size();
+    const auto &other_size = other.source_size();
+    double fi = static_cast<double>(this_size.height) / other_size.height;
+    double fj = static_cast<double>(this_size.width) / other_size.width;
+
+    for (int i = 0; i < this_size.height; ++i) {
+        for (int j = 0; j < this_size.width; ++j) {
+            if (m_source.is_globally_masked(i, j)) continue;
+
+            int ilow = static_cast<int>(std::min(i / fi, static_cast<double>(other_size.height - 1)));
+            int jlow = static_cast<int>(std::min(j / fj, static_cast<double>(other_size.width - 1)));
+            auto this_value = mutable_ptr(i, j);
+            auto other_value = other.ptr(ilow, jlow);
+
+            this_value[0] = static_cast<int>(other_value[0] * fi);
+            this_value[1] = static_cast<int>(other_value[1] * fj);
+            this_value[2] = _distance(i, j, this_value[0], this_value[1]);
+        }
+    }
+
+    _randomize_field(max_retry, false);
+}
+
+void NearestNeighborField::minimize(int nr_pass) {
+    const auto &this_size = source_size();
+    while (nr_pass--) {
+        for (int i = 0; i < this_size.height; ++i)
+            for (int j = 0; j < this_size.width; ++j) {
+                if (m_source.is_globally_masked(i, j)) continue;
+                if (at(i, j, 2) > 0) _minimize_link(i, j, +1);
+            }
+        for (int i = this_size.height - 1; i >= 0; --i)
+            for (int j = this_size.width - 1; j >= 0; --j) {
+                if (m_source.is_globally_masked(i, j)) continue;
+                if (at(i, j, 2) > 0) _minimize_link(i, j, -1);
+            }
+    }
+}
+
+void NearestNeighborField::_minimize_link(int y, int x, int direction) {
+    const auto &this_size = source_size();
+    const auto &this_target_size = target_size();
+    auto this_ptr = mutable_ptr(y, x);
+
+    // propagation along the y direction.
+    if (y - direction >= 0 && y - direction < this_size.height && !m_source.is_globally_masked(y - direction, x)) {
+        int yp = at(y - direction, x, 0) + direction;
+        int xp = at(y - direction, x, 1);
+        int dp = _distance(y, x, yp, xp);
+        if (dp < at(y, x, 2)) {
+            this_ptr[0] = yp, this_ptr[1] = xp, this_ptr[2] = dp;
+        }
+    }
+
+    // propagation along the x direction.
+    if (x - direction >= 0 && x - direction < this_size.width && !m_source.is_globally_masked(y, x - direction)) {
+        int yp = at(y, x - direction, 0);
+        int xp = at(y, x - direction, 1) + direction;
+        int dp = _distance(y, x, yp, xp);
+        if (dp < at(y, x, 2)) {
+            this_ptr[0] = yp, this_ptr[1] = xp, this_ptr[2] = dp;
+        }
+    }
+
+    // random search with a progressive step size.
+    int random_scale = (std::min(this_target_size.height, this_target_size.width) - 1) / 2;
+    while (random_scale > 0) {
+        int yp = this_ptr[0] + (rand() % (2 * random_scale + 1) - random_scale);
+        int xp = this_ptr[1] + (rand() % (2 * random_scale + 1) - random_scale);
+        yp = clamp(yp, 0, target_size().height - 1);
+        xp = clamp(xp, 0, target_size().width - 1);
+
+        if (m_target.is_globally_masked(yp, xp)) {
+            random_scale /= 2;
+        }
+
+        int dp = _distance(y, x, yp, xp);
+        if (dp < at(y, x, 2)) {
+            this_ptr[0] = yp, this_ptr[1] = xp, this_ptr[2] = dp;
+        }
+        random_scale /= 2;
+    }
+}
+
+const int PatchDistanceMetric::kDistanceScale = 65535;
+const int PatchSSDDistanceMetric::kSSDScale = 9 * 255 * 255;
+
+namespace {
+
+inline int pow2(int i) {
+    return i * i;
+}
+
+int distance_masked_images(
+    const MaskedImage &source, int ys, int xs,
+    const MaskedImage &target, int yt, int xt,
+    int patch_size
+) {
+    long double distance = 0;
+    long double wsum = 0;
+
+    source.compute_image_gradients();
+    target.compute_image_gradients();
+
+    auto source_size = source.size();
+    auto target_size = target.size();
+
+    for (int dy = -patch_size; dy <= patch_size; ++dy) {
+        const int yys = ys + dy, yyt = yt + dy;
+
+        if (yys <= 0 || yys >= source_size.height - 1 || yyt <= 0 || yyt >= target_size.height - 1) {
+            distance += (long double)(PatchSSDDistanceMetric::kSSDScale) * (2 * patch_size + 1);
+            wsum += 2 * patch_size + 1;
+            continue;
+        }
+
+        const auto *p_si = source.image().ptr<unsigned char>(yys, 0);
+        const auto *p_ti = target.image().ptr<unsigned char>(yyt, 0);
+        const auto *p_sm = source.mask().ptr<unsigned char>(yys, 0);
+        const auto *p_tm = target.mask().ptr<unsigned char>(yyt, 0);
+
+        const unsigned char *p_sgm = nullptr;
+        const unsigned char *p_tgm = nullptr;
+        if (!source.global_mask().empty()) {
+            p_sgm = source.global_mask().ptr<unsigned char>(yys, 0);
+            p_tgm = target.global_mask().ptr<unsigned char>(yyt, 0);
+        }
+
+        const auto *p_sgy = source.grady().ptr<unsigned char>(yys, 0);
+        const auto *p_tgy = target.grady().ptr<unsigned char>(yyt, 0);
+        const auto *p_sgx = source.gradx().ptr<unsigned char>(yys, 0);
+        const auto *p_tgx = target.gradx().ptr<unsigned char>(yyt, 0);
+
+        for (int dx = -patch_size; dx <= patch_size; ++dx) {
+            int xxs = xs + dx, xxt = xt + dx;
+            wsum += 1;
+
+            if (xxs <= 0 || xxs >= source_size.width - 1 || xxt <= 0 || xxt >= source_size.width - 1) {
+                distance += PatchSSDDistanceMetric::kSSDScale;
+                continue;
+            }
+
+            if (p_sm[xxs] || p_tm[xxt] || (p_sgm && p_sgm[xxs]) || (p_tgm && p_tgm[xxt]) ) {
+                distance += PatchSSDDistanceMetric::kSSDScale;
+                continue;
+            }
+
+            int ssd = 0;
+            for (int c = 0; c < 3; ++c) {
+                int s_value = p_si[xxs * 3 + c];
+                int t_value = p_ti[xxt * 3 + c];
+                int s_gy = p_sgy[xxs * 3 + c];
+                int t_gy = p_tgy[xxt * 3 + c];
+                int s_gx = p_sgx[xxs * 3 + c];
+                int t_gx = p_tgx[xxt * 3 + c];
+
+                ssd += pow2(static_cast<int>(s_value) - t_value);
+                ssd += pow2(static_cast<int>(s_gx) - t_gx);
+                ssd += pow2(static_cast<int>(s_gy) - t_gy);
+            }
+            distance += ssd;
+        }
+    }
+
+    distance /= (long double)(PatchSSDDistanceMetric::kSSDScale);
+
+    int res = int(PatchDistanceMetric::kDistanceScale * distance / wsum);
+    if (res < 0 || res > PatchDistanceMetric::kDistanceScale) return PatchDistanceMetric::kDistanceScale;
+    return res;
+}
+
+}
+
+int PatchSSDDistanceMetric::operator ()(const MaskedImage &source, int source_y, int source_x, const MaskedImage &target, int target_y, int target_x) const {
+    return distance_masked_images(source, source_y, source_x, target, target_y, target_x, m_patch_size);
+}
+
+int DebugPatchSSDDistanceMetric::operator ()(const MaskedImage &source, int source_y, int source_x, const MaskedImage &target, int target_y, int target_x) const {
+    fprintf(stderr, "DebugPatchSSDDistanceMetric: %d %d %d %d\n", source.size().width, source.size().height, m_width, m_height);
+    return distance_masked_images(source, source_y, source_x, target, target_y, target_x, m_patch_size);
+}
+
+int RegularityGuidedPatchDistanceMetricV1::operator ()(const MaskedImage &source, int source_y, int source_x, const MaskedImage &target, int target_y, int target_x) const {
+    double dx = remainder(double(source_x - target_x) / source.size().width, m_dx1);
+    double dy = remainder(double(source_y - target_y) / source.size().height, m_dy2);
+
+    double score1 = sqrt(dx * dx + dy *dy) / m_scale;
+    if (score1 < 0 || score1 > 1) score1 = 1;
+    score1 *= PatchDistanceMetric::kDistanceScale;
+
+    double score2 = distance_masked_images(source, source_y, source_x, target, target_y, target_x, m_patch_size);
+    double score = score1 * m_weight + score2 / (1 + m_weight);
+    return static_cast<int>(score / (1 + m_weight));
+}
+
+int RegularityGuidedPatchDistanceMetricV2::operator ()(const MaskedImage &source, int source_y, int source_x, const MaskedImage &target, int target_y, int target_x) const {
+    if (target_y < 0 || target_y >= target.size().height || target_x < 0 || target_x >= target.size().width)
+        return PatchDistanceMetric::kDistanceScale;
+
+    int source_scale = m_ijmap.size().height / source.size().height;
+    int target_scale = m_ijmap.size().height / target.size().height;
+
+    // fprintf(stderr, "RegularityGuidedPatchDistanceMetricV2 %d %d %d %d\n", source_y * source_scale, m_ijmap.size().height, source_x * source_scale, m_ijmap.size().width);
+
+    double score1 = PatchDistanceMetric::kDistanceScale;
+    if (!source.is_globally_masked(source_y, source_x) && !target.is_globally_masked(target_y, target_x)) {
+        auto source_ij = m_ijmap.ptr<float>(source_y * source_scale, source_x * source_scale);
+        auto target_ij = m_ijmap.ptr<float>(target_y * target_scale, target_x * target_scale);
+
+        float di = fabs(source_ij[0] - target_ij[0]); if (di > 0.5) di = 1 - di;
+        float dj = fabs(source_ij[1] - target_ij[1]); if (dj > 0.5) dj = 1 - dj;
+        score1 = sqrt(di * di + dj *dj) / 0.707;
+        if (score1 < 0 || score1 > 1) score1 = 1;
+        score1 *= PatchDistanceMetric::kDistanceScale;
+    }
+
+    double score2 = distance_masked_images(source, source_y, source_x, target, target_y, target_x, m_patch_size);
+    double score = score1 * m_weight + score2;
+    return int(score / (1 + m_weight));
+}
+

PyPatchMatch/csrc/nnf.h

@@ -1,133 +1,133 @@
-#pragma once
-
-#include <opencv2/core.hpp>
-#include "masked_image.h"
-
-class PatchDistanceMetric {
-public:
-    PatchDistanceMetric(int patch_size) : m_patch_size(patch_size) {}
-    virtual ~PatchDistanceMetric() = default;
-
-    inline int patch_size() const { return m_patch_size; }
-    virtual int operator()(const MaskedImage &source, int source_y, int source_x, const MaskedImage &target, int target_y, int target_x) const = 0;
-    static const int kDistanceScale;
-
-protected:
-    int m_patch_size;
-};
-
-class NearestNeighborField {
-public:
-    NearestNeighborField() : m_source(), m_target(), m_field(), m_distance_metric(nullptr) {
-        // pass
-    }
-    NearestNeighborField(const MaskedImage &source, const MaskedImage &target, const PatchDistanceMetric *metric, int max_retry = 20)
-        : m_source(source), m_target(target), m_distance_metric(metric) {
-        m_field = cv::Mat(m_source.size(), CV_32SC3);
-        _randomize_field(max_retry);
-    }
-    NearestNeighborField(const MaskedImage &source, const MaskedImage &target, const PatchDistanceMetric *metric, const NearestNeighborField &other, int max_retry = 20)
-            : m_source(source), m_target(target), m_distance_metric(metric) {
-        m_field = cv::Mat(m_source.size(), CV_32SC3);
-        _initialize_field_from(other, max_retry);
-    }
-
-    const MaskedImage &source() const {
-        return m_source;
-    }
-    const MaskedImage &target() const {
-        return m_target;
-    }
-    inline cv::Size source_size() const {
-        return m_source.size();
-    }
-    inline cv::Size target_size() const {
-        return m_target.size();
-    }
-    inline void set_source(const MaskedImage &source) {
-        m_source = source;
-    }
-    inline void set_target(const MaskedImage &target) {
-        m_target = target;
-    }
-
-    inline int *mutable_ptr(int y, int x) {
-        return m_field.ptr<int>(y, x);
-    }
-    inline const int *ptr(int y, int x) const {
-        return m_field.ptr<int>(y, x);
-    }
-
-    inline int at(int y, int x, int c) const {
-        return m_field.ptr<int>(y, x)[c];
-    }
-    inline int &at(int y, int x, int c) {
-        return m_field.ptr<int>(y, x)[c];
-    }
-    inline void set_identity(int y, int x) {
-        auto ptr = mutable_ptr(y, x);
-        ptr[0] = y, ptr[1] = x, ptr[2] = 0;
-    }
-
-    void minimize(int nr_pass);
-
-private:
-    inline int _distance(int source_y, int source_x, int target_y, int target_x) {
-        return (*m_distance_metric)(m_source, source_y, source_x, m_target, target_y, target_x);
-    }
-
-    void _randomize_field(int max_retry = 20, bool reset = true);
-    void _initialize_field_from(const NearestNeighborField &other, int max_retry);
-    void _minimize_link(int y, int x, int direction);
-
-    MaskedImage m_source;
-    MaskedImage m_target;
-    cv::Mat m_field;  // { y_target, x_target, distance_scaled }
-    const PatchDistanceMetric *m_distance_metric;
-};
-
-
-class PatchSSDDistanceMetric : public PatchDistanceMetric {
-public:
-    using PatchDistanceMetric::PatchDistanceMetric;
-    virtual int operator ()(const MaskedImage &source, int source_y, int source_x, const MaskedImage &target, int target_y, int target_x) const;
-    static const int kSSDScale;
-};
-
-class DebugPatchSSDDistanceMetric : public PatchDistanceMetric {
-public:
-    DebugPatchSSDDistanceMetric(int patch_size, int width, int height) : PatchDistanceMetric(patch_size), m_width(width), m_height(height) {}
-    virtual int operator ()(const MaskedImage &source, int source_y, int source_x, const MaskedImage &target, int target_y, int target_x) const;
-protected:
-    int m_width, m_height;
-};
-
-class RegularityGuidedPatchDistanceMetricV1 : public PatchDistanceMetric {
-public:
-    RegularityGuidedPatchDistanceMetricV1(int patch_size, double dx1, double dy1, double dx2, double dy2, double weight)
-        : PatchDistanceMetric(patch_size), m_dx1(dx1), m_dy1(dy1), m_dx2(dx2), m_dy2(dy2), m_weight(weight) {
-
-        assert(m_dy1 == 0);
-        assert(m_dx2 == 0);
-        m_scale = sqrt(m_dx1 * m_dx1 + m_dy2 * m_dy2) / 4;
-    }
-    virtual int operator ()(const MaskedImage &source, int source_y, int source_x, const MaskedImage &target, int target_y, int target_x) const;
-
-protected:
-    double m_dx1, m_dy1, m_dx2, m_dy2;
-    double m_scale, m_weight;
-};
-
-class RegularityGuidedPatchDistanceMetricV2 : public PatchDistanceMetric {
-public:
-    RegularityGuidedPatchDistanceMetricV2(int patch_size, cv::Mat ijmap, double weight)
-        : PatchDistanceMetric(patch_size), m_ijmap(ijmap), m_weight(weight) {
-
-    }
-    virtual int operator ()(const MaskedImage &source, int source_y, int source_x, const MaskedImage &target, int target_y, int target_x) const;
-
-protected:
-    cv::Mat m_ijmap;
-    double m_width, m_height, m_weight;
-};
-
+#pragma once
+
+#include <opencv2/core.hpp>
+#include "masked_image.h"
+
+class PatchDistanceMetric {
+public:
+    PatchDistanceMetric(int patch_size) : m_patch_size(patch_size) {}
+    virtual ~PatchDistanceMetric() = default;
+
+    inline int patch_size() const { return m_patch_size; }
+    virtual int operator()(const MaskedImage &source, int source_y, int source_x, const MaskedImage &target, int target_y, int target_x) const = 0;
+    static const int kDistanceScale;
+
+protected:
+    int m_patch_size;
+};
+
+class NearestNeighborField {
+public:
+    NearestNeighborField() : m_source(), m_target(), m_field(), m_distance_metric(nullptr) {
+        // pass
+    }
+    NearestNeighborField(const MaskedImage &source, const MaskedImage &target, const PatchDistanceMetric *metric, int max_retry = 20)
+        : m_source(source), m_target(target), m_distance_metric(metric) {
+        m_field = cv::Mat(m_source.size(), CV_32SC3);
+        _randomize_field(max_retry);
+    }
+    NearestNeighborField(const MaskedImage &source, const MaskedImage &target, const PatchDistanceMetric *metric, const NearestNeighborField &other, int max_retry = 20)
+            : m_source(source), m_target(target), m_distance_metric(metric) {
+        m_field = cv::Mat(m_source.size(), CV_32SC3);
+        _initialize_field_from(other, max_retry);
+    }
+
+    const MaskedImage &source() const {
+        return m_source;
+    }
+    const MaskedImage &target() const {
+        return m_target;
+    }
+    inline cv::Size source_size() const {
+        return m_source.size();
+    }
+    inline cv::Size target_size() const {
+        return m_target.size();
+    }
+    inline void set_source(const MaskedImage &source) {
+        m_source = source;
+    }
+    inline void set_target(const MaskedImage &target) {
+        m_target = target;
+    }
+
+    inline int *mutable_ptr(int y, int x) {
+        return m_field.ptr<int>(y, x);
+    }
+    inline const int *ptr(int y, int x) const {
+        return m_field.ptr<int>(y, x);
+    }
+
+    inline int at(int y, int x, int c) const {
+        return m_field.ptr<int>(y, x)[c];
+    }
+    inline int &at(int y, int x, int c) {
+        return m_field.ptr<int>(y, x)[c];
+    }
+    inline void set_identity(int y, int x) {
+        auto ptr = mutable_ptr(y, x);
+        ptr[0] = y, ptr[1] = x, ptr[2] = 0;
+    }
+
+    void minimize(int nr_pass);
+
+private:
+    inline int _distance(int source_y, int source_x, int target_y, int target_x) {
+        return (*m_distance_metric)(m_source, source_y, source_x, m_target, target_y, target_x);
+    }
+
+    void _randomize_field(int max_retry = 20, bool reset = true);
+    void _initialize_field_from(const NearestNeighborField &other, int max_retry);
+    void _minimize_link(int y, int x, int direction);
+
+    MaskedImage m_source;
+    MaskedImage m_target;
+    cv::Mat m_field;  // { y_target, x_target, distance_scaled }
+    const PatchDistanceMetric *m_distance_metric;
+};
+
+
+class PatchSSDDistanceMetric : public PatchDistanceMetric {
+public:
+    using PatchDistanceMetric::PatchDistanceMetric;
+    virtual int operator ()(const MaskedImage &source, int source_y, int source_x, const MaskedImage &target, int target_y, int target_x) const;
+    static const int kSSDScale;
+};
+
+class DebugPatchSSDDistanceMetric : public PatchDistanceMetric {
+public:
+    DebugPatchSSDDistanceMetric(int patch_size, int width, int height) : PatchDistanceMetric(patch_size), m_width(width), m_height(height) {}
+    virtual int operator ()(const MaskedImage &source, int source_y, int source_x, const MaskedImage &target, int target_y, int target_x) const;
+protected:
+    int m_width, m_height;
+};
+
+class RegularityGuidedPatchDistanceMetricV1 : public PatchDistanceMetric {
+public:
+    RegularityGuidedPatchDistanceMetricV1(int patch_size, double dx1, double dy1, double dx2, double dy2, double weight)
+        : PatchDistanceMetric(patch_size), m_dx1(dx1), m_dy1(dy1), m_dx2(dx2), m_dy2(dy2), m_weight(weight) {
+
+        assert(m_dy1 == 0);
+        assert(m_dx2 == 0);
+        m_scale = sqrt(m_dx1 * m_dx1 + m_dy2 * m_dy2) / 4;
+    }
+    virtual int operator ()(const MaskedImage &source, int source_y, int source_x, const MaskedImage &target, int target_y, int target_x) const;
+
+protected:
+    double m_dx1, m_dy1, m_dx2, m_dy2;
+    double m_scale, m_weight;
+};
+
+class RegularityGuidedPatchDistanceMetricV2 : public PatchDistanceMetric {
+public:
+    RegularityGuidedPatchDistanceMetricV2(int patch_size, cv::Mat ijmap, double weight)
+        : PatchDistanceMetric(patch_size), m_ijmap(ijmap), m_weight(weight) {
+
+    }
+    virtual int operator ()(const MaskedImage &source, int source_y, int source_x, const MaskedImage &target, int target_y, int target_x) const;
+
+protected:
+    cv::Mat m_ijmap;
+    double m_width, m_height, m_weight;
+};
+

PyPatchMatch/csrc/pyinterface.cpp

@@ -1,107 +1,107 @@
-#include "pyinterface.h"
-#include "inpaint.h"
-
-static unsigned int PM_seed = 1212;
-static bool PM_verbose = false;
-
-int _dtype_py_to_cv(int dtype_py);
-int _dtype_cv_to_py(int dtype_cv);
-cv::Mat _py_to_cv2(PM_mat_t pymat);
-PM_mat_t _cv2_to_py(cv::Mat cvmat);
-
-void PM_set_random_seed(unsigned int seed) {
-    PM_seed = seed;
-}
-
-void PM_set_verbose(int value) {
-    PM_verbose = static_cast<bool>(value);
-}
-
-void PM_free_pymat(PM_mat_t pymat) {
-    free(pymat.data_ptr);
-}
-
-PM_mat_t PM_inpaint(PM_mat_t source_py, PM_mat_t mask_py, int patch_size) {
-    cv::Mat source = _py_to_cv2(source_py);
-    cv::Mat mask = _py_to_cv2(mask_py);
-    auto metric = PatchSSDDistanceMetric(patch_size);
-    cv::Mat result = Inpainting(source, mask, &metric).run(PM_verbose, false, PM_seed);
-    return _cv2_to_py(result);
-}
-
-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) {
-    cv::Mat source = _py_to_cv2(source_py);
-    cv::Mat mask = _py_to_cv2(mask_py);
-    cv::Mat ijmap = _py_to_cv2(ijmap_py);
-
-    auto metric = RegularityGuidedPatchDistanceMetricV2(patch_size, ijmap, guide_weight);
-    cv::Mat result = Inpainting(source, mask, &metric).run(PM_verbose, false, PM_seed);
-    return _cv2_to_py(result);
-}
-
-PM_mat_t PM_inpaint2(PM_mat_t source_py, PM_mat_t mask_py, PM_mat_t global_mask_py, int patch_size) {
-    cv::Mat source = _py_to_cv2(source_py);
-    cv::Mat mask = _py_to_cv2(mask_py);
-    cv::Mat global_mask = _py_to_cv2(global_mask_py);
-
-    auto metric = PatchSSDDistanceMetric(patch_size);
-    cv::Mat result = Inpainting(source, mask, global_mask, &metric).run(PM_verbose, false, PM_seed);
-    return _cv2_to_py(result);
-}
-
-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) {
-    cv::Mat source = _py_to_cv2(source_py);
-    cv::Mat mask = _py_to_cv2(mask_py);
-    cv::Mat global_mask = _py_to_cv2(global_mask_py);
-    cv::Mat ijmap = _py_to_cv2(ijmap_py);
-
-    auto metric = RegularityGuidedPatchDistanceMetricV2(patch_size, ijmap, guide_weight);
-    cv::Mat result = Inpainting(source, mask, global_mask, &metric).run(PM_verbose, false, PM_seed);
-    return _cv2_to_py(result);
-}
-
-int _dtype_py_to_cv(int dtype_py) {
-    switch (dtype_py) {
-        case PM_UINT8: return CV_8U;
-        case PM_INT8: return CV_8S;
-        case PM_UINT16: return CV_16U;
-        case PM_INT16: return CV_16S;
-        case PM_INT32: return CV_32S;
-        case PM_FLOAT32: return CV_32F;
-        case PM_FLOAT64: return CV_64F;
-    }
-
-    return CV_8U;
-}
-
-int _dtype_cv_to_py(int dtype_cv) {
-    switch (dtype_cv) {
-        case CV_8U: return PM_UINT8;
-        case CV_8S: return PM_INT8;
-        case CV_16U: return PM_UINT16;
-        case CV_16S: return PM_INT16;
-        case CV_32S: return PM_INT32;
-        case CV_32F: return PM_FLOAT32;
-        case CV_64F: return PM_FLOAT64;
-    }
-
-    return PM_UINT8;
-}
-
-cv::Mat _py_to_cv2(PM_mat_t pymat) {
-    int dtype = _dtype_py_to_cv(pymat.dtype);
-    dtype = CV_MAKETYPE(pymat.dtype, pymat.shape.channels);
-    return cv::Mat(cv::Size(pymat.shape.width, pymat.shape.height), dtype, pymat.data_ptr).clone();
-}
-
-PM_mat_t _cv2_to_py(cv::Mat cvmat) {
-    PM_shape_t shape = {cvmat.size().width, cvmat.size().height, cvmat.channels()};
-    int dtype = _dtype_cv_to_py(cvmat.depth());
-    size_t dsize = cvmat.total() * cvmat.elemSize();
-
-    void *data_ptr = reinterpret_cast<void *>(malloc(dsize));
-    memcpy(data_ptr, reinterpret_cast<void *>(cvmat.data), dsize);
-
-    return PM_mat_t {data_ptr, shape, dtype};
-}
-
+#include "pyinterface.h"
+#include "inpaint.h"
+
+static unsigned int PM_seed = 1212;
+static bool PM_verbose = false;
+
+int _dtype_py_to_cv(int dtype_py);
+int _dtype_cv_to_py(int dtype_cv);
+cv::Mat _py_to_cv2(PM_mat_t pymat);
+PM_mat_t _cv2_to_py(cv::Mat cvmat);
+
+void PM_set_random_seed(unsigned int seed) {
+    PM_seed = seed;
+}
+
+void PM_set_verbose(int value) {
+    PM_verbose = static_cast<bool>(value);
+}
+
+void PM_free_pymat(PM_mat_t pymat) {
+    free(pymat.data_ptr);
+}
+
+PM_mat_t PM_inpaint(PM_mat_t source_py, PM_mat_t mask_py, int patch_size) {
+    cv::Mat source = _py_to_cv2(source_py);
+    cv::Mat mask = _py_to_cv2(mask_py);
+    auto metric = PatchSSDDistanceMetric(patch_size);
+    cv::Mat result = Inpainting(source, mask, &metric).run(PM_verbose, false, PM_seed);
+    return _cv2_to_py(result);
+}
+
+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) {
+    cv::Mat source = _py_to_cv2(source_py);
+    cv::Mat mask = _py_to_cv2(mask_py);
+    cv::Mat ijmap = _py_to_cv2(ijmap_py);
+
+    auto metric = RegularityGuidedPatchDistanceMetricV2(patch_size, ijmap, guide_weight);
+    cv::Mat result = Inpainting(source, mask, &metric).run(PM_verbose, false, PM_seed);
+    return _cv2_to_py(result);
+}
+
+PM_mat_t PM_inpaint2(PM_mat_t source_py, PM_mat_t mask_py, PM_mat_t global_mask_py, int patch_size) {
+    cv::Mat source = _py_to_cv2(source_py);
+    cv::Mat mask = _py_to_cv2(mask_py);
+    cv::Mat global_mask = _py_to_cv2(global_mask_py);
+
+    auto metric = PatchSSDDistanceMetric(patch_size);
+    cv::Mat result = Inpainting(source, mask, global_mask, &metric).run(PM_verbose, false, PM_seed);
+    return _cv2_to_py(result);
+}
+
+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) {
+    cv::Mat source = _py_to_cv2(source_py);
+    cv::Mat mask = _py_to_cv2(mask_py);
+    cv::Mat global_mask = _py_to_cv2(global_mask_py);
+    cv::Mat ijmap = _py_to_cv2(ijmap_py);
+
+    auto metric = RegularityGuidedPatchDistanceMetricV2(patch_size, ijmap, guide_weight);
+    cv::Mat result = Inpainting(source, mask, global_mask, &metric).run(PM_verbose, false, PM_seed);
+    return _cv2_to_py(result);
+}
+
+int _dtype_py_to_cv(int dtype_py) {
+    switch (dtype_py) {
+        case PM_UINT8: return CV_8U;
+        case PM_INT8: return CV_8S;
+        case PM_UINT16: return CV_16U;
+        case PM_INT16: return CV_16S;
+        case PM_INT32: return CV_32S;
+        case PM_FLOAT32: return CV_32F;
+        case PM_FLOAT64: return CV_64F;
+    }
+
+    return CV_8U;
+}
+
+int _dtype_cv_to_py(int dtype_cv) {
+    switch (dtype_cv) {
+        case CV_8U: return PM_UINT8;
+        case CV_8S: return PM_INT8;
+        case CV_16U: return PM_UINT16;
+        case CV_16S: return PM_INT16;
+        case CV_32S: return PM_INT32;
+        case CV_32F: return PM_FLOAT32;
+        case CV_64F: return PM_FLOAT64;
+    }
+
+    return PM_UINT8;
+}
+
+cv::Mat _py_to_cv2(PM_mat_t pymat) {
+    int dtype = _dtype_py_to_cv(pymat.dtype);
+    dtype = CV_MAKETYPE(pymat.dtype, pymat.shape.channels);
+    return cv::Mat(cv::Size(pymat.shape.width, pymat.shape.height), dtype, pymat.data_ptr).clone();
+}
+
+PM_mat_t _cv2_to_py(cv::Mat cvmat) {
+    PM_shape_t shape = {cvmat.size().width, cvmat.size().height, cvmat.channels()};
+    int dtype = _dtype_cv_to_py(cvmat.depth());
+    size_t dsize = cvmat.total() * cvmat.elemSize();
+
+    void *data_ptr = reinterpret_cast<void *>(malloc(dsize));
+    memcpy(data_ptr, reinterpret_cast<void *>(cvmat.data), dsize);
+
+    return PM_mat_t {data_ptr, shape, dtype};
+}
+

PyPatchMatch/csrc/pyinterface.h

@@ -1,38 +1,38 @@
-#include <opencv2/core.hpp>
-#include <cstdlib>
-#include <cstdio>
-#include <cstring>
-
-extern "C" {
-
-struct PM_shape_t {
-    int width, height, channels;
-};
-
-enum PM_dtype_e {
-    PM_UINT8,
-    PM_INT8,
-    PM_UINT16,
-    PM_INT16,
-    PM_INT32,
-    PM_FLOAT32,
-    PM_FLOAT64,
-};
-
-struct PM_mat_t {
-    void *data_ptr;
-    PM_shape_t shape;
-    int dtype;
-};
-
-void PM_set_random_seed(unsigned int seed);
-void PM_set_verbose(int value);
-
-void PM_free_pymat(PM_mat_t pymat);
-PM_mat_t PM_inpaint(PM_mat_t image, PM_mat_t mask, int patch_size);
-PM_mat_t PM_inpaint_regularity(PM_mat_t image, PM_mat_t mask, PM_mat_t ijmap, int patch_size, float guide_weight);
-PM_mat_t PM_inpaint2(PM_mat_t image, PM_mat_t mask, PM_mat_t global_mask, int patch_size);
-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);
-
-} /*  extern "C" */
-
+#include <opencv2/core.hpp>
+#include <cstdlib>
+#include <cstdio>
+#include <cstring>
+
+extern "C" {
+
+struct PM_shape_t {
+    int width, height, channels;
+};
+
+enum PM_dtype_e {
+    PM_UINT8,
+    PM_INT8,
+    PM_UINT16,
+    PM_INT16,
+    PM_INT32,
+    PM_FLOAT32,
+    PM_FLOAT64,
+};
+
+struct PM_mat_t {
+    void *data_ptr;
+    PM_shape_t shape;
+    int dtype;
+};
+
+void PM_set_random_seed(unsigned int seed);
+void PM_set_verbose(int value);
+
+void PM_free_pymat(PM_mat_t pymat);
+PM_mat_t PM_inpaint(PM_mat_t image, PM_mat_t mask, int patch_size);
+PM_mat_t PM_inpaint_regularity(PM_mat_t image, PM_mat_t mask, PM_mat_t ijmap, int patch_size, float guide_weight);
+PM_mat_t PM_inpaint2(PM_mat_t image, PM_mat_t mask, PM_mat_t global_mask, int patch_size);
+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);
+
+} /*  extern "C" */
+

PyPatchMatch/examples/.gitignore

@@ -1,2 +1,2 @@
-/cpp_example.exe
-/images/*recovered.bmp
+/cpp_example.exe
+/images/*recovered.bmp

PyPatchMatch/examples/cpp_example.cpp

@@ -1,31 +1,31 @@
-#include <iostream>
-#include <opencv2/imgcodecs.hpp>
-#include <opencv2/highgui.hpp>
-
-#include "masked_image.h"
-#include "nnf.h"
-#include "inpaint.h"
-
-int main() {
-    auto source = cv::imread("./images/forest_pruned.bmp", cv::IMREAD_COLOR);
-
-    auto mask = cv::Mat(source.size(), CV_8UC1);
-    mask = cv::Scalar::all(0);
-    for (int i = 0; i < source.size().height; ++i) {
-        for (int j = 0; j < source.size().width; ++j) {
-            auto source_ptr = source.ptr<unsigned char>(i, j);
-            if (source_ptr[0] == 255 && source_ptr[1] == 255 && source_ptr[2] == 255) {
-                mask.at<unsigned char>(i, j) = 1;
-            }
-        }
-    }
-
-    auto metric = PatchSSDDistanceMetric(3);
-    auto result = Inpainting(source, mask, &metric).run(true, true);
-    // cv::imwrite("./images/forest_recovered.bmp", result);
-    // cv::imshow("Result", result);
-    // cv::waitKey();
-
-    return 0;
-}
-
+#include <iostream>
+#include <opencv2/imgcodecs.hpp>
+#include <opencv2/highgui.hpp>
+
+#include "masked_image.h"
+#include "nnf.h"
+#include "inpaint.h"
+
+int main() {
+    auto source = cv::imread("./images/forest_pruned.bmp", cv::IMREAD_COLOR);
+
+    auto mask = cv::Mat(source.size(), CV_8UC1);
+    mask = cv::Scalar::all(0);
+    for (int i = 0; i < source.size().height; ++i) {
+        for (int j = 0; j < source.size().width; ++j) {
+            auto source_ptr = source.ptr<unsigned char>(i, j);
+            if (source_ptr[0] == 255 && source_ptr[1] == 255 && source_ptr[2] == 255) {
+                mask.at<unsigned char>(i, j) = 1;
+            }
+        }
+    }
+
+    auto metric = PatchSSDDistanceMetric(3);
+    auto result = Inpainting(source, mask, &metric).run(true, true);
+    // cv::imwrite("./images/forest_recovered.bmp", result);
+    // cv::imshow("Result", result);
+    // cv::waitKey();
+
+    return 0;
+}
+

PyPatchMatch/examples/cpp_example_run.sh

@@ -1,18 +1,18 @@
-#! /bin/bash
-#
-# cpp_example_run.sh
-# Copyright (C) 2020 Jiayuan Mao <maojiayuan@gmail.com>
-#
-# Distributed under terms of the MIT license.
-#
-
-set -x
-
-CFLAGS="-std=c++14 -O2 $(pkg-config --cflags opencv)"
-LDFLAGS="$(pkg-config --libs opencv)"
-g++ $CFLAGS cpp_example.cpp -I../csrc/ -L../ -lpatchmatch $LDFLAGS -o cpp_example.exe
-
-export DYLD_LIBRARY_PATH=../:$DYLD_LIBRARY_PATH  # For macOS
-export LD_LIBRARY_PATH=../:$LD_LIBRARY_PATH  # For Linux
-time ./cpp_example.exe
-
+#! /bin/bash
+#
+# cpp_example_run.sh
+# Copyright (C) 2020 Jiayuan Mao <maojiayuan@gmail.com>
+#
+# Distributed under terms of the MIT license.
+#
+
+set -x
+
+CFLAGS="-std=c++14 -O2 $(pkg-config --cflags opencv)"
+LDFLAGS="$(pkg-config --libs opencv)"
+g++ $CFLAGS cpp_example.cpp -I../csrc/ -L../ -lpatchmatch $LDFLAGS -o cpp_example.exe
+
+export DYLD_LIBRARY_PATH=../:$DYLD_LIBRARY_PATH  # For macOS
+export LD_LIBRARY_PATH=../:$LD_LIBRARY_PATH  # For Linux
+time ./cpp_example.exe
+

PyPatchMatch/examples/py_example.py

@@ -1,21 +1,21 @@
-#! /usr/bin/env python3
-# -*- coding: utf-8 -*-
-# File   : test.py
-# Author : Jiayuan Mao
-# Email  : maojiayuan@gmail.com
-# Date   : 01/09/2020
-#
-# Distributed under terms of the MIT license.
-
-from PIL import Image
-
-import sys
-sys.path.insert(0, '../')
-import patch_match
-
-
-if __name__ == '__main__':
-    source = Image.open('./images/forest_pruned.bmp')
-    result = patch_match.inpaint(source, patch_size=3)
-    Image.fromarray(result).save('./images/forest_recovered.bmp')
-
+#! /usr/bin/env python3
+# -*- coding: utf-8 -*-
+# File   : test.py
+# Author : Jiayuan Mao
+# Email  : maojiayuan@gmail.com
+# Date   : 01/09/2020
+#
+# Distributed under terms of the MIT license.
+
+from PIL import Image
+
+import sys
+sys.path.insert(0, '../')
+import patch_match
+
+
+if __name__ == '__main__':
+    source = Image.open('./images/forest_pruned.bmp')
+    result = patch_match.inpaint(source, patch_size=3)
+    Image.fromarray(result).save('./images/forest_recovered.bmp')
+

PyPatchMatch/examples/py_example_global_mask.py

@@ -1,27 +1,27 @@
-#! /usr/bin/env python3
-# -*- coding: utf-8 -*-
-# File   : test.py
-# Author : Jiayuan Mao
-# Email  : maojiayuan@gmail.com
-# Date   : 01/09/2020
-#
-# Distributed under terms of the MIT license.
-
-import numpy as np
-from PIL import Image
-
-import sys
-sys.path.insert(0, '../')
-import patch_match
-
-
-if __name__ == '__main__':
-    patch_match.set_verbose(True)
-    source = Image.open('./images/forest_pruned.bmp')
-    source = np.array(source)
-    source[:100, :100] = 255
-    global_mask = np.zeros_like(source[..., 0])
-    global_mask[:100, :100] = 1
-    result = patch_match.inpaint(source, global_mask=global_mask, patch_size=3)
-    Image.fromarray(result).save('./images/forest_recovered.bmp')
-
+#! /usr/bin/env python3
+# -*- coding: utf-8 -*-
+# File   : test.py
+# Author : Jiayuan Mao
+# Email  : maojiayuan@gmail.com
+# Date   : 01/09/2020
+#
+# Distributed under terms of the MIT license.
+
+import numpy as np
+from PIL import Image
+
+import sys
+sys.path.insert(0, '../')
+import patch_match
+
+
+if __name__ == '__main__':
+    patch_match.set_verbose(True)
+    source = Image.open('./images/forest_pruned.bmp')
+    source = np.array(source)
+    source[:100, :100] = 255
+    global_mask = np.zeros_like(source[..., 0])
+    global_mask[:100, :100] = 1
+    result = patch_match.inpaint(source, global_mask=global_mask, patch_size=3)
+    Image.fromarray(result).save('./images/forest_recovered.bmp')
+

PyPatchMatch/patch_match.py

@@ -1,201 +1,263 @@
-#! /usr/bin/env python3
-# -*- coding: utf-8 -*-
-# File   : patch_match.py
-# Author : Jiayuan Mao
-# Email  : maojiayuan@gmail.com
-# Date   : 01/09/2020
-#
-# Distributed under terms of the MIT license.
-
-import ctypes
-import os.path as osp
-from typing import Optional, Union
-
-import numpy as np
-from PIL import Image
-
-try:
-    # If the Jacinle library (https://github.com/vacancy/Jacinle) is present, use its auto_travis feature.
-    from jacinle.jit.cext import auto_travis
-    auto_travis(__file__, required_files=['*.so'])
-except ImportError as e:
-    # Otherwise, fall back to the subprocess.
-    import subprocess
-    print('Compiling and loading c extensions from "{}".'.format(osp.realpath(osp.dirname(__file__))))
-    subprocess.check_call("make clean && make", cwd=osp.dirname(__file__), shell=True)
-
-
-__all__ = ['set_random_seed', 'set_verbose', 'inpaint', 'inpaint_regularity']
-
-
-class CShapeT(ctypes.Structure):
-    _fields_ = [
-        ('width', ctypes.c_int),
-        ('height', ctypes.c_int),
-        ('channels', ctypes.c_int),
-    ]
-
-
-class CMatT(ctypes.Structure):
-    _fields_ = [
-        ('data_ptr', ctypes.c_void_p),
-        ('shape', CShapeT),
-        ('dtype', ctypes.c_int)
-    ]
-
-
-PMLIB = ctypes.CDLL(osp.join(osp.dirname(__file__), 'libpatchmatch.so'))
-
-PMLIB.PM_set_random_seed.argtypes = [ctypes.c_uint]
-PMLIB.PM_set_verbose.argtypes = [ctypes.c_int]
-PMLIB.PM_free_pymat.argtypes = [CMatT]
-PMLIB.PM_inpaint.argtypes = [CMatT, CMatT, ctypes.c_int]
-PMLIB.PM_inpaint.restype = CMatT
-PMLIB.PM_inpaint_regularity.argtypes = [CMatT, CMatT, CMatT, ctypes.c_int, ctypes.c_float]
-PMLIB.PM_inpaint_regularity.restype = CMatT
-PMLIB.PM_inpaint2.argtypes = [CMatT, CMatT, CMatT, ctypes.c_int]
-PMLIB.PM_inpaint2.restype = CMatT
-PMLIB.PM_inpaint2_regularity.argtypes = [CMatT, CMatT, CMatT, CMatT, ctypes.c_int, ctypes.c_float]
-PMLIB.PM_inpaint2_regularity.restype = CMatT
-
-
-def set_random_seed(seed: int):
-    PMLIB.PM_set_random_seed(ctypes.c_uint(seed))
-
-
-def set_verbose(verbose: bool):
-    PMLIB.PM_set_verbose(ctypes.c_int(verbose))
-
-
-def inpaint(
-    image: Union[np.ndarray, Image.Image],
-    mask: Optional[Union[np.ndarray, Image.Image]] = None,
-    *,
-    global_mask: Optional[Union[np.ndarray, Image.Image]] = None,
-    patch_size: int = 15
-) -> np.ndarray:
-    """
-    PatchMatch based inpainting proposed in:
-
-        PatchMatch : A Randomized Correspondence Algorithm for Structural Image Editing
-        C.Barnes, E.Shechtman, A.Finkelstein and Dan B.Goldman
-        SIGGRAPH 2009
-
-    Args:
-        image (Union[np.ndarray, Image.Image]): the input image, should be 3-channel RGB/BGR.
-        mask (Union[np.array, Image.Image], optional): the mask of the hole(s) to be filled, should be 1-channel.
-        If not provided (None), the algorithm will treat all purely white pixels as the holes (255, 255, 255).
-        global_mask (Union[np.array, Image.Image], optional): the target mask of the output image.
-        patch_size (int): the patch size for the inpainting algorithm.
-
-    Return:
-        result (np.ndarray): the repaired image, of the same size as the input image.
-    """
-
-    if isinstance(image, Image.Image):
-        image = np.array(image)
-    image = np.ascontiguousarray(image)
-    assert image.ndim == 3 and image.shape[2] == 3 and image.dtype == 'uint8'
-
-    if mask is None:
-        mask = (image == (255, 255, 255)).all(axis=2, keepdims=True).astype('uint8')
-        mask = np.ascontiguousarray(mask)
-    else:
-        mask = _canonize_mask_array(mask)
-
-    if global_mask is None:
-        ret_pymat = PMLIB.PM_inpaint(np_to_pymat(image), np_to_pymat(mask), ctypes.c_int(patch_size))
-    else:
-        global_mask = _canonize_mask_array(global_mask)
-        ret_pymat = PMLIB.PM_inpaint2(np_to_pymat(image), np_to_pymat(mask), np_to_pymat(global_mask), ctypes.c_int(patch_size))
-
-    ret_npmat = pymat_to_np(ret_pymat)
-    PMLIB.PM_free_pymat(ret_pymat)
-
-    return ret_npmat
-
-
-def inpaint_regularity(
-    image: Union[np.ndarray, Image.Image],
-    mask: Optional[Union[np.ndarray, Image.Image]],
-    ijmap: np.ndarray,
-    *,
-    global_mask: Optional[Union[np.ndarray, Image.Image]] = None,
-    patch_size: int = 15, guide_weight: float = 0.25
-) -> np.ndarray:
-    if isinstance(image, Image.Image):
-        image = np.array(image)
-    image = np.ascontiguousarray(image)
-
-    assert isinstance(ijmap, np.ndarray) and ijmap.ndim == 3 and ijmap.shape[2] == 3 and ijmap.dtype == 'float32'
-    ijmap = np.ascontiguousarray(ijmap)
-
-    assert image.ndim == 3 and image.shape[2] == 3 and image.dtype == 'uint8'
-    if mask is None:
-        mask = (image == (255, 255, 255)).all(axis=2, keepdims=True).astype('uint8')
-        mask = np.ascontiguousarray(mask)
-    else:
-        mask = _canonize_mask_array(mask)
-
-
-    if global_mask is None:
-        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))
-    else:
-        global_mask = _canonize_mask_array(global_mask)
-        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))
-
-    ret_npmat = pymat_to_np(ret_pymat)
-    PMLIB.PM_free_pymat(ret_pymat)
-
-    return ret_npmat
-
-
-def _canonize_mask_array(mask):
-    if isinstance(mask, Image.Image):
-        mask = np.array(mask)
-    if mask.ndim == 2 and mask.dtype == 'uint8':
-        mask = mask[..., np.newaxis]
-    assert mask.ndim == 3 and mask.shape[2] == 1 and mask.dtype == 'uint8'
-    return np.ascontiguousarray(mask)
-
-
-dtype_pymat_to_ctypes = [
-    ctypes.c_uint8,
-    ctypes.c_int8,
-    ctypes.c_uint16,
-    ctypes.c_int16,
-    ctypes.c_int32,
-    ctypes.c_float,
-    ctypes.c_double,
-]
-
-
-dtype_np_to_pymat = {
-    'uint8': 0,
-    'int8': 1,
-    'uint16': 2,
-    'int16': 3,
-    'int32': 4,
-    'float32': 5,
-    'float64': 6,
-}
-
-
-def np_to_pymat(npmat):
-    assert npmat.ndim == 3
-    return CMatT(
-        ctypes.cast(npmat.ctypes.data, ctypes.c_void_p),
-        CShapeT(npmat.shape[1], npmat.shape[0], npmat.shape[2]),
-        dtype_np_to_pymat[str(npmat.dtype)]
-    )
-
-
-def pymat_to_np(pymat):
-    npmat = np.ctypeslib.as_array(
-        ctypes.cast(pymat.data_ptr, ctypes.POINTER(dtype_pymat_to_ctypes[pymat.dtype])),
-        (pymat.shape.height, pymat.shape.width, pymat.shape.channels)
-    )
-    ret = np.empty(npmat.shape, npmat.dtype)
-    ret[:] = npmat
-    return ret
-
+#! /usr/bin/env python3
+# -*- coding: utf-8 -*-
+# File   : patch_match.py
+# Author : Jiayuan Mao
+# Email  : maojiayuan@gmail.com
+# Date   : 01/09/2020
+#
+# Distributed under terms of the MIT license.
+
+import ctypes
+import os.path as osp
+from typing import Optional, Union
+
+import numpy as np
+from PIL import Image
+
+
+import os
+if os.name!="nt":
+    # Otherwise, fall back to the subprocess.
+    import subprocess
+    print('Compiling and loading c extensions from "{}".'.format(osp.realpath(osp.dirname(__file__))))
+    # subprocess.check_call(['./travis.sh'], cwd=osp.dirname(__file__))
+    subprocess.check_call("make clean && make", cwd=osp.dirname(__file__), shell=True)
+
+
+__all__ = ['set_random_seed', 'set_verbose', 'inpaint', 'inpaint_regularity']
+
+
+class CShapeT(ctypes.Structure):
+    _fields_ = [
+        ('width', ctypes.c_int),
+        ('height', ctypes.c_int),
+        ('channels', ctypes.c_int),
+    ]
+
+
+class CMatT(ctypes.Structure):
+    _fields_ = [
+        ('data_ptr', ctypes.c_void_p),
+        ('shape', CShapeT),
+        ('dtype', ctypes.c_int)
+    ]
+
+import tempfile
+from urllib.request import urlopen, Request
+import shutil
+from pathlib import Path
+from tqdm import tqdm
+
+def download_url_to_file(url, dst, hash_prefix=None, progress=True):
+    r"""Download object at the given URL to a local path.
+
+    Args:
+        url (string): URL of the object to download
+        dst (string): Full path where object will be saved, e.g. ``/tmp/temporary_file``
+        hash_prefix (string, optional): If not None, the SHA256 downloaded file should start with ``hash_prefix``.
+            Default: None
+        progress (bool, optional): whether or not to display a progress bar to stderr
+            Default: True
+    https://pytorch.org/docs/stable/_modules/torch/hub.html#load_state_dict_from_url
+    """
+    file_size = None
+    req = Request(url)
+    u = urlopen(req)
+    meta = u.info()
+    if hasattr(meta, 'getheaders'):
+        content_length = meta.getheaders("Content-Length")
+    else:
+        content_length = meta.get_all("Content-Length")
+    if content_length is not None and len(content_length) > 0:
+        file_size = int(content_length[0])
+
+    # We deliberately save it in a temp file and move it after
+    # download is complete. This prevents a local working checkpoint
+    # being overridden by a broken download.
+    dst = os.path.expanduser(dst)
+    dst_dir = os.path.dirname(dst)
+    f = tempfile.NamedTemporaryFile(delete=False, dir=dst_dir)
+
+    try:
+        with tqdm(total=file_size, disable=not progress,
+                  unit='B', unit_scale=True, unit_divisor=1024) as pbar:
+            while True:
+                buffer = u.read(8192)
+                if len(buffer) == 0:
+                    break
+                f.write(buffer)
+                pbar.update(len(buffer))
+
+        f.close()
+        shutil.move(f.name, dst)
+    finally:
+        f.close()
+        if os.path.exists(f.name):
+            os.remove(f.name)
+
+if os.name!="nt":
+    PMLIB = ctypes.CDLL(osp.join(osp.dirname(__file__), 'libpatchmatch.so'))
+else:
+    if not os.path.exists(osp.join(osp.dirname(__file__), 'libpatchmatch.dll')):
+        download_url_to_file(url="https://github.com/lkwq007/PyPatchMatch/releases/download/v0.1/libpatchmatch.dll",dst=osp.join(osp.dirname(__file__), 'libpatchmatch.dll'))
+    if not os.path.exists(osp.join(osp.dirname(__file__), 'opencv_world460.dll')):
+        download_url_to_file(url="https://github.com/lkwq007/PyPatchMatch/releases/download/v0.1/opencv_world460.dll",dst=osp.join(osp.dirname(__file__), 'opencv_world460.dll'))
+    if not os.path.exists(osp.join(osp.dirname(__file__), 'libpatchmatch.dll')):
+        print("[Dependency Missing] Please download https://github.com/lkwq007/PyPatchMatch/releases/download/v0.1/libpatchmatch.dll and put it into the PyPatchMatch folder")
+    if not os.path.exists(osp.join(osp.dirname(__file__), 'opencv_world460.dll')):
+        print("[Dependency Missing] Please download https://github.com/lkwq007/PyPatchMatch/releases/download/v0.1/opencv_world460.dll and put it into the PyPatchMatch folder")
+    PMLIB = ctypes.CDLL(osp.join(osp.dirname(__file__), 'libpatchmatch.dll'))
+
+PMLIB.PM_set_random_seed.argtypes = [ctypes.c_uint]
+PMLIB.PM_set_verbose.argtypes = [ctypes.c_int]
+PMLIB.PM_free_pymat.argtypes = [CMatT]
+PMLIB.PM_inpaint.argtypes = [CMatT, CMatT, ctypes.c_int]
+PMLIB.PM_inpaint.restype = CMatT
+PMLIB.PM_inpaint_regularity.argtypes = [CMatT, CMatT, CMatT, ctypes.c_int, ctypes.c_float]
+PMLIB.PM_inpaint_regularity.restype = CMatT
+PMLIB.PM_inpaint2.argtypes = [CMatT, CMatT, CMatT, ctypes.c_int]
+PMLIB.PM_inpaint2.restype = CMatT
+PMLIB.PM_inpaint2_regularity.argtypes = [CMatT, CMatT, CMatT, CMatT, ctypes.c_int, ctypes.c_float]
+PMLIB.PM_inpaint2_regularity.restype = CMatT
+
+
+def set_random_seed(seed: int):
+    PMLIB.PM_set_random_seed(ctypes.c_uint(seed))
+
+
+def set_verbose(verbose: bool):
+    PMLIB.PM_set_verbose(ctypes.c_int(verbose))
+
+
+def inpaint(
+    image: Union[np.ndarray, Image.Image],
+    mask: Optional[Union[np.ndarray, Image.Image]] = None,
+    *,
+    global_mask: Optional[Union[np.ndarray, Image.Image]] = None,
+    patch_size: int = 15
+) -> np.ndarray:
+    """
+    PatchMatch based inpainting proposed in:
+
+        PatchMatch : A Randomized Correspondence Algorithm for Structural Image Editing
+        C.Barnes, E.Shechtman, A.Finkelstein and Dan B.Goldman
+        SIGGRAPH 2009
+
+    Args:
+        image (Union[np.ndarray, Image.Image]): the input image, should be 3-channel RGB/BGR.
+        mask (Union[np.array, Image.Image], optional): the mask of the hole(s) to be filled, should be 1-channel.
+        If not provided (None), the algorithm will treat all purely white pixels as the holes (255, 255, 255).
+        global_mask (Union[np.array, Image.Image], optional): the target mask of the output image.
+        patch_size (int): the patch size for the inpainting algorithm.
+
+    Return:
+        result (np.ndarray): the repaired image, of the same size as the input image.
+    """
+
+    if isinstance(image, Image.Image):
+        image = np.array(image)
+    image = np.ascontiguousarray(image)
+    assert image.ndim == 3 and image.shape[2] == 3 and image.dtype == 'uint8'
+
+    if mask is None:
+        mask = (image == (255, 255, 255)).all(axis=2, keepdims=True).astype('uint8')
+        mask = np.ascontiguousarray(mask)
+    else:
+        mask = _canonize_mask_array(mask)
+
+    if global_mask is None:
+        ret_pymat = PMLIB.PM_inpaint(np_to_pymat(image), np_to_pymat(mask), ctypes.c_int(patch_size))
+    else:
+        global_mask = _canonize_mask_array(global_mask)
+        ret_pymat = PMLIB.PM_inpaint2(np_to_pymat(image), np_to_pymat(mask), np_to_pymat(global_mask), ctypes.c_int(patch_size))
+
+    ret_npmat = pymat_to_np(ret_pymat)
+    PMLIB.PM_free_pymat(ret_pymat)
+
+    return ret_npmat
+
+
+def inpaint_regularity(
+    image: Union[np.ndarray, Image.Image],
+    mask: Optional[Union[np.ndarray, Image.Image]],
+    ijmap: np.ndarray,
+    *,
+    global_mask: Optional[Union[np.ndarray, Image.Image]] = None,
+    patch_size: int = 15, guide_weight: float = 0.25
+) -> np.ndarray:
+    if isinstance(image, Image.Image):
+        image = np.array(image)
+    image = np.ascontiguousarray(image)
+
+    assert isinstance(ijmap, np.ndarray) and ijmap.ndim == 3 and ijmap.shape[2] == 3 and ijmap.dtype == 'float32'
+    ijmap = np.ascontiguousarray(ijmap)
+
+    assert image.ndim == 3 and image.shape[2] == 3 and image.dtype == 'uint8'
+    if mask is None:
+        mask = (image == (255, 255, 255)).all(axis=2, keepdims=True).astype('uint8')
+        mask = np.ascontiguousarray(mask)
+    else:
+        mask = _canonize_mask_array(mask)
+
+
+    if global_mask is None:
+        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))
+    else:
+        global_mask = _canonize_mask_array(global_mask)
+        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))
+
+    ret_npmat = pymat_to_np(ret_pymat)
+    PMLIB.PM_free_pymat(ret_pymat)
+
+    return ret_npmat
+
+
+def _canonize_mask_array(mask):
+    if isinstance(mask, Image.Image):
+        mask = np.array(mask)
+    if mask.ndim == 2 and mask.dtype == 'uint8':
+        mask = mask[..., np.newaxis]
+    assert mask.ndim == 3 and mask.shape[2] == 1 and mask.dtype == 'uint8'
+    return np.ascontiguousarray(mask)
+
+
+dtype_pymat_to_ctypes = [
+    ctypes.c_uint8,
+    ctypes.c_int8,
+    ctypes.c_uint16,
+    ctypes.c_int16,
+    ctypes.c_int32,
+    ctypes.c_float,
+    ctypes.c_double,
+]
+
+
+dtype_np_to_pymat = {
+    'uint8': 0,
+    'int8': 1,
+    'uint16': 2,
+    'int16': 3,
+    'int32': 4,
+    'float32': 5,
+    'float64': 6,
+}
+
+
+def np_to_pymat(npmat):
+    assert npmat.ndim == 3
+    return CMatT(
+        ctypes.cast(npmat.ctypes.data, ctypes.c_void_p),
+        CShapeT(npmat.shape[1], npmat.shape[0], npmat.shape[2]),
+        dtype_np_to_pymat[str(npmat.dtype)]
+    )
+
+
+def pymat_to_np(pymat):
+    npmat = np.ctypeslib.as_array(
+        ctypes.cast(pymat.data_ptr, ctypes.POINTER(dtype_pymat_to_ctypes[pymat.dtype])),
+        (pymat.shape.height, pymat.shape.width, pymat.shape.channels)
+    )
+    ret = np.empty(npmat.shape, npmat.dtype)
+    ret[:] = npmat
+    return ret
+

PyPatchMatch/travis.sh

@@ -1,9 +1,9 @@
-#! /bin/bash
-#
-# travis.sh
-# Copyright (C) 2020 Jiayuan Mao <maojiayuan@gmail.com>
-#
-# Distributed under terms of the MIT license.
-#
-
-make clean && make
+#! /bin/bash
+#
+# travis.sh
+# Copyright (C) 2020 Jiayuan Mao <maojiayuan@gmail.com>
+#
+# Distributed under terms of the MIT license.
+#
+
+make clean && make

config.yaml

@@ -0,0 +1,18 @@
+shortcut:
+  clear: Escape
+  load: Ctrl+o
+  save: Ctrl+s
+  export: Ctrl+e
+  upload: Ctrl+u
+  selection: 1
+  canvas: 2
+  eraser: 3
+  outpaint: d
+  accept: a
+  cancel: c
+  retry: r
+  prev: q
+  next: e
+  zoom_in: z
+  zoom_out: x
+  random_seed: s

convert_checkpoint.py

@@ -0,0 +1,706 @@
+# coding=utf-8
+# Copyright 2022 The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# https://github.com/huggingface/diffusers/blob/main/scripts/convert_original_stable_diffusion_to_diffusers.py
+""" Conversion script for the LDM checkpoints. """
+
+import argparse
+import os
+
+import torch
+
+
+try:
+    from omegaconf import OmegaConf
+except ImportError:
+    raise ImportError(
+        "OmegaConf is required to convert the LDM checkpoints. Please install it with `pip install OmegaConf`."
+    )
+
+from diffusers import (
+    AutoencoderKL,
+    DDIMScheduler,
+    LDMTextToImagePipeline,
+    LMSDiscreteScheduler,
+    PNDMScheduler,
+    StableDiffusionPipeline,
+    UNet2DConditionModel,
+)
+from diffusers.pipelines.latent_diffusion.pipeline_latent_diffusion import LDMBertConfig, LDMBertModel
+from diffusers.pipelines.stable_diffusion import StableDiffusionSafetyChecker
+from transformers import AutoFeatureExtractor, BertTokenizerFast, CLIPTextModel, CLIPTokenizer
+
+
+def shave_segments(path, n_shave_prefix_segments=1):
+    """
+    Removes segments. Positive values shave the first segments, negative shave the last segments.
+    """
+    if n_shave_prefix_segments >= 0:
+        return ".".join(path.split(".")[n_shave_prefix_segments:])
+    else:
+        return ".".join(path.split(".")[:n_shave_prefix_segments])
+
+
+def renew_resnet_paths(old_list, n_shave_prefix_segments=0):
+    """
+    Updates paths inside resnets to the new naming scheme (local renaming)
+    """
+    mapping = []
+    for old_item in old_list:
+        new_item = old_item.replace("in_layers.0", "norm1")
+        new_item = new_item.replace("in_layers.2", "conv1")
+
+        new_item = new_item.replace("out_layers.0", "norm2")
+        new_item = new_item.replace("out_layers.3", "conv2")
+
+        new_item = new_item.replace("emb_layers.1", "time_emb_proj")
+        new_item = new_item.replace("skip_connection", "conv_shortcut")
+
+        new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments)
+
+        mapping.append({"old": old_item, "new": new_item})
+
+    return mapping
+
+
+def renew_vae_resnet_paths(old_list, n_shave_prefix_segments=0):
+    """
+    Updates paths inside resnets to the new naming scheme (local renaming)
+    """
+    mapping = []
+    for old_item in old_list:
+        new_item = old_item
+
+        new_item = new_item.replace("nin_shortcut", "conv_shortcut")
+        new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments)
+
+        mapping.append({"old": old_item, "new": new_item})
+
+    return mapping
+
+
+def renew_attention_paths(old_list, n_shave_prefix_segments=0):
+    """
+    Updates paths inside attentions to the new naming scheme (local renaming)
+    """
+    mapping = []
+    for old_item in old_list:
+        new_item = old_item
+
+        #         new_item = new_item.replace('norm.weight', 'group_norm.weight')
+        #         new_item = new_item.replace('norm.bias', 'group_norm.bias')
+
+        #         new_item = new_item.replace('proj_out.weight', 'proj_attn.weight')
+        #         new_item = new_item.replace('proj_out.bias', 'proj_attn.bias')
+
+        #         new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments)
+
+        mapping.append({"old": old_item, "new": new_item})
+
+    return mapping
+
+
+def renew_vae_attention_paths(old_list, n_shave_prefix_segments=0):
+    """
+    Updates paths inside attentions to the new naming scheme (local renaming)
+    """
+    mapping = []
+    for old_item in old_list:
+        new_item = old_item
+
+        new_item = new_item.replace("norm.weight", "group_norm.weight")
+        new_item = new_item.replace("norm.bias", "group_norm.bias")
+
+        new_item = new_item.replace("q.weight", "query.weight")
+        new_item = new_item.replace("q.bias", "query.bias")
+
+        new_item = new_item.replace("k.weight", "key.weight")
+        new_item = new_item.replace("k.bias", "key.bias")
+
+        new_item = new_item.replace("v.weight", "value.weight")
+        new_item = new_item.replace("v.bias", "value.bias")
+
+        new_item = new_item.replace("proj_out.weight", "proj_attn.weight")
+        new_item = new_item.replace("proj_out.bias", "proj_attn.bias")
+
+        new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments)
+
+        mapping.append({"old": old_item, "new": new_item})
+
+    return mapping
+
+
+def assign_to_checkpoint(
+    paths, checkpoint, old_checkpoint, attention_paths_to_split=None, additional_replacements=None, config=None
+):
+    """
+    This does the final conversion step: take locally converted weights and apply a global renaming
+    to them. It splits attention layers, and takes into account additional replacements
+    that may arise.
+
+    Assigns the weights to the new checkpoint.
+    """
+    assert isinstance(paths, list), "Paths should be a list of dicts containing 'old' and 'new' keys."
+
+    # Splits the attention layers into three variables.
+    if attention_paths_to_split is not None:
+        for path, path_map in attention_paths_to_split.items():
+            old_tensor = old_checkpoint[path]
+            channels = old_tensor.shape[0] // 3
+
+            target_shape = (-1, channels) if len(old_tensor.shape) == 3 else (-1)
+
+            num_heads = old_tensor.shape[0] // config["num_head_channels"] // 3
+
+            old_tensor = old_tensor.reshape((num_heads, 3 * channels // num_heads) + old_tensor.shape[1:])
+            query, key, value = old_tensor.split(channels // num_heads, dim=1)
+
+            checkpoint[path_map["query"]] = query.reshape(target_shape)
+            checkpoint[path_map["key"]] = key.reshape(target_shape)
+            checkpoint[path_map["value"]] = value.reshape(target_shape)
+
+    for path in paths:
+        new_path = path["new"]
+
+        # These have already been assigned
+        if attention_paths_to_split is not None and new_path in attention_paths_to_split:
+            continue
+
+        # Global renaming happens here
+        new_path = new_path.replace("middle_block.0", "mid_block.resnets.0")
+        new_path = new_path.replace("middle_block.1", "mid_block.attentions.0")
+        new_path = new_path.replace("middle_block.2", "mid_block.resnets.1")
+
+        if additional_replacements is not None:
+            for replacement in additional_replacements:
+                new_path = new_path.replace(replacement["old"], replacement["new"])
+
+        # proj_attn.weight has to be converted from conv 1D to linear
+        if "proj_attn.weight" in new_path:
+            checkpoint[new_path] = old_checkpoint[path["old"]][:, :, 0]
+        else:
+            checkpoint[new_path] = old_checkpoint[path["old"]]
+
+
+def conv_attn_to_linear(checkpoint):
+    keys = list(checkpoint.keys())
+    attn_keys = ["query.weight", "key.weight", "value.weight"]
+    for key in keys:
+        if ".".join(key.split(".")[-2:]) in attn_keys:
+            if checkpoint[key].ndim > 2:
+                checkpoint[key] = checkpoint[key][:, :, 0, 0]
+        elif "proj_attn.weight" in key:
+            if checkpoint[key].ndim > 2:
+                checkpoint[key] = checkpoint[key][:, :, 0]
+
+
+def create_unet_diffusers_config(original_config):
+    """
+    Creates a config for the diffusers based on the config of the LDM model.
+    """
+    unet_params = original_config.model.params.unet_config.params
+
+    block_out_channels = [unet_params.model_channels * mult for mult in unet_params.channel_mult]
+
+    down_block_types = []
+    resolution = 1
+    for i in range(len(block_out_channels)):
+        block_type = "CrossAttnDownBlock2D" if resolution in unet_params.attention_resolutions else "DownBlock2D"
+        down_block_types.append(block_type)
+        if i != len(block_out_channels) - 1:
+            resolution *= 2
+
+    up_block_types = []
+    for i in range(len(block_out_channels)):
+        block_type = "CrossAttnUpBlock2D" if resolution in unet_params.attention_resolutions else "UpBlock2D"
+        up_block_types.append(block_type)
+        resolution //= 2
+
+    config = dict(
+        sample_size=unet_params.image_size,
+        in_channels=unet_params.in_channels,
+        out_channels=unet_params.out_channels,
+        down_block_types=tuple(down_block_types),
+        up_block_types=tuple(up_block_types),
+        block_out_channels=tuple(block_out_channels),
+        layers_per_block=unet_params.num_res_blocks,
+        cross_attention_dim=unet_params.context_dim,
+        attention_head_dim=unet_params.num_heads,
+    )
+
+    return config
+
+
+def create_vae_diffusers_config(original_config):
+    """
+    Creates a config for the diffusers based on the config of the LDM model.
+    """
+    vae_params = original_config.model.params.first_stage_config.params.ddconfig
+    _ = original_config.model.params.first_stage_config.params.embed_dim
+
+    block_out_channels = [vae_params.ch * mult for mult in vae_params.ch_mult]
+    down_block_types = ["DownEncoderBlock2D"] * len(block_out_channels)
+    up_block_types = ["UpDecoderBlock2D"] * len(block_out_channels)
+
+    config = dict(
+        sample_size=vae_params.resolution,
+        in_channels=vae_params.in_channels,
+        out_channels=vae_params.out_ch,
+        down_block_types=tuple(down_block_types),
+        up_block_types=tuple(up_block_types),
+        block_out_channels=tuple(block_out_channels),
+        latent_channels=vae_params.z_channels,
+        layers_per_block=vae_params.num_res_blocks,
+    )
+    return config
+
+
+def create_diffusers_schedular(original_config):
+    schedular = DDIMScheduler(
+        num_train_timesteps=original_config.model.params.timesteps,
+        beta_start=original_config.model.params.linear_start,
+        beta_end=original_config.model.params.linear_end,
+        beta_schedule="scaled_linear",
+    )
+    return schedular
+
+
+def create_ldm_bert_config(original_config):
+    bert_params = original_config.model.parms.cond_stage_config.params
+    config = LDMBertConfig(
+        d_model=bert_params.n_embed,
+        encoder_layers=bert_params.n_layer,
+        encoder_ffn_dim=bert_params.n_embed * 4,
+    )
+    return config
+
+
+def convert_ldm_unet_checkpoint(checkpoint, config):
+    """
+    Takes a state dict and a config, and returns a converted checkpoint.
+    """
+
+    # extract state_dict for UNet
+    unet_state_dict = {}
+    unet_key = "model.diffusion_model."
+    keys = list(checkpoint.keys())
+    for key in keys:
+        if key.startswith(unet_key):
+            unet_state_dict[key.replace(unet_key, "")] = checkpoint.pop(key)
+
+    new_checkpoint = {}
+
+    new_checkpoint["time_embedding.linear_1.weight"] = unet_state_dict["time_embed.0.weight"]
+    new_checkpoint["time_embedding.linear_1.bias"] = unet_state_dict["time_embed.0.bias"]
+    new_checkpoint["time_embedding.linear_2.weight"] = unet_state_dict["time_embed.2.weight"]
+    new_checkpoint["time_embedding.linear_2.bias"] = unet_state_dict["time_embed.2.bias"]
+
+    new_checkpoint["conv_in.weight"] = unet_state_dict["input_blocks.0.0.weight"]
+    new_checkpoint["conv_in.bias"] = unet_state_dict["input_blocks.0.0.bias"]
+
+    new_checkpoint["conv_norm_out.weight"] = unet_state_dict["out.0.weight"]
+    new_checkpoint["conv_norm_out.bias"] = unet_state_dict["out.0.bias"]
+    new_checkpoint["conv_out.weight"] = unet_state_dict["out.2.weight"]
+    new_checkpoint["conv_out.bias"] = unet_state_dict["out.2.bias"]
+
+    # Retrieves the keys for the input blocks only
+    num_input_blocks = len({".".join(layer.split(".")[:2]) for layer in unet_state_dict if "input_blocks" in layer})
+    input_blocks = {
+        layer_id: [key for key in unet_state_dict if f"input_blocks.{layer_id}" in key]
+        for layer_id in range(num_input_blocks)
+    }
+
+    # Retrieves the keys for the middle blocks only
+    num_middle_blocks = len({".".join(layer.split(".")[:2]) for layer in unet_state_dict if "middle_block" in layer})
+    middle_blocks = {
+        layer_id: [key for key in unet_state_dict if f"middle_block.{layer_id}" in key]
+        for layer_id in range(num_middle_blocks)
+    }
+
+    # Retrieves the keys for the output blocks only
+    num_output_blocks = len({".".join(layer.split(".")[:2]) for layer in unet_state_dict if "output_blocks" in layer})
+    output_blocks = {
+        layer_id: [key for key in unet_state_dict if f"output_blocks.{layer_id}" in key]
+        for layer_id in range(num_output_blocks)
+    }
+
+    for i in range(1, num_input_blocks):
+        block_id = (i - 1) // (config["layers_per_block"] + 1)
+        layer_in_block_id = (i - 1) % (config["layers_per_block"] + 1)
+
+        resnets = [
+            key for key in input_blocks[i] if f"input_blocks.{i}.0" in key and f"input_blocks.{i}.0.op" not in key
+        ]
+        attentions = [key for key in input_blocks[i] if f"input_blocks.{i}.1" in key]
+
+        if f"input_blocks.{i}.0.op.weight" in unet_state_dict:
+            new_checkpoint[f"down_blocks.{block_id}.downsamplers.0.conv.weight"] = unet_state_dict.pop(
+                f"input_blocks.{i}.0.op.weight"
+            )
+            new_checkpoint[f"down_blocks.{block_id}.downsamplers.0.conv.bias"] = unet_state_dict.pop(
+                f"input_blocks.{i}.0.op.bias"
+            )
+
+        paths = renew_resnet_paths(resnets)
+        meta_path = {"old": f"input_blocks.{i}.0", "new": f"down_blocks.{block_id}.resnets.{layer_in_block_id}"}
+        assign_to_checkpoint(
+            paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config
+        )
+
+        if len(attentions):
+            paths = renew_attention_paths(attentions)
+            meta_path = {"old": f"input_blocks.{i}.1", "new": f"down_blocks.{block_id}.attentions.{layer_in_block_id}"}
+            assign_to_checkpoint(
+                paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config
+            )
+
+    resnet_0 = middle_blocks[0]
+    attentions = middle_blocks[1]
+    resnet_1 = middle_blocks[2]
+
+    resnet_0_paths = renew_resnet_paths(resnet_0)
+    assign_to_checkpoint(resnet_0_paths, new_checkpoint, unet_state_dict, config=config)
+
+    resnet_1_paths = renew_resnet_paths(resnet_1)
+    assign_to_checkpoint(resnet_1_paths, new_checkpoint, unet_state_dict, config=config)
+
+    attentions_paths = renew_attention_paths(attentions)
+    meta_path = {"old": "middle_block.1", "new": "mid_block.attentions.0"}
+    assign_to_checkpoint(
+        attentions_paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config
+    )
+
+    for i in range(num_output_blocks):
+        block_id = i // (config["layers_per_block"] + 1)
+        layer_in_block_id = i % (config["layers_per_block"] + 1)
+        output_block_layers = [shave_segments(name, 2) for name in output_blocks[i]]
+        output_block_list = {}
+
+        for layer in output_block_layers:
+            layer_id, layer_name = layer.split(".")[0], shave_segments(layer, 1)
+            if layer_id in output_block_list:
+                output_block_list[layer_id].append(layer_name)
+            else:
+                output_block_list[layer_id] = [layer_name]
+
+        if len(output_block_list) > 1:
+            resnets = [key for key in output_blocks[i] if f"output_blocks.{i}.0" in key]
+            attentions = [key for key in output_blocks[i] if f"output_blocks.{i}.1" in key]
+
+            resnet_0_paths = renew_resnet_paths(resnets)
+            paths = renew_resnet_paths(resnets)
+
+            meta_path = {"old": f"output_blocks.{i}.0", "new": f"up_blocks.{block_id}.resnets.{layer_in_block_id}"}
+            assign_to_checkpoint(
+                paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config
+            )
+
+            if ["conv.weight", "conv.bias"] in output_block_list.values():
+                index = list(output_block_list.values()).index(["conv.weight", "conv.bias"])
+                new_checkpoint[f"up_blocks.{block_id}.upsamplers.0.conv.weight"] = unet_state_dict[
+                    f"output_blocks.{i}.{index}.conv.weight"
+                ]
+                new_checkpoint[f"up_blocks.{block_id}.upsamplers.0.conv.bias"] = unet_state_dict[
+                    f"output_blocks.{i}.{index}.conv.bias"
+                ]
+
+                # Clear attentions as they have been attributed above.
+                if len(attentions) == 2:
+                    attentions = []
+
+            if len(attentions):
+                paths = renew_attention_paths(attentions)
+                meta_path = {
+                    "old": f"output_blocks.{i}.1",
+                    "new": f"up_blocks.{block_id}.attentions.{layer_in_block_id}",
+                }
+                assign_to_checkpoint(
+                    paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config
+                )
+        else:
+            resnet_0_paths = renew_resnet_paths(output_block_layers, n_shave_prefix_segments=1)
+            for path in resnet_0_paths:
+                old_path = ".".join(["output_blocks", str(i), path["old"]])
+                new_path = ".".join(["up_blocks", str(block_id), "resnets", str(layer_in_block_id), path["new"]])
+
+                new_checkpoint[new_path] = unet_state_dict[old_path]
+
+    return new_checkpoint
+
+
+def convert_ldm_vae_checkpoint(checkpoint, config):
+    # extract state dict for VAE
+    vae_state_dict = {}
+    vae_key = "first_stage_model."
+    keys = list(checkpoint.keys())
+    for key in keys:
+        if key.startswith(vae_key):
+            vae_state_dict[key.replace(vae_key, "")] = checkpoint.get(key)
+
+    new_checkpoint = {}
+
+    new_checkpoint["encoder.conv_in.weight"] = vae_state_dict["encoder.conv_in.weight"]
+    new_checkpoint["encoder.conv_in.bias"] = vae_state_dict["encoder.conv_in.bias"]
+    new_checkpoint["encoder.conv_out.weight"] = vae_state_dict["encoder.conv_out.weight"]
+    new_checkpoint["encoder.conv_out.bias"] = vae_state_dict["encoder.conv_out.bias"]
+    new_checkpoint["encoder.conv_norm_out.weight"] = vae_state_dict["encoder.norm_out.weight"]
+    new_checkpoint["encoder.conv_norm_out.bias"] = vae_state_dict["encoder.norm_out.bias"]
+
+    new_checkpoint["decoder.conv_in.weight"] = vae_state_dict["decoder.conv_in.weight"]
+    new_checkpoint["decoder.conv_in.bias"] = vae_state_dict["decoder.conv_in.bias"]
+    new_checkpoint["decoder.conv_out.weight"] = vae_state_dict["decoder.conv_out.weight"]
+    new_checkpoint["decoder.conv_out.bias"] = vae_state_dict["decoder.conv_out.bias"]
+    new_checkpoint["decoder.conv_norm_out.weight"] = vae_state_dict["decoder.norm_out.weight"]
+    new_checkpoint["decoder.conv_norm_out.bias"] = vae_state_dict["decoder.norm_out.bias"]
+
+    new_checkpoint["quant_conv.weight"] = vae_state_dict["quant_conv.weight"]
+    new_checkpoint["quant_conv.bias"] = vae_state_dict["quant_conv.bias"]
+    new_checkpoint["post_quant_conv.weight"] = vae_state_dict["post_quant_conv.weight"]
+    new_checkpoint["post_quant_conv.bias"] = vae_state_dict["post_quant_conv.bias"]
+
+    # Retrieves the keys for the encoder down blocks only
+    num_down_blocks = len({".".join(layer.split(".")[:3]) for layer in vae_state_dict if "encoder.down" in layer})
+    down_blocks = {
+        layer_id: [key for key in vae_state_dict if f"down.{layer_id}" in key] for layer_id in range(num_down_blocks)
+    }
+
+    # Retrieves the keys for the decoder up blocks only
+    num_up_blocks = len({".".join(layer.split(".")[:3]) for layer in vae_state_dict if "decoder.up" in layer})
+    up_blocks = {
+        layer_id: [key for key in vae_state_dict if f"up.{layer_id}" in key] for layer_id in range(num_up_blocks)
+    }
+
+    for i in range(num_down_blocks):
+        resnets = [key for key in down_blocks[i] if f"down.{i}" in key and f"down.{i}.downsample" not in key]
+
+        if f"encoder.down.{i}.downsample.conv.weight" in vae_state_dict:
+            new_checkpoint[f"encoder.down_blocks.{i}.downsamplers.0.conv.weight"] = vae_state_dict.pop(
+                f"encoder.down.{i}.downsample.conv.weight"
+            )
+            new_checkpoint[f"encoder.down_blocks.{i}.downsamplers.0.conv.bias"] = vae_state_dict.pop(
+                f"encoder.down.{i}.downsample.conv.bias"
+            )
+
+        paths = renew_vae_resnet_paths(resnets)
+        meta_path = {"old": f"down.{i}.block", "new": f"down_blocks.{i}.resnets"}
+        assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config)
+
+    mid_resnets = [key for key in vae_state_dict if "encoder.mid.block" in key]
+    num_mid_res_blocks = 2
+    for i in range(1, num_mid_res_blocks + 1):
+        resnets = [key for key in mid_resnets if f"encoder.mid.block_{i}" in key]
+
+        paths = renew_vae_resnet_paths(resnets)
+        meta_path = {"old": f"mid.block_{i}", "new": f"mid_block.resnets.{i - 1}"}
+        assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config)
+
+    mid_attentions = [key for key in vae_state_dict if "encoder.mid.attn" in key]
+    paths = renew_vae_attention_paths(mid_attentions)
+    meta_path = {"old": "mid.attn_1", "new": "mid_block.attentions.0"}
+    assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config)
+    conv_attn_to_linear(new_checkpoint)
+
+    for i in range(num_up_blocks):
+        block_id = num_up_blocks - 1 - i
+        resnets = [
+            key for key in up_blocks[block_id] if f"up.{block_id}" in key and f"up.{block_id}.upsample" not in key
+        ]
+
+        if f"decoder.up.{block_id}.upsample.conv.weight" in vae_state_dict:
+            new_checkpoint[f"decoder.up_blocks.{i}.upsamplers.0.conv.weight"] = vae_state_dict[
+                f"decoder.up.{block_id}.upsample.conv.weight"
+            ]
+            new_checkpoint[f"decoder.up_blocks.{i}.upsamplers.0.conv.bias"] = vae_state_dict[
+                f"decoder.up.{block_id}.upsample.conv.bias"
+            ]
+
+        paths = renew_vae_resnet_paths(resnets)
+        meta_path = {"old": f"up.{block_id}.block", "new": f"up_blocks.{i}.resnets"}
+        assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config)
+
+    mid_resnets = [key for key in vae_state_dict if "decoder.mid.block" in key]
+    num_mid_res_blocks = 2
+    for i in range(1, num_mid_res_blocks + 1):
+        resnets = [key for key in mid_resnets if f"decoder.mid.block_{i}" in key]
+
+        paths = renew_vae_resnet_paths(resnets)
+        meta_path = {"old": f"mid.block_{i}", "new": f"mid_block.resnets.{i - 1}"}
+        assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config)
+
+    mid_attentions = [key for key in vae_state_dict if "decoder.mid.attn" in key]
+    paths = renew_vae_attention_paths(mid_attentions)
+    meta_path = {"old": "mid.attn_1", "new": "mid_block.attentions.0"}
+    assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config)
+    conv_attn_to_linear(new_checkpoint)
+    return new_checkpoint
+
+
+def convert_ldm_bert_checkpoint(checkpoint, config):
+    def _copy_attn_layer(hf_attn_layer, pt_attn_layer):
+        hf_attn_layer.q_proj.weight.data = pt_attn_layer.to_q.weight
+        hf_attn_layer.k_proj.weight.data = pt_attn_layer.to_k.weight
+        hf_attn_layer.v_proj.weight.data = pt_attn_layer.to_v.weight
+
+        hf_attn_layer.out_proj.weight = pt_attn_layer.to_out.weight
+        hf_attn_layer.out_proj.bias = pt_attn_layer.to_out.bias
+
+    def _copy_linear(hf_linear, pt_linear):
+        hf_linear.weight = pt_linear.weight
+        hf_linear.bias = pt_linear.bias
+
+    def _copy_layer(hf_layer, pt_layer):
+        # copy layer norms
+        _copy_linear(hf_layer.self_attn_layer_norm, pt_layer[0][0])
+        _copy_linear(hf_layer.final_layer_norm, pt_layer[1][0])
+
+        # copy attn
+        _copy_attn_layer(hf_layer.self_attn, pt_layer[0][1])
+
+        # copy MLP
+        pt_mlp = pt_layer[1][1]
+        _copy_linear(hf_layer.fc1, pt_mlp.net[0][0])
+        _copy_linear(hf_layer.fc2, pt_mlp.net[2])
+
+    def _copy_layers(hf_layers, pt_layers):
+        for i, hf_layer in enumerate(hf_layers):
+            if i != 0:
+                i += i
+            pt_layer = pt_layers[i : i + 2]
+            _copy_layer(hf_layer, pt_layer)
+
+    hf_model = LDMBertModel(config).eval()
+
+    # copy  embeds
+    hf_model.model.embed_tokens.weight = checkpoint.transformer.token_emb.weight
+    hf_model.model.embed_positions.weight.data = checkpoint.transformer.pos_emb.emb.weight
+
+    # copy layer norm
+    _copy_linear(hf_model.model.layer_norm, checkpoint.transformer.norm)
+
+    # copy hidden layers
+    _copy_layers(hf_model.model.layers, checkpoint.transformer.attn_layers.layers)
+
+    _copy_linear(hf_model.to_logits, checkpoint.transformer.to_logits)
+
+    return hf_model
+
+
+def convert_ldm_clip_checkpoint(checkpoint):
+    text_model = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14")
+
+    keys = list(checkpoint.keys())
+
+    text_model_dict = {}
+
+    for key in keys:
+        if key.startswith("cond_stage_model.transformer"):
+            text_model_dict[key[len("cond_stage_model.transformer.") :]] = checkpoint[key]
+
+    text_model.load_state_dict(text_model_dict)
+
+    return text_model
+
+import os
+def convert_checkpoint(checkpoint_path, inpainting=False):
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument(
+        "--checkpoint_path", default=checkpoint_path, type=str, help="Path to the checkpoint to convert."
+    )
+    # !wget https://raw.githubusercontent.com/CompVis/stable-diffusion/main/configs/stable-diffusion/v1-inference.yaml
+    parser.add_argument(
+        "--original_config_file",
+        default=None,
+        type=str,
+        help="The YAML config file corresponding to the original architecture.",
+    )
+    parser.add_argument(
+        "--scheduler_type",
+        default="pndm",
+        type=str,
+        help="Type of scheduler to use. Should be one of ['pndm', 'lms', 'ddim']",
+    )
+    parser.add_argument("--dump_path", default=None, type=str, help="Path to the output model.")
+
+    args = parser.parse_args([])
+    if args.original_config_file is None:
+        if inpainting:
+            args.original_config_file = "./models/v1-inpainting-inference.yaml"
+        else:
+            args.original_config_file = "./models/v1-inference.yaml"
+
+    original_config = OmegaConf.load(args.original_config_file)
+    checkpoint = torch.load(args.checkpoint_path)["state_dict"]
+
+    num_train_timesteps = original_config.model.params.timesteps
+    beta_start = original_config.model.params.linear_start
+    beta_end = original_config.model.params.linear_end
+    if args.scheduler_type == "pndm":
+        scheduler = PNDMScheduler(
+            beta_end=beta_end,
+            beta_schedule="scaled_linear",
+            beta_start=beta_start,
+            num_train_timesteps=num_train_timesteps,
+            skip_prk_steps=True,
+        )
+    elif args.scheduler_type == "lms":
+        scheduler = LMSDiscreteScheduler(beta_start=beta_start, beta_end=beta_end, beta_schedule="scaled_linear")
+    elif args.scheduler_type == "ddim":
+        scheduler = DDIMScheduler(
+            beta_start=beta_start,
+            beta_end=beta_end,
+            beta_schedule="scaled_linear",
+            clip_sample=False,
+            set_alpha_to_one=False,
+        )
+    else:
+        raise ValueError(f"Scheduler of type {args.scheduler_type} doesn't exist!")
+
+    # Convert the UNet2DConditionModel model.
+    unet_config = create_unet_diffusers_config(original_config)
+    converted_unet_checkpoint = convert_ldm_unet_checkpoint(checkpoint, unet_config)
+
+    unet = UNet2DConditionModel(**unet_config)
+    unet.load_state_dict(converted_unet_checkpoint)
+
+    # Convert the VAE model.
+    vae_config = create_vae_diffusers_config(original_config)
+    converted_vae_checkpoint = convert_ldm_vae_checkpoint(checkpoint, vae_config)
+
+    vae = AutoencoderKL(**vae_config)
+    vae.load_state_dict(converted_vae_checkpoint)
+
+    # Convert the text model.
+    text_model_type = original_config.model.params.cond_stage_config.target.split(".")[-1]
+    if text_model_type == "FrozenCLIPEmbedder":
+        text_model = convert_ldm_clip_checkpoint(checkpoint)
+        tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14")
+        safety_checker = StableDiffusionSafetyChecker.from_pretrained("CompVis/stable-diffusion-safety-checker")
+        feature_extractor = AutoFeatureExtractor.from_pretrained("CompVis/stable-diffusion-safety-checker")
+        pipe = StableDiffusionPipeline(
+            vae=vae,
+            text_encoder=text_model,
+            tokenizer=tokenizer,
+            unet=unet,
+            scheduler=scheduler,
+            safety_checker=safety_checker,
+            feature_extractor=feature_extractor,
+        )
+    else:
+        text_config = create_ldm_bert_config(original_config)
+        text_model = convert_ldm_bert_checkpoint(checkpoint, text_config)
+        tokenizer = BertTokenizerFast.from_pretrained("bert-base-uncased")
+        pipe = LDMTextToImagePipeline(vqvae=vae, bert=text_model, tokenizer=tokenizer, unet=unet, scheduler=scheduler)
+
+    return pipe

js/keyboard.js

@@ -0,0 +1,37 @@
+
+window.my_setup_keyboard=setInterval(function(){
+    let app=document.querySelector("gradio-app");
+    app=app.shadowRoot??app;
+    let frame=app.querySelector("#sdinfframe").contentWindow;
+    console.log("Check iframe...");
+    if(frame.setup_shortcut)
+    {
+        frame.setup_shortcut(json);
+        clearInterval(window.my_setup_keyboard);
+    }
+}, 1000);
+var config=JSON.parse(json);
+var key_map={};
+Object.keys(config.shortcut).forEach(k=>{
+    key_map[config.shortcut[k]]=k;
+});
+document.addEventListener("keydown", e => {
+    if(e.target.tagName!="INPUT"&&e.target.tagName!="GRADIO-APP"&&e.target.tagName!="TEXTAREA")
+    {
+        let key=e.key;
+        if(e.ctrlKey)
+        {
+            key="Ctrl+"+e.key;
+            if(key in key_map)
+            {
+                e.preventDefault();
+            }
+        }
+        let app=document.querySelector("gradio-app");
+        app=app.shadowRoot??app;
+        let frame=app.querySelector("#sdinfframe").contentDocument;
+        frame.dispatchEvent(
+          new KeyboardEvent("keydown", {key: e.key, ctrlKey: e.ctrlKey})
+        );
+    }
+})

js/mode.js

@@ -1,6 +1,6 @@
-function(mode){
-    let app=document.querySelector("gradio-app");
-    let frame=app.querySelector("#sdinfframe").contentWindow;
-    frame.postMessage(["mode", mode], "*");
-    return mode;
+function(mode){
+    let app=document.querySelector("gradio-app").shadowRoot;
+    let frame=app.querySelector("#sdinfframe").contentWindow.document;
+    frame.querySelector("#mode").value=mode;
+    return mode;
 }

js/outpaint.js

@@ -1,31 +1,23 @@
-function(a){
-    if(!window.my_observe_outpaint)
-    {
-        console.log("setup outpaint here");
-        window.my_observe_outpaint = new MutationObserver(function (event) {
-            console.log(event);
-            let app=document.querySelector("gradio-app");
-            let frame=app.querySelector("#sdinfframe").contentWindow;
-            var str=document.querySelector("gradio-app").querySelector("#output textarea").value;
-            frame.postMessage(["outpaint", str], "*");
-        });
-        window.my_observe_outpaint_target=document.querySelector("gradio-app").querySelector("#output span")
-        window.my_observe_outpaint.observe(window.my_observe_outpaint_target, {
-            attributes: false, 
-            subtree: true,
-            childList: true, 
-            characterData: true
-        });
-        window.addEventListener("message", function(e){
-            if(e.data[0]=="transfer")
-            {
-                document.querySelector("gradio-app").querySelector("#input textarea").value=e.data[1];
-                document.querySelector("gradio-app").querySelector("#proceed").click();
-            }
-        });
-    }
-    let app=document.querySelector("gradio-app");
-    let frame=app.querySelector("#sdinfframe").contentWindow;
-    frame.postMessage(["transfer"],"*")
-    return a;
+function(a){
+    if(!window.my_observe_outpaint)
+    {
+        console.log("setup outpaint here");
+        window.my_observe_outpaint = new MutationObserver(function (event) {
+            console.log(event);
+            let app=document.querySelector("gradio-app");
+            app=app.shadowRoot??app;
+            let frame=app.querySelector("#sdinfframe").contentWindow;
+            frame.postMessage(["outpaint", ""], "*");
+        });
+        var app=document.querySelector("gradio-app");
+        app=app.shadowRoot??app;
+        window.my_observe_outpaint_target=app.querySelector("#output span");
+        window.my_observe_outpaint.observe(window.my_observe_outpaint_target, {
+            attributes: false, 
+            subtree: true,
+            childList: true, 
+            characterData: true
+        });
+    }
+    return a;
 }

js/proceed.js

@@ -1,22 +1,42 @@
-function(sel_buffer_str,
-    prompt_text,
-    strength,
-    guidance,
-    step,
-    resize_check,
-    fill_mode,
-    enable_safety,
-    state){
-    sel_buffer = document.querySelector("gradio-app").querySelector("#input textarea").value;
-    return [
-        sel_buffer,
-        prompt_text,
-        strength,
-        guidance,
-        step,
-        resize_check,
-        fill_mode,
-        enable_safety,
-        state
-    ]
+function(sel_buffer_str,
+    prompt_text,
+    negative_prompt_text,
+    strength,
+    guidance,
+    step,
+    resize_check,
+    fill_mode,
+    enable_safety,
+    use_correction,
+    enable_img2img,
+    use_seed,
+    seed_val,
+    generate_num,
+    scheduler,
+    scheduler_eta,
+    state){
+    let app=document.querySelector("gradio-app");
+    app=app.shadowRoot??app;
+    sel_buffer=app.querySelector("#input textarea").value;
+    let use_correction_bak=false;
+    ({resize_check,enable_safety,use_correction_bak,enable_img2img,use_seed,seed_val}=window.config_obj);
+    return [
+        sel_buffer,
+        prompt_text,
+        negative_prompt_text,
+        strength,
+        guidance,
+        step,
+        resize_check,
+        fill_mode,
+        enable_safety,
+        use_correction,
+        enable_img2img,
+        use_seed,
+        seed_val,
+        generate_num,
+        scheduler,
+        scheduler_eta,
+        state,
+    ]
 }

js/setup.js

@@ -1,22 +1,28 @@
-function(token_val, width, height, size){
-    let app=document.querySelector("gradio-app");
-    app.querySelector("#sdinfframe").style.height=height+"px";
-    let frame=app.querySelector("#sdinfframe").contentWindow.document;
-    if(frame.querySelector("#setup").value=="0")
-    {
-        window.my_setup=setInterval(function(){
-            let frame=document.querySelector("gradio-app").querySelector("#sdinfframe").contentWindow.document;
-            console.log("Check PyScript...")
-            if(frame.querySelector("#setup").value=="1")
-            {
-                frame.querySelector("#draw").click();
-                clearInterval(window.my_setup);
-            }
-        },100)
-    }
-    else
-    {
-        frame.querySelector("#draw").click();
-    }
-    return [token_val, width, height, size];
+function(token_val, width, height, size, model_choice, model_path){
+    let app=document.querySelector("gradio-app");
+    app=app.shadowRoot??app;
+    app.querySelector("#sdinfframe").style.height=80+Number(height)+"px";
+    // app.querySelector("#setup_row").style.display="none";
+    app.querySelector("#model_path_input").style.display="none";
+    let frame=app.querySelector("#sdinfframe").contentWindow.document;
+
+    if(frame.querySelector("#setup").value=="0")
+    {
+        window.my_setup=setInterval(function(){
+            let app=document.querySelector("gradio-app");
+            app=app.shadowRoot??app;
+            let frame=app.querySelector("#sdinfframe").contentWindow.document;
+            console.log("Check PyScript...")
+            if(frame.querySelector("#setup").value=="1")
+            {
+                frame.querySelector("#draw").click();
+                clearInterval(window.my_setup);
+            }
+        }, 100)
+    }
+    else
+    {
+        frame.querySelector("#draw").click();
+    }
+    return [token_val, width, height, size, model_choice, model_path];
 }

js/toolbar.js

@@ -0,0 +1,581 @@
+// import { w2ui,w2toolbar,w2field,query,w2alert, w2utils,w2confirm} from "https://rawgit.com/vitmalina/w2ui/master/dist/w2ui.es6.min.js"
+// import { w2ui,w2toolbar,w2field,query,w2alert, w2utils,w2confirm} from "https://cdn.jsdelivr.net/gh/vitmalina/w2ui@master/dist/w2ui.es6.min.js"
+
+// https://stackoverflow.com/questions/36280818/how-to-convert-file-to-base64-in-javascript
+function getBase64(file) {
+   var reader = new FileReader();
+   reader.readAsDataURL(file);
+   reader.onload = function () {
+    add_image(reader.result);
+    //  console.log(reader.result);
+   };
+   reader.onerror = function (error) {
+     console.log("Error: ", error);
+   };
+}
+
+function getText(file) {
+   var reader = new FileReader();
+   reader.readAsText(file);
+   reader.onload = function () {
+    window.postMessage(["load",reader.result],"*")
+    //  console.log(reader.result);
+   };
+   reader.onerror = function (error) {
+     console.log("Error: ", error);
+   };
+}
+
+document.querySelector("#upload_file").addEventListener("change", (event)=>{
+    console.log(event);
+    let file = document.querySelector("#upload_file").files[0];
+    getBase64(file);
+})
+
+document.querySelector("#upload_state").addEventListener("change", (event)=>{
+    console.log(event);
+    let file = document.querySelector("#upload_state").files[0];
+    getText(file);
+})
+
+open_setting = function() {
+    if (!w2ui.foo) {
+        new w2form({
+            name: "foo",
+            style: "border: 0px; background-color: transparent;",
+            fields: [{
+                    field: "canvas_width",
+                    type: "int",
+                    required: true,
+                    html: {
+                        label: "Canvas Width"
+                    }
+                },
+                {
+                    field: "canvas_height",
+                    type: "int",
+                    required: true,
+                    html: {
+                        label: "Canvas Height"
+                    }
+                },
+            ],
+            record: {
+                canvas_width: 1200,
+                canvas_height: 600,
+            },
+            actions: {
+                Save() {
+                    this.validate();
+                    let record = this.getCleanRecord();
+                    window.postMessage(["resize",record.canvas_width,record.canvas_height],"*");
+                    w2popup.close();
+                },
+                custom: {
+                    text: "Cancel",
+                    style: "text-transform: uppercase",
+                    onClick(event) {
+                        w2popup.close();
+                    }
+                }
+            }
+        });
+    }
+    w2popup.open({
+            title: "Form in a Popup",
+            body: "<div id='form' style='width: 100%; height: 100%;''></div>",
+            style: "padding: 15px 0px 0px 0px",
+            width: 500,
+            height: 280,
+            showMax: true,
+            async onToggle(event) {
+                await event.complete
+                w2ui.foo.resize();
+            }
+        })
+        .then((event) => {
+            w2ui.foo.render("#form")
+        });
+}
+
+var button_lst=["clear", "load", "save", "export", "upload", "selection", "canvas", "eraser", "outpaint", "accept", "cancel", "retry", "prev", "current", "next", "eraser_size_btn", "eraser_size", "resize_selection", "scale", "zoom_in", "zoom_out", "help"];
+var upload_button_lst=['clear', 'load', 'save', "upload", 'export', 'outpaint', 'resize_selection', 'help', "setting"];
+var resize_button_lst=['clear', 'load', 'save', "upload", 'export', "selection", "canvas", "eraser", 'outpaint', 'resize_selection',"zoom_in", "zoom_out", 'help', "setting"];
+var outpaint_button_lst=['clear', 'load', 'save', "canvas", "eraser", "upload", 'export', 'resize_selection', "zoom_in", "zoom_out",'help', "setting"];
+var outpaint_result_lst=["accept", "cancel", "retry", "prev", "current", "next"];
+var outpaint_result_func_lst=["accept", "retry", "prev", "current", "next"];
+
+function check_button(id,text="",checked=true,tooltip="")
+{
+    return { type: "check",  id: id, text: text, icon: checked?"fa-solid fa-square-check":"fa-regular fa-square", checked: checked, tooltip: tooltip };
+}
+
+var toolbar=new w2toolbar({
+    box: "#toolbar",
+    name: "toolbar",
+    tooltip: "top",
+    items: [
+        { type: "button", id: "clear", text: "Reset", tooltip: "Reset Canvas", icon: "fa-solid fa-rectangle-xmark" },
+        { type: "break" },
+        { type: "button", id: "load", tooltip: "Load Canvas", icon: "fa-solid fa-file-import" },
+        { type: "button", id: "save", tooltip: "Save Canvas", icon: "fa-solid fa-file-export" },
+        { type: "button", id: "export", tooltip: "Export Image", icon: "fa-solid fa-floppy-disk" },
+        { type: "break" },
+        { type: "button", id: "upload", text: "Upload Image", icon: "fa-solid fa-upload" },
+        { type: "break" },
+        { type: "radio", id: "selection", group: "1", tooltip: "Selection", icon: "fa-solid fa-arrows-up-down-left-right", checked: true },
+        { type: "radio", id: "canvas", group: "1", tooltip: "Canvas", icon: "fa-solid fa-image" },
+        { type: "radio", id: "eraser", group: "1", tooltip: "Eraser", icon: "fa-solid fa-eraser" },
+        { type: "break" },
+        { type: "button", id: "outpaint", text: "Outpaint", tooltip: "Run Outpainting", icon: "fa-solid fa-brush" },
+        { type: "break" },
+        { type: "button", id: "accept", text: "Accept", tooltip: "Accept current result", icon: "fa-solid fa-check", hidden: true, disable:true,},
+        { type: "button", id: "cancel", text: "Cancel", tooltip: "Cancel current outpainting/error", icon: "fa-solid fa-ban", hidden: true},
+        { type: "button", id: "retry", text: "Retry", tooltip: "Retry", icon: "fa-solid fa-rotate", hidden: true, disable:true,},
+        { type: "button", id: "prev", tooltip: "Prev Result", icon: "fa-solid fa-caret-left", hidden: true, disable:true,},
+        { type: "html", id: "current", hidden: true, disable:true,
+            async onRefresh(event) {
+                await event.complete
+                let fragment = query.html(`
+                <div class="w2ui-tb-text">
+                <div class="w2ui-tb-count">
+                    <span>${this.sel_value ?? "1/1"}</span>
+                </div> </div>`)
+                query(this.box).find("#tb_toolbar_item_current").append(fragment)
+            }
+        },
+        { type: "button", id: "next", tooltip: "Next Result", icon: "fa-solid fa-caret-right", hidden: true,disable:true,},
+        { type: "button", id: "add_image", text: "Add Image", icon: "fa-solid fa-file-circle-plus", hidden: true,disable:true,},
+        { type: "button", id: "delete_image", text: "Delete Image", icon: "fa-solid fa-trash-can", hidden: true,disable:true,},
+        { type: "button", id: "confirm", text: "Confirm", icon: "fa-solid fa-check", hidden: true,disable:true,},
+        { type: "button", id: "cancel_overlay", text: "Cancel", icon: "fa-solid fa-ban", hidden: true,disable:true,},
+        { type: "break" },
+        { type: "spacer" },
+        { type: "break" },
+        { type: "button", id: "eraser_size_btn", tooltip: "Eraser Size", text:"Size", icon: "fa-solid fa-eraser", hidden: true, count: 32},
+        { type: "html", id: "eraser_size", hidden: true,
+            async onRefresh(event) {
+                await event.complete
+                // let fragment = query.html(`
+                //     <input type="number" size="${this.eraser_size ? this.eraser_size.length:"2"}" style="margin: 0px 3px; padding: 4px;" min="8" max="${this.eraser_max ?? "256"}" value="${this.eraser_size ?? "32"}">
+                //     <input type="range" style="margin: 0px 3px; padding: 4px;" min="8" max="${this.eraser_max ?? "256"}" value="${this.eraser_size ?? "32"}">`)
+                let fragment = query.html(`
+                    <input type="range" style="margin: 0px 3px; padding: 4px;" min="8" max="${this.eraser_max ?? "256"}" value="${this.eraser_size ?? "32"}">
+                    `)
+                fragment.filter("input").on("change", event => {
+                    this.eraser_size = event.target.value;
+                    window.overlay.freeDrawingBrush.width=this.eraser_size;
+                    this.setCount("eraser_size_btn", event.target.value);
+                    window.postMessage(["eraser_size", event.target.value],"*")
+                    this.refresh();
+                })
+                query(this.box).find("#tb_toolbar_item_eraser_size").append(fragment)
+            }
+        },
+        // { type: "button", id: "resize_eraser", tooltip: "Resize Eraser", icon: "fa-solid fa-sliders" },
+        { type: "button", id: "resize_selection", text: "Resize Selection", tooltip: "Resize Selection", icon: "fa-solid fa-expand" },
+        { type: "break" },
+        { type: "html", id: "scale",
+            async onRefresh(event) {
+                await event.complete
+                let fragment = query.html(`
+                <div class="">
+                <div style="padding: 4px; border: 1px solid silver">
+                    <span>${this.scale_value ?? "100%"}</span>
+                </div></div>`)
+                query(this.box).find("#tb_toolbar_item_scale").append(fragment)
+            }
+        },
+        { type: "button", id: "zoom_in", tooltip: "Zoom In", icon: "fa-solid fa-magnifying-glass-plus" },
+        { type: "button", id: "zoom_out", tooltip: "Zoom Out", icon: "fa-solid fa-magnifying-glass-minus" },
+        { type: "break" },
+        { type: "button", id: "help", tooltip: "Help", icon: "fa-solid fa-circle-info" },
+        { type: "new-line"},
+        { type: "button", id: "setting", text: "Canvas Setting", tooltip: "Resize Canvas Here", icon: "fa-solid fa-sliders" },
+        { type: "break" },
+        check_button("enable_img2img","Enable Img2Img",false),
+        // check_button("use_correction","Photometric Correction",false),
+        check_button("resize_check","Resize Small Input",true),
+        check_button("enable_safety","Enable Safety Checker",true),
+        check_button("square_selection","Square Selection Only",false),
+        {type: "break"},
+        check_button("use_seed","Use Seed:",false),
+        { type: "html", id: "seed_val",
+            async onRefresh(event) {
+                await event.complete
+                let fragment = query.html(`
+                    <input type="number" style="margin: 0px 3px; padding: 4px; width:100px;" value="${this.config_obj.seed_val ?? "0"}">`)
+                fragment.filter("input").on("change", event => {
+                    this.config_obj.seed_val = event.target.value;
+                    parent.config_obj=this.config_obj;
+                    this.refresh();
+                })
+                query(this.box).find("#tb_toolbar_item_seed_val").append(fragment)
+            }
+        },
+        { type: "button", id: "random_seed", tooltip: "Set a random seed", icon: "fa-solid fa-dice" },
+    ],
+    onClick(event) {
+        switch(event.target){
+            case "setting":
+                open_setting();
+                break;
+            case "upload":
+                this.upload_mode=true
+                document.querySelector("#overlay_container").style.pointerEvents="auto";
+                this.click("canvas");
+                this.click("selection");
+                this.show("confirm","cancel_overlay","add_image","delete_image");
+                this.enable("confirm","cancel_overlay","add_image","delete_image");
+                this.disable(...upload_button_lst);
+                query("#upload_file").click();
+                if(this.upload_tip)
+                {
+                    this.upload_tip=false;
+                    w2utils.notify("Note that only visible images will be added to canvas",{timeout:10000,where:query("#container")})
+                }
+                break;
+            case "resize_selection":
+                this.resize_mode=true;
+                this.disable(...resize_button_lst);
+                this.enable("confirm","cancel_overlay");
+                this.show("confirm","cancel_overlay");
+                window.postMessage(["resize_selection",""],"*");
+                document.querySelector("#overlay_container").style.pointerEvents="auto";
+                break;
+            case "confirm":
+                if(this.upload_mode)
+                {
+                    export_image();
+                }
+                else
+                {
+                    let sel_box=this.selection_box;
+                    window.postMessage(["resize_selection",sel_box.x,sel_box.y,sel_box.width,sel_box.height],"*");
+                }
+            case "cancel_overlay":
+                end_overlay();
+                this.hide("confirm","cancel_overlay","add_image","delete_image");
+                if(this.upload_mode){
+                    this.enable(...upload_button_lst);
+                }
+                else
+                {
+                    this.enable(...resize_button_lst);
+                    window.postMessage(["resize_selection","",""],"*");
+                    if(event.target=="cancel_overlay")
+                    {
+                        this.selection_box=this.selection_box_bak;
+                    }
+                }
+                if(this.selection_box)
+                {
+                    this.setCount("resize_selection",`${Math.floor(this.selection_box.width/8)*8}x${Math.floor(this.selection_box.height/8)*8}`);
+                }
+                this.disable("confirm","cancel_overlay","add_image","delete_image");
+                this.upload_mode=false;
+                this.resize_mode=false;
+                this.click("selection");
+                break;
+            case "add_image":
+                query("#upload_file").click();
+                break;
+            case "delete_image":
+                let active_obj = window.overlay.getActiveObject();
+                if(active_obj)
+                {
+                    window.overlay.remove(active_obj);
+                    window.overlay.renderAll();
+                }
+                else
+                {
+                    w2utils.notify("You need to select an image first",{error:true,timeout:2000,where:query("#container")})
+                }
+                break;
+            case "load":
+                query("#upload_state").click();
+                this.selection_box=null;
+                this.setCount("resize_selection","");
+                break;
+            case "next":
+            case "prev":
+                window.postMessage(["outpaint", "", event.target], "*");
+                break;
+            case "outpaint":
+                this.click("selection");
+                this.disable(...outpaint_button_lst);
+                this.show(...outpaint_result_lst);
+                if(this.outpaint_tip)
+                {
+                    this.outpaint_tip=false;
+                    w2utils.notify("The canvas stays locked until you accept/cancel current outpainting",{timeout:10000,where:query("#container")})
+                }
+                document.querySelector("#container").style.pointerEvents="none";
+            case "retry":
+                this.disable(...outpaint_result_func_lst);
+                window.postMessage(["transfer",""],"*")
+                break;
+            case "accept":
+            case "cancel":
+                this.hide(...outpaint_result_lst);
+                this.disable(...outpaint_result_func_lst);
+                this.enable(...outpaint_button_lst);
+                document.querySelector("#container").style.pointerEvents="auto";
+                window.postMessage(["click", event.target],"*");
+                let app=parent.document.querySelector("gradio-app");
+                app=app.shadowRoot??app;
+                app.querySelector("#cancel").click();
+                break;
+            case "eraser":
+            case "selection":
+            case "canvas":
+                if(event.target=="eraser")
+                {
+                    this.show("eraser_size","eraser_size_btn");
+                    window.overlay.freeDrawingBrush.width=this.eraser_size;
+                    window.overlay.isDrawingMode = true;
+                }
+                else
+                {
+                    this.hide("eraser_size","eraser_size_btn");
+                    window.overlay.isDrawingMode = false;
+                }
+                if(this.upload_mode)
+                {
+                    if(event.target=="canvas")
+                    {
+                        window.postMessage(["mode", event.target],"*")
+                        document.querySelector("#overlay_container").style.pointerEvents="none";
+                        document.querySelector("#overlay_container").style.opacity = 0.5;
+                    }
+                    else
+                    {
+                        document.querySelector("#overlay_container").style.pointerEvents="auto";
+                        document.querySelector("#overlay_container").style.opacity = 1.0;
+                    }
+                }
+                else
+                {
+                    window.postMessage(["mode", event.target],"*")
+                }
+                break;
+            case "help":
+                w2popup.open({
+                    title: "Document",
+                    body: "Usage: <a href='https://github.com/lkwq007/stablediffusion-infinity/blob/master/docs/usage.md'  target='_blank'>https://github.com/lkwq007/stablediffusion-infinity/blob/master/docs/usage.md</a>"
+                })
+                break;
+            case "clear":
+                w2confirm("Reset canvas?").yes(() => {
+                    window.postMessage(["click", event.target],"*");
+                }).no(() => {})
+                break;
+            case "random_seed":
+                this.config_obj.seed_val=Math.floor(Math.random() * 3000000000);
+                parent.config_obj=this.config_obj;
+                this.refresh();
+                break;
+            case "enable_img2img":
+            case "use_correction":
+            case "resize_check":
+            case "enable_safety":
+            case "use_seed":
+            case "square_selection":
+                let target=this.get(event.target);
+                target.icon=target.checked?"fa-regular fa-square":"fa-solid fa-square-check";
+                this.config_obj[event.target]=!target.checked;
+                parent.config_obj=this.config_obj;
+                this.refresh();
+                break;
+            case "save":
+            case "export":
+                ask_filename(event.target);
+                break;
+            default:
+                // clear, save, export, outpaint, retry
+                // break, save, export, accept, retry, outpaint
+                window.postMessage(["click", event.target],"*")
+        }
+        console.log("Target: "+ event.target, event)
+    }
+})
+window.w2ui=w2ui;
+w2ui.toolbar.config_obj={
+    resize_check: true,
+    enable_safety: true,
+    use_correction: false,
+    enable_img2img: false,
+    use_seed: false,
+    seed_val: 0,
+    square_selection: false,
+};
+w2ui.toolbar.outpaint_tip=true;
+w2ui.toolbar.upload_tip=true;
+window.update_count=function(cur,total){
+  w2ui.toolbar.sel_value=`${cur}/${total}`;
+  w2ui.toolbar.refresh();
+}
+window.update_eraser=function(val,max_val){
+  w2ui.toolbar.eraser_size=`${val}`;
+  w2ui.toolbar.eraser_max=`${max_val}`;
+  w2ui.toolbar.setCount("eraser_size_btn", `${val}`);
+  w2ui.toolbar.refresh();
+}
+window.update_scale=function(val){
+  w2ui.toolbar.scale_value=`${val}`;
+  w2ui.toolbar.refresh();
+}
+window.enable_result_lst=function(){
+  w2ui.toolbar.enable(...outpaint_result_lst);
+}
+function onObjectScaled(e)
+{
+    let object = e.target;
+    if(object.isType("rect"))
+    {
+        let width=object.getScaledWidth();
+        let height=object.getScaledHeight();
+        object.scale(1);
+        width=Math.max(Math.min(width,window.overlay.width-object.left),256);
+        height=Math.max(Math.min(height,window.overlay.height-object.top),256);
+        let l=Math.max(Math.min(object.left,window.overlay.width-width-object.strokeWidth),0);
+        let t=Math.max(Math.min(object.top,window.overlay.height-height-object.strokeWidth),0);
+        if(window.w2ui.toolbar.config_obj.square_selection)
+        {
+            let max_val = Math.min(Math.max(width,height),window.overlay.width,window.overlay.height);
+            width=max_val;
+            height=max_val;
+        }
+        object.set({ width: width, height: height, left:l,top:t})
+        window.w2ui.toolbar.selection_box={width: width, height: height, x:object.left, y:object.top};
+        window.w2ui.toolbar.setCount("resize_selection",`${Math.floor(width/8)*8}x${Math.floor(height/8)*8}`);
+        window.w2ui.toolbar.refresh();
+    }
+}
+function onObjectMoved(e)
+{
+    let object = e.target;
+    if(object.isType("rect"))
+    {   
+        let l=Math.max(Math.min(object.left,window.overlay.width-object.width-object.strokeWidth),0);
+        let t=Math.max(Math.min(object.top,window.overlay.height-object.height-object.strokeWidth),0);
+        object.set({left:l,top:t});
+        window.w2ui.toolbar.selection_box={width: object.width, height: object.height, x:object.left, y:object.top};
+    }
+}
+window.setup_overlay=function(width,height)
+{
+    if(window.overlay)
+    {
+        window.overlay.setDimensions({width:width,height:height});
+        let app=parent.document.querySelector("gradio-app");
+        app=app.shadowRoot??app;
+        app.querySelector("#sdinfframe").style.height=80+Number(height)+"px";
+        document.querySelector("#container").style.height= height+"px";
+        document.querySelector("#container").style.width = width+"px";
+    }
+    else
+    {
+        canvas=new fabric.Canvas("overlay_canvas");
+        canvas.setDimensions({width:width,height:height});
+        let app=parent.document.querySelector("gradio-app");
+        app=app.shadowRoot??app;
+        app.querySelector("#sdinfframe").style.height=80+Number(height)+"px";
+        canvas.freeDrawingBrush = new fabric.EraserBrush(canvas);
+        canvas.on("object:scaling", onObjectScaled);
+        canvas.on("object:moving", onObjectMoved);
+        window.overlay=canvas;
+    }
+    document.querySelector("#overlay_container").style.pointerEvents="none";
+}
+window.update_overlay=function(width,height)
+{
+    window.overlay.setDimensions({width:width,height:height},{backstoreOnly:true});
+    // document.querySelector("#overlay_container").style.pointerEvents="none";
+}
+window.adjust_selection=function(x,y,width,height)
+{
+    var rect = new fabric.Rect({
+        left: x,
+        top: y,
+        fill: "rgba(0,0,0,0)",
+        strokeWidth: 3, 
+        stroke: "rgba(0,0,0,0.7)",
+        cornerColor: "red",
+        cornerStrokeColor: "red",
+        borderColor: "rgba(255, 0, 0, 1.0)",
+        width: width,
+        height: height,
+        lockRotation: true,
+    });
+    rect.setControlsVisibility({ mtr: false });
+    window.overlay.add(rect);
+    window.overlay.setActiveObject(window.overlay.item(0));
+    window.w2ui.toolbar.selection_box={width: width, height: height, x:x, y:y};
+    window.w2ui.toolbar.selection_box_bak={width: width, height: height, x:x, y:y};
+}
+function add_image(url)
+{
+    fabric.Image.fromURL(url,function(img){
+        window.overlay.add(img);
+        window.overlay.setActiveObject(img);
+    },{left:100,top:100});
+}
+function export_image()
+{
+    data=window.overlay.toDataURL();
+    document.querySelector("#upload_content").value=data;
+    window.postMessage(["upload",""],"*");
+    end_overlay();
+}
+function end_overlay()
+{
+    window.overlay.clear();
+    document.querySelector("#overlay_container").style.opacity = 1.0;
+    document.querySelector("#overlay_container").style.pointerEvents="none";
+}
+function ask_filename(target)
+{
+    w2prompt({
+        label: "Enter filename",
+        value: `outpaint_${((new Date(Date.now() -(new Date()).getTimezoneOffset() * 60000))).toISOString().replace("T","_").replace(/[^0-9_]/g, "").substring(0,15)}`,
+    })
+    .change((event) => {
+        console.log("change", event.detail.originalEvent.target.value);
+    })
+    .ok((event) => {
+        console.log("value=", event.detail.value);
+        window.postMessage(["click",target,event.detail.value],"*");
+    })
+    .cancel((event) => {
+        console.log("cancel");
+    });
+}
+
+document.querySelector("#container").addEventListener("wheel",(e)=>{e.preventDefault()})
+window.setup_shortcut=function(json)
+{
+    var config=JSON.parse(json);
+    var key_map={};
+    Object.keys(config.shortcut).forEach(k=>{
+        key_map[config.shortcut[k]]=k;
+    })
+    document.addEventListener("keydown",(e)=>{
+        if(e.target.tagName!="INPUT")
+        {
+            let key=e.key;
+            if(e.ctrlKey)
+            {
+                key="Ctrl+"+e.key;
+                if(key in key_map)
+                {
+                    e.preventDefault();
+                }
+            }
+            if(key in key_map)
+            {
+                w2ui.toolbar.click(key_map[key]);
+            }
+        }
+    })
+}

js/upload.js

@@ -1,20 +1,19 @@
-function(a,b){
-    if(!window.my_observe_upload)
-    {
-        console.log("setup upload here");
-        window.my_observe_upload = new MutationObserver(function (event) {
-            console.log(event);
-            var frame=document.querySelector("gradio-app").querySelector("#sdinfframe").contentWindow;
-            var str=document.querySelector("gradio-app").querySelector("#upload textarea").value;
-            frame.postMessage(["upload", str], "*");
-        });
-        window.my_observe_upload_target = document.querySelector("gradio-app").querySelector("#upload span");
-        window.my_observe_upload.observe(window.my_observe_upload_target, {
-            attributes: false, 
-            subtree: true,
-            childList: true, 
-            characterData: true
-        });
-    }
-    return [a,b];
+function(a,b){
+    if(!window.my_observe_upload)
+    {
+        console.log("setup upload here");
+        window.my_observe_upload = new MutationObserver(function (event) {
+            console.log(event);
+            var frame=document.querySelector("gradio-app").shadowRoot.querySelector("#sdinfframe").contentWindow.document;
+            frame.querySelector("#upload").click();
+        });
+        window.my_observe_upload_target = document.querySelector("gradio-app").shadowRoot.querySelector("#upload span");
+        window.my_observe_upload.observe(window.my_observe_upload_target, {
+            attributes: false, 
+            subtree: true,
+            childList: true, 
+            characterData: true
+        });
+    }
+    return [a,b];
 }

js/xss.js

@@ -0,0 +1,31 @@
+var setup_outpaint=function(){
+    if(!window.my_observe_outpaint)
+    {
+        console.log("setup outpaint here");
+        window.my_observe_outpaint = new MutationObserver(function (event) {
+            console.log(event);
+            let app=document.querySelector("gradio-app");
+            app=app.shadowRoot??app;
+            let frame=app.querySelector("#sdinfframe").contentWindow;
+            frame.postMessage(["outpaint", ""], "*");
+        });
+        var app=document.querySelector("gradio-app");
+        app=app.shadowRoot??app;
+        window.my_observe_outpaint_target=app.querySelector("#output span");
+        window.my_observe_outpaint.observe(window.my_observe_outpaint_target, {
+            attributes: false, 
+            subtree: true,
+            childList: true, 
+            characterData: true
+        });
+    }
+}; 
+window.config_obj={
+    resize_check: true,
+    enable_safety: true,
+    use_correction: false,
+    enable_img2img: false,
+    use_seed: false,
+    seed_val: 0,
+};
+setup_outpaint();

models/v1-inference.yaml

@@ -0,0 +1,70 @@
+model:
+  base_learning_rate: 1.0e-04
+  target: ldm.models.diffusion.ddpm.LatentDiffusion
+  params:
+    linear_start: 0.00085
+    linear_end: 0.0120
+    num_timesteps_cond: 1
+    log_every_t: 200
+    timesteps: 1000
+    first_stage_key: "jpg"
+    cond_stage_key: "txt"
+    image_size: 64
+    channels: 4
+    cond_stage_trainable: false   # Note: different from the one we trained before
+    conditioning_key: crossattn
+    monitor: val/loss_simple_ema
+    scale_factor: 0.18215
+    use_ema: False
+
+    scheduler_config: # 10000 warmup steps
+      target: ldm.lr_scheduler.LambdaLinearScheduler
+      params:
+        warm_up_steps: [ 10000 ]
+        cycle_lengths: [ 10000000000000 ] # incredibly large number to prevent corner cases
+        f_start: [ 1.e-6 ]
+        f_max: [ 1. ]
+        f_min: [ 1. ]
+
+    unet_config:
+      target: ldm.modules.diffusionmodules.openaimodel.UNetModel
+      params:
+        image_size: 32 # unused
+        in_channels: 4
+        out_channels: 4
+        model_channels: 320
+        attention_resolutions: [ 4, 2, 1 ]
+        num_res_blocks: 2
+        channel_mult: [ 1, 2, 4, 4 ]
+        num_heads: 8
+        use_spatial_transformer: True
+        transformer_depth: 1
+        context_dim: 768
+        use_checkpoint: True
+        legacy: False
+
+    first_stage_config:
+      target: ldm.models.autoencoder.AutoencoderKL
+      params:
+        embed_dim: 4
+        monitor: val/rec_loss
+        ddconfig:
+          double_z: true
+          z_channels: 4
+          resolution: 256
+          in_channels: 3
+          out_ch: 3
+          ch: 128
+          ch_mult:
+          - 1
+          - 2
+          - 4
+          - 4
+          num_res_blocks: 2
+          attn_resolutions: []
+          dropout: 0.0
+        lossconfig:
+          target: torch.nn.Identity
+
+    cond_stage_config:
+      target: ldm.modules.encoders.modules.FrozenCLIPEmbedder

models/v1-inpainting-inference.yaml

@@ -0,0 +1,70 @@
+model:
+  base_learning_rate: 7.5e-05
+  target: ldm.models.diffusion.ddpm.LatentInpaintDiffusion
+  params:
+    linear_start: 0.00085
+    linear_end: 0.0120
+    num_timesteps_cond: 1
+    log_every_t: 200
+    timesteps: 1000
+    first_stage_key: "jpg"
+    cond_stage_key: "txt"
+    image_size: 64
+    channels: 4
+    cond_stage_trainable: false   # Note: different from the one we trained before
+    conditioning_key: hybrid   # important
+    monitor: val/loss_simple_ema
+    scale_factor: 0.18215
+    finetune_keys: null
+
+    scheduler_config: # 10000 warmup steps
+      target: ldm.lr_scheduler.LambdaLinearScheduler
+      params:
+        warm_up_steps: [ 2500 ] # NOTE for resuming. use 10000 if starting from scratch
+        cycle_lengths: [ 10000000000000 ] # incredibly large number to prevent corner cases
+        f_start: [ 1.e-6 ]
+        f_max: [ 1. ]
+        f_min: [ 1. ]
+
+    unet_config:
+      target: ldm.modules.diffusionmodules.openaimodel.UNetModel
+      params:
+        image_size: 32 # unused
+        in_channels: 9  # 4 data + 4 downscaled image + 1 mask
+        out_channels: 4
+        model_channels: 320
+        attention_resolutions: [ 4, 2, 1 ]
+        num_res_blocks: 2
+        channel_mult: [ 1, 2, 4, 4 ]
+        num_heads: 8
+        use_spatial_transformer: True
+        transformer_depth: 1
+        context_dim: 768
+        use_checkpoint: True
+        legacy: False
+
+    first_stage_config:
+      target: ldm.models.autoencoder.AutoencoderKL
+      params:
+        embed_dim: 4
+        monitor: val/rec_loss
+        ddconfig:
+          double_z: true
+          z_channels: 4
+          resolution: 256
+          in_channels: 3
+          out_ch: 3
+          ch: 128
+          ch_mult:
+          - 1
+          - 2
+          - 4
+          - 4
+          num_res_blocks: 2
+          attn_resolutions: []
+          dropout: 0.0
+        lossconfig:
+          target: torch.nn.Identity
+
+    cond_stage_config:
+      target: ldm.modules.encoders.modules.FrozenCLIPEmbedder