lnyan/stablediffusion-infinity

1#include <algorithm>2#include <iostream>3#include <opencv2/imgcodecs.hpp>4#include <opencv2/imgproc.hpp>5#include <opencv2/highgui.hpp>6 7#include "inpaint.h"8 9namespace {10    static std::vector<double> kDistance2Similarity;11 12    void init_kDistance2Similarity() {13        double base[11] = {1.0, 0.99, 0.96, 0.83, 0.38, 0.11, 0.02, 0.005, 0.0006, 0.0001, 0};14        int length = (PatchDistanceMetric::kDistanceScale + 1);15        kDistance2Similarity.resize(length);16        for (int i = 0; i < length; ++i) {17            double t = (double) i / length;18            int j = (int) (100 * t);19            int k = j + 1;20            double vj = (j < 11) ? base[j] : 0;21            double vk = (k < 11) ? base[k] : 0;22            kDistance2Similarity[i] = vj + (100 * t - j) * (vk - vj);23        }24    }25 26 27    inline void _weighted_copy(const MaskedImage &source, int ys, int xs, cv::Mat &target, int yt, int xt, double weight) {28        if (source.is_masked(ys, xs)) return;29        if (source.is_globally_masked(ys, xs)) return;30 31        auto source_ptr = source.get_image(ys, xs);32        auto target_ptr = target.ptr<double>(yt, xt);33 34#pragma unroll35        for (int c = 0; c < 3; ++c)36            target_ptr[c] += static_cast<double>(source_ptr[c]) * weight;37        target_ptr[3] += weight;38    }39}40 41/**42 * This algorithme uses a version proposed by Xavier Philippeau.43 */44 45Inpainting::Inpainting(cv::Mat image, cv::Mat mask, const PatchDistanceMetric *metric)46    : m_initial(image, mask), m_distance_metric(metric), m_pyramid(), m_source2target(), m_target2source() {47    _initialize_pyramid();48}49 50Inpainting::Inpainting(cv::Mat image, cv::Mat mask, cv::Mat global_mask, const PatchDistanceMetric *metric)51    : m_initial(image, mask, global_mask), m_distance_metric(metric), m_pyramid(), m_source2target(), m_target2source() {52    _initialize_pyramid();53}54 55void Inpainting::_initialize_pyramid() {56    auto source = m_initial;57    m_pyramid.push_back(source);58    while (source.size().height > m_distance_metric->patch_size() && source.size().width > m_distance_metric->patch_size()) {59        source = source.downsample();60        m_pyramid.push_back(source);61    }62 63    if (kDistance2Similarity.size() == 0) {64        init_kDistance2Similarity();65    }66}67 68cv::Mat Inpainting::run(bool verbose, bool verbose_visualize, unsigned int random_seed) {69    srand(random_seed);70    const int nr_levels = m_pyramid.size();71 72    MaskedImage source, target;73    for (int level = nr_levels - 1; level >= 0; --level) {74        if (verbose) std::cerr << "Inpainting level: " << level << std::endl;75 76        source = m_pyramid[level];77 78        if (level == nr_levels - 1) {79            target = source.clone();80            target.clear_mask();81            m_source2target = NearestNeighborField(source, target, m_distance_metric);82            m_target2source = NearestNeighborField(target, source, m_distance_metric);83        } else {84            m_source2target = NearestNeighborField(source, target, m_distance_metric, m_source2target);85            m_target2source = NearestNeighborField(target, source, m_distance_metric, m_target2source);86        }87 88        if (verbose) std::cerr << "Initialization done." << std::endl;89 90        if (verbose_visualize) {91            auto visualize_size = m_initial.size();92            cv::Mat source_visualize(visualize_size, m_initial.image().type());93            cv::resize(source.image(), source_visualize, visualize_size);94            cv::imshow("Source", source_visualize);95            cv::Mat target_visualize(visualize_size, m_initial.image().type());96            cv::resize(target.image(), target_visualize, visualize_size);97            cv::imshow("Target", target_visualize);98            cv::waitKey(0);99        }100 101        target = _expectation_maximization(source, target, level, verbose);102    }103 104    return target.image();105}106 107// EM-Like algorithm (see "PatchMatch" - page 6).108// Returns a double sized target image (unless level = 0).109MaskedImage Inpainting::_expectation_maximization(MaskedImage source, MaskedImage target, int level, bool verbose) {110    const int nr_iters_em = 1 + 2 * level;111    const int nr_iters_nnf = static_cast<int>(std::min(7, 1 + level));112    const int patch_size = m_distance_metric->patch_size();113 114    MaskedImage new_source, new_target;115 116    for (int iter_em = 0; iter_em < nr_iters_em; ++iter_em) {117        if (iter_em != 0) {118            m_source2target.set_target(new_target);119            m_target2source.set_source(new_target);120            target = new_target;121        }122 123        if (verbose) std::cerr << "EM Iteration: " << iter_em << std::endl;124 125        auto size = source.size();126        for (int i = 0; i < size.height; ++i) {127            for (int j = 0; j < size.width; ++j) {128                if (!source.contains_mask(i, j, patch_size)) {129                    m_source2target.set_identity(i, j);130                    m_target2source.set_identity(i, j);131                }132            }133        }134        if (verbose) std::cerr << "  NNF minimization started." << std::endl;135        m_source2target.minimize(nr_iters_nnf);136        m_target2source.minimize(nr_iters_nnf);137        if (verbose) std::cerr << "  NNF minimization finished." << std::endl;138 139        // Instead of upsizing the final target, we build the last target from the next level source image.140        // Thus, the final target is less blurry (see "Space-Time Video Completion" - page 5).141        bool upscaled = false;142        if (level >= 1 && iter_em == nr_iters_em - 1) {143            new_source = m_pyramid[level - 1];144            new_target = target.upsample(new_source.size().width, new_source.size().height, m_pyramid[level - 1].global_mask());145            upscaled = true;146        } else {147            new_source = m_pyramid[level];148            new_target = target.clone();149        }150 151        auto vote = cv::Mat(new_target.size(), CV_64FC4);152        vote.setTo(cv::Scalar::all(0));153 154        // Votes for best patch from NNF Source->Target (completeness) and Target->Source (coherence).155        _expectation_step(m_source2target, 1, vote, new_source, upscaled);156        if (verbose) std::cerr << "  Expectation source to target finished." << std::endl;157        _expectation_step(m_target2source, 0, vote, new_source, upscaled);158        if (verbose) std::cerr << "  Expectation target to source finished." << std::endl;159 160        // Compile votes and update pixel values.161        _maximization_step(new_target, vote);162        if (verbose) std::cerr << "  Minimization step finished." << std::endl;163    }164 165    return new_target;166}167 168// Expectation step: vote for best estimations of each pixel.169void Inpainting::_expectation_step(170    const NearestNeighborField &nnf, bool source2target,171    cv::Mat &vote, const MaskedImage &source, bool upscaled172) {173    auto source_size = nnf.source_size();174    auto target_size = nnf.target_size();175    const int patch_size = m_distance_metric->patch_size();176 177    for (int i = 0; i < source_size.height; ++i) {178        for (int j = 0; j < source_size.width; ++j) {179            if (nnf.source().is_globally_masked(i, j)) continue;180            int yp = nnf.at(i, j, 0), xp = nnf.at(i, j, 1), dp = nnf.at(i, j, 2);181            double w = kDistance2Similarity[dp];182 183            for (int di = -patch_size; di <= patch_size; ++di) {184                for (int dj = -patch_size; dj <= patch_size; ++dj) {185                    int ys = i + di, xs = j + dj, yt = yp + di, xt = xp + dj;186                    if (!(ys >= 0 && ys < source_size.height && xs >= 0 && xs < source_size.width)) continue;187                    if (nnf.source().is_globally_masked(ys, xs)) continue;188                    if (!(yt >= 0 && yt < target_size.height && xt >= 0 && xt < target_size.width)) continue;189                    if (nnf.target().is_globally_masked(yt, xt)) continue;190 191                    if (!source2target) {192                        std::swap(ys, yt);193                        std::swap(xs, xt);194                    }195 196                    if (upscaled) {197                        for (int uy = 0; uy < 2; ++uy) {198                            for (int ux = 0; ux < 2; ++ux) {199                                _weighted_copy(source, 2 * ys + uy, 2 * xs + ux, vote, 2 * yt + uy, 2 * xt + ux, w);200                            }201                        }202                    } else {203                        _weighted_copy(source, ys, xs, vote, yt, xt, w);204                    }205                }206            }207        }208    }209}210 211// Maximization Step: maximum likelihood of target pixel.212void Inpainting::_maximization_step(MaskedImage &target, const cv::Mat &vote) {213    auto target_size = target.size();214    for (int i = 0; i < target_size.height; ++i) {215        for (int j = 0; j < target_size.width; ++j) {216            const double *source_ptr = vote.ptr<double>(i, j);217            unsigned char *target_ptr = target.get_mutable_image(i, j);218 219            if (target.is_globally_masked(i, j)) {220                continue;221            }222 223            if (source_ptr[3] > 0) {224                unsigned char r = cv::saturate_cast<unsigned char>(source_ptr[0] / source_ptr[3]);225                unsigned char g = cv::saturate_cast<unsigned char>(source_ptr[1] / source_ptr[3]);226                unsigned char b = cv::saturate_cast<unsigned char>(source_ptr[2] / source_ptr[3]);227                target_ptr[0] = r, target_ptr[1] = g, target_ptr[2] = b;228            } else {229                target.set_mask(i, j, 0);230            }231        }232    }233}234 235