混合ガウス分布を用いたNPCのためのEMアルゴリズム

npc.c

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>

#include <cv.h>
#include <highgui.h>


#define DIM 2 // 次元
#define LEVEL 32 // 階調
#define M 512 // 混合するガウス分布の数
#define MAX_ITR 100 // 繰り返し回数


// グローバル変数
int N; // 入力ベクトル数
double (*input)[DIM]; // 入力 input[N][dim] (x,y)
double mean[M][DIM]; // ガウス分布の平均 
double cov[M][DIM*DIM]; // ガウス分布の分散 (xx, xy // yx, yy)
double (*gamma)[M]; // 寄与率 gamma[n][m]
int H, W; // 画像サイズ

// 関数
void initParams( const IplImage* inputImage , const IplImage* gaussian);
void saveCurrentImage( char append[] );
void saveGaussianImage( char append[] );
double noramlDistr( const double x[DIM], const double m[DIM], const double c[DIM*DIM] );
void updateGamma(); 
void updateOther();

int main( int argc, char* argv[]){
  
  IplImage *inputImage = cvLoadImage( argv[1], CV_LOAD_IMAGE_GRAYSCALE );
  IplImage *inputGaussian = cvLoadImage( argv[2], CV_LOAD_IMAGE_GRAYSCALE );

  if( inputImage == NULL  || inputGaussian == NULL){
    printf("input argument[1] : input image\n");
    printf("input argument[2] : initial gauusian\n");
    return 1;
  }

  IplImage* inputMin = cvCreateImage( cvSize( inputImage->width/2, inputImage->height/2), IPL_DEPTH_8U, 1);
  cvResize( inputImage, inputMin, CV_INTER_AREA);

  initParams( inputImage, inputGaussian);

  saveCurrentImage("initial");

  for(int itr = 0; itr < MAX_ITR; ++itr){

    printf("itr = %d\n", itr);

    // E-step
    updateGamma();
    printf("E-step done\n");

    // M-step
    updateOther();
    printf("M-step done\n");

    char append[256];
    sprintf(append, "%02d", itr);
    saveCurrentImage(append);
    saveGaussianImage(append);
  }


}

void initParams( const IplImage* inputImage , const IplImage* gaussian)
{
  // count up N
  H = inputImage->height;
  W = inputImage->width;
  N = 0;
  int step = 256/LEVEL;
  for( int h = 0; h < H; ++h){
    for( int w = 0; w < W; ++w){
      N += (CV_IMAGE_ELEM(inputImage, uchar, h, w)) / step;
    }
  }
  printf("N = %d\n", N);
  
  // set input
  input = (double(*)[DIM])malloc(sizeof(double[DIM])*N);

  int idx = 0;
  for( int h = 0; h < H; ++h){
    for( int w = 0; w < W; ++w){
      int l = (CV_IMAGE_ELEM(inputImage, uchar, h, w)) / step;
      for(int i = 0; i <l; ++i){
	input[idx][0] = (double)w ;
	input[idx][1] = (double)h ;
	idx++;
      }
    }
  }  

  // set mean
  idx = 0;
  for(int h = 0; h < gaussian->height; ++h){
    for(int w = 0; w < gaussian->width; ++w){
      if( CV_IMAGE_ELEM( gaussian, uchar, h, w) == 0){
	mean[idx][0] = (double)w ;
	mean[idx][1] = (double)h ;
	printf("%04d mean : %lf, %lf\n", idx, mean[idx][0], mean[idx][1]);
	idx++;
      }
    }
  }

  int count[M];
  for( int k = 0; k < M; ++k){
    count[k] = 0;
    for( int c = 0; c < 4; ++c)
      cov[k][c] = 0;
  }
  for( int n  = 0; n < N; ++n ){
    double min = DBL_MAX;
    int minIdx;
    for( int k = 0; k < M; ++k){
      double r = ( input[n][0] - mean[k][0]) * ( input[n][0] - mean[k][0]);
      r +=( input[n][1] - mean[k][1]) * ( input[n][1] - mean[k][1]);
      if( min > r ){ min = r; minIdx = k;}
    }
    
    cov[minIdx][0] += sqrt(min);
    cov[minIdx][3] += sqrt(min);
    count[minIdx] ++;

    if( n % (W*10) == 0 )
      printf("%6d outof %6d done\n", n, N);

  }

  for( int k = 0; k < M; ++k){
    if(count[k] == 0 || cov[k][0] == 0.0){
      cov[k][0] = 0.0001;
      cov[k][3] = 0.0001;
    }else{
      double val = cov[k][0] / (double)count[k];
      cov[k][0] = val*val;
      cov[k][3] = val*val;
    }
    printf("%04d cov : %lf %lf %lf %lf count = %d\n", k,cov[k][0],cov[k][1],cov[k][2],cov[k][3], count[k]);
  }


  gamma = (double(*)[M])malloc( sizeof(double[M])*N);

}



