cv::Mat_ generalizedJointBilateralFilter()

in source/depth_estimation/TemporalBilateralFilter.h [39:123]
76 lines of code
10 McCabe index (conditional complexity)

cv::Mat_<TPixel> generalizedJointBilateralFilter(
    const cv::Mat_<TPixel>& image, // Either float, Vec2f or Vec3f
    const cv::Mat_<TGuide>& guide,
    const cv::Mat_<TGuide>& neighborTGuide,
    const cv::Mat_<bool>& mask,
    const int radius,
    const float sigma,
    const float weight0 = 1.0f,
    const float weight1 = 1.0f,
    const float weight2 = 1.0f,
    const int numThreads = -1) {
  CHECK_EQ(guide.size(), neighborTGuide.size());
  CHECK_EQ(image.size(), guide.size());
  CHECK_EQ(guide.size(), mask.size());

  const TPixel zero = 0.0;

  cv::Mat_<TPixel> dest(image.size());
  ThreadPool threadPool(numThreads);
  const int edgeX = image.cols;
  const int edgeY = 1;
  for (int yBegin = 0; yBegin < image.rows; yBegin += edgeY) {
    for (int xBegin = 0; xBegin < image.cols; xBegin += edgeX) {
      const int xEnd = std::min(xBegin + edgeX, image.cols);
      const int yEnd = std::min(yBegin + edgeY, image.rows);
      threadPool.spawn([&, xBegin, yBegin, xEnd, yEnd] {
        for (int y = yBegin; y < yEnd; ++y) {
          for (int x = xBegin; x < xEnd; ++x) {
            if (!mask(y, x)) {
              dest(y, x) = image(y, x);
              continue;
            }

            const auto guideColor = guide(y, x);
            float sumWeight = 0.0f;
            TPixel weightedAvg = zero;

            const float guideFactor = 1 / cv_util::maxPixelValue(guide);
            const float neighborTGuideFactor = 1 / cv_util::maxPixelValue(neighborTGuide);

            for (int v = -radius; v <= radius; ++v) {
              for (int u = -radius; u <= radius; ++u) {
                const int sampleX = math_util::clamp(x + u, 0, image.cols - 1);
                const int sampleY = math_util::clamp(y + v, 0, image.rows - 1);

                if (!mask(sampleY, sampleX)) {
                  continue;
                }

                const TGuide& neighborTGuideColor = neighborTGuide(sampleY, sampleX);

                const float colorDiffSq = // BGR
                    weight0 *
                        math_util::square(
                            (guideColor[0] * guideFactor) -
                            (neighborTGuideColor[0] * neighborTGuideFactor)) +
                    weight1 *
                        math_util::square(
                            (guideColor[1] * guideFactor) -
                            (neighborTGuideColor[1] * neighborTGuideFactor)) +
                    weight2 *
                        math_util::square(
                            (guideColor[2] * guideFactor) -
                            (neighborTGuideColor[2] * neighborTGuideFactor));
                const float weight =
                    expf((-colorDiffSq / 3.0f) / (2.0f * math_util::square(sigma)));

                sumWeight += weight;
                weightedAvg += weight * image(sampleY, sampleX);
              }
            }
            if (sumWeight != 0.0f) {
              weightedAvg /= sumWeight;
              dest(y, x) = weightedAvg;
            } else {
              dest(y, x) = image(y, x);
            }
          }
        }
      });
    }
  }
  threadPool.join();
  return dest;
}