ROS機器人Diego 1#整合Tensorflow object_detection，圖像識別

google最近又公佈了物體識別的Api，使得圖像識別變得更加方便，並提供了一個預訓練模型，及示例代碼，官方文檔請見https://github.com/tensorflow/models/blob/master/object_detection/g3doc/installation.md
從官方提供的效果圖來看效果還是很不錯的，這篇文章就基於官方提供的示例代碼，製作一個ROS節點，訂閱Image主題，然後調用Object detection api來識別，再將識別的結果，通過CompressedImage消息發送出去。

1.安裝Object_detection

這裏假設已經安裝好了tensorflow，object_detection只要按照官方的安裝說明安裝即可，如下安裝的腳步

sudo apt-get install protobuf-compiler python-pil python-lxml
sudo pip install jupyter
sudo pip install matplotlib

2.創建diego_tensorflow 包

catkin_create_pkg diego_tensorflow std_msgs rospy roscpp cv_bridge

在diego_tensorflow目錄下創建scripts和launch目錄

scripts目錄用於存放Python的源代碼
launch目錄用於存放ROS launch文件

下載object_detection代碼到scripts目錄

3.ROS節點

代碼請見ObjectDetectionDemo.py

#!/usr/bin/env python

import rospy
from sensor_msgs.msg import Image as ROSImage
from sensor_msgs.msg import CompressedImage as ROSImage_C
from cv_bridge import CvBridge
import cv2
import matplotlib
import numpy as np
import os
import six.moves.urllib as urllib
import sys
import tarfile
import tensorflow as tf
import zipfile

from collections import defaultdict
from io import StringIO
from PIL import Image

# This is needed since the notebook is stored in the object_detection folder.
from object_detection.utils import label_map_util
from object_detection.utils import visualization_utils as vis_util

class ObjectDetectionDemo():
    def __init__(self):
    rospy.init_node('object_detection_demo')

    # Set the shutdown function (stop the robot)
        rospy.on_shutdown(self.shutdown)

        model_path = rospy.get_param("~model_path", "")
        image_topic = rospy.get_param("~image_topic", "")

        self._cv_bridge = CvBridge()

        # What model to download.
    MODEL_NAME = 'ssd_mobilenet_v1_coco_11_06_2017'
    MODEL_FILE = MODEL_NAME + '.tar.gz'
    DOWNLOAD_BASE = 'http://download.tensorflow.org/models/object_detection/'

    # Path to frozen detection graph. This is the actual model that is used for the object detection.
    PATH_TO_CKPT = MODEL_NAME + '/frozen_inference_graph.pb'

    # List of the strings that is used to add correct label for each box.
    PATH_TO_LABELS = os.path.join(model_path+'/data', 'mscoco_label_map.pbtxt')
    #PATH_TO_LABELS = os.path.join('data', 'mscoco_label_map.pbtxt')

    NUM_CLASSES = 90

    opener = urllib.request.URLopener()
    opener.retrieve(DOWNLOAD_BASE + MODEL_FILE, MODEL_FILE)
    tar_file = tarfile.open(MODEL_FILE)
    for file in tar_file.getmembers():
        file_name = os.path.basename(file.name)
        if 'frozen_inference_graph.pb' in file_name:
                tar_file.extract(file, os.getcwd())

    self.detection_graph = tf.Graph()
    with self.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='')

    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)
    self.category_index = label_map_util.create_category_index(categories)

    self._sub = rospy.Subscriber(image_topic, ROSImage, self.callback, queue_size=1)    
    self._pub = rospy.Publisher('object_detection', ROSImage_C, queue_size=1)


    def callback(self, image_msg):
    with self.detection_graph.as_default():
        with tf.Session(graph=self.detection_graph) as sess:
             cv_image = self._cv_bridge.imgmsg_to_cv2(image_msg, "bgr8")
             pil_img = Image.fromarray(cv_image)             
             (im_width, im_height) = pil_img.size            
             # the array based representation of the image will be used later in order to prepare the
             # result image with boxes and labels on it.
             image_np =np.array(pil_img.getdata()).reshape((im_height, im_width, 3)).astype(np.uint8)
             # Expand dimensions since the model expects images to have shape: [1, None, None, 3]
             image_np_expanded = np.expand_dims(image_np, axis=0)
             image_tensor = self.detection_graph.get_tensor_by_name('image_tensor:0')
             # Each box represents a part of the image where a particular object was detected.
             boxes = self.detection_graph.get_tensor_by_name('detection_boxes:0')
             # Each score represent how level of confidence for each of the objects.
             # Score is shown on the result image, together with the class label.
             scores = self.detection_graph.get_tensor_by_name('detection_scores:0')
             classes = self.detection_graph.get_tensor_by_name('detection_classes:0')
             num_detections = self.detection_graph.get_tensor_by_name('num_detections:0')
             # Actual detection.
             (boxes, scores, classes, num_detections) = sess.run(
                [boxes, scores, classes, num_detections],
                feed_dict={image_tensor: image_np_expanded})
             # Visualization of the results of a detection.
             vis_util.visualize_boxes_and_labels_on_image_array(
                image_np,
                np.squeeze(boxes),
                np.squeeze(classes).astype(np.int32),
                np.squeeze(scores),
                self.category_index,
                use_normalized_coordinates=True,
                line_thickness=8)

             ros_compressed_image=self._cv_bridge.cv2_to_compressed_imgmsg(image_np)
             self._pub.publish(ros_compressed_image)


    def shutdown(self):
        rospy.loginfo("Stopping the tensorflow object detection...")
        rospy.sleep(1) 

