Tensorflow之实现物体检测

你开到的不止于此

项目背景

产品看到竞品可以标记物体的功能，秉承一贯的他有我也要有，他没有我更要有的作风，丢过来一网站，说这个功能很简单，一定可以实现

这时候万能的谷歌发挥了作用，在茫茫的数据大海中发现了Tensorflow机器学习框架，也就是目前非常火爆的的深度学习（人工智能），既然方案已有，就差一个程序员了

Tensorflow介绍

百科介绍：TensorFlow是谷歌基于DistBelief进行研发的第二代人工智能学习系统，可被用于语音识别或图像识别等多项机器学习和深度学习领域。

翻译成大白话：是一个深度学习和神经网络的框架，底层C++，通过Python进行控制，当然，也是支持Go,Java等语言

环境搭建

Unix/Linux(笔者使用Mac)
Python3.6
protoc 3.5.1
tensorflow 1.7.0

1、克隆文件

git clone https://github.com/guandeng/tensorflow.git

文件目录格式如下

└── tensorflow
    ├── Dockerfile
    ├── README.md
    ├── data
    │   ├── models
    │   ├── pbtxt
    │   └── tf_models
    ├── object_detection_api.py
    ├── server.py
    ├── sh
    │   ├── download_data.sh
    │   └── ods.sh
    ├── static
    ├── templates
    └── upload

data/models 存放
data/pbtxt 物体标识名称
data/tf_models 存放tensorflow/models数据

2、安装依赖库

pip3 install -r requirements.txt

3、下载模型

sh sh/download_data.sh

4、添加环境变量PYTHONPATH

echo 'PYTHONPATH=$PYTHONPATH:pwd/data/tf_models/models/research'>> ~/.bashrc && source ~/.bashrc

5、启动服务

python3 server.py

没有报错，说明你已成功搭建环境，使用过程是不是非常简单，下面介绍代码调用逻辑过程

模型选用

我从谷歌提供几种模型选出来对比

Model name	Speed	mAP[^1]
ssd_mobilenet_v1_coco	30	21
ssd_mobilenet_v2_coco	31	22
ssd_inception_v2_coco	42	24
faster_rcnn_inception_resnet_v2_atrous_coco	620	37

Speed 是识别物体速度，值越小，识别越快
mAP（平均准确率）是精度和检测边界盒的乘积，值越高神经网络的识别精确度越高，对应Speed越大

为了测试方便，笔者选用轻量级（ssd_mobilenet）作为本次识别物体模型

引入Python库

import numpy as np
import os
import tensorflow as tf
import json
import time
from PIL import Image
# 兼容Python2.7版本
try:
    import urllib.request as ulib
except Exception as e:
    import urllib as ulib
import re
from object_detection.utils import label_map_util

载入模型

MODEL_NAME = 'data/models/ssd_mobilenet_v2_coco_2018_03_29'
PATH_TO_CKPT = MODEL_NAME + '/frozen_inference_graph.pb'
PATH_TO_LABELS = os.path.join('data/pbtxt','mscoco_label_map.pbtxt')  # CWH: Add object_detection path
# data下mscoco_label_map.pbtxt最大item.id
NUM_CLASSES = 90
detection_graph = tf.Graph()
with detection_graph.as_default():
  od_graph_def = tf.GraphDef()
  # 加载模型
  with tf.gfile.GFile(PATH_TO_CKPT, 'rb') as fid:
    serialized_graph = fid.read()
    od_graph_def.ParseFromString(serialized_graph)
    tf.import_graph_def(od_graph_def, name='')

载入标签映射，内置函数返回整数会映射到pbtxt字符标签

mscoco_label_map.pbtxt格式如下

item {
  name: "/m/01g317"
  id: 1
  display_name: "person"
}
item {
  name: "/m/0199g"
  id: 2
  display_name: "bicycle"
}

# 加载标签
label_map = label_map_util.load_labelmap(PATH_TO_LABELS)
categories = label_map_util.convert_label_map_to_categories(
    label_map, max_num_classes=NUM_CLASSES, use_display_name=True)
category_index = label_map_util.create_category_index(categories)

with detection_graph.as_default():
  config = tf.ConfigProto()
  config.gpu_options.allow_growth = True
  with tf.Session(graph=detection_graph,config=config) as sess:
    image_tensor = detection_graph.get_tensor_by_name('image_tensor:0')
    # 物体坐标
    detection_boxes = detection_graph.get_tensor_by_name('detection_boxes:0')
    # 检测到物体的准确度
    detection_scores = detection_graph.get_tensor_by_name('detection_scores:0')
    detection_classes = detection_graph.get_tensor_by_name('detection_classes:0')
    num_detections = detection_graph.get_tensor_by_name('num_detections:0')

def get_objects(file_name, threshold=0.5):
    image = Image.open(file_name)
    # 判断文件是否是jpeg格式
    if not image.format=='JPEG':
        result['status'] = 0
        result['msg'] = file_name+ ' is ' + image.format + ' ods system allow jpeg or jpg'
        return result
    image_np = load_image_into_numpy_array(image)
    # 扩展维度
    image_np_expanded = np.expand_dims(image_np, axis=0)
    output = []
    # 获取运算结果
    (boxes, scores, classes, num) = sess.run(
        [detection_boxes, detection_scores, detection_classes, num_detections],
        feed_dict={image_tensor: image_np_expanded})
    # 去掉纬度为1的数组
    classes = np.squeeze(classes).astype(np.int32)
    scores = np.squeeze(scores)
    boxes = np.squeeze(boxes)
    for c in range(0, len(classes)):
        if scores[c] >= threshold:
            item = Object()
            item.class_name = category_index[classes[c]]['name'] # 物体名称
            item.score = float(scores[c]) # 准确率
            # 物体坐标轴百分比
            item.y1 = float(boxes[c][0])
            item.x1 = float(boxes[c][1])
            item.y2 = float(boxes[c][2])
            item.x2 = float(boxes[c][3])
            output.append(item)
    # 返回JSON格式
    outputJson = json.dumps([ob.__dict__ for ob in output])
    return outputJson

运行路由

server.py下的逻辑

def image():
    startTime = time.time()
    if request.method=='POST':
        image_file = request.files['file']
        base_path = os.path.abspath(os.path.dirname(__file__))
        upload_path = os.path.join(base_path,'static/upload/')
        # 保存上传图片文件
        file_name = upload_path + image_file.filename
        image_file.save(file_name)
        # 准确率过滤值
        threshold = request.form.get('threshold',0.5)
        # 调用Api服务
        objects = object_detection_api.get_objects(file_name, threshold)
        # 模板显示
        return render_template('index.html',json_data = objects,img=image_file.filename)

curl http://localhost:5000 | python -m json.tool

[
    {
        "y2": 0.9886252284049988,
        "class_name": "bed",
        "x2": 0.4297400414943695,
        "score": 0.9562674164772034,
        "y1": 0.5202791094779968,
        "x1": 0
    },
    {
        "y2": 0.9805927872657776,
        "class_name": "couch",
        "x2": 0.4395904541015625,
        "score": 0.6422878503799438,
        "y1": 0.5051193833351135,
        "x1": 0.00021047890186309814
    }
]

在浏览器访问网址体验

http://localhost:5000/upload

小结

Tensorflow使用GPU效率提升几个数量级
可以尝试不同的模型比较速度和准确度
本案例也是支持python2,为了跟上时代步伐，建议使用python3

大家肯定很好奇，怎么训练自己需要检测的物体，可以期待下一篇文章