【手撕算法】Criminisi图像修复算法

int main()
{
  std::chrono::steady_clock clk;
  auto begin = clk.now();//开始计时

  cv::Mat srcImage = cv::imread("image/image1.jpg");//读取图片
  cv::Mat mask = cv::imread("image/mask1.jpg", 0);//读取掩膜  8位单通道的掩膜
  if (srcImage.empty())
  {
    printf_s("读取图片失败");
    return -1;
  }
  cv::Mat lab_img;
  cv::cvtColor(srcImage, lab_img, COLOR_BGR2Lab);//转换颜色空间为LAB

  Criminisi criminisi(lab_img);

  criminisi.mask(mask);//赋值给Criminisi类的_mask变量
  const auto& res_lab = criminisi.generate(); //开启criminisi函数迭代

  cv::Mat resImage;
  cv::cvtColor(res_lab, resImage, COLOR_Lab2BGR);

  cv::imshow("src", srcImage);
  cv::imshow("mask", mask);
  cv::imshow("res", resImage);

  //结束计时
  auto time = std::chrono::duration_cast<std::chrono::milliseconds> (clk.now() - begin);
  std::cerr << "time : " << time.count() << std::endl;

  std::cout << "Press 'q' to exit" << std::endl;
  while ('q' != cv::waitKey(1));

  return 0;
}

/*********************************
函数声明：generate函数
参数：NONE
注释：Criminsi算法主程序体
测试：NONE
**********************************/
cv::Mat Criminisi::generate(void)
{
  generate_contour();//生成轮廓
  //把矩阵_mask中元素不为0的点全部变为value值
  _modified = _original.clone().setTo(cv::Vec3b(0, 0, 0),  _mask);//得到待修复的原图

  generate_priority();//生成优先队列，得到了对应点坐标的优先级

  cv::Mat resSSD;//用来存储匹配度的矩阵
  cv::Mat pMask;
  cv::Mat pInvMask;
  cv::Mat templateMask;
  cv::Mat mergeArrays[3];

  cv::Mat dilatedMask;

  while (_pq.size()) {//遍历_pq set容器

    const std::pair<int, int>& point = _pq.rbegin()->second;
    const cv::Point p(point.first, point.second);//获取待修复轮廓点

    const auto& phi_p = patch(p, _modified);//提取一小块待修复区域
    const int radius = (phi_p.rows - 1) / 2;//得到修复半径

    pMask = patch(p, _mask, radius);//得到p点待修复区域的_mask的掩膜块，待修复区域为1，已知区域为0
    pInvMask = ~pMask;//取反

    templateMask = (pInvMask);
    //单通道掩膜块合并成三通道掩膜块
    for (int i = 0; i < 3; ++i)
      mergeArrays[i] = templateMask;
    cv::merge(mergeArrays, 3, templateMask);//合并，得到三通道的样本掩膜块

    cv::matchTemplate(_modified, phi_p, resSSD, cv::TM_SQDIFF, templateMask); //在_modified中搜索与模板phi_p的templateMask区域相匹配的块并得到相似度矩阵resSSD

    cv::dilate(_mask, dilatedMask, cv::Mat(), cv::Point(-1, -1), radius);//将_mask膨胀为dilatedMask

    std::cerr << "Points in contour : " << _pq.size() << std::endl;

    //将待修复区域对应的resSSD区域设为匹配值最大，排除其干扰
    resSSD.setTo(std::numeric_limits<float>::max(), //将resSSD的所有或部分数组元素设置为指定的值max()
      dilatedMask(cv::Range(radius, _rows - radius),
        cv::Range(radius, _cols - radius)));

    cv::Point q;
    cv::minMaxLoc(resSSD, NULL, NULL, &q);//resSSD的最小值位置，即最佳匹配位置

    q = q + cv::Point(radius, radius);//得到该位置定位点

    const auto& phi_q = patch(q, _modified, radius);//提取出最佳匹配块

    phi_q.copyTo(phi_p, pMask);//复制填充

    cv::Mat confPatch = patch(p, _confidence);
    const double confidence = cv::sum(confPatch)[0] /confPatch.total();
    confPatch.setTo(confidence, pMask);

    pMask.setTo(0);

    update_contour(p);//修复p点后，更新边界信息
  }

  std::cerr << "Completed" << std::endl;
  return _modified;
}

