Image segmentation semantic - prediction using a single picture DeeplabV3 +

When the training is good DeeplabV3 + model, generated .ckpt file, the next hope to use the model to predict the real scene, common practice to generate .pb file. The benefit of this is: 1. variable into a constant, reduced model, to facilitate 2 other languages call (.pb files can be directly C / C ++ / Java / NodeJS read, etc.).

Run export_model.py generation model

Export_model.py generated code files using the official frozen_inference_graph.pb file, use the file to predict. It should be noted that: must know the input and output model, which can be obtained by looking at the code.

python export_model.py \
  --checkpoint_path="./checkpoint_1/model.ckpt-518495" \  # 训练得到的ckpt文件
  --export_path="./output_model/frozen_inference_graph.pb" \  # 需要导出的模型名称
  --model_variant="xception_65" \
  --atrous_rates=6 \
  --atrous_rates=12 \
  --atrous_rates=18 \
  --output_stride=16 \
  --decoder_output_stride=4 \
  --num_classes=3 \
  --crop_size=1440 \  # 需要预测图片的大小，如果预测的图像大小比该值大，将报错
  --crop_size=1440 \
  --inference_scales=1.0

Source clearly shows: input_name is 'ImageTensor', shape is [1, None, None, 3], the data type is tf.uint8, you can change the type of data here, output_name is 'SemanticPredictions'. Know the input and outp, it can be predicted.

# export_model.py部分代码

# Input name of the exported model.
_INPUT_NAME = 'ImageTensor'

# Output name of the exported model.
_OUTPUT_NAME = 'SemanticPredictions'


def _create_input_tensors():
  """Creates and prepares input tensors for DeepLab model.

  This method creates a 4-D uint8 image tensor 'ImageTensor' with shape
  [1, None, None, 3]. The actual input tensor name to use during inference is
  'ImageTensor:0'.
  """
  # input_preprocess takes 4-D image tensor as input.
  input_image = tf.placeholder(tf.uint8, [1, None, None, 3], name=_INPUT_NAME)

Predict a single image

.Pb file using the generated predicted a new picture:

1. The image read and converted to uint8, shape is [1, None, None, 3] format;

2. Read .pb file, specifies the input and output;

3. The demand output, label output is 0, 1, 2 ..., so that the black spotted;

4. Results After treatment, this step will be personal

import tensorflow as tf
from keras.preprocessing.image import load_img, img_to_array

img = load_img(img_path)  # 输入预测图片的url
img = img_to_array(img)  
img = np.expand_dims(img, axis=0).astype(np.uint8)  # uint8是之前导出模型时定义的

# 加载模型
sess = tf.Session()
with open("frozen_inference_graph.pb", "rb") as f:
    graph_def = tf.GraphDef()
    graph_def.ParseFromString(f.read())
    output = tf.import_graph_def(graph_def, input_map={"ImageTensor:0": img},
                                     return_elements=["SemanticPredictions:0"])
    # input_map 就是指明 输入是什么；
    # return_elements 就是指明输出是什么；两者在前面已介绍

result = sess.run(output) 
print(result[0].shape)  # (1, height, width)

The results show:

My condition is a three-classification problem, enter the picture 1040X868, on a personal notebook, forecasting slower: 40s, deployed on the server, 0.4s.

Tips- Note TF version:

When forecasting, problems often arise out of memory, but the model has only 157MB, the memory is 16GB, may not have been the solution. Original Tensorflow is 1.8.0, download from a local Github, CUDA with version 9.1. Later, download the official Tensorflow1.8.0, CUDA 9.0 support, I had to re-install CUDA, memory overflow problem goes away. So: Please download the software from the formal channels.