レベルセット法の練習

levelset_method.c

/*  -*- mode: c; coding: utf-8-unix -*-  */
/* 
   Level set method
   Build:
   gcc `pkg-config --libs --cflags opencv` -pg levelset_method.c -o levelset_method 
   Usage:
   ./levelset_method imagefile

   TODO:
   ・高速化
   ・リファクタリング
*/
#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <cv.h>
#include <highgui.h>

#define WINDOW_NAME_1 "Window1"
#define WINDOW_NAME_2 "Window2"

/** 2次元行列matの(x, y)要素を取得*/
#define GET_2D(mat, x, y, step) (mat[((step) * (y)) + (x)])

/** statusのラベル */
enum { FARAWAY, BAND, RESET_BAND, FRONT };

/** 点の配列 */
typedef struct PointArray {
  /** 点 */
  CvPoint *points;
  /** 点の最大数 */
  int max_size;
  /** 現在のデータ数 */
  int count;
} PointArray;


/** レベルセット法で使う画像データの型 */
typedef struct LevelSetImageType {
  /** 各ピクセルの輝度(8ビットグレイスケール) */
  unsigned char *image_gray;
  /** 各ピクセルの状態 */
  unsigned char *status;
  /** 各ピクセルの補助関数の値 */
  double *phi;
  /** 各ピクセルの補助関数の変化量 */
  double *dphi;
  /** 各ピクセルの成長速度 */
  double *f;
  /** 画像の幅 */
  uint width;
  /** 画像の高さ */
  uint height;
} LevelSetImage;

/**
   画像データを作成
   @param src 元となるIplImageのポインタ
   @return 作成した画像データ
*/
LevelSetImage* new_levelset_image(const IplImage *src) {
  int size = src->width * src->height;
  LevelSetImage *img = (LevelSetImage*)malloc(sizeof(LevelSetImage));
  if (img == NULL)
    return NULL;
  
  img->width = src->width;
  img->height = src->height;

  img->image_gray = (unsigned char *)malloc(sizeof(unsigned char) * size);
  img->status = (unsigned char *)malloc(sizeof(unsigned char) * size);
  img->phi = (double *)malloc(sizeof(double) * size);
  img->dphi = (double *)malloc(sizeof(double) * size);
  img->f = (double *)malloc(sizeof(double) * size);

  if (img->image_gray == NULL || img->status == NULL ||
      img->phi == NULL || img->dphi == NULL || img->f == NULL) {
    if (img->image_gray != NULL)
      free(img->image_gray);
    if (img->status != NULL)
      free(img->status);
    if (img->phi != NULL)
      free(img->phi);
    if (img->dphi != NULL)
      free(img->dphi);
    if (img->f != NULL)
      free(img->f);
    return NULL;
  }
  memcpy(img->image_gray, src->imageData, sizeof(unsigned char) * size);
  
  return img;
}

/**
   画像データのメモリを開放
   @param img 画像データのポインタ
*/
void free_levelset_image(LevelSetImage* img) {
  free(img->image_gray);
  free(img->status);
  free(img->phi);
  free(img->dphi);
  free(img->f);

  free(img);
}

/**
   PointArrayを作成
   @param max_size 作成するPointArrayの最大サイズ
   @return 作成したデータ
*/
PointArray* new_point_array(int max_size) {
  PointArray* p = (PointArray *)malloc(sizeof(PointArray));
  if (p == NULL)
    return NULL;
  
  p->points = (CvPoint *)malloc(sizeof(CvPoint) * max_size);
  if (p->points == NULL) {
    free(p);
    return NULL;
  }
  
  p->max_size = max_size;
  p->count = 0;
  return p;
}

/**
   PointArrayのメモリを開放
   @param p 開放するPointArrayのポインタ
*/
void free_point_array(PointArray *p) {
  free(p->points);
  free(p);
}


