******************************************************************************
*
* @Author : xm
* @Data : 2019-9-8
* @file : ssd_inference.cpp
* @Desctiption: the member function SSD_Inference class
*
*****************************************************************************/
#include "ssd_inference.h"
#include <time.h>
//setenv("CUDA_VISIBLE_DEVICES", "", 1)
SSD_Inference::SSD_Inference()
{
// 构造函数
session_ = nullptr;
confid_thres_ = 0.5;
}
SSD_Inference::~SSD_Inference()
{
// 析构函数
}
int SSD_Inference::loadModelFile(const std::string& model_file_path)
{
Status status = tensorflow::NewSession(SessionOptions(), &session_);
GraphDef graph_def;
tensorflow::SessionOptions sess_options;
status = ReadBinaryProto(Env::Default(), model_file_path, &graph_def);
if (!status.ok())
return -1;
status = session_->Create(graph_def);
if (!status.ok())
return -2;
return 0;
}
bool SSD_Inference::setClassInfo(const std::vector<std::string>& class_name_ls, int image_width, int image_height)
{
class_name_ls_ = class_name_ls;
image_width_ = image_width;
image_height_ = image_height;
return true;
}
bool SSD_Inference::setConfidenceThreshold(float confid_thres)
{
confid_thres_ = confid_thres;
return 0;
}
int SSD_Inference::predict(const cv::Mat& src_mat)
{
src_image_ = src_mat;
if (src_image_.rows * src_image_.cols == 0)
{
return -2;
}
cv::Mat test_image;
cv::cvtColor(src_mat, test_image, COLOR_BGR2RGB);
Tensor resized_tensor(DT_FLOAT, TensorShape({1, image_height_, image_width_, 3}));
cvmatToTensor(test_image, &resized_tensor, image_height_, image_width_);
output_tensor_ls_.clear();
std::string input_tensor_name = "input_1";
std::string output_tensor_name = "ssd_decoded_predictions/loop_over_batch/TensorArrayStack/TensorArrayGatherV3";
clock_t start1 = clock();
Status status_run = session_->Run({ { input_tensor_name, resized_tensor } }, { output_tensor_name }, {}, &output_tensor_ls_);
clock_t end1 = clock();
cout <<"run time is :" << double (end1 - start1) / 1000 << endl;
if (!status_run.ok())
return -1;
return 0;
}
void SSD_Inference::getDetectResult(std::vector<DetectResult>& detect_result_ls)
{
detect_result_ls.clear();
Tensor output_tensor = output_tensor_ls_[0];
auto result_tensor = output_tensor.tensor<float, 3>();
float x_ratio = src_image_.cols / float(image_width_);
float y_ratio = src_image_.rows / float(image_height_);
for (int idx_1 = 0, idx_cnt_1 = result_tensor.dimension(1); idx_1 != idx_cnt_1; ++idx_1) {
if (result_tensor(0, idx_1, 1) >= confid_thres_) {
DetectResult cur_result;
cur_result.class_idx = int(result_tensor(0, idx_1, 0));
cur_result.confidence = result_tensor(0, idx_1, 1);
float x_min = result_tensor(0, idx_1, 2) * x_ratio;
float y_min = result_tensor(0, idx_1, 3) * y_ratio;
float x_max = result_tensor(0, idx_1, 4) * x_ratio;
float y_max = result_tensor(0, idx_1, 5) * y_ratio;
cur_result.top_left = cv::Point(int(x_min), int(y_min));
cur_result.bottom_right = cv::Point(int(x_max), int(y_max));
detect_result_ls.push_back(cur_result);
}
}
}
void SSD_Inference::cvmatToTensor(cv::Mat img, Tensor* output_tensor, int input_rows, int input_cols)
{
resize(img, img, cv::Size(input_cols, input_rows));
img.convertTo(img, CV_32FC3);
float* p = output_tensor->flat<float>().data();
cv::Mat tempMat(input_rows, input_cols, CV_32FC3, p);
img.convertTo(tempMat, CV_32FC3);
}
#ifndef SSD_INFERENCE_H
#define SSD_INFERENCE_H
#include <fstream>
#include <utility>
#include <iostream>
#include <string>
#include <vector>
#include <Eigen/Core>
#include <Eigen/Dense>
#include <opencv4/opencv2/opencv.hpp>
#include "tensorflow/cc/ops/const_op.h"
#include "tensorflow/cc/ops/image_ops.h"
#include "tensorflow/cc/ops/standard_ops.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/graph/default_device.h"
#include "tensorflow/core/graph/graph_def_builder.h"
#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/lib/core/stringpiece.h"
#include "tensorflow/core/lib/core/threadpool.h"
#include "tensorflow/core/lib/io/path.h"
#include "tensorflow/core/lib/strings/stringprintf.h"
#include "tensorflow/core/public/session.h"
#include "tensorflow/core/util/command_line_flags.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/init_main.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/platform/types.h"
using namespace std;
using namespace cv;
using namespace tensorflow::ops;
using namespace tensorflow;
using tensorflow::Flag;
using tensorflow::int32;
using tensorflow::Status;
using tensorflow::string;
using tensorflow::Tensor;
// the object detect result information
typedef struct
{
// class index
int class_idx;
// the object confidence
float confidence;
// top left point
Point top_left;
// bottom right point
Point bottom_right;
}DetectResult;
class SSD_Inference
{
public:
/*
* @brief : structor function
*/
SSD_Inference();
/*
* @brief : destructor function
*/
~SSD_Inference();
/*
* @brief : load the ssd network model file in format .pb
* @param : model_file_path: the model file path [in]
* @return : int: if load the file successfully, return 0
*/
int loadModelFile(const string& modle_file_path);
/*
* @brief : set the class information
* @param : class_name_ls: set class name list [in]
* @param : image_width: set input image width of network [in]
* @param : image_height: set input image height of network [in]
* @return : bool: if set the class information successfully, return true
*/
bool setClassInfo(const vector<string>& class_name_ls, int image_width, int image_geight);
/*
* @brief : set the detect confidence threshold
* @param : confid_thres: the detection confidence threshold [in]
* @return : bool: if set the parameter successfully, return true
*/
bool setConfidenceThreshold(float confid_thres);
/*
* @brief : detect the input image
* @param : src_mat: the image to be detected [in]
* @return : int: if detect image successfully, return 0
*/
int predict(const Mat& src_mat);
/*
* @brief : get the detection result
* @param : detect_result_ls: the detection result [out]
*/
void getDetectResult(vector<DetectResult>& detect_result_ls);
/*
* @brief : convert the cv image to tf tensor
* @param : img: the image in cv::Mat [in]
* @param : output_tensor: the tensor in tensorflow [out]
* @param : input_rows: the input image width of network [in]
* @param : input_cols: the input image height of network [in]
*/
void cvmatToTensor(Mat img, Tensor* output_tensor, int input_rows, int input_cols);
private:
// tensorflow session pointer
Session* session_;
// the output tensor list of network detection
vector<Tensor> output_tensor_ls_;
// the detection result list
vector<DetectResult> detect_result_ls_;
// the class name list, the first element must be the "background"
vector<string> class_name_ls_;
// the detection confidence threshlod
float confid_thres_;
// the input image width of network
int image_width_;
// the input image height of network
int image_height_;
// the input image to be detected
Mat src_image_;
};
#endif //SSD_INFENCE_H
/******************************************************************************
*
* @Author : xm
* @Data : 2019-9-8
* @file : main.cpp
* @Desctiption: SSD 模型的运行 main 函数
*
*****************************************************************************/
#include <iostream>
#include <opencv4/opencv2/opencv.hpp>
#include "ssd_inference.h"
#include <time.h>
using namespace std;
using namespace cv;
/*
* @brief : SSD model c++ test
* @param : model_path : saved the tensorflow format xxxx.pb model
* @param : image_path : test image file
* @return : DetectResult: Detece box and class name
*/
int main()
{
Mat Image;
// load the pb model
string model_path = "./dataset_2019-10-11-circle-square_color_2.pb";
// load image_path
string image_path = "./20.bmp";
Image = cv::imread(image_path);
if (Image.empty())
{
cout << "could not load image ..." << endl;
return -1;
}
namedWindow("test image",WINDOW_AUTOSIZE);
imshow("test image", Image);
// image exchange color channels RGB to BGR
//cvtColor(Image, Image, COLOR_RGB2BGR);
//imshow("test image", Image);
// the ssd network inference
SSD_Inference ssd_inference;
ssd_inference.loadModelFile(model_path);
// load class name
std::vector<std::string> class_name_ls{"background", "circle", "square"};
//std::vector<std::string> class_name_ls{"background", "Oreo", "PacificBiscuit"};
// set class information
ssd_inference.setClassInfo(class_name_ls, 300, 300);
// set confidence threshold
ssd_inference.setConfidenceThreshold(0.7);
// set detect result
std::vector<DetectResult> detect_result_ls;
clock_t start, end;
//for (int i=0;i<10;i++)
//{
start = clock();
// set prdict image
ssd_inference.predict(Image);
// get detect result
//
ssd_inference.getDetectResult(detect_result_ls);
end = clock();
cout<<"total time is : "<< double (end-start) / 1000 << endl;
//}
// draw the result
Mat result_image = Image;
for (const auto& ele_detect : detect_result_ls)
{
auto class_info = class_name_ls.at(ele_detect.class_idx);
class_info += ":" + std::to_string(ele_detect.confidence);
cout << class_info << endl;
rectangle(result_image, ele_detect.top_left, ele_detect.bottom_right, cv::Scalar(0, 0, 255), 5);
Point text_point(ele_detect.top_left.x, (ele_detect.top_left.y + ele_detect.bottom_right.y) / 2.0);
putText(result_image, class_info, text_point, FONT_HERSHEY_PLAIN, 2, cv::Scalar(0, 0, 255), 2);
}
// image exchange color channels RGB to BGR
//cvtColor(result_image, result_image, COLOR_BGR2RGB);
imshow("detect result", result_image);
waitKey(0);
imwrite("result.bmp", result_image);
return 0;
}
对应的 python 推理 pb 的过程
# -*- coding : utf-8 -*-
# @Data : 2020-12-28
# @Author : Ming Xu
# @Email : 920972751@qq.com
# @File : inferencePb.py
# Desctiption: MobileNetV2SSDlite
import os
import cv2
import time
import numpy as np
import tensorflow as tf
#os.environ["CUDA_VISIBLE_DEVICES"]="-1"
###############################################################################
#- 定义识别函数
def recognize(src_image):
"""
MobileNetV2-SSDLite
:param src_image: 输入视频流或图像
:param pb_file_path: the model file path
:return:
"""
with tf.Graph().as_default():
output_graph_def = tf.GraphDef()
#---------------------------------------
# 打开 .pb 模型
pb_file = "ssd300_pascal_07+12_epoch-86_loss-1.2568_val_loss-0.5428.pb"
with open(pb_file, "rb") as f:
output_graph_def.ParseFromString(f.read())
tensors = tf.import_graph_def(output_graph_def, name="")
print("tensors:",tensors)
with tf.Session() as sess:
# init = tf.global_variables_initializer()
# sess.run(init)
#---------------------------------------
# 打开图中所有的操作
#op = sess.graph.get_operations()
#for i,m in enumerate(op):
# print('op{}:'.format(i),m.values())
#---------------------------------------
# 模型的的输入和输出名称
#--------------------------------------
# 遍历某目录下的图像
input_x = sess.graph.get_tensor_by_name("input_1:0")
#print("input_X:",input_x)
output_tensor = sess.graph.get_tensor_by_name("ssd_decoded_predictions/loop_over_batch/TensorArrayStack/TensorArrayGatherV3:0")
#print("Output:",output_tensor)
#--------------------------------------
# 计算时间, 持续加载同一张图像
src_image = cv2.imread(src_image)
org_img = src_image.copy()
img=cv2.resize(src_image,(300,300))
img=img.astype(np.float32)
y_pred = sess.run([output_tensor], feed_dict={input_x:np.reshape(img,(1,300,300,3))})
confidence_threshold = 0.8
y_pred_array = np.array(y_pred[0])
y_pred_thresh = [y_pred_array[k][y_pred_array[k,:,1] > confidence_threshold] for k in range(y_pred_array.shape[0])]
print(y_pred_thresh)
classes = ['background', 'tank']
image_size = (300, 300, 3)
for box in y_pred_thresh[0]:
xmin = box[2] * org_img.shape[1] / image_size[0]
ymin = box[3] * org_img.shape[0] / image_size[1]
xmax = box[4] * org_img.shape[1] / image_size[1]
ymax = box[5] * org_img.shape[0] / image_size[0]
label = '{}: {:.2f}'.format(classes[int(box[0])], box[1])
print("label", label)
def main():
src_image = "./002394.jpg"
#pb_file = "ssd300_pascal_07+12_epoch-86_loss-1.2568_val_loss-0.5428.pb"
recognize(src_image)
if __name__ == '__main__':
main()