Ubuntu上运行tensorflow C++的完整例子

个人博客原文：http://www.bearoom.xyz/2019/08/25/ubuntu-tensorflow-cc-example/

之前记录的运行Tensorflow的C++接口的例子都是零散的，现在写一个完整的例子。

一、模型文件转换

首先是需要有训练好的模型文件，然后将其转化为tensorflow的C++接口能够读取的.pb文件，这个前面也有记录，现在贴下完整的代码：

###################################################
#
#   Script to
#   - Loading a pre-trained model(.h5)
#   - Generate a .pb model file from .h5
#
##################################################

#Python
import numpy as np
import configparser
import time
from matplotlib import pyplot as plt
#Keras
from keras.models import model_from_json
from keras.models import Model
#scikit learn
from sklearn.metrics import roc_curve
from sklearn.metrics import roc_auc_score
from sklearn.metrics import confusion_matrix
from sklearn.metrics import precision_recall_curve
from sklearn.metrics import jaccard_similarity_score
from sklearn.metrics import f1_score
import sys

from keras.models import load_model
import tensorflow as tf
from keras import backend as K
from tensorflow.python.framework import graph_io
from tensorflow.python.framework.graph_util import convert_variables_to_constants

def freeze_session(session, keep_var_names=None, output_names=None, clear_devices=True):
    graph = session.graph
    with graph.as_default():
        freeze_var_names = list(set(v.op.name for v in tf.global_variables()).difference(keep_var_names or []))
        output_names = output_names or []
        output_names += [v.op.name for v in tf.global_variables()]
        input_graph_def = graph.as_graph_def()
        if clear_devices:
            for node in input_graph_def.node:
                node.device = ""
        frozen_graph = convert_variables_to_constants(session, input_graph_def, output_names, freeze_var_names)
        return frozen_graph


model = model_from_json(open('CDNet_v4_architecture.json').read())
model.load_weights('CDNet_v4_best_weights.h5')

model.summary()

print('input is :', model.input.name)
print ('output is:', model.output.name)

"""------------------------保存为.pb格式------------------------"""
sess = K.get_session()
frozen_graph = freeze_session(K.get_session(), output_names=[model.output.op.name])
graph_io.write_graph(frozen_graph, './', 'CDNet_v4_best_weights.pb', as_text=False)

然后会输出模型的每层信息，以及输入输出信息：

二、CMakeLists.txt文件

CMakeLists.txt文件里面会添加opencv和tensorflow的路径：

#设置cmake的最小版本
cmake_minimum_required(VERSION 3.0)
#项目名称
project(demo)
#设置c++编译器
set(CMAKE_CXX_STANDARD 11)
#设置TENSORFLOW_DIR变量，变量内容为安装的tensorflow文件夹路径
set(TENSORFLOW_DIR /home/zcx/tensorflow-r1.12)
#项目中的include路径
include_directories(${TENSORFLOW_DIR})
include_directories(${TENSORFLOW_DIR}/tensorflow/contrib/makefile/downloads/absl)
include_directories(${TENSORFLOW_DIR}/bazel-genfiles)
include_directories(${TENSORFLOW_DIR}/tensorflow/contrib/makefile/downloads/eigen)

# opencv
include_directories(/usr/local/opencv347/include /usr/local/opencv347/include/opencv /usr/local/opencv347/include/opencv2)

#项目中lib路径
link_directories(${TENSORFLOW_DIR}/tensorflow/contrib/makefile/gen/lib)
link_directories(${TENSORFLOW_DIR}/bazel-bin/tensorflow)

link_directories(/usr/local/opencv347/lib)

add_executable(demo main.cpp)
#连接libtensorflow_cc.so和libtensorflow_framework库。
target_link_libraries(demo tensorflow_cc tensorflow_framework opencv_core  opencv_highgui opencv_imgproc opencv_imgcodecs)

 

三、C++文件：

因为我这个模型是进行视网膜血管图像分割，最终的输出层是一个2维tensor，第一维图像长宽的乘积，即数据大小，第二维是通道数，因为有背景和血管两个类，所以分了两个通道。最后得到输出张量的时候也要通过赋值或者变形等操作将数据转为图像数据来显示。

#include <tensorflow/core/platform/env.h>
#include <tensorflow/core/public/session.h>

