lnyan/stablediffusion-infinity

1#include "masked_image.h"2#include <algorithm>3#include <iostream>4 5const cv::Size MaskedImage::kDownsampleKernelSize = cv::Size(6, 6);6const int MaskedImage::kDownsampleKernel[6] = {1, 5, 10, 10, 5, 1};7 8bool MaskedImage::contains_mask(int y, int x, int patch_size) const {9    auto mask_size = size();10    for (int dy = -patch_size; dy <= patch_size; ++dy) {11        for (int dx = -patch_size; dx <= patch_size; ++dx) {12            int yy = y + dy, xx = x + dx;13            if (yy >= 0 && yy < mask_size.height && xx >= 0 && xx < mask_size.width) {14                if (is_masked(yy, xx) && !is_globally_masked(yy, xx)) return true;15            }16        }17    }18    return false;19}20 21MaskedImage MaskedImage::downsample() const {22    const auto &kernel_size = MaskedImage::kDownsampleKernelSize;23    const auto &kernel = MaskedImage::kDownsampleKernel;24 25    const auto size = this->size();26    const auto new_size = cv::Size(size.width / 2, size.height / 2);27 28    auto ret = MaskedImage(new_size.width, new_size.height);29    if (!m_global_mask.empty()) ret.init_global_mask_mat();30    for (int y = 0; y < size.height - 1; y += 2) {31        for (int x = 0; x < size.width - 1; x += 2) {32            int r = 0, g = 0, b = 0, ksum = 0;33            bool is_gmasked = true;34 35            for (int dy = -kernel_size.height / 2 + 1; dy <= kernel_size.height / 2; ++dy) {36                for (int dx = -kernel_size.width / 2 + 1; dx <= kernel_size.width / 2; ++dx) {37                    int yy = y + dy, xx = x + dx;38                    if (yy >= 0 && yy < size.height && xx >= 0 && xx < size.width) {39                        if (!is_globally_masked(yy, xx)) {40                            is_gmasked = false;41                        }42                        if (!is_masked(yy, xx)) {43                            auto source_ptr = get_image(yy, xx);44                            int k = kernel[kernel_size.height / 2 - 1 + dy] * kernel[kernel_size.width / 2 - 1 + dx];45                            r += source_ptr[0] * k, g += source_ptr[1] * k, b += source_ptr[2] * k;46                            ksum += k;47                        }48                    }49                }50            }51 52            if (ksum > 0) r /= ksum, g /= ksum, b /= ksum;53 54            if (!m_global_mask.empty()) {55                ret.set_global_mask(y / 2, x / 2, is_gmasked);56            }57            if (ksum > 0) {58                auto target_ptr = ret.get_mutable_image(y / 2, x / 2);59                target_ptr[0] = r, target_ptr[1] = g, target_ptr[2] = b;60                ret.set_mask(y / 2, x / 2, 0);61            } else {62                ret.set_mask(y / 2, x / 2, 1);63            }64        }65    }66 67    return ret;68}69 70MaskedImage MaskedImage::upsample(int new_w, int new_h) const {71    const auto size = this->size();72    auto ret = MaskedImage(new_w, new_h);73    if (!m_global_mask.empty()) ret.init_global_mask_mat();74    for (int y = 0; y < new_h; ++y) {75        for (int x = 0; x < new_w; ++x) {76            int yy = y * size.height / new_h;77            int xx = x * size.width / new_w;78 79            if (is_globally_masked(yy, xx)) {80                ret.set_global_mask(y, x, 1);81                ret.set_mask(y, x, 1);82            } else {83                if (!m_global_mask.empty()) ret.set_global_mask(y, x, 0);84 85                if (is_masked(yy, xx)) {86                    ret.set_mask(y, x, 1);87                } else {88                    auto source_ptr = get_image(yy, xx);89                    auto target_ptr = ret.get_mutable_image(y, x);90                    for (int c = 0; c < 3; ++c)91                        target_ptr[c] = source_ptr[c];92                    ret.set_mask(y, x, 0);93                }94            }95        }96    }97 98    return ret;99}100 101MaskedImage MaskedImage::upsample(int new_w, int new_h, const cv::Mat &new_global_mask) const {102    auto ret = upsample(new_w, new_h);103    ret.set_global_mask_mat(new_global_mask);104    return ret;105}106 107void MaskedImage::compute_image_gradients() {108    if (m_image_grad_computed) {109        return;110    }111 112    const auto size = m_image.size();113    m_image_grady = cv::Mat(size, CV_8UC3);114    m_image_gradx = cv::Mat(size, CV_8UC3);115    m_image_grady = cv::Scalar::all(0);116    m_image_gradx = cv::Scalar::all(0);117 118    for (int i = 1; i < size.height - 1; ++i) {119        const auto *ptr = m_image.ptr<unsigned char>(i, 0);120        const auto *ptry1 = m_image.ptr<unsigned char>(i + 1, 0);121        const auto *ptry2 = m_image.ptr<unsigned char>(i - 1, 0);122        const auto *ptrx1 = m_image.ptr<unsigned char>(i, 0) + 3;123        const auto *ptrx2 = m_image.ptr<unsigned char>(i, 0) - 3;124        auto *mptry = m_image_grady.ptr<unsigned char>(i, 0);125        auto *mptrx = m_image_gradx.ptr<unsigned char>(i, 0);126        for (int j = 3; j < size.width * 3 - 3; ++j) {127            mptry[j] = (ptry1[j] / 2 - ptry2[j] / 2) + 128;128            mptrx[j] = (ptrx1[j] / 2 - ptrx2[j] / 2) + 128;129        }130    }131 132    m_image_grad_computed = true;133}134 135void MaskedImage::compute_image_gradients() const {136    const_cast<MaskedImage *>(this)->compute_image_gradients();137}138 139