#include <opencv/highgui.h>
#include <opencv2/opencv.hpp>
#include <cmath>

//OSTU算法
int OSTU(cv::Mat& src)
{
	int threshold=0;	//阈值
	float sigma_temp=0.0f;	
	float sigma_max = 0.0f;	
	int pixel_num = src.rows * src.cols;	//像素点个数
	int pixel_count[256];	//各像素个数
	float pixel_pro[256];	//各像素概率

	//初始化
	for(int i=0; i<256; i++)
	{
		pixel_count[i] = 0;
	}

	// 统计像素个数
	for(int i=0; i < src.rows; i++)
	{	
		uchar* data = src.ptr<uchar>(i);

		for(int j=0; j < src.cols; j++)
		{
			pixel_count[data[j]]++;
		}
	}
		
	//计算像素概率
	for(int i=0; i<256; i++)
	{
		pixel_pro[i] = (float)(pixel_count[i]) / (float)(pixel_num);
	}

	//计算阈值
	for(int i=0; i<256; i++)
	{
		float w0=0, w1=0, u0=0, u1=0, u=0;
		//遍历像素
		for(int j=0; j<256; j++)
		{	
			//背景像素
			if(j < i)
			{
				w0 += pixel_pro[j];
				u0 += pixel_pro[j] * j;	
			}
			else	//前景像素
			{
				u1 += pixel_pro[j] * j;
			}
		}

		//全图像素均值
		//u = w0 * u0 + (1-w0) * u1;
		
		//方差
		sigma_temp = w0 * (1-w0) * pow((u0 - u1), 2);
		
		//更新阈值
		if(sigma_temp > sigma_max)
		{
			sigma_max = sigma_temp;
			threshold = i;
		}
	}

	return threshold;
}

//二值化图像
void BinaryImage(cv::Mat& src, cv::Mat& des, int threshold)
{
	for(int i=0; i < src.rows; i++)
	{	
		uchar* src_data = src.ptr<uchar>(i);
		uchar* des_data = des.ptr<uchar>(i);

		for(int j=0; j < src.cols; j++)
		{
			if(src_data[j] < threshold)
			{
				des_data[j] = 0;
			}
			else
			{
				des_data[j] = 255;
			}
		}
	}
}

int main()
{	
	cv::Mat src, gray;

	cv::namedWindow("src", cv::WINDOW_NORMAL);
	cv::namedWindow("gray", cv::WINDOW_NORMAL);
	cv::namedWindow("binary", cv::WINDOW_NORMAL);


	src = cv::imread("../wallhaven-698240.jpg", cv::IMREAD_COLOR);
	cv::cvtColor(src, gray, cv::COLOR_RGB2GRAY);	//转换成灰度图

	int threshold = OSTU(gray);		//计算阈值

	cv::Mat binary = cv::Mat(gray.rows, gray.cols, CV_8U);		//初始化二值化图像
	BinaryImage(gray, binary, threshold);

	printf("threshold: %d\n", threshold);

	//显示RGB图像, 灰度图像, 二值化图像
	cv::imshow("src", src);
	cv::imshow("gray", gray);
	cv::imshow("binary", binary);

	cv::waitKey(0);

	return 0;
}