void cleanup(){}
void saveCurrentImage( char append[] ){

  IplImage* img = cvCreateImage( cvSize(W,H), IPL_DEPTH_8U, 1);
  IplImage* distrub = cvCreateImage( cvSize(W,H), IPL_DEPTH_8U, 1);
  char filename[256];
  cvSetZero( img );
  
  for(int k = 0; k < M; ++k){
    cvSetZero(distrub);
    CvPoint pt = cvPoint( 1.0*mean[k][0], 1.0*mean[k][1] );
    CvScalar color = cvScalarAll( LEVEL  );
    CvSize axis = cvSize( 2.0*sqrt(cov[k][0]), 2.0*sqrt(cov[k][3]));
    double angle = acos(cov[k][1] / sqrt( cov[k][0]*cov[k][3] ) ) * 180.0 / M_PI;
    printf("k = %4d, ", k);
    printf("pt = %03d, %03d ", pt.x, pt.y);
    printf("Color = %lf ", color.val[0]);
    printf("axis = %03d, %03d ", axis.width, axis.height);
    printf("angle = %lf\n", angle);
    if( axis.width > W/8 ) axis.width = 1.0;
    if( axis.height > H/8 ) axis.height = 1.0;
    cvEllipse( distrub, pt, axis, angle, 0, 360.0, color, -1, 8, 0);

    for(int h = 0; h < H; ++h){
      for( int w = 0; w < W; ++w){
	CV_IMAGE_ELEM( img, uchar, h, w) += CV_IMAGE_ELEM( distrub, uchar, h, w);
      }
    }

  }
  
  sprintf( filename, "imgs/result-%s.png",append);
  cvSaveImage( filename, img, NULL);


  cvSetZero(img);
  for( int n = 0; n < N; ++n ){
    int h = input[n][1] * H;
    int w = input[n][0] * W;
    CV_IMAGE_ELEM( img, uchar, h, w ) += 256.0/LEVEL;
  }
  cvSaveImage("imgs/input.png", img, NULL);

}

void saveGaussianImage( char append[] )
{
  IplImage *dst = cvCreateImage( cvSize(W, H), IPL_DEPTH_64F, 1);
  cvSetZero( dst );
  for( int k = 0; k < M; ++k){
    for( int h = 0; h < H; ++h){
      for( int w = 0; w < W; ++w){
	double vec[2] = {w, h };
	CV_IMAGE_ELEM( dst, double, h, w )
	  += noramlDistr( vec, mean[k], cov[k] );
      }
    }
  }

  double min, max;
  cvMinMaxLoc( dst, &min, &max, NULL, NULL, NULL);
  IplImage *img = cvCreateImage( cvSize(W,H), IPL_DEPTH_8U, 1);
  cvConvertScale( dst, img, 255.0/(max-min), -min*255.0/(max-min));
  char filename[256];
  sprintf( filename, "imgs/gaussian-%s.png", append);
  cvSaveImage( filename, img, NULL );
  cvReleaseImage( &img );
  cvReleaseImage( &dst );

}


double noramlDistr( const double x[DIM], const double m[DIM], const double c[DIM*DIM] )
{
  
  double det = c[0]*c[3] - c[1]*c[2];
  double v[2];
  v[0] = x[0] - m[0];
  v[1] = x[1] - m[1];

  double u[2];
  u[0] = v[0] * c[3] - v[1] * c[2];
  u[1] = -v[0]* c[1] + v[1] * c[0];

  double e = exp( -(u[0]*v[0] + u[1]*v[1]) / ( 2.0 * det ));
  return e / (2.0 * M_PI * sqrt(det));
}

void updateGamma(){
  
  double w = 1.0 / (double)M;

  for(int n = 0;n < N; ++n){
    double sum = 0.0;
    double tmp[2] = { input[n][0], input[n][1] };
    for(int j = 0; j < M; ++j)
      sum += w * noramlDistr( tmp, mean[j], cov[j] );

    for(int k = 0; k < M; ++k)
      gamma[n][k] = w * noramlDistr(tmp, mean[k], cov[k] ) / sum;

    if( n % (W*H) == 0 )
      printf("n = %d : out of %d\n", n, N);

  }
}



void updateOther(){
  int n, k;
  
  double Nk[M];
  for( k = 0; k < M; ++k) Nk[k] = 0.0;
  
  for( n = 0; n < N; ++n){
    for( k = 0; k < M; ++k){ Nk[k] += gamma[n][k]; }
    if( n %(N/1000) == 0 ) printf("calc Nk %d out of %d done\n", n, N);
  }


  // update mean
  for( k = 0; k < M; ++k){
    mean[k][0] = mean[k][1] = 0.0;
  }

  for( n = 0; n < N ; ++n){
    for( k = 0; k < M; ++k ){
      mean[k][0] += input[n][0] * gamma[n][k];
      mean[k][1] += input[n][1] * gamma[n][k];
    }

  }

  for( k = 0; k < M; ++k){
    if( Nk[k] ==  0.0 ) continue;
    mean[k][0] /= Nk[k];
    mean[k][1] /= Nk[k];
    printf("mean %03d : %lf, %lf\n", k, mean[k][0], mean[k][1]);
  }

  // update cov
  for( k = 0; k < M; ++k){
    cov[k][0] = cov[k][1] = cov[k][2] = cov[k][3] = 0.0;
  }
  for( n = 0; n < N ; ++n){
    for( k = 0; k < M; ++k){
      double v[2];
      v[0] = input[n][0] - mean[k][0];
      v[1] = input[n][1] - mean[k][1];
      cov[k][0] += v[0]*v[0] * gamma[n][k];
      cov[k][1] += v[0]*v[1] * gamma[n][k];
      cov[k][2] += v[1]*v[0] * gamma[n][k];
      cov[k][3] += v[1]*v[1] * gamma[n][k];
    }
    if( n%(W*H) == 0 )
      printf("%d out of %d done\n", n, N);
    
  }   

  for( k = 0; k < M; ++k){
    if( Nk[k] == 0.0 ) continue;
    cov[k][0] /= Nk[k];
    cov[k][1] /= Nk[k];
    cov[k][2] /= Nk[k];
    cov[k][3] /= Nk[k];
    printf("cov %03d : %lf, %lf, %lf, %lf\n", k, cov[k][0], cov[k][1], cov[k][2], cov[k][3]);
  }


}