charrin / retinaface-cpp Goto Github PK

#include <stdio.h>
#include <vector>
#include <opencv2/core/core.hpp>
#include <opencv2/highgui/highgui.hpp>
#include <opencv2/imgproc/imgproc.hpp>

#include <iostream>
#include <string>
#include <fstream>
#include <sstream>
#include <stdlib.h>
#include <cstdio>

#include "platform.h"
#include "net.h"
#include<windows.h>
#include "anchor_generator.h"
#include "config.h"
#include "tools.h"

using namespace std;

const int target_size = 300;
//------------------------------------------------------------------------------------------
struct Object
{
	cv::Rect_<float> rect;
	int label;
	float prob;
};

static int detect_mobilenet(const cv::Mat& bgr, std::vector<Anchor>& result, const char* paramPath, const char* modelPath)
{
	ncnn::Net retina;

	retina.load_param(paramPath);
	retina.load_model(modelPath);

	int img_w = bgr.cols;
	int img_h = bgr.rows;

	ncnn::Mat in = ncnn::Mat::from_pixels_resize(bgr.data, ncnn::Mat::PIXEL_BGR2RGB, bgr.cols, bgr.rows, img_w, img_h);
	//ncnn::Mat in = ncnn::Mat::from_pixels_resize(bgr.data, ncnn::Mat::PIXEL_BGR, bgr.cols, bgr.rows, target_size, target_size);
	
	//const float mean_vals[3] = { 103.939f, 116.779f, 123.68f }; //retina
	//const float norm_vals[3] = { 1.0 / mean_vals[0],1.0 / mean_vals[1],1.0 / mean_vals[2] };
	const float mean_vals[3] = { 0,0,0 }; //retina
	const float norm_vals[3] = { 1,1,1 };
	in.substract_mean_normalize(mean_vals, norm_vals);

	ncnn::Extractor ex = retina.create_extractor();

	ex.input("data", in);

	//--------------------------------------------------------
	std::vector<AnchorGenerator> ac(_feat_stride_fpn.size());
	for (int i = 0; i < _feat_stride_fpn.size(); ++i) {
		int stride = _feat_stride_fpn[i];
		ac[i].Init(stride, anchor_cfg[stride], false);
	}
	std::vector<Anchor> proposals;
	proposals.clear();

	for (int i = 0; i < _feat_stride_fpn.size(); ++i) {
		ncnn::Mat cls;
		ncnn::Mat reg;
		ncnn::Mat pts;

		char clsname[100]; sprintf(clsname, "face_rpn_cls_prob_reshape_stride%d", _feat_stride_fpn[i]);
		char regname[100]; sprintf(regname, "face_rpn_bbox_pred_stride%d", _feat_stride_fpn[i]);
		char ptsname[100]; sprintf(ptsname, "face_rpn_landmark_pred_stride%d", _feat_stride_fpn[i]);

		ex.extract(clsname, cls);
		ex.extract(regname, reg);
		ex.extract(ptsname, pts);

		ac[i].FilterAnchor(cls, reg, pts, proposals);

	}
	
	nms_cpu(proposals, nms_threshold, result);

	retina.clear();
	ac.clear();
	proposals.clear();

	return 0;
}

void crop_objects(cv::Mat& bgr, std::vector<Anchor> result, const char* imagepath) {
	//resize(bgr, bgr, cv::Size(target_size, target_size));
	cv::Mat image = bgr.clone();
	
	if (result.size() == 0) {
		cout << "[ " << imagepath << " ] crop fail!\n" << endl;
		system("pause");
		exit(0);
	}
	Anchor res = result[0];
	for (int i = 1; i < result.size(); i++)
	{	
		cout << res.score << endl;
		cout << result[i].score << endl;

		if (res.score < result[i].score) {
			res = result[i];
			cout << "!" << endl;
		}

	}
    
	cout << res.score << endl << endl;
	cout << res.finalbox.x << " " << res.finalbox.y << " " << res.finalbox.width << " " << res.finalbox.height << endl;

	float w,h;
	w = h = 0;
	w = res.finalbox.width - res.finalbox.x + 1;
	h = res.finalbox.height - res.finalbox.y + 1;

	if (res.finalbox.x < 0) {
		res.finalbox.x = 0;
	}
	if ((res.finalbox.x + w) > image.cols) {
		w = image.cols - res.finalbox.x;
	}

	if (res.finalbox.y < 0) {
		res.finalbox.y = 0;
	}
	if ((res.finalbox.y + h) > image.rows) {
		h = image.rows - res.finalbox.y;
	}

	cout << res.finalbox.x << " " << res.finalbox.y << " " << w << " " << h << endl;
	image = bgr(cvRect(res.finalbox.x, res.finalbox.y, w, h));

	cv::imwrite(imagepath, image);
	//cv::imshow("t", image);
	//cv::waitKey(true);
	//Sleep(200);
	cout << "[ " << imagepath << " ] crop complete!\n" << endl;
}

int main(int argc, char** argv)
{
	string cropListPathStr = "C:/Users/SEARECLUSE/Desktop/RetinaFaceDemo/x64/Release/dataList.txt";
	string paramPathStr = "C:/Users/SEARECLUSE/Desktop/RetinaFaceDemo/x64/Release/model/mnet.ncnn.param";
	string modelPathStr = "C:/Users/SEARECLUSE/Desktop/RetinaFaceDemo/x64/Release/model/mnet.ncnn.bin";

	if (argc > 1) {
		cropListPathStr = argv[1];
		paramPathStr = argv[2];
		modelPathStr = argv[3];
	}

	cout << "确认裁剪列表: " << cropListPathStr << endl;
	cout << "确认模型结构: " << paramPathStr << endl;
	cout << "确认模型文件：" << modelPathStr << endl;
	cout << "准备就绪。" << endl;
	system("pause");

	const char* cropListPath = cropListPathStr.data();
	const char* paramPath = paramPathStr.data();
	const char* modelPath = modelPathStr.data();

	ifstream infile;
	infile.open(cropListPath, ios::in);
	while (!infile.eof()) {
		string str;
		getline(infile, str);

		const char* imagepath = str.data();
		cv::Mat m = cv::imread(imagepath, 1);

		cout << imagepath << endl;

		if (m.empty())
		{
			fprintf(stderr, "cv::imread %s failed\n", imagepath);
			return -1;
		}

		std::vector<Anchor> result;
		detect_mobilenet(m, result, paramPath, modelPath);
		crop_objects(m, result, imagepath);
		result.clear();
	}

	system("pause");
	return 0;
}