/**
   距離マップを作成
   TODO: ちゃんとNULLチェックする
   @param band_width Narrow Bandの幅
   @return 距離マップ(要素数: band_width + 1)
*/
PointArray** init_circle_map(int band_width) {
  int x, y, i, d;
  int num_points = 7 * (band_width + 1);
  
  PointArray **circle_maps = (PointArray **)malloc(sizeof(PointArray *) * (band_width + 1));
  if (circle_maps == NULL)
    return NULL;
  
  for (i = 0; i < band_width + 1; i++) {
    if ((circle_maps[i] = new_point_array(num_points)) == NULL) {
      while (i >= 0)
	free_point_array(circle_maps[i--]);
      return NULL;
    }
  }
  
  for (y = -band_width; y <= band_width; y++) {
    for (x = -band_width; x <= band_width; x++) {
      d = (int)sqrt(x * x + y * y);
      if (d <= band_width) {
	i = circle_maps[d]->count;
	circle_maps[d]->points[i].x = x;
	circle_maps[d]->points[i].y = y;
	circle_maps[d]->count++;
      }
    }
  }
  return circle_maps;
}


/**
   距離マップのメモリを開放
   @param 距離マップ
   @param band_width Narrow bandの幅
*/
void free_circle_maps(PointArray** circle_maps, int band_width) {
  int i;
  for (i = 0; i < band_width + 1; i++)
    free_point_array(circle_maps[i]);
  free(circle_maps);
}

/**
   Zero level setの初期化 (Step.2 符号付き距離場の構築)
   @param front Zero level set
   @param img 画像データ
   @param init_offset 初期曲面のオフセット
   @param band_width Narrow bandの幅
*/
void init_front(PointArray* front, LevelSetImage* img,
		int init_offset, int band_width) {
  int x, y, i;
  int width, height;
  width = img->width;
  height = img->height;

  memset(img->status, FARAWAY, width * height);

  // Frontグリッドを設定
  i = 0;
  for (x = init_offset; x < width - init_offset; x++) {
    GET_2D(img->status, x, init_offset, width) = FRONT;
    GET_2D(img->phi, x, init_offset, width) = 0.0;
    front->points[i].x = x;
    front->points[i].y = init_offset;
    i++;

    GET_2D(img->status, x, height - init_offset - 1, width) = FRONT;
    GET_2D(img->phi, x, height - init_offset - 1, width) = 0.0;
    front->points[i].x = x;
    front->points[i].y = height - init_offset - 1;
    i++;
  }

  for (y = init_offset; y < height - init_offset; y++) {
    GET_2D(img->status, init_offset, y, width) = FRONT;
    GET_2D(img->phi, init_offset, y, width) = 0.0;
    front->points[i].x = init_offset;
    front->points[i].y = y;
    i++;

    GET_2D(img->status, width - init_offset - 1, y, width) = FRONT;
    GET_2D(img->phi, width - init_offset - 1, y, width) = 0.0;
    front->points[i].x = width - init_offset - 1;
    front->points[i].y = y;
    i++;
  }

  front->count = i;

  // その他の場所を初期化
  for (y = 0; y < height; y++) {
    for (x = 0; x < width; x++) {
      if (GET_2D(img->status, x, y, width) != FRONT) {
	if (x > init_offset && x < width - init_offset - 1 &&
	    y > init_offset && y < height - init_offset - 1) {
	  GET_2D(img->phi, x, y, width) = -band_width;
	}
	else
	  GET_2D(img->phi, x, y, width) = band_width;
      }
    }
  }
}

