Jetson + ROS

JetsonTX2 + ROS + Segnet

2017/7/19のROS勉強会にて[「Jetson + ROS」]というタイトルでLTを行いました．
スライドは以下で公開してあります．
https://docs.google.com/presentation/d/1edOO_PCEk5WQyNRetSBnmB5Bjew2wPAttXtYuYsHLhQ/edit#slide=id.p

内容は，TX2でSegnetをROSに繋げて動かしてみました，というお話です．

ハードウェア

動画は車にwebcameraを貼り付けまして，

助手席にPCを置いてrosbagをとりました．

動画ではこんな感じです．
https://drive.google.com/open?id=0B5oEuTcLOvlDaEIyemQyWEJ5bDg

jetsonを車に乗せた事もありましたが，PCの方が手軽です．

ソースコード

SegnetのROS対応コードを作ったので参考に公開しておきます．

Sensor_msgsのimageトピックを受け取り，セグメンテーション後のimageを出力します．

インストールの注意点として，Segnetをビルドする際にcaffeが必要なのですが，Alexさんの公開しているcaffe-segnetでないとエラーが出たのでこちらを入れる必要があります．

OTLさんの本にある「ROS」を使って「分散画像処理」の章を参考にしながら，Segnet_Tutorialのwebcam_demo.pyを書き換えました．

ros_segnet.py

#!/usr/bin/env python

#$ python Scripts/ros_segnet.py --model Example_Models/segnet_model_driving_webdemo.prototxt --weights Example_Models/segnet_weights_driving_webdemo.caffemodel --colours Scripts/camvid12.png

import numpy as np
import matplotlib.pyplot as plt
import os.path
import scipy
import argparse
import math
import cv2
import sys
import time

#ros
import rospy
from sensor_msgs.msg import Image
from cv_bridge import CvBridge,CvBridgeError

class segnet:
    def __init__(self):
        self.image_in_pub = rospy.Publisher("segnet_input", Image)
        self.image_out_pub = rospy.Publisher("segnet_output", Image)
        self.bridge = CvBridge()
        self.image_sub = rospy.Subscriber("/image_topic",Image,self.callback)



    def callback(self,message):

        start = time.time() 
        try:
            cv2_imag = self.bridge.imgmsg_to_cv2(message,"bgr8")
        except CvBridgeError,e:
            rospy.loginfo(e)
        end = time.time()
        print '%30s' % 'image_input ', str((end - start)*1000), 'ms'

        start = time.time()
        cv2_imag = cv2_imag[0:380,0:640]    
        cv2_imag = cv2.resize(cv2_imag, (input_shape[3],input_shape[2]))
        input_image = cv2_imag.transpose((2,0,1))

        input_image = np.asarray([input_image])
        end = time.time()
        print '%30s' % 'Resized image in ', str((end - start)*1000), 'ms'

        start = time.time()
        out = net.forward_all(data=input_image)
        end = time.time()
        print '%30s' % 'Executed SegNet in ', str((end - start)*1000), 'ms'

        start = time.time()
        segmentation_ind = np.squeeze(net.blobs['argmax'].data)
        segmentation_ind_3ch = np.resize(segmentation_ind,(3,input_shape[2],input_shape[3]))
        segmentation_ind_3ch = segmentation_ind_3ch.transpose(1,2,0).astype(np.uint8)
        segmentation_rgb = np.zeros(segmentation_ind_3ch.shape, dtype=np.uint8)

        cv2.LUT(segmentation_ind_3ch,label_colours,segmentation_rgb)
        segmentation_rgb = segmentation_rgb.astype(float)/255

        end = time.time()
        print '%30s' % 'Processed results in ', str((end - start)*1000), 'ms\n'


        #cv2.imshow("Input", cv2_imag)  
        #cv2.imshow("SegNet", segmentation_rgb)
        #cv2.waitKey(3)

        segmentation_rgb = (segmentation_rgb*256).astype(np.uint8)

        try:
            self.image_in_pub.publish(self.bridge.cv2_to_imgmsg(cv2_imag, "bgr8"))
            self.image_out_pub.publish(self.bridge.cv2_to_imgmsg(segmentation_rgb, "bgr8"))
        except CvBridgeError as e:
            print(e)

def main(args):
  ic = image_converter()
  rospy.init_node('image_converter', anonymous=True)
  try:
    rospy.spin()
  except KeyboardInterrupt:
    print("Shutting down")
  cv2.destroyAllWindows()

def main():
    sys.path.append('/usr/local/lib/python2.7/site-packages')

    # Make sure that caffe is on the python path:
    caffe_root = '${HOME}/caffe-segnet/'
    sys.path.insert(0, caffe_root + 'python')
    import caffe

    # Import arguments
    parser = argparse.ArgumentParser()
    parser.add_argument('--model', type=str, required=True)
    parser.add_argument('--weights', type=str, required=True)
    parser.add_argument('--colours', type=str, required=True)
    args = parser.parse_args()

    global net
    net = caffe.Net(args.model,
                args.weights,
                caffe.TEST)

    caffe.set_mode_gpu()

    global  input_shape
    global output_label_coloursshape
    input_shape = net.blobs['data'].data.shape
    output_shape = net.blobs['argmax'].data.shape

    global label_colours
    label_colours = cv2.imread(args.colours).astype(np.uint8)

    sn = segnet()

    rospy.init_node('segnet_node',anonymous=True)
    rospy.loginfo("start")
    try:
        rospy.spin()
    except KeyboardInterrupt:
        pass
    cv2.destroyAllWindows()

if __name__ == '__main__':
    main()