tensorflow 的c++部署--NIMA

NIMA.h

/*****************************************************************************
* @brief : NIMA 美学评价，keras版本实现，模型已经由Keras的.h5转到tensorflow的.pb格式
* @author : liuwei
* @date : 2018/12/21 13:49
* @version : ver 1.0 
*****************************************************************************/
#pragma once
#include "../common/interface.h"
#include <tensorflow/core/public/session.h>
#include <tensorflow/core/platform/env.h>
#include <opencv2/opencv.hpp>

using namespace tensorflow;
using namespace cv;

///  NIMA 美学评价
///     
class NIMA : public IAssessment
{
public:
	NIMA();
	~NIMA();

	void initial(std::string filename = "");	//初始化
	float assessment(std::string path);		//预测得分

private:
	Session* m_session;

	float NIMAMean(const tensorflow::TTypes<float, 1>::Tensor prediction);
	tensorflow::Tensor Mat2Tensor(cv::Mat &img);
};

NIMA.cpp

#include "NIMA.h"

NIMA::NIMA()
{
	m_session = NULL;
}

NIMA::~NIMA()
{
	
}

void NIMA::initial(std::string filename /* = "" */)
{
	//创建session
	Status status = NewSession(SessionOptions(), &m_session);
	if (!status.ok())
	{
		MYLOG::MYLOG(MYLOG::ERROR, status.ToString());
		exit(-1);
	}
	else
	{
		MYLOG::MYLOG(MYLOG::INFO, "NIMA: Session created successfully");
	}

	//导入模型
	GraphDef graph_def;
	status = ReadBinaryProto(Env::Default(), filename, &graph_def);
	if (!status.ok()) 
	{
		MYLOG::MYLOG(MYLOG::ERROR, status.ToString());
		exit(-1);
	}
	else
	{
		MYLOG::MYLOG(MYLOG::INFO, "NIMA: Load graph protobuf successfully");
	}

	//将模型导入到session
	status = m_session->Create(graph_def);
	if (!status.ok())
	{
		MYLOG::MYLOG(MYLOG::ERROR, status.ToString());
		exit(-1);
	}
	else
	{
		MYLOG::MYLOG(MYLOG::INFO, "NIMA: Add graph to session successfully");
	}
}

Tensor NIMA::Mat2Tensor(cv::Mat &img) {

	tensorflow::Tensor image_input = tensorflow::Tensor(tensorflow::DT_FLOAT, tensorflow::TensorShape(
		{ 1, img.size().height, img.size().width, img.channels() }));

	float *tensor_data_ptr = image_input.flat<float>().data();
	cv::Mat fake_mat(img.rows, img.cols, CV_32FC(img.channels()), tensor_data_ptr);
	img.convertTo(fake_mat, CV_32FC(img.channels()));

	fake_mat /= 255.0;	
	fake_mat -= Scalar(0.5, 0.5, 0.5);
	fake_mat *= 2.0;

	return image_input;
}

float NIMA::NIMAMean(const tensorflow::TTypes<float, 1>::Tensor prediction)
{
	float m = 0;
	for (int i = 1; i <= 10; i++)
		m += i * prediction(i - 1);
	return m;
}

float NIMA::assessment(std::string path)
{
	Tensor phase(DT_BOOL, TensorShape()); // input phase
	phase.scalar<bool>()() = 0;

	Mat img = imread(path, IMREAD_COLOR | IMREAD_IGNORE_ORIENTATION);
	if (img.empty())
	{
		MYLOG::MYLOG(MYLOG::ERROR, "Reading image failed");
		exit(-1);
	}
	cvtColor(img, img, CV_BGR2RGB);
	resize(img, img, Size(224, 224), 0, 0, INTER_LINEAR);
	
	Tensor input_img = Mat2Tensor(img);
	std::vector<std::pair<std::string, Tensor> > inputs = {
		{"input_1:0",input_img},{"conv1_bn/keras_learning_phase:0",phase}
	};
	std::vector<Tensor> outputs;
	Status status = m_session->Run(inputs, { "dense_1/Softmax:0" }, {}, &outputs);
	if (!status.ok())
	{
		MYLOG::MYLOG(MYLOG::ERROR, status.ToString());
	}

	const tensorflow::TTypes<float, 1>::Tensor& prediction = outputs[0].flat_inner_dims<float, 1>();
	float m = NIMAMean(prediction);

	return m;
}