/**
   Narrow band上の成長速度を設定+再初期化処理(Step.3, Step.6)
   @param narrow_band Narrow band
   @param img 画像データ
   @param front Zero level set
   @param circle_maps 参照マップ
   @param band_width Narrow bandの幅
   @param reset 0以外ならnarrow_bandを再設定する(Step.6)、0ならしない
   @return 成長速度の合計値
*/
double set_growing_rate(PointArray* narrow_band, LevelSetImage* img,
			PointArray* front, PointArray** circle_maps,
			int band_width, int reset_band_width, int reset) {
  double sum_f = 0.0;
  int i, j, d;
  int x, y, xf, yf;
  double ki;
  double dx, dy;
  double phi_x, phi_y;
  double phi_xx, phi_yy, phi_xy;
  double df;
  double kappa;
  double phi00, phi01, phi02;
  double phi10, phi11, phi12;
  double phi20, phi21, phi22;
  int width, height;
  width = img->width;
  height = img->height;
  
  // Narrow band上の成長速度をクリア
  for (i = 0; i < narrow_band->count; i++) {
    x = narrow_band->points[i].x;
    y = narrow_band->points[i].y;
    GET_2D(img->f, x, y, width) = 0.0;
    GET_2D(img->dphi, x, y, width) = 0.0;
    // Narrow bandの再初期化
    if (reset) {
      if (GET_2D(img->status, x, y, width) != FRONT)
	GET_2D(img->status, x, y, width) = FARAWAY;
    }
  }
  
  // まずはZero level setの成長速度を更新
  for (i = 0; i < front->count; i++) {
    x = front->points[i].x;
    y = front->points[i].y;
    
    if (x < 1 || x >= width - 1 ||
	y < 1 || y >= height - 1)
      continue;

    phi00 = GET_2D(img->phi, x - 1, y - 1, width);
    phi01 = GET_2D(img->phi, x, y - 1, width);
    phi02 = GET_2D(img->phi, x + 1, y - 1, width);
    phi10 = GET_2D(img->phi, x - 1, y, width);
    phi11 = GET_2D(img->phi, x, y, width);
    phi12 = GET_2D(img->phi, x + 1, y, width);
    phi20 = GET_2D(img->phi, x - 1, y + 1, width);
    phi21 = GET_2D(img->phi, x, y + 1, width);
    phi22 = GET_2D(img->phi, x + 1, y + 1, width);
    
    phi_x = ((phi22 - phi20) + 2.0 * (phi12 - phi10) + (phi02 - phi20)) / 4.0 / 2.0;
    phi_y = ((phi22 - phi02) + 2.0 * (phi21 - phi01) + (phi20 - phi20)) / 4.0 / 2.0;
    phi_xy = ((phi22 - phi20) - (phi02 - phi20)) / 4.0 / 2.0;
    phi_xx = ((phi22 + phi20 - 2.0 * phi21) + 2.0 * (phi12 - phi10 - 2.0 * phi11)
	      + (phi02 - phi00 - 2.0 * phi01)) / 4.0;
    phi_yy = ((phi22 + phi02 - 2.0 * phi12) + 2.0 * (phi21 + phi01 - 2.0 * phi11)
	      + (phi20 - phi00 - 2.0 * phi10)) / 4.0;
    
    df = sqrt(phi_x * phi_x + phi_y * phi_y);
    if (df * df - 0.001 > 0) // df != 0
      kappa = (phi_xx * phi_y * phi_y - 2.0 * phi_x * phi_y * phi_xy
	       + phi_yy * phi_x * phi_x) / (df * df * df); // (22)
    else
      kappa = 0.0;

    dx = GET_2D(img->image_gray, x + 1, y, width) - GET_2D(img->image_gray, x, y, width);
    dy = GET_2D(img->image_gray, x, y + 1, width) - GET_2D(img->image_gray, x, y, width);

    ki = 1.0 / (1.0 + sqrt(dx * dx + dy * dy)); // (21)
    GET_2D(img->f, x, y, width) = ki * (-1.0 - 0.1 * kappa); // (20) (a = -1.0, b = 0.1)
    sum_f += GET_2D(img->f, x, y, width);
  }

  // 次にNarrow bandの成長速度を更新
  for (d = band_width; d >= 0; d--) { // d == 0のときはFRONTだから更新しなくておｋ?
    for (i = 0; i < front->count; i++) {
      xf = front->points[i].x;
      yf = front->points[i].y;
      if (reset)
	GET_2D(img->phi, xf, yf, width) = 0.0;

      for (j = 0; j < circle_maps[d]->count; j++) {
	x = xf + circle_maps[d]->points[j].x;
	y = yf + circle_maps[d]->points[j].y;
	if (x < 0 || x >= width ||
	    y < 0 || y >= height)
	  continue;
	if (GET_2D(img->status, x, y, width) == FRONT)
	  continue;
	if (reset) {
	  // 境界がNarrow bandの端に近づいたらRESET_BANDにする
	  if (band_width - d < reset_band_width)
	    GET_2D(img->status, x, y, width) = RESET_BAND;
	  else
	    GET_2D(img->status, x, y, width) = BAND;
	  GET_2D(img->phi, x, y, width) = (GET_2D(img->phi, x, y, width) < 0) ? -d : d; // 再初期化(Step.6)
	}
	// Frontでなければ上書き
	GET_2D(img->f, x, y, width) = GET_2D(img->f, xf, yf, width);
      }
    }
  }
  
  if (reset) {
    // Narrow bandに登録
    i = 0;
    for (y = 0; y < height; y++) {
      for (x = 0; x < width; x++) {
	if (GET_2D(img->status, x, y, width) != FARAWAY) {
	  narrow_band->points[i].x = x;
	  narrow_band->points[i].y = y;
	  i++;
	  if (i >= narrow_band->max_size) {
	    // FIXME: 足りなくなったら増やすようにする
	    fprintf(stderr, "Error: Too many narrow band points (%d)\n", i);
	    narrow_band->count = i;
	    return 0.0;
	  }
	}
      }
    }
    narrow_band->count = i;
  }

  return sum_f;
}

