// face-predict.cpp : 定义控制台应用程序的入口点。
//

#include "stdafx.h"
#include "opencv.hpp"
#include "face.hpp"
#include <vector>

using namespace std;
using namespace cv;
using namespace face;

static int load_data(vector<Mat>& data, vector<int>& label);
static void train();
static void test();
static void generate_facedata();
static void realtime_detection_face();

vector<Mat> g_Faces;		// 存储人脸数据
vector<int> g_Labels;		// 存储人脸对应的标签

//load_data(g_Faces, g_Labels);
//imshow("", g_Faces.at(0));
//waitKey(0);


/* 

1. 人脸检测与数据保存：generate_facedata();

2. 训练人脸识别模型：load_data(g_Faces, g_Labels); train();

3. 测试人脸识别模型： load_data(g_Faces, g_Labels); test();

4. 人脸实时识别：realtime_detection_face();

*/

int main()
{
	generate_facedata();
	realtime_detection_face();
	load_data(g_Faces, g_Labels);
	train();
	test();
	system("pause");
    return 0;
}

//////////////////////////////////////////////////////////////////////////
//		测试人脸识别， 输出测试精度
//
void test()
{
	printf("------ Eigen Face Recognizer Test ! \r\n");
	printf("------ create model...\r\n");
	Ptr<FaceRecognizer> model = createEigenFaceRecognizer();

	printf("------ load model... \r\n");
	model->load("model.xml");

	printf("------ start predict...\r\n");
	int count = 0;
	for (int i = 1; i != g_Faces.size(); i++)
	{
		int label = -1;
		double accuracy = 0;
		model->predict(g_Faces.at(i), label, accuracy);
		if (label == g_Labels.at(i))
			count++;
	}
	double accuracy = ((double)count) / ((double)g_Faces.size());
	printf("------ predict end, accuracy = %f. \r\n", accuracy);
}

//////////////////////////////////////////////////////////////////////////
//    训练人脸识别模型，并保存模型
//
void train()
{
	printf("------ Eigin Face Recognizer Train ! \r\n");
	printf("------ create model...\r\n");
	Ptr<FaceRecognizer> model = createEigenFaceRecognizer();

	printf("------ start train, it will take some time, please wait... \r\n");
	model->train(g_Faces, g_Labels);

	printf("------ train end, save model.xml. \r\n");
	model->save("model.xml");
}

//////////////////////////////////////////////////////////////////////////
//    加载本地人脸数据集
//
int load_data(vector<Mat>& data, vector<int>& label)
{
	char filename[128];
	data.clear();
	label.clear();

	for (int i = 1; i != 41; i++)
	{
		for (int k = 1; k != 11; k++)
		{
			sprintf_s(filename, "./data/s%d/%d.pgm", i, k);
			Mat image = imread(filename, IMREAD_GRAYSCALE);
			data.push_back(image);
			label.push_back(i);
		}
	}

	for (int i = 1; i != 11; i++)
	{
		sprintf_s(filename, "./data/s%d/%d.jpg", 41, i);
		Mat image = imread(filename, IMREAD_GRAYSCALE);
		data.push_back(image);
		label.push_back(41);
	}

	printf("------ load data success, size = %ld \r\n", data.size());
	return 0;
}

//////////////////////////////////////////////////////////////////////////
//  打开摄像头检测人脸， 按下S或s保存人脸数据
//
static void generate_facedata()
{
	VideoCapture cap(0);
	CascadeClassifier ccface;

	// 加载人脸分类器
	if (!ccface.load("haarcascade_frontalface_alt.xml"))
	{
		fprintf(stderr, "------ ERROR: load model failed. %s \r\n", __FUNCTION__);
		return;
	}

	// 检测摄像头是否打开成功
	if (!cap.isOpened())
	{
		fprintf(stderr, "------ ERROR: open video capture failed. %s \r\n", __FUNCTION__);
		return;
	}

	while (true)
	{
		// 用于存储当前一帧的图像
		Mat frame;
		Mat show;

		// 从摄像头获取一帧图像
		cap >> frame;
		show = frame;

		// 转换为灰度图
		cvtColor(frame, frame, CV_BGR2GRAY);

		// 检测人脸
		vector<Rect> faces;
		ccface.detectMultiScale(frame, faces);

		Mat face;
		if (!faces.empty())
		{
			// 框出人脸
			for (int i = 0; i != faces.size(); i++)
			{
				rectangle(show, faces.at(i), Scalar(0, 255, 0));
			}
			face = frame(faces.at(0));
		}

		imshow("", show);

		char key = waitKey(32);
		if(key == 'q' || key == 'Q')
			break;

		if (key == 's' || key == 'S')
		{
			if (!face.empty())
			{
				static int  i = 0;
				char fmt[32];
				resize(face, face, Size(92, 112));
				sprintf_s(fmt, "%d.jpg", i);
				imwrite(fmt, face);
				i++;
			}
		}
	}

	// 关闭摄像头并销毁所有显示窗口
	cap.release();
	destroyAllWindows();
}


void realtime_detection_face()
{
	VideoCapture cap(0);
	CascadeClassifier ccface;

	// 加载人脸检测器
	if (!ccface.load("haarcascade_frontalface_alt.xml"))
	{
		fprintf(stderr, "------ ERROR: load model failed. %s \r\n", __FUNCTION__);
		return;
	}

	// 加载人脸识别器
	Ptr<FaceRecognizer> model = createEigenFaceRecognizer();
	model->load("model.xml");


	// 检测摄像头是否打开成功
	if (!cap.isOpened())
	{
		fprintf(stderr, "------ ERROR: open video capture failed. %s \r\n", __FUNCTION__);
		return;
	}

	while (true)
	{
		// 用于存储当前一帧的图像
		Mat frame;
		Mat show;

		// 从摄像头获取一帧图像
		cap >> frame;
		show = frame;

		// 转换为灰度图
		cvtColor(frame, frame, CV_BGR2GRAY);

		// 检测人脸
		vector<Rect> faces;
		ccface.detectMultiScale(frame, faces);		// 人脸多尺度检测

		if (!faces.empty())
		{
			// 框出人脸
			for (int i = 0; i != faces.size(); i++)
			{
				rectangle(show, faces.at(i), Scalar(0, 255, 0));
			}
			Mat face = frame(faces.at(0));
			resize(face, face, Size(92, 112));

			int label = -1;
			double accuracy = 0.0;
			model->predict(face, label, accuracy);
			printf("------ label:%d        accuracy: %0.2f \r\n", label, accuracy);
		}

		imshow("", show);
		char key = waitKey(30);
		if (key == 'q' || key == 'Q')
			break;
	}

	cap.release();
	destroyAllWindows();
}