//---------------------------------------------------------
// 概要      : 正規分布モデルを用いた背景差分法
// File Name : sgm.cpp
// Library   : OpenCV for MS-Windows 1.0
//---------------------------------------------------------

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

#define	NUM_OF_BACKGROUND_FRAMES	50	//	背景モデルを生成するのに使用する画像の枚数
#define	THRESHOLD_COEFFICIENT		5.0	//	閾値を決める際に使用する標準偏差にかける数

#define	CONSTANT_THRESHOLD	20	//	画像で1つの閾値を使う際の閾値
#define CONSTANT_MODE		0	//	画像全体で1つの閾値を用いるフラグ値
#define DYNAMIC_MODE		1	//	画素ごとに閾値を設けるフラグ値

CvCapture *capture = NULL;	//	カメラキャプチャ用の構造体

IplImage *backgroundAverageImage = NULL;	//	背景の平均値保存用IplImage
IplImage *backgroundThresholdImage = NULL;	//	背景の閾値保存用IplImage

//
//	画像を上下反転して表示する
//
//	引数:
//		windowName : 画像を表示するウィンドウの名前
//		image      : 上下反転して表示する画像用IplImage
//
void showFlipImage( char *windowName, IplImage *image ) {
	if ( image->origin == 0 ) {
		cvFlip( image, image, 0 );
		cvShowImage( windowName, image );
		cvFlip( image, image, 0 );
	}
}

//
//	背景モデルを初期化する
//
//	引数:
//		num  : 背景モデルを生成するのに使用する画像の枚数
//		size : 画像サイズ
//
void initializeBackgroundModel( int num, CvSize size ){
	int i;

	// 以前の背景情報があれば破棄
	if( backgroundAverageImage != NULL ){
		cvReleaseImage( &backgroundAverageImage );
	}
	if( backgroundThresholdImage != NULL ){
		cvReleaseImage( &backgroundThresholdImage );
	}

	//	画像情報蓄積用バッファを確保する
	IplImage *acc = cvCreateImage( size, IPL_DEPTH_32F, 3 );
	IplImage *acc2 = cvCreateImage( size, IPL_DEPTH_32F, 3 );
	
	//	画像の初期化を行う
	cvSetZero( acc );
	cvSetZero( acc2 );

	//	画像情報の蓄積
	printf( "背景取得中...\n" );
	IplImage *frameImage;
	for( i = 0; i < num; i++ ){
		frameImage = cvQueryFrame( capture );
		cvAcc( frameImage, acc );
		cvSquareAcc( frameImage, acc2 );
		printf( "%d 枚中 %d 枚目\n", num, i + 1 );
	}
	printf( "背景取得完了\n" );

	//	cvAddS, cvSubS はあるが cvMulS はないので、cvConvertScale を使う
	cvConvertScale( acc, acc, 1.0 / num );		// 平均
	cvConvertScale( acc2, acc2, 1.0 / num );	// 二乗和の平均

	//	平均が求まったので backgroundAverageImage に格納する。
	backgroundAverageImage = cvCreateImage( size, IPL_DEPTH_8U, 3 );
	cvConvert( acc, backgroundAverageImage );

	//	分散を計算する
	IplImage *dispersion = cvCreateImage( size, IPL_DEPTH_32F, 3 );
	cvMul( acc, acc, acc );
	cvSub( acc2, acc, dispersion );

	//	標準偏差を計算する
	IplImage *sd = cvCreateImage( size, IPL_DEPTH_32F, 3 );
	cvPow( dispersion, sd, 0.5 );

	//	閾値を計算する
	backgroundThresholdImage = cvCreateImage( size, IPL_DEPTH_8U, 3 );
	cvConvertScale( sd, backgroundThresholdImage, THRESHOLD_COEFFICIENT );

	//	メモリを解放する
	cvReleaseImage( &acc );
	cvReleaseImage( &acc2 );
	cvReleaseImage( &dispersion );
	cvReleaseImage( &sd );
}

