void DepthVideoProcessor::flowGuidedFilter()

in lib/Processor.cpp [315:590]


void DepthVideoProcessor::flowGuidedFilter(const Params& params) {
  LOG(INFO) << "Applying flow guided filter...";

  if (!params.frameRange.isConsecutive()) {
    throw std::runtime_error("Frame range must be consecutive.");
  }

  const ColorStream& cs = video_->colorStream("down");
  const int w = cs.width();
  const int h = cs.height();

  std::vector<std::pair<int, int>> flowPairs;
  std::string flowPath = fmt::format("{:s}/flow", video_->path());
  for (auto& entry :
      boost::make_iterator_range(fs::directory_iterator(flowPath), {})) {
    std::string fileName = entry.path().stem().string();
    if (fileName.size() != 18 ||
        fileName.substr(0, 5) != "flow_") {
      continue;
    }
    int f0 = std::stoi(fileName.substr(5, 6));
    int f1 = std::stoi(fileName.substr(12, 6));
    flowPairs.emplace_back(f0, f1);
  }

  auto loadFlow = [&](const int frame0, const int frame1)
      -> std::unique_ptr<Mat2f> {
    std::string flowFile = fmt::format(
        "{:s}/flow/flow_{:06d}_{:06d}.raw",
        video_->path(), frame0, frame1);
    if (!fs::exists(flowFile)) {
      return nullptr;
    }
    std::unique_ptr<Mat2f> flow = std::make_unique<Mat2f>();
    freadim(flowFile, *flow);
    if (flow->cols != w || flow->rows != h) {
      return nullptr;
    }

    return flow;
  };

  auto loadFlowMask = [&](const int frame0, const int frame1)
      -> std::unique_ptr<Mat1b> {
    std::string maskFile = fmt::format(
         "{:s}/flow_mask/mask_{:06d}_{:06d}.png",
         video_->path(), frame0, frame1);
    if (!fs::exists(maskFile)) {
      return nullptr;
    }
    std::unique_ptr<Mat1b> mask = std::make_unique<Mat1b>();
    *mask = imread(maskFile, IMREAD_GRAYSCALE);
    if (mask->cols != w || mask->rows != h) {
      return nullptr;
    }

    return mask;
  };

  const DepthStream& srcDs = video_->depthStream(params.sourceDepthStream);
  DepthStream& dstDs = video_->depthStream(params.depthStream);

  for (int frame : params.frameRange) {
    LOG(INFO) << "Processing frame " << frame <<
        " of range " << params.frameRange.toString() << "...";

    const DepthFrame& refDf = srcDs.frame(frame);
    const Vector3f refPosition = refDf.extrinsics.position;
    const Vector3f refForward = refDf.extrinsics.forward();

    Mat1f filteredDepth(h, w);

    // Fetching a stack of (2 * frameRadius + 1) *guidance images*, temporally
    // centered around the current frame.
    const int f0 = std::max(0, frame - params.frameRadius);
    const int f1 =
        std::min(params.frameRange.lastFrame(), frame + params.frameRadius);
    const int temporalWindowSize = f1 - f0 + 1;
    const int middleFrame = frame - f0;
    std::vector<std::unique_ptr<Mat2f>> forwardFlowStack(temporalWindowSize);
    std::vector<std::unique_ptr<Mat1b>> forwardMaskStack(temporalWindowSize);
    std::vector<std::unique_ptr<Mat2f>> backwardFlowStack(temporalWindowSize);
    std::vector<std::unique_ptr<Mat1b>> backwardMaskStack(temporalWindowSize);
    for (int i = 0; i < temporalWindowSize; ++i) {
      if (i >= middleFrame && i < temporalWindowSize - 1) {
        forwardFlowStack[i] = loadFlow(f0 + i, f0 + i + 1);
        forwardMaskStack[i] = loadFlowMask(f0 + i, f0 + i + 1);
        CHECK(forwardFlowStack[i]);
        CHECK(forwardMaskStack[i]);
      }

      if (i <= middleFrame && i > 0) {
        backwardFlowStack[i] = loadFlow(f0 + i, f0 + i - 1);
        backwardMaskStack[i] = loadFlowMask(f0 + i, f0 + i - 1);
        CHECK(backwardFlowStack[i]);
        CHECK(backwardMaskStack[i]);
      }
    }

    std::vector<std::unique_ptr<Mat2f>> farFlowStack;
    std::vector<std::unique_ptr<Mat1b>> farMaskStack;
    std::vector<int> farIndex;
    if (params.farConnections) {
      for (const auto& pair : flowPairs) {
        const int fi = pair.second;
        if (pair.first == frame && (fi < f0 || fi > f1)) {
          farIndex.push_back(fi);
          farFlowStack.push_back(loadFlow(frame, fi));
          farMaskStack.push_back(loadFlowMask(frame, fi));
        }
      }
    }

    struct SampleInfo {
      int frame;
      Vector2f loc;
      float depth;
      float weight;
    };
    std::vector<SampleInfo> samples;

    auto addSample = [&](const Vector2f& loc, const int fi) {
      samples.emplace_back();
      SampleInfo& s = samples.back();
      s.frame = fi;
      s.loc = loc;
      Vector2f ndc(loc.x() / w, loc.y() / h * video_->invAspect());
      Vector3f position =
          video_->project(params.sourceDepthStream, fi, ndc, false);
      s.depth = (position - refPosition).dot(refForward);
    };

    // Looping over the output image pixels.
    for (int y = 0; y < h; ++y) {
      const int y0 = std::max(0, y - params.spatialRadius);
      const int y1 = std::min(h - 1, y + params.spatialRadius);

      float* const dstPtr = (float* const)filteredDepth.ptr(y);
      for (int x = 0; x < w; ++x) {
        const int x0 = std::max(0, x - params.spatialRadius);
        const int x1 = std::min(w - 1, x + params.spatialRadius);

        samples.clear();

        float referenceDepth = FLT_MAX;
        for (int wy = y0; wy <= y1; ++wy) {
          for (int wx = x0; wx <= x1; ++wx) {
            addSample(Vector2f(wx, wy), frame);

            if (wx == x && wy == y) {
              assert(!samples.empty()); // For lint.
              referenceDepth = samples.back().depth;
            }

            // Forward pass
            Vector2f loc = Vector2f(wx, wy);
            for (int fi = frame + 1; fi <= f1; ++fi) {
              const int i = fi - f0;
              Mat2f& flow = *forwardFlowStack[i - 1];
              Mat1b& mask = *forwardMaskStack[i - 1];

              int ix = std::min(int(loc.x() + 0.5f), w - 1);
              int iy = std::min(int(loc.y() + 0.5f), h - 1);

              if (!mask(iy, ix)) {
                break;
              }

              const Vec2f& f = flow(iy, ix);
              loc.x() += f(0);
              loc.y() += f(1);
              ix = loc.x() + 0.5f;
              iy = loc.y() + 0.5f;
              if (ix < 0 || ix >= w || iy < 0 || iy >= h) {
                break;
              }

              addSample(loc, fi);
            }

            // Backward pass
            loc = Vector2f(wx, wy);
            for (int fi = frame - 1; fi >= f0; --fi) {
              const int i = fi - f0;
              Mat2f& flow = *backwardFlowStack[i + 1];
              Mat1b& mask = *backwardMaskStack[i + 1];

              int ix = std::min(int(loc.x() + 0.5f), w - 1);
              int iy = std::min(int(loc.y() + 0.5f), h - 1);

              if (!mask(iy, ix)) {
                break;
              }

              const Vec2f& f = flow(iy, ix);
              loc.x() += f(0);
              loc.y() += f(1);
              ix = loc.x() + 0.5f;
              iy = loc.y() + 0.5f;
              if (ix < 0 || ix >= w || iy < 0 || iy >= h) {
                break;
              }

              addSample(loc, fi);
            }

            for (int i = 0; i < farIndex.size(); ++i) {
              assert(!farFlowStack.empty()); // For lint.
              assert(!farMaskStack.empty()); // For lint.
              Mat2f& flow = *farFlowStack[i];
              Mat1b& mask = *farMaskStack[i];

              loc = Vector2f(wx, wy);
              int ix = std::min(int(loc.x() + 0.5f), w - 1);
              int iy = std::min(int(loc.y() + 0.5f), h - 1);

              if (!mask(iy, ix)) {
                break;
              }

              const Vec2f& f = flow(iy, ix);
              loc.x() += f(0);
              loc.y() += f(1);
              ix = loc.x() + 0.5f;
              iy = loc.y() + 0.5f;
              if (ix < 0 || ix >= w || iy < 0 || iy >= h) {
                break;
              }

              addSample(loc, farIndex[i]);
            }
          }
        }

        // Get depth samples
        float depthSum = 0.f;
        float weightSum = 0.f;
        for (SampleInfo& s : samples) {
          const float value =
              std::max(s.depth, referenceDepth) /
              std::min(s.depth, referenceDepth);

          s.weight = expf(-value * 3.f);

          depthSum += s.depth * s.weight;

          weightSum += s.weight;
        }

        if (params.median) {
          float halfWeight = weightSum / 2.f;
          std::sort(samples.begin(), samples.end(),
              [&](SampleInfo& lhs, SampleInfo& rhs) {
            return (lhs.depth < rhs.depth);
          });
          float cumWeight = 0.f;
          for (const SampleInfo& s : samples) {
            cumWeight += s.weight;
            if (cumWeight >= halfWeight) {
              dstPtr[x] = s.depth;
              break;
            }
          }
        } else {
          if (weightSum > 0.f) {
            dstPtr[x] = depthSum / weightSum;
          } else {
            dstPtr[x] = 0.f;
          }
        }
      }
    }

    dstDs.frame(frame).setDepth(filteredDepth);
  }
}