/**
   補助関数の更新(Step.4)
   @param narrow_band Narrow band
   @param img 画像データ
*/
void update_phi(const PointArray* narrow_band, LevelSetImage *img) {
  int x, y;
  int i;
  double fdxp, fdxm;
  double fdyp, fdym;
  int width, height;
  width = img->width;
  height = img->height;

  for (i = 0; i < narrow_band->count; i++) {
    x = narrow_band->points[i].x;
    y = narrow_band->points[i].y;
    if (x < 1 || x >= width - 1 ||
	y < 1 || y >= height - 1)
      continue;

    // Upwind scheme
    if (GET_2D(img->f, x, y, width) > 0.0) {
      fdxm = MAX(GET_2D(img->phi, x, y, width) - GET_2D(img->phi, x - 1, y, width), 0);
      fdxp = MIN(GET_2D(img->phi, x + 1, y, width) - GET_2D(img->phi, x, y, width), 0);
      fdym = MAX(GET_2D(img->phi, x, y, width) - GET_2D(img->phi, x, y - 1, width), 0);
      fdyp = MIN(GET_2D(img->phi, x, y + 1, width) - GET_2D(img->phi, x, y, width), 0);
    }
    else {
      fdxm = MIN(GET_2D(img->phi, x, y, width) - GET_2D(img->phi, x - 1, y, width), 0);
      fdxp = MAX(GET_2D(img->phi, x + 1, y, width) - GET_2D(img->phi, x, y, width), 0);
      fdym = MIN(GET_2D(img->phi, x, y, width) - GET_2D(img->phi, x, y - 1, width), 0);
      fdyp = MAX(GET_2D(img->phi, x, y + 1, width) - GET_2D(img->phi, x, y, width), 0);
    }
    GET_2D(img->dphi, x, y, width) = sqrt(fdxm * fdxm + fdxp * fdxp
					  + fdym * fdym + fdyp * fdyp); // (31) or (32)
  }

  for (i = 0; i < narrow_band->count; i++) {
    x = narrow_band->points[i].x;
    y = narrow_band->points[i].y;
    GET_2D(img->phi, x, y, width) -= GET_2D(img->f, x, y, width) * GET_2D(img->dphi, x, y, width);
    // (28) ただしdt=1なので省略
  }
}

/**
   Frontのラベルを再設定
   @param front Zero level set
   @param img 画像データ
   @param narrow_band Narrow band
   @return > 0なら再初期化が必要、 ==0なら不要、< 0ならエラー(Zero level setの容量不足)
*/
int relabeling(PointArray *front, LevelSetImage *img, const PointArray *narrow_band) {
  int update_required = 0;
  int i, n;
  int x, y;
  int width, height;
  width = img->width;
  height = img->height;
  
  n = 0;
  for (i = 0; i < narrow_band->count; i++) {
    x = narrow_band->points[i].x;
    y = narrow_band->points[i].y;

    if (x < 1 || x >= width - 1 ||
	y < 1 || y >= height - 1)
      continue;

    // Step.5 Zero level setの検出
    if (GET_2D(img->phi, x, y, width) >= 0.0 &&
	((GET_2D(img->phi, x + 1, y, width) + GET_2D(img->phi, x, y, width) <= 0.0) ||
	 (GET_2D(img->phi, x - 1, y, width) + GET_2D(img->phi, x, y, width) <= 0.0) ||
	 (GET_2D(img->phi, x, y + 1, width) + GET_2D(img->phi, x, y, width) <= 0.0) ||
	 (GET_2D(img->phi, x, y - 1, width) + GET_2D(img->phi, x, y, width) <= 0.0)) ||
	GET_2D(img->phi, x, y, width) <= 0.0 &&
	((GET_2D(img->phi, x + 1, y, width) + GET_2D(img->phi, x, y, width) >= 0.0) ||
	 (GET_2D(img->phi, x - 1, y, width) + GET_2D(img->phi, x, y, width) >= 0.0) ||
	 (GET_2D(img->phi, x, y + 1, width) + GET_2D(img->phi, x, y, width) >= 0.0) ||
	 (GET_2D(img->phi, x, y - 1, width) + GET_2D(img->phi, x, y, width) >= 0.0))) {

      if (GET_2D(img->status, x, y, width) == RESET_BAND)
	update_required = 1;

      GET_2D(img->status, x, y, width) = FRONT;
      front->points[n].x = x;
      front->points[n].y = y;
      n++;

      if (i >= front->max_size) {
	fprintf(stderr, "Too many front points (%d)\n", n);
	front->count = n;
	return -1;
      }
    }
    else if (GET_2D(img->status, x, y, width) == FRONT)
      GET_2D(img->status, x, y, width) = BAND;
      
  }
  front->count = n;
  
  return update_required;
}


