#include "FaceQuality.h"

FaceQuality* FaceQuality::m_instance;

const static unsigned char _rfhgdfger_pose_onnx_model[] = {
    #include "algorithm_module/models/pose_onnx.dat"
};

FaceQuality* FaceQuality::GetInstance()
{
	if (m_instance == nullptr)
	{
		m_instance = new FaceQuality;
	}
	return m_instance;
}

FaceQuality::FaceQuality()
{
	m_pose_net = cv::dnn::readNetFromONNX((const char*)_rfhgdfger_pose_onnx_model, sizeof(_rfhgdfger_pose_onnx_model));
	m_input_size.width = 64;
	m_input_size.height = 64;
	m_mean_vals = cv::Scalar(127.5, 127.5, 127.5);
}

int FaceQuality::DetectPose(const cv::Mat &img, cv::Rect &bbox, cv::Mat &pose)
{
	std::lock_guard<std::mutex> lock(m_mt);
    cv::Rect fbox = bbox;;
	if(fbox.x < 0) fbox.x = 0;
	if(fbox.y < 0) fbox.y = 0;
	if(fbox.width  + fbox.x > img.cols) fbox.width  = img.cols - fbox.x;
	if(fbox.height + fbox.y > img.rows) fbox.height = img.rows - fbox.y;

    if(bbox.width > bbox.height)
    {
        fbox.y = fbox.y - (fbox.width - fbox.height) / 2;
        if(fbox.y < 0) fbox.y = 0;

        fbox.height = fbox.height + (fbox.width - fbox.height);
        if(fbox.height + fbox.y > img.rows) fbox.height = img.rows - fbox.y;

    }
    else if(bbox.width < bbox.height)
    {
        fbox.x = fbox.x - (fbox.height - fbox.width) / 2;
        if(fbox.x < 0) fbox.x = 0;

        fbox.width = fbox.width + (fbox.height - fbox.width);
        if(fbox.width + fbox.x > img.cols) fbox.width = img.cols - fbox.x;
    }
    cv::Mat face = img(fbox).clone();
	cv::cvtColor(face, face, cv::COLOR_BGR2RGB);
    cv::Mat inputBlob = cv::dnn::blobFromImage(face, 0.007843, m_input_size, m_mean_vals, false, false);
    m_pose_net.setInput(inputBlob, "data");
    pose = m_pose_net.forward("fc1") * 100;

	return 0;
}

int FaceQuality::DetectQuality(const cv::Mat &img, cv::Rect& face_rect)
{
	cv::Mat pose;
	DetectPose(img, face_rect, pose);
	float angle_yaw = pose.at<float>(0, 0);
	float angle_pitch = pose.at<float>(0, 1);
	float angle_roll = pose.at<float>(0, 2);

	if (angle_yaw > 90) angle_yaw = 90;
	if (angle_pitch > 90) angle_pitch = 90;

	float yaw, pitch;
	yaw = 100 - abs(angle_yaw);
	pitch = 100 - abs(angle_pitch);

	if (yaw < 0 || pitch < 0) {
		yaw = 0;
		pitch = 0;
	}

	std::cout << "yaw: " << yaw << "\t\t";
	std::cout << "pitch: " << pitch << "\t\t";
	std::cout << "--------------------" << std::endl;
	int size = sqrt(face_rect.width * face_rect.width + face_rect.height * face_rect.height) / 3;
	if (size > 100) size = 100;

	int quality = yaw *  0.7 + 0.3 * pitch + 0.05 * size;
	if (quality > 100) quality = 100;
	if ((face_rect.x - 20 < 0) || (face_rect.y - 20 < 0) ||
		(face_rect.x + face_rect.width + 20 > img.cols) ||
		(face_rect.y + face_rect.height + 20 > img.rows)) quality = quality - 10;

	return quality;
}