poisson.cpp

#include <opencv2/opencv.hpp>

#include <iostream>
#include <vector>
#include <cmath>
#include <assert.h>

using namespace std;
using namespace cv;

#define MAX_ITERATION 300

inline double clampValue (double val) {
  if (val > 1) return 1;
  else if (val < 0) return 0;
  else return (double)val;
}

void montage (Mat &srcImage, Mat &destImage, Mat &maskImage) {
  int w = srcImage.cols;
  int h = srcImage.rows;
  
  if (w != maskImage.cols || h != maskImage.rows) {
    cerr << "mask size doesn't match src size" << endl;
    return;
  }

  vector<Point2i> destPoints;
  vector<double> constants;
  
  int size=0;
  for (int y=0; y<h; y++) {
    for (int x=0; x<w; x++) {
            
      uchar c = maskImage.at<Vec3b>(y,x)[0];
      
      if (c > 0) {
        //printf("input (%d,%d)\n", x, y);
        destPoints.push_back(Point2i(x,y));
            
        int constant[3] = {0};
        double sum1[3]={0};
        double sum2[3]={0};
        // right
        if (x < srcImage.cols-1) {
          for (int i=0; i<3; i++) {
            double val1 = destImage.at<Vec3d>(y,x+1)[i] - destImage.at<Vec3d>(y,x)[i];
            double val2 = srcImage.at<Vec3d>(y,x+1)[i] - srcImage.at<Vec3d>(y,x)[i];
            sum1[i]+=val1;
            sum2[i]+=val2;
          }
        }
        // left
        if (x > 0) {
          for (int i=0; i<3; i++) {
            double val1 = destImage.at<Vec3d>(y,x-1)[i] - destImage.at<Vec3d>(y,x)[i];
            double val2 = srcImage.at<Vec3d>(y,x-1)[i] - srcImage.at<Vec3d>(y,x)[i];
            sum1[i]+=val1;
            sum2[i]+=val2;
          }
        }
        // top
        if (y > 0) {
          for (int i=0; i<3; i++) {
            double val1 = destImage.at<Vec3d>(y-1,x)[i] - destImage.at<Vec3d>(y,x)[i];
            double val2 = srcImage.at<Vec3d>(y-1,x)[i] - srcImage.at<Vec3d>(y,x)[i];
            sum1[i]+=val1;
            sum2[i]+=val2;
          }
        }
        // bottom
        if (y < srcImage.rows-1) {
          for (int i=0; i<3; i++) {
            double val1 = destImage.at<Vec3d>(y+1,x)[i] - destImage.at<Vec3d>(y,x)[i];
            double val2 = srcImage.at<Vec3d>(y+1,x)[i] - srcImage.at<Vec3d>(y,x)[i];
            sum1[i]+=val1;
            sum2[i]+=val2;
          }
        }
        for (int i=0; i<3; i++) {
          constants.push_back(sum2[i]);
        }
      }
      size++;
    } 
  }
  
  printf("destPoints size=%d\n", (int)destPoints.size());
  printf("constants size=%d\n", (int)constants.size());

  Mat final(destImage.cols, destImage.rows, CV_64FC3);
  
  for (int x=0; x<destImage.cols; x++) 
    for (int y=0; y<destImage.rows; y++) 
        final.at<Vec3d>(y,x) = destImage.at<Vec3d>(y,x);
  
  // ヤコビ法で連立一次方程式を解く
  // "I solve a system of linear equations in the Jacobi method"
  for (int loop=0; loop<MAX_ITERATION; loop++) {
    int n = destPoints.size();
    for (int p=0; p<n; p++) {
      //int destIndex = destPoints[p];
      int y = destPoints[p].y;
      int x = destPoints[p].x;
      //printf("check (%d,%d)\n", x, y);
      double values[3] = {0};
      // right
      int count = 0;
      if (x < destImage.cols-1) {
        count++;
        for (int i=0; i<3; i++) {
          values[i] += final.at<Vec3d>(y, x+1)[i];
        }
      }
      // left
      if (x > 0) {
        count++;
        for (int i=0; i<3; i++) {
          values[i] += final.at<Vec3d>(y, x-1)[i];
        }
      }
      // top
      if (y > 0) {
        count++;
        for (int i=0; i<3; i++) {
          values[i] += final.at<Vec3d>(y-1, x)[i];
        }
      }
      // bottom
      if (y < destImage.rows-1) {
        count++;
        for (int i=0; i<3; i++) {
          values[i] += final.at<Vec3d>(y+1, x)[i];
        }
      }

      // 更新
      for (int j=0; j<3; j++) {
        double newval = (values[j] - constants[p*3+j]) / count;
        double oldval = final.at<Vec3d>(y,x)[j];
        final.at<Vec3d>(y,x)[j] = newval;
      }
    }
  }
  
  {
    int n = destPoints.size();
    for (int p=0; p<n; p++) {
      int y = destPoints[p].y;
      int x = destPoints[p].x;
      for (int i=0; i<3; i++) {
        destImage.at<Vec3d>(y,x)[i] = clampValue(final.at<Vec3d>(y,x)[i]);
      }
    }
  }

  return;
}


int main(int argc, char** argv) {
  const char* srcFile = argc == 4 ? argv[1] : "black1.jpg";
  const char* destFile = argc == 4 ? argv[2] : "takahiro.jpg";
  const char* maskFile = argc == 4 ? argv[3] : "mask_black.jpg";

  Mat srcImage = imread(srcFile);
  Mat destImage = imread(destFile);
  Mat maskImage = imread(maskFile);

  srcImage.convertTo(srcImage, CV_64FC3, 1.0/255, 0);
  destImage.convertTo(destImage, CV_64FC3, 1.0/255, 0);
  maskImage.convertTo(maskImage, CV_8UC3); // i know 3 mask channels arent needed...

  
  montage(srcImage, destImage, maskImage);

  imshow("Poisson Image Editing", destImage);
  imshow("poisson.jpeg", destImage);
  
  waitKey(0);
  

  return 0;
}