art.cpp

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

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

#include <iostream>
#include <queue>

using namespace std;

#define SELECTED 255
#define ABORT 100
#define UNCHECKED 0

int threshhold = 10;
int patch_radius;

IplImage *input = NULL;
IplImage *mask = NULL;
IplImage *res = NULL;

int width;
int height;
int widthstep;
int gwidthstep;
int nch;

CvPoint center_pos;

struct Color{
	int r;
	int g;
	int b;
};

queue <CvPoint> p; // 队列p中所存的是当前待传播的点的信息

int colordiff(Color *cur, Color *center)
{
	double diff;
	diff = (cur->b - center->b) * (cur->b - center->b)
		+ (cur->g - center->g) * (cur->g - center->g)
		+ (cur->r - center->r) * (cur->r - center->r);

	diff = sqrt(diff);

	if (diff <= threshhold)
		return 1;
	else
		return 0;
}

// 以某一点为中心，周围3*3的块进行判断颜色相似性
void point_grow(int i, int j, Color center, int radius)
{
	Color cur;

	int st_row = i - radius, ed_row = i + radius;
	int st_col = j - radius, ed_col = j + radius;

	if (st_row <= center_pos.y - patch_radius)
		st_row = center_pos.y - patch_radius;
	if (ed_row > center_pos.y + patch_radius)
		ed_row = center_pos.y + patch_radius;
	if (st_col <= center_pos.x - patch_radius)
		st_col = center_pos.x - patch_radius;
	if (ed_col > center_pos.x + patch_radius)
		ed_col = center_pos.x + patch_radius;

	if (st_row <= 0)
		st_row = 1;
	if (ed_row > height)
		ed_row = height;
	if (st_col <= 0)
		st_col = 1;
	if (ed_col > width)
		ed_col = width;

	int diff_flag;

	CvPoint cur_pos;

	for (int m = st_row; m <= ed_row; m++)
	{
		for (int n = st_col; n <= ed_col; n++)
		{
			if (!((m == i) && (n == j)))
			{
				int mask_value = *(uchar *)(mask->imageData + (m - 1) * gwidthstep + (n - 1));
				if ((mask_value != ABORT) && (mask_value != SELECTED))
				{
					cur.b = *(uchar *)(input->imageData + (m - 1) * widthstep + (n - 1) * nch);
					cur.g = *(uchar *)(input->imageData + (m - 1) * widthstep + (n - 1) * nch + 1);
					cur.r = *(uchar *)(input->imageData + (m - 1) * widthstep + (n - 1) * nch + 2);

					diff_flag = colordiff(&cur, &center);

					if (diff_flag == 1) // 说明颜色比较接近，应当被选择
					{
						*(uchar *)(mask->imageData + (m - 1) * gwidthstep + (n - 1)) = SELECTED;
						cur_pos.x = j;
						cur_pos.y = i;

						p.push(cur_pos);

					}
					else // 说明颜色相差较大，不应该被选择，标记为128，防止以后折回生长
					{
						*(uchar *)(mask->imageData + (m - 1) * gwidthstep + (n - 1)) = ABORT;
					} // End if diff_flag
				}
			}
		}
	} // End for loop m, n
}

// Patch中心点为当前要计算DarkChannel的patch的中心点
void rgrow(IplImage *res, IplImage *input, IplImage *mask, int threshhold, int patch_radius)
{
	int radius; // 当前传播的窗口半径
	radius = 1;

	Color center;

	for (int i = 1; i <= height; i++)
	{
		for (int j = 1; j <= width; j++)
		{
			center.b = *(uchar *)(input->imageData + (i - 1) * widthstep + (j - 1) * nch);
			center.g = *(uchar *)(input->imageData + (i - 1) * widthstep + (j - 1) * nch + 1);
			center.r = *(uchar *)(input->imageData + (i - 1) * widthstep + (j - 1) * nch + 2);

			//CvPoint center_pos;
			center_pos.x = j;
			center_pos.y = i;

			p.push(center_pos);

			CvPoint cur_pos;

			while (p.empty() != true) // 即队列非空
			{
				cur_pos.x = p.front().x;
				cur_pos.y = p.front().y;

				p.pop();

				point_grow(cur_pos.y, cur_pos.x, center, radius);
			}
		}
	} // End for loop i, j
}

int main(int argc, char **argv)
{
	if (argc < 2)
		input = cvLoadImage("test.jpg");
	else
		input = cvLoadImage(argv[1]);

	width = input->width;
	height = input->height;
	widthstep = input->widthStep;
	nch = input->nChannels;

	cvNamedWindow("Input");
	cvShowImage("Input", input);
	cvWaitKey();

	cvNamedWindow("Mask");
	cvResizeWindow("Mask", width, height);

	mask = cvCreateImage(cvSize(width, height), IPL_DEPTH_8U, 1);
	cvZero(mask);
	gwidthstep = mask->widthStep;

	res = cvCreateImage(cvSize(width, height), IPL_DEPTH_8U, 3);

	patch_radius = (int)(MIN(width, height) * 0.02);

	clock_t begin, end;
	double time_spent;
	begin = clock();
	rgrow(res, input, mask, threshhold, patch_radius);
	end = clock();
	time_spent = (double)(end - begin) / CLOCKS_PER_SEC;
	printf("Region Grow %.0lfs\n", time_spent);

	cvNamedWindow("Mask");
	cvShowImage("Mask", mask);

	cvWaitKey();
	
	cvDestroyAllWindows();
	cvReleaseImage(&input);
	cvReleaseImage(&mask);
	cvReleaseImage(&res);

	return 0;
}