tensorflow c++ mobilenetv2 ssd

******************************************************************************
 *
 * @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()