tensorflow C++接口調用圖像分類pb模型代碼

#include <fstream>
#include <utility>
#include <Eigen/Core>
#include <Eigen/Dense>
#include <iostream>
 
#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"
 
#include "opencv2/opencv.hpp"
 
using namespace tensorflow::ops;
using namespace tensorflow;
using namespace std;
using namespace cv;
using tensorflow::Flag;
using tensorflow::Tensor;
using tensorflow::Status;
using tensorflow::string;
using tensorflow::int32 ;
 
// 定義一個函數講OpenCV的Mat數據轉化為tensor，python里面只要對cv2.read讀進來的矩陣進行np.reshape之后，
// 數據類型就成了一個tensor，即tensor與矩陣一樣，然后就可以輸入到網絡的入口了，但是C++版本，我們網絡開放的入口
// 也需要將輸入圖片轉化成一個tensor，所以如果用OpenCV讀取圖片的話，就是一個Mat，然后就要考慮怎么將Mat轉化為
// Tensor了
void CVMat_to_Tensor(Mat img,Tensor* output_tensor,int input_rows,int input_cols)
{
    //imshow("input image",img);
    //圖像進行resize處理
    resize(img,img,cv::Size(input_cols,input_rows));
    //imshow("resized image",img);
 
    //歸一化
    img.convertTo(img,CV_32FC1);
    img=1-img/255;
 
    //創建一個指向tensor的內容的指針
    float *p = output_tensor->flat<float>().data();
 
    //創建一個Mat，與tensor的指針綁定,改變這個Mat的值，就相當於改變tensor的值
    cv::Mat tempMat(input_rows, input_cols, CV_32FC1, p);
    img.convertTo(tempMat,CV_32FC1);
 
//    waitKey(0);
 
}
 
int main(int argc, char** argv )
{
    /*--------------------------------配置關鍵信息------------------------------*/
    string model_path="../inception_v3_2016_08_28_frozen.pb";
    string image_path="../test.jpg";
    int input_height =299;
    int input_width=299;
    string input_tensor_name="input";
    string output_tensor_name="InceptionV3/Predictions/Reshape_1";
 
    /*--------------------------------創建session------------------------------*/
    Session* session;
    Status status = NewSession(SessionOptions(), &session);//創建新會話Session
 
    /*--------------------------------從pb文件中讀取模型--------------------------------*/
    GraphDef graphdef; //Graph Definition for current model
 
    Status status_load = ReadBinaryProto(Env::Default(), model_path, &graphdef); //從pb文件中讀取圖模型;
    if (!status_load.ok()) {
        cout << "ERROR: Loading model failed..." << model_path << std::endl;
        cout << status_load.ToString() << "\n";
        return -1;
    }
    Status status_create = session->Create(graphdef); //將模型導入會話Session中;
    if (!status_create.ok()) {
        cout << "ERROR: Creating graph in session failed..." << status_create.ToString() << std::endl;
        return -1;
    }
    cout << "<----Successfully created session and load graph.------->"<< endl;
 
    /*---------------------------------載入測試圖片-------------------------------------*/
    cout<<endl<<"<------------loading test_image-------------->"<<endl;
    Mat img=imread(image_path,0);
    if(img.empty())
    {
        cout<<"can't open the image!!!!!!!"<<endl;
        return -1;
    }
 
    //創建一個tensor作為輸入網絡的接口
    Tensor resized_tensor(DT_FLOAT, TensorShape({1,input_height,input_width,3}));
 
    //將Opencv的Mat格式的圖片存入tensor
    CVMat_to_Tensor(img,&resized_tensor,input_height,input_width);
 
    cout << resized_tensor.DebugString()<<endl;
 
    /*-----------------------------------用網絡進行測試-----------------------------------------*/
    cout<<endl<<"<-------------Running the model with test_image--------------->"<<endl;
    //前向運行，輸出結果一定是一個tensor的vector
    vector<tensorflow::Tensor> outputs;
    string output_node = output_tensor_name;
    Status status_run = session->Run({{input_tensor_name, resized_tensor}}, {output_node}, {}, &outputs);
 
    if (!status_run.ok()) {
        cout << "ERROR: RUN failed..."  << std::endl;
        cout << status_run.ToString() << "\n";
        return -1;
    }
    //把輸出值給提取出來
    cout << "Output tensor size:" << outputs.size() << std::endl;
    for (std::size_t i = 0; i < outputs.size(); i++) {
        cout << outputs[i].DebugString()<<endl;
    }
 
    Tensor t = outputs[0];                   // Fetch the first tensor
    auto tmap = t.tensor<float, 2>();        // Tensor Shape: [batch_size, target_class_num]
    int output_dim = t.shape().dim_size(1);  // Get the target_class_num from 1st dimension
 
    // Argmax: Get Final Prediction Label and Probability
    int output_class_id = -1;
    double output_prob = 0.0;
    for (int j = 0; j < output_dim; j++)
    {
        cout << "Class " << j << " prob:" << tmap(0, j) << "," << std::endl;
        if (tmap(0, j) >= output_prob) {
            output_class_id = j;
            output_prob = tmap(0, j);
        }
    }
 
    // 輸出結果
    cout << "Final class id: " << output_class_id << std::endl;
    cout << "Final class prob: " << output_prob << std::endl;
 
    return 0;
}