#include <iostream>

#include <opencv.hpp>

using namespace std;
using namespace tensorflow;

void CVMat_to_Tensor(cv::Mat &img, Tensor* output_tensor, int input_rows,int input_cols)    
{
    //imshow("input image",img);
    //图像进行resize处理
    cv::resize(img,img,cv::Size(input_cols,input_rows));
    //imshow("resized image",img);

    //归一化
    img.convertTo(img,CV_32FC1);
    img=img/255;

    float *p = output_tensor->flat<float>().data();

    cv::Mat tempMat(input_rows, input_cols, CV_32FC1, p);
    img.convertTo(tempMat,CV_32FC1);

}

int main()
{
    Session* session;
    Status status = NewSession(SessionOptions(), &session);
    if (!status.ok())
    {
        cout << status.ToString() << "\n";
        return -1;
    }
    cout << "Session successfully created.\n";

    string model_path = "./CDNet_v4_best_weights.pb";

    GraphDef graphdef;
    Status status_load = ReadBinaryProto(Env::Default(), model_path, &graphdef);
    if (!status_load.ok()) 
    {
        cout << status_load.ToString() << "\n";
        return -1;
    }
    Status status_create = session->Create(graphdef);
    if (!status_create.ok()) 
    {
        cout << status_create.ToString() << "\n";
        return -1;
    }

    //输出每一曾的名字
    for (int i=0; i < graphdef.node_size(); i++)
    {
        std::string name = graphdef.node(i).name();
        std::cout << name << std::endl;
    }

    cv::Mat src = cv::imread("21_training.tif", 0);
    Tensor src_tensor = Tensor(DT_FLOAT, TensorShape({ 1, 1, 544, 544}));
    CVMat_to_Tensor(src, &src_tensor, 544, 544);

    string input_tensor_name="input_1:0";
    string output_tensor_name="activation_1/truediv:0";
    vector<tensorflow::Tensor> outputs;
    string output_node = output_tensor_name;
    
    Status status_run = session->Run({{input_tensor_name, src_tensor}}, {output_node}, {}, &outputs);
    if (!status_run.ok()) 
    {
        cout << status_run.ToString() << "\n";
        return -1;
    }

    std::cout << outputs[0].shape().dim_size(1) << std::endl;
    std::cout << outputs[0].shape().dim_size(2) << std::endl;

    float *pDstTensor = outputs[0].flat<float>().data();
    cv::Mat tmp = cv::Mat(544, 544, CV_32FC2, pDstTensor);
    cv::Mat tmp1 = cv::Mat(cv::Size(544, 544), CV_32FC1);
    cv::Mat dst0 = cv::Mat(cv::Size(544, 544), CV_8UC1);
    cv::Mat dst1 = cv::Mat(cv::Size(544, 544), CV_8UC1);

    for (int i = 0; i < 544; i++)
    {
        float *ptrTmp = tmp.ptr<float>(i);
        float *ptrTmp1 = tmp1.ptr<float>(i);
        for (int j = 0; j < 544; j++)
        {
            *(ptrTmp1 + j) = *(ptrTmp + j*2)*255;
        }
    }
    
    tmp1.convertTo(dst0, CV_8UC1);
    cv::imwrite("dst0.jpg", dst0);

    for (int i = 0; i < 544; i++)
    {
        float *ptrTmp = tmp.ptr<float>(i);
        float *ptrTmp1 = tmp1.ptr<float>(i);
        for (int j = 0; j < 544; j++)
        {
            *(ptrTmp1 + j) = *(ptrTmp + j*2 + 1)*255;
        }
    }
    
    tmp1.convertTo(dst1, CV_8UC1);
    cv::imwrite("dst1.jpg", dst1);

    return 0;
}

输入图像：

输出图像我可视化了背景和血管两个通道，转为图像来显示，当然每个像素是乘以255转出来的，毕竟输出的是概率，在0-1之间：

背景(白的背景，黑的为血管)：

血管(白的为血管，黑的为背景)：

完结，撒花...

风鬟雨鬓，偏是来无准。

倦倚玉兰看月晕，容易语低香近。

软风吹遍窗纱，心期便隔天涯。

从此伤春伤别，黄昏只对梨花。

  -- 纳兰性德 《清平乐·风鬟雨鬓》