if __name__ == '__main__':
    try:
        ObjectDetectionDemo()
        rospy.spin()
    except rospy.ROSInterruptException:
        rospy.loginfo("RosTensorFlow_ObjectDetectionDemo has started.")

這段代碼主要是將官方提供的示例代碼，封裝爲ROS節點，訂閱image主題，結果作爲CompressedImage發佈，具體的識別原理可以查看官方的說明，這裏只介紹ROS封裝部分的代碼

    def __init__(self):
    rospy.init_node('object_detection_demo')

    # Set the shutdown function (stop the robot)
        rospy.on_shutdown(self.shutdown)

        model_path = rospy.get_param("~model_path", "")
        image_topic = rospy.get_param("~image_topic", "")

以上代碼是標準的ROS初始化話代碼，和參數讀取代碼

    self._sub = rospy.Subscriber(image_topic, ROSImage, self.callback, queue_size=1)    
    self._pub = rospy.Publisher('object_detection', ROSImage_C, queue_size=1)

以上代碼訂閱了Image主題，並定義回調函數callback。定義發佈的主題爲object_detection，後續我們可以訂閱object_detection主題來顯示識別結果

    def callback(self, image_msg):
    with self.detection_graph.as_default():
        with tf.Session(graph=self.detection_graph) as sess:
             cv_image = self._cv_bridge.imgmsg_to_cv2(image_msg, "bgr8")
             pil_img = Image.fromarray(cv_image)             
             (im_width, im_height) = pil_img.size            
             # the array based representation of the image will be used later in order to prepare the
             # result image with boxes and labels on it.
             image_np =np.array(pil_img.getdata()).reshape((im_height, im_width, 3)).astype(np.uint8)
             # Expand dimensions since the model expects images to have shape: [1, None, None, 3]
             image_np_expanded = np.expand_dims(image_np, axis=0)
             image_tensor = self.detection_graph.get_tensor_by_name('image_tensor:0')
             # Each box represents a part of the image where a particular object was detected.
             boxes = self.detection_graph.get_tensor_by_name('detection_boxes:0')
             # Each score represent how level of confidence for each of the objects.
             # Score is shown on the result image, together with the class label.
             scores = self.detection_graph.get_tensor_by_name('detection_scores:0')
             classes = self.detection_graph.get_tensor_by_name('detection_classes:0')
             num_detections = self.detection_graph.get_tensor_by_name('num_detections:0')
             # Actual detection.
             (boxes, scores, classes, num_detections) = sess.run(
                [boxes, scores, classes, num_detections],
                feed_dict={image_tensor: image_np_expanded})
             # Visualization of the results of a detection.
             vis_util.visualize_boxes_and_labels_on_image_array(
                image_np,
                np.squeeze(boxes),
                np.squeeze(classes).astype(np.int32),
                np.squeeze(scores),
                self.category_index,
                use_normalized_coordinates=True,
                line_thickness=8)

             ros_compressed_image=self._cv_bridge.cv2_to_compressed_imgmsg(image_np)
             self._pub.publish(ros_compressed_image)

callback函數是主要的處理函數，將攝像頭捕捉到的圖片，經過識別處理後發佈爲object_detection主題

4.launch文件

在launch文件目錄下創建object_detection_demo.launch，文件內容如下

<launch>

    <node pkg="diego_tensorflow" name="ObjectDetectionDemo" type="ObjectDetectionDemo.py" output="screen">
        <param name="image_topic" value="/usb_cam/image_raw" />
        <param name="model_path" value="$(find diego_tensorflow)/scripts/object_detection" />
    </node> 
</launch>

5.啓動

roscore
roslaunch usb_cam usb_cam-test.launch
roslaunch diego_tensorflow object_detection_demo.launch

6.通過手機APP訂閱識別處理結果

我們可以用手機APP訂閱處理結果，只需要將Image的主題設置爲/object_detection就可以了

下圖爲增加深度信息的效果圖，深度信息是利用深度相機xtion計算出來，可以利用此信息來判斷識別物體與機器人之間的距離。

從測試結果開看，只要是模型中包含的，識別效果還是非常好的，只是我diego1# 雖然使用了I7的處理器，但還是有明顯的滯後，4核的I7cpu使用率已經非常高了，8g的內存也使用到了一半，看來還是需要專門的GPU來處理

Tensorflow object detection提供瞭如下所示的5種模型，這裏只是使用了ssd_mobilenet_v1_coco，這是一個精簡模型，適合在移動設備上使用，而更加準確的模型應該是faster_rcnn_inception_resnet_v2_atrous_coco，但運行此模型需要專門的GPU，在i7的mini pc上如果針對每個視頻幀都處理的話，處理速度基本上和蝸牛差不多，但識別率，明顯高很多。