#include <iostream>
#include <opencv2\opencv.hpp>

using namespace std;
using namespace cv;

Mat srcImage, srcImage_, maskImage;
Mat maskWaterShed;  // watershed()函式的引數
Point clickPoint;	// 滑鼠點下去的位置

void on_Mouse(int event, int x, int y, int flags, void*);
void helpText();

int main(int argc, char** argv)
{
	/* 操作提示 */
	helpText();

	srcImage = imread("fly.jpg");
	srcImage_ = srcImage.clone();  // 程式中srcImage會被改變，所以這裡做備份
	maskImage = Mat(srcImage.size(), CV_8UC1);  // 掩模，在上面做標記，然後傳給findContours
	maskImage = Scalar::all(0);

	int areaCount = 1;  // 計數，在按【0】時繪製每個區域

	imshow("在影象中做標記", srcImage);

	setMouseCallback("在影象中做標記", on_Mouse, 0);

	while (true)
	{
		int c = waitKey(0);

		if ((char)c == 27)	// 按【ESC】鍵退出
			break;

		if ((char)c == '2')  // 按【2】恢復原圖
		{
			maskImage = Scalar::all(0);
			srcImage = srcImage_.clone();
			imshow("在影象中做標記", srcImage);
		}

		if ((char)c == '1')  // 按【1】處理圖片
		{
			vector<vector<Point>> contours;
			vector<Vec4i> hierarchy;

			findContours(maskImage, contours, hierarchy, RETR_CCOMP, CHAIN_APPROX_SIMPLE);

			if (contours.size() == 0)  // 如果沒有做標記，即沒有輪廓，則退出該if語句
				break;
			cout << contours.size() << "個輪廓" << endl;

			maskWaterShed = Mat(maskImage.size(), CV_32S);
			maskWaterShed = Scalar::all(0);

			/* 在maskWaterShed上繪製輪廓 */
			for (int index = 0; index < contours.size(); index++)
				drawContours(maskWaterShed, contours, index, Scalar::all(index + 1), -1, 8, hierarchy, INT_MAX);
			
			/* 如果imshow這個maskWaterShed，我們會發現它是一片黑，原因是在上面我們只給它賦了1,2,3這樣的值，通過程式碼80行的處理我們才能清楚的看出結果 */
			watershed(srcImage_, maskWaterShed);  // 註釋一

			vector<Vec3b> colorTab;  // 隨機生成幾種顏色
			for (int i = 0; i < contours.size(); i++)
			{
				int b = theRNG().uniform(0, 255);
				int g = theRNG().uniform(0, 255);
				int r = theRNG().uniform(0, 255);

				colorTab.push_back(Vec3b((uchar)b, (uchar)g, (uchar)r));
			}

			Mat resImage = Mat(srcImage.size(), CV_8UC3);  // 宣告一個最後要顯示的影象
			for (int i = 0; i < maskImage.rows; i++)
			{
				for (int j = 0; j < maskImage.cols; j++)
				{	// 根據經過watershed處理過的maskWaterShed來繪製每個區域的顏色
					int index = maskWaterShed.at<int>(i, j);  // 這裡的maskWaterShed是經過watershed處理的
					if (index == -1)  // 區域間的值被置為-1（邊界）
						resImage.at<Vec3b>(i, j) = Vec3b(255, 255, 255);
					else if (index <= 0 || index > contours.size())  // 沒有標記清楚的區域被置為0
						resImage.at<Vec3b>(i, j) = Vec3b(0, 0, 0);
					else  // 其他每個區域的值保持不變：1，2，...，contours.size()
						resImage.at<Vec3b>(i, j) = colorTab[index - 1];  // 然後把這些區域繪製成不同顏色
				}
			}
			imshow("resImage", resImage);
			addWeighted(resImage, 0.3, srcImage_, 0.7, 0, resImage);
			imshow("分水嶺結果", resImage);
		}

		if ((char)c == '0')  // 多次點按【0】依次顯示每個被分割的區域，需要先按【1】處理影象
		{
			Mat resImage = Mat(srcImage.size(), CV_8UC3);  // 宣告一個最後要顯示的影象
			for (int i = 0; i < maskImage.rows; i++)
			{
				for (int j = 0; j < maskImage.cols; j++)
				{
					int index = maskWaterShed.at<int>(i, j);
					if (index == areaCount)
						resImage.at<Vec3b>(i, j) = srcImage_.at<Vec3b>(i, j);
					else
						resImage.at<Vec3b>(i, j) = Vec3b(0, 0, 0);
				}
			}
			imshow("分水嶺結果", resImage);
			areaCount++;
			if (areaCount == 4)
				areaCount = 1;
		}
	}

	return 0;
}

void on_Mouse(int event, int x, int y, int flags, void*)
{
	// 如果滑鼠不在視窗中則返回
	if (x < 0 || x >= srcImage.cols || y < 0 || y >= srcImage.rows)
		return;

	// 如果滑鼠左鍵被按下，獲取滑鼠當前位置；當滑鼠左鍵按下並且移動時，繪製白線；
	if (event == EVENT_LBUTTONDOWN)
	{
		clickPoint = Point(x, y);
	}		
	else if (event == EVENT_MOUSEMOVE && (flags & EVENT_FLAG_LBUTTON))
	{
		Point point(x, y);
		line(maskImage, clickPoint, point, Scalar::all(255), 5, 8, 0);
		line(srcImage, clickPoint, point, Scalar::all(255), 5, 8, 0);
		clickPoint = point;
		imshow("在影象中做標記", srcImage);
	}
}

void helpText()
{
	cout << "先用滑鼠在圖片視窗中標記出大致的區域" << endl;
	cout << "如果想把圖片分割為N個區域，就要做N個標記" << endl;
	cout << "鍵盤按鍵【1】	- 執行的分水嶺分割演算法" << endl;
	cout << "鍵盤按鍵【2】	- 恢復原始圖片" << endl;
	cout << "鍵盤按鍵【0】	- 依次分割每個區域（必須先按【1】）" << endl;
	cout << "鍵盤按鍵【ESC】	- 退出程式" << endl << endl;
}

/* 註釋一：watershed(srcImage_, maskWaterShed);
 * 注意它的兩個引數
 * srcImage_是沒做任何修改的原圖，CV_8UC3型別
 * maskWaterShed宣告為CV_32S型別（32位單通道），且全部元素為0
 *		然後作為drawContours的第一個引數傳入，在上面繪製輪廓
 *		最後作為watershed的引數
 * 另外，watershed的第二個引數maskWaterShed是InputOutputArray型別
 *		即作為輸入，也作為輸出儲存函式呼叫的結果
 */