void show_phi(const LevelSetImage *img, const PointArray *front, int band_width, IplImage *dst) {
  int i;
  
  memset(dst->imageData, 255, dst->width * dst->height);
  for (i = 0; i < front->count; i++)
    cvSet2D(dst, front->points[i].y, front->points[i].x, CV_RGB(0, 0, 0));

  cvShowImage(WINDOW_NAME_1, dst);
}

int main(int argc, char *argv[]) {
  int i, max_count;
  IplImage *src, *dst;
  LevelSetImage *img;
  PointArray* front;
  int band_width = 10;
  int reset_band_width = 3;
  int init_offset = 1;
  PointArray** circle_maps;
  PointArray *narrow_band;
  int update_required = 1;
  double sum_f = 0.0;
  
  if (argc < 2) {
    fprintf(stderr, "Usage: %s imagefile\n", argv[0]);
    return -1;
  }
  if ((src = cvLoadImage(argv[1], CV_LOAD_IMAGE_GRAYSCALE)) == 0) {
    fprintf(stderr, "Error: Failed to load %s\n", argv[1]);
    return -1;
  }

  // Step.1 変数の準備
  img = new_levelset_image(src);
  front = new_point_array(img->width * img->height * 10);
  circle_maps = init_circle_map(band_width);

  // Step.2 符号付き距離場の構築
  narrow_band = new_point_array(img->width * img->height * 10);
  init_front(front, img, init_offset, band_width);
  set_growing_rate(narrow_band, img, front, circle_maps,
		   band_width, reset_band_width, 1);

  dst = cvCloneImage(src);
  cvNamedWindow(WINDOW_NAME_1, CV_WINDOW_AUTOSIZE);
  cvNamedWindow(WINDOW_NAME_2, CV_WINDOW_AUTOSIZE);

  max_count = 300;
  for (i = 0; i < max_count; i++) {
    // Step.3 成長速度の計算
    sum_f =  set_growing_rate(narrow_band, img, front, circle_maps,
			      band_width, reset_band_width, update_required);
    /*
      show_phi(img, front, band_width, dst);
      if (cvWaitKey(10) == 'q')
      break;

      if (i % 50 == 0)
      printf("[%d] sum_f = %.4f\n", i, sum_f);
    */
    if (sum_f * sum_f < 200.0)
      break;

    // Step.4 補助関数の更新
    update_phi(narrow_band, img);
    update_required = relabeling(front, img, narrow_band);

    if (update_required < 0) // Error!!
      break;
  }

  printf("End at [%d] sum_f = %.4f\n", i, sum_f);
  memset(dst->imageData, 255, dst->width * dst->height);
  for (i = 0; i < front->count; i++)
    cvSet2D(dst, front->points[i].y, front->points[i].x, CV_RGB(100, 100, 100));
  cvShowImage(WINDOW_NAME_1, src);
  cvShowImage(WINDOW_NAME_2, dst);
  cvWaitKey(0);

  cvDestroyWindow(WINDOW_NAME_1);
  cvDestroyWindow(WINDOW_NAME_2);
  cvReleaseImage(&src);
  cvReleaseImage(&dst);

  free_point_array(front);
  free_point_array(narrow_band);
  free_levelset_image(img);
  free_circle_maps(circle_maps, band_width);
  
  return 0;
}