int main( void ){ 
	int key;		//	キー入力用の変数
	int	mode = 0;	// 0: 画素ごとに異なる閾値 / 1: 画像全体で一つの閾値

	char *windowNameCapture = "Capture";		//	キャプチャした画像を表示するウィンドウの名前
	char *windowNameResult  = "Result";			//	背景差分結果を表示するウィンドウの名前
	char *windowNameBackground = "Background";	//	背景画像を表示するウィンドウの名前
	char *windowNameThreshold = "Threshold";	//	閾値を表示するウィンドウの名前

	IplImage *frameImage;		//	キャプチャ画像用IplImage
	IplImage *differenceImage;	//	差分画像用IplImage
	IplImage *resultImage;		//	結果画像用IplImage

	char *mode_str[2] = {
		"画素単位",
		"画像で一つ"
	};

	//	カメラを初期化する
	if ( ( capture = cvCreateCameraCapture( -1 ) ) == NULL ){
		//	カメラが見つからなかった場合
		printf( "カメラが見つかりません\n" );
		return -1;
	}
	
	//	ウィンドウを生成する
	cvNamedWindow( windowNameCapture, CV_WINDOW_AUTOSIZE );
	cvNamedWindow( windowNameResult, CV_WINDOW_AUTOSIZE );
	cvNamedWindow( windowNameBackground, CV_WINDOW_AUTOSIZE );
	cvNamedWindow( windowNameThreshold, CV_WINDOW_AUTOSIZE );

	//	カメラの現フレーム画像をキャプチャする
	frameImage = cvQueryFrame( capture );

	//	画像サイズを保存する
	CvSize imageSize = cvSize( frameImage->width, frameImage->height );

	//	2値化結果格納用画像を確保する
	differenceImage = cvCloneImage( frameImage );
	resultImage = cvCreateImage( imageSize, IPL_DEPTH_8U, 1 );
								 
	//	背景モデルを初期化する
	initializeBackgroundModel( NUM_OF_BACKGROUND_FRAMES, imageSize );

	while ( 1 ){
		frameImage = cvQueryFrame( capture );

		//	現在の背景との差の絶対値を成分ごとに取る
		cvAbsDiff( frameImage, backgroundAverageImage, differenceImage );

		//	Sub はマイナスになったら0に切り詰めてくれる
		if( mode == 0 ){
			cvSub( differenceImage, backgroundThresholdImage, differenceImage );
		} else{
			cvSubS( differenceImage, cvScalarAll( CONSTANT_THRESHOLD ), differenceImage );
		}

		//	differenceImage の要素が1つでも0以上だったら前景
		cvCvtColor( differenceImage, resultImage, CV_BGR2GRAY );
        cvThreshold( resultImage, resultImage, 0, 255, CV_THRESH_BINARY );
		
		//	メディアンフィルタでノイズを除去する
		cvSmooth( resultImage, resultImage, CV_MEDIAN );

		//	画像を表示する
		cvShowImage( windowNameCapture, frameImage );
		showFlipImage( windowNameResult, resultImage );
		showFlipImage( windowNameBackground, backgroundAverageImage );

		IplImage *tmp = cvCloneImage( backgroundThresholdImage );
		cvConvertScale( tmp, tmp, 3 );
		showFlipImage( windowNameThreshold, tmp );
		cvReleaseImage( &tmp );

		//	キー入力判定
		key = cvWaitKey( 1 );
		if( key == 'q' ){
			//	'q'キーが押されたらループを抜ける
			break;
		} else if( key == 'b' ){
			//	'b'キーが押されたらその時点での画像を背景画像とする
			initializeBackgroundModel( NUM_OF_BACKGROUND_FRAMES, imageSize );
			printf( "背景情報更新\n" );
		} else if( key == 'm' ){
			//	'm'キーが押されたら閾値の設定方法を変更する
			mode = 1 - mode;
			printf( "閾値: %s\n", mode_str[mode] );
		}
	}
	
	//	キャプチャを解放する
	cvReleaseCapture( &capture );
	//	メモリを解放する
	cvReleaseImage( &differenceImage );
	cvReleaseImage( &resultImage );
	cvReleaseImage( &backgroundAverageImage );
	cvReleaseImage( &backgroundThresholdImage );
	//	ウィンドウを破棄する
	cvDestroyWindow( windowNameCapture );
	cvDestroyWindow( windowNameResult );
	cvDestroyWindow( windowNameBackground );
	cvDestroyWindow( windowNameThreshold );

	return 0;
}