/*********************************
函数声明：generate_contour生成轮廓函数
参数：NONE
注释：遍历整幅图像得到待修复边界
测试：NONE
**********************************/
void Criminisi::generate_contour(void)
{
  constexpr int NUM_N = 8;

  const int dx8[NUM_N] = { -1, -1,  0,  1, 1, 1, 0, -1 };//用来访问八邻域的数组
  const int dy8[NUM_N] = { 0, -1, -1, -1, 0, 1, 1,  1 };

  _contour.clear();//清除变量_contour

  for (int i = 0; i < _cols; ++i) {
    for (int j = 0; j < _rows; ++j) {//遍历图像
      for (int k = 0; k < NUM_N; ++k) { //遍历八邻域

        if (!_mask.at<uchar>(j, i))//如果当前像素不属于_mask，即不属于待修复区域，则退出
          continue;

        const int x = i + dx8[k];
        const int y = j + dy8[k];

        if (x >= 0 && x < _cols && y >= 0 && y < _rows) {//八邻域没有出界

          if (!_mask.at<uchar>(y, x)) {//八邻域像素不属于_mask，即代表该点属于边缘

            _contour.emplace(i, j);//将pair(i,j)存到set：_contour中
            break;//只需有一个就可以退出该for循环
          }
        }
      }
    }
  }
}

/*********************************
函数声明：生成优先级队列 函数
参数：NONE
注释：遍历所有点，根据轮廓点坐标计算其优先级
测试：NONE
**********************************/
void Criminisi::generate_priority(void)
{
  _confidence = cv::Mat::zeros(_rows, _cols, CV_64FC1); //置信 矩阵，单通道全为零的矩阵，代表了当前位置点被修复的优先级

  for (int i = 0; i < _cols; ++i)
    for (int j = 0; j < _rows; ++j)//遍历置信矩阵
      _confidence.at<double>(j, i) = !_mask.at<uchar>(j, i);//_mask为0的点置为1，_mask为1的点置为0

  _pq.clear(); //存储优先级与轮廓点
  _map.clear();//轮廓点到优先级的映射

  for (const auto& c : _contour) {
    //根据点位置c计算优先级pri
    const double pri = priority(c);//c为轮廓点set集  pri为该轮廓点的优先级priority
    _pq.emplace(pri, c);//建立轮廓点与优先级的对应关系
    _map[c] = pri;
  }

}

void cv::matchTemplate(
    cv::InputArray image, // 用于搜索的输入图像, 8U 或 32F, 大小 W-H
    cv::InputArray templ, // 用于匹配的模板，和image类型相同， 大小 w-h
    cv::OutputArray result, // 匹配结果图像, 类型 32F, 大小 (W-w+1)-(H-h+1)
    int method // 用于比较的方法
    InputArray mask //搜索模板templ的掩码。它必须与templ具有相同的数据类型和大小
);

/*********************************
函数声明：更新轮廓函数
参数：NONE
注释：NONE
测试：NONE
**********************************/
void Criminisi::update_contour(const cv::Point& p)
{
  constexpr int NUM_N = 8;

  const int dx8[NUM_N] = { -1, -1,  0,  1, 1, 1, 0, -1 };//用来访问八邻域的数组
  const int dy8[NUM_N] = { 0, -1, -1, -1, 0, 1, 1,  1 };

  const int x_begin = std::max(p.x - 2 * _radius, 0);
  const int y_begin = std::max(p.y - 2 * _radius, 0);
  const int x_end = std::min(p.x + 2 * _radius, _cols - 1);
  const int y_end = std::min(p.y + 2 * _radius, _rows - 1);

  for (int i = x_begin; i <= x_end; ++i) {
    for (int j = y_begin; j <= y_end; ++j) {

      const std::pair<int, int> p = std::make_pair(i, j);

      if (_contour.count(p)) {
        const double pri = _map[p];
        _contour.erase(p);
        _pq.erase(std::make_pair(pri, p));
        _map.erase(p);
      }
    }
  }

  std::set<std::pair<int, int>> new_contour;

  for (int i = x_begin; i <= x_end; ++i) {
    for (int j = y_begin; j <= y_end; ++j) {
      for (int k = 0; k < NUM_N; ++k) {

        if (!_mask.at<uchar>(j, i))
          continue;

        const int x = i + dx8[k];
        const int y = j + dy8[k];

        if (x >= 0 && x < _cols && y >= 0 && y < _rows) {

          if (!_mask.at<uchar>(y, x)) {

            new_contour.emplace(i, j);
            break;
          }
        }
      }
    }
  }

  for (const auto& nc : new_contour)
    _contour.emplace(nc);

  for (const auto& nc : new_contour) {

    const double pri = priority(nc);
    _pq.emplace(std::make_pair(pri, nc));
    _map[nc] = pri;
  }
}