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()