使用SSD模型检测自定目标

SSD简介

SSD（Single Shot MultiBox Detector）是深度学习领域一种新型的目标检测算法。在过去的几次国际比赛中，SSD在速度和准确性方面均取得优异成绩，与其他检测算法一度拉开很大差距。

SSD的算法流程大体可以概括为产生候选区域、框选、判断、过滤几个步骤。其中，产生候选区域、框选和过滤的算法是固定的，而针对给定的候选区域，判断区域中的图像是否是待检测目标，需要使用不同的模型。常用的识别模型有VGG、Inception、ResNet、MobileNet等。其中，基于MobileNet构建的SSD模型具有最快的检测速度。

本次训练使用的是 ssd_mobilenet_v2_coco 这一实现方式，该模型使用新版本的mobile_net，在Coco数据集上进行训练。

模型的大体工作流程

SSD模型细节非常复杂，但如果只是使用，可以无需了解其实现细节。在TensorFlow中，使用SSD模型需要以下几个步骤：

1. 使用GFile工具类导入二进制模型文件；
2. 将二进制模型文件转换成TensorFlow中的可计算图（Graph）；
3. 分析图的结构（或查看文档）找到图的数据入口和结果出口；
4. 将数据入口保存到 tf.Placeholder 类中；
5. 将结果出口保存到对应的 Tensor 中；
6. 在 Session 中运行 Tensor 对应的 Operation，将真实数据 feed 到对应的 Placeholder 中，然后取出结果数据；
7. 对结果数据进行清洗，可视化等操作。

模型的训练方法

标注和准备数据

将待标注的图片放置在同一个文件夹，随后下载安装 labelImg 工具，在工具中打开待标注文件夹，并选择结果保存文件夹，然后开始手动标注图片。每张图片对应一个名字相同的XML文件。

标注完成后，运行 PrepareData.bat，即可生成 TfRecord 文件。TfRecord 文件是 TensorFlow 中数据传输的标准格式，使用 TfRecord 文件可以避免小文件的多次读取，减少 IO 操作次数。

准备预训练模型和配置文件

前往GitHub上下载预训练模型，然后解压。

压缩包中除了预训练模型，还有和模型配套的配置文件。配置文件中大多数内容不需要修改，需要修改的内容如下：

model {
  ssd {
    num_classes: 1
    # 检测的类别，集装箱号固定为1类
    image_resizer {
      fixed_shape_resizer {
        height: 300
        width: 300
        # 检测时图片大小，默认300*300
      }
    }
  }
}
train_config {
  batch_size: 8
  # 训练时，每一步的图片数，内存允许的情况下越大越好
  data_augmentation_options {
    random_rgb_to_gray {
      probability : 0.5
    }
  }
  data_augmentation_options {
    random_adjust_brightness {
      max_delta: 0.5
    }
  }
  data_augmentation_options {
    random_adjust_hue {
      max_delta: 0.25
    }
  }
  data_augmentation_options {
    random_adjust_saturation {
    }
  }
  # 上面4个都是对图片进行随机变换，提高样本数
  fine_tune_checkpoint: "object_detection/data/model/model.ckpt"
  # 预训练模型路径
  num_steps: 100000
  # 训练次数
  fine_tune_checkpoint_type: "detection"
}
train_input_reader {
  label_map_path: "object_detection/data/LabelMap.pbtxt"
  # 类别和名称的对应关系文件
  tf_record_input_reader {
    input_path: "object_detection/data/mscoco_train.record"
    # 输入的训练数据
  }
}

修改完毕后，需要指定类别和名称的对应关系文件。文件内容如下：

item {
  id: 1
  name: 'conNum'
}

随后即可开始训练。训练调用 train.py，也可以直接打开 StartTrain.bat。

导出模型

训练过程中会自动保存模型，但是保存的是上下文模型，要得到二进制模型（更精简，更快速），需要调用 export_inference_graph.py。可以点击 ConvertModel.bat 完成操作，也可以直接输入命令：

python object_detection\export_inference_graph.py \
--input_type image_tensor \
--pipeline_config_path object_detection\data\ModelConfig.config \
--trained_checkpoint_prefix object_detection\data\output\model.ckpt-20000 \
--output_directory object_detection\data\output

模型的使用

代码介绍如下：

import tensorflow as tf
from PIL import ImageDraw
from PIL import ImageFont
from PIL import Image
import numpy as np
import time
import os


INPUT_DIR = 'E:\\TrainData\\TestImage'
OUTPUT_DIR = 'E:\\TrainData\\TestResult2'


with tf.gfile.GFile('object_detection/data/output/frozen_inference_graph.pb', 'rb') as file:
    graph_def = tf.GraphDef()
    graph_def.ParseFromString(file.read())
with tf.Graph().as_default() as graph:
    tf.import_graph_def(graph_def)
    image_holder = graph.get_tensor_by_name('import/image_tensor:0')
    boxes = graph.get_tensor_by_name('import/detection_boxes:0')
    scores = graph.get_tensor_by_name('import/detection_scores:0')
    num_detections = graph.get_tensor_by_name('import/num_detections:0')
    # classes = graph.get_tensor_by_name('detection_classes:0')
begin = time.time()
with tf.Session(graph=graph) as sess:
    font_en = ImageFont.truetype('C:\\Windows\\Fonts\\batang.ttc', 48)
    for filename in os.listdir(INPUT_DIR):
        image_raw = Image.open(os.path.join(INPUT_DIR, filename))
        image_in = image_raw.resize((300, 300))
        image_in = np.array(image_in, np.uint8)
        if np.shape(image_in)[-1] == 4:
            image_raw = image_raw.convert('RGB')
            image_in = image_in[:, :, :3]
        image_in = np.expand_dims(image_in, axis=0)
        boxes_out, scores_out, nums_out = sess.run([boxes, scores, num_detections], feed_dict={image_holder: image_in})
        height, width, draw = image_raw.height, image_raw.width, ImageDraw.Draw(image_raw)
        for idx in range(int(nums_out)):
            bbox = boxes_out[0][idx]
            draw.rectangle([(bbox[1] * width, bbox[0] * height), (bbox[3] * width, bbox[2] * height)], fill=(0, 0, 0))
            draw.text((bbox[1] * width, bbox[0] * height), str(scores_out[0][idx]), font=font_en, fill=(255, 255, 255))
        if int(nums_out) > 0:
            image_raw.save(os.path.join(OUTPUT_DIR, filename))
            print('Detect: %s' % filename)
        else:
            print('No Result: %s' % filename)
        image_raw.close()
print('Cost: %.2f' % (time.time() - begin))

运行输出的结果是一个目标区域被画上黑框的图片。在训练一小时以后，检测率达到70%左右。预计延长训练时间、加大样本数量可以大幅度提高准确度。