nnf.cpp

#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);
}

unsigned long NearestNeighborField::minimize(int nr_pass, bool is_source2target, bool conditional_skip) {
    const auto& this_size = source_size();
    unsigned long total_distance = 0;
    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;

                bool can_skip = false;
                if (conditional_skip)
                    if (is_source2target)
                        can_skip = !m_source.is_masked(i, j);
                    else
                        can_skip = m_source.is_masked(i, j);
                    if (can_skip)
                        continue;
                     if (at(i, j, 2) > 0){
                        _minimize_link(i, j, +1);
                        auto this_ptr = mutable_ptr(i, j);
                        total_distance += this_ptr[2];
                    }
            }
        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;

                bool can_skip = false;
                if (conditional_skip)
                    if (is_source2target)
                        can_skip = !m_source.is_masked(i, j);
                    else
                        can_skip = m_source.is_masked(i, j);
                if (can_skip)
                    continue;


                if (at(i, j, 2) > 0) {
                    _minimize_link(i, j, -1);
                    auto this_ptr = mutable_ptr(i, j);
                    total_distance += this_ptr[2];
                }
            }
    }
    return total_distance;
}

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) {
                unsigned int xxs = xs + dx, xxt = xt + dx;
                wsum += 1;
                //std::cerr << xxt << std::endl;
                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));
}