Image Inpainting (Poisson iteration)

bitmap.h

#pragma once
#include <vector>

typedef unsigned short WORD;
typedef unsigned long DWORD;
typedef long           LONG;
typedef unsigned char  BYTE;
#pragma pack(push, 1)
struct BITMAPFILEHEADER {
  WORD  bfType;
  DWORD bfSize;
  WORD  bfReserved1;
  WORD  bfReserved2;
  DWORD bfOffBits;
};
struct BITMAPINFOHEADER {
  DWORD biSize; 
  LONG  biWidth;
  LONG  biHeight;
  WORD  biPlanes;
  WORD  biBitCount;
  DWORD biCompression;
  DWORD biSizeImage;
  LONG  biXPelsPerMeter;
  LONG  biYPelsPerMeter;
  DWORD biClrUsed;
  DWORD biClrImportant;
};
#pragma pop
const DWORD BI_RGB       = 0x00;
const DWORD BI_RLE8      = 0x01;
const DWORD BI_RLE4      = 0x02;
const DWORD BI_BITFIELDS = 0x03;

struct $FOPEN {
  FILE *fp;
  $FOPEN(const char *file, const char *sw) {
    fp = fopen(file, sw);
  }
  ~$FOPEN() { 
    fclose(fp); 
  }
};

struct RGB { double R, G, B; };
typedef std::vector< std::vector<RGB> > Image;
Image readBMP(const char *file) {
  $FOPEN FOPEN(file, "rb");
  Image img;

  FILE *fp = FOPEN.fp;
  if (!fp) {
    fprintf(stderr, "Cannot open %s\n", file);
    return img;
  }
  BITMAPFILEHEADER head;
  fread(&head, 1, sizeof(head), fp);
  if (head.bfType != ('B' | ('M' << 8))) {
    fprintf(stderr, "%s is not BMP format\n", file);
    return img;
  }
  BITMAPINFOHEADER info;
  fread(&info, 1, sizeof(BITMAPINFOHEADER), fp);
  if (info.biBitCount != 24) {
    fprintf(stderr, "%s is not 24bit color BMP\n", file);
    return img;
  }
  fseek(fp, head.bfOffBits, SEEK_SET);

  int width  = info.biWidth;
  int height = info.biHeight;
  int size   = head.bfSize - head.bfOffBits;

  std::vector<BYTE> data(size);
  fread(&(data[0]), 1, size, fp);

  Image bmp(height, std::vector<RGB>(width));
  int k = 0;
  for (int i = 0; i < height; ++i) {
    for (int j = 0; j < width; ++j) {
      bmp[i][j].B = data[k++];
      bmp[i][j].G = data[k++];
      bmp[i][j].R = data[k++];
    }
    for (int j = width; j % 4 > 0; ++j) {
      ++k;
    }
  }
  return bmp;
}

void saveBMP(Image bmp, const char *file) {
  $FOPEN FOPEN(file, "wb");
  FILE *fp = FOPEN.fp;
  std::vector<BYTE> data;

  int height = bmp.size();
  int width  = bmp[0].size();
  for (int i = 0; i < height; ++i) {
    for (int j = 0; j < width; ++j) {
      data.push_back(bmp[i][j].B);
      data.push_back(bmp[i][j].G);
      data.push_back(bmp[i][j].R);
    }
    while (data.size() % 4) {
      data.push_back(0);
    }
  }
  BITMAPINFOHEADER info = {
    /* biSize          = */ sizeof(info),
    /* biWidth         = */ bmp[0].size(),
    /* biHeight        = */ bmp.size(),
    /* biPlanes        = */ 1,
    /* biBitCount      = */ 24, 
    /* biCompression   = */ BI_RGB,
    /* biSizeImage     = */ data.size(),
    /* biXPelsPerMeter = */ 0,
    /* biYPelsPerMeter = */ 0,
    /* biClrUsed       = */ 0,
    /* biClrImportant  = */ 0,
  };
  BITMAPFILEHEADER head = {
    /* bfType          = */ ('B' | ('M' << 8)),
    /* bfSize          = */ data.size()+sizeof(head)+sizeof(info),
    /* bfReserved1     = */ 0,
    /* bfReserved2     = */ 0,
    /* bfOffBits       = */ sizeof(head)+sizeof(info),
  };
  fwrite(&head, sizeof(head), 1, fp);
  fwrite(&info, sizeof(info), 1, fp);
  fwrite(&data[0], 1, data.size(), fp);
}

inpaint.cc

//
// Image inpainting (Poisson iteration)
// 
// Reference: 
//   M. M. Oliveira, B. Bowen, R. McKenna, and Y-S. Chang (2001):
//     "Fast digital image inpainting".
//     Proceedings of the International Conference of Visualization, Imaging, and Image Processing, 
//     pp. 261-266. 
//   ( http://elynxsdk.free.fr/ext-docs/Inpainting/todo/Fast_Digital_Image_Inpainting.pdf )
//
#include <iostream>
#include <vector>
#include <cstdio>
#include <cstdlib>
#include <map>
#include <cmath>
#include <cstring>
#include <functional>
#include <algorithm>
#include <cassert>
#include "bitmap.h"

using namespace std;

#define fst first
#define snd second

void imageReconstruction(Image &bmp) {
  const double a = 0.073235, b = 0.176765,
             w[] = { b, a, b, a, b, a, b, a};
  const int di[] = { 0, 1, 1, 1, 0,-1,-1,-1}, 
            dj[] = { 1, 1, 0,-1,-1,-1, 0, 1};

  int height = bmp.size();
  int width  = bmp[0].size();

  vector< pair<int,int> > inpaint;
  for (int i = 0; i < height; ++i) 
    for (int j = 0; j < width; ++j) 
      if (!bmp[i][j].R && !bmp[i][j].G && !bmp[i][j].B) {
        //bmp[i][j].R = 255;
        inpaint.push_back( make_pair(i, j) );
      }

  Image tmp(bmp);
  for (int iter = 0; iter < 1000; ++iter) {
    for (int k = 0; k < inpaint.size(); ++k) {
      int i = inpaint[k].fst;
      int j = inpaint[k].snd;

      double r = 0;
      tmp[i][j].R = tmp[i][j].G = tmp[i][j].B = 0;
      for (int t = 0; t < 8; ++t) {
        if (i+di[t] >= 0 && i+di[t] < height &&
            j+dj[t] >= 0 && j+dj[t] < width  &&
            (bmp[i+di[t]][j+dj[t]].R ||
             bmp[i+di[t]][j+dj[t]].G ||
             bmp[i+di[t]][j+dj[t]].B)) {
          r += w[t];
          tmp[i][j].R += w[t]*bmp[i+di[t]][j+dj[t]].R;
          tmp[i][j].G += w[t]*bmp[i+di[t]][j+dj[t]].G;
          tmp[i][j].B += w[t]*bmp[i+di[t]][j+dj[t]].B;
        }
      }
      if (r > 0) {
        tmp[i][j].R = min(tmp[i][j].R/r, 255.);
        tmp[i][j].G = min(tmp[i][j].G/r, 255.);
        tmp[i][j].B = min(tmp[i][j].B/r, 255.);
      }
    }
    swap(bmp, tmp);
  }
}


int main() { 
  Image bmp = readBMP("hoge.bmp");
  imageReconstruction(bmp);
  saveBMP(bmp, "hoge.save.bmp");
}