Tensorflow之实现物体检测

项目背景

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

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

Tensorflow介绍

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

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

环境搭建

Linux/Unix(笔者使用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 'export PYTHONPATH=$PYTHONPATH:pwd/data/tf_models/models/research'>> ~/.bashrc && source ~/.bashrc

5、启动服务

python3 server.py

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

模型选用

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

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/pbtxt下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

    }

]

class_name表示物体标签名称
score 可信度值
x1,y1表示对象所在最左上点位置
x2,y2表示对象最右下点位置

在浏览器访问网址体验

http://localhost:5000/upload

小结

Tensorflow使用GPU效率提升几个数量级
可以尝试不同的模型比较速度和准确度
本案例也是支持python2,为了跟上时代步伐，建议使用python3
案例有个摄像头演示，需要https支持，且使用安卓系统

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

巴特西