LoginSignup
13
7

More than 5 years have passed since last update.

@shunchan0677

Tensorflow_in_ROSのコードを解説する

Last updated at Posted at 2016-11-26

まえがき

  • 以前, Tensorflowの学習済みMNISTが動くROSのノードを作った.
  • 紹介記事をQiitaに, コードをGitHubにあげた.
  • 今回は上のQiita記事で書ききれなかったコード解説の部分を書く.

結果

細かいコード解説に入る前に実行するとどんなことができるかだけを先に示しておく.Screenshot from 2016-11-23 15_44_07.png

こんな感じで右がカメラから入力された手書きの9
左が学習済みのCNNからの推定結果であり, 9を見てちゃんと9を返している.
今回はこんな感じのCNNがROS上で動かせるようになるコードを書いた.
実行方法や準備等の詳細は以前のQiita記事参照のこと.

処理の概要

処理の概要を復習しておくと

  1. CNNを構成して学習済みファイルの読み出し
  2. カメラノードから画像情報をSubscribe
  3. 画像情報をMNISTのCNNに入るよう, 28*28に圧縮＆白黒反転2値化
  4. CNNに画像を見せて数字を推定
  5. 結果をPublish

という感じである.

コード全体

コード全体は以下のようになる.

tensorflow_in_ros_mnist.py
import rospy
from sensor_msgs.msg import Image
from std_msgs.msg import Int16
from cv_bridge import CvBridge
import cv2
import numpy as np
import tensorflow as tf


def weight_variable(shape):
  initial = tf.truncated_normal(shape, stddev=0.1)
  return tf.Variable(initial)

def bias_variable(shape):
  initial = tf.constant(0.1, shape=shape)
  return tf.Variable(initial)

def conv2d(x, W):
  return tf.nn.conv2d(x, W, strides=[1, 1, 1, 1], 
                      padding='SAME')

def max_pool_2x2(x):
  return tf.nn.max_pool(x, ksize=[1, 2, 2, 1],
                        strides=[1, 2, 2, 1], padding='SAME')


def makeCNN(x,keep_prob):
    # --- define CNN model
    W_conv1 = weight_variable([5, 5, 1, 32])
    b_conv1 = bias_variable([32])
    h_conv1 = tf.nn.relu(conv2d(x, W_conv1) + b_conv1)

    h_pool1 = max_pool_2x2(h_conv1)

    W_conv2 = weight_variable([3, 3, 32, 64])
    b_conv2 = bias_variable([64])
    h_conv2 = tf.nn.relu(conv2d(h_pool1, W_conv2) + b_conv2)

    h_pool2 = max_pool_2x2(h_conv2)

    W_fc1 = weight_variable([7 * 7 * 64, 1024])
    b_fc1 = bias_variable([1024])
    h_pool2_flat = tf.reshape(h_pool2, [-1, 7 * 7 * 64])
    h_fc1 = tf.nn.relu(tf.matmul(h_pool2_flat, W_fc1) + b_fc1)

    h_fc1_drop = tf.nn.dropout(h_fc1, keep_prob)

    W_fc2 = weight_variable([1024, 10])
    b_fc2 = bias_variable([10])

    y_conv = tf.nn.softmax(tf.matmul(h_fc1_drop, W_fc2) + b_fc2)

    return y_conv



class RosTensorFlow():
    def __init__(self):
        self._cv_bridge = CvBridge()
        self._sub = rospy.Subscriber('image', Image, self.callback, queue_size=1)
        self._pub = rospy.Publisher('result', Int16, queue_size=1)

        self.x = tf.placeholder(tf.float32, [None,28,28,1], name="x")
        self.keep_prob = tf.placeholder("float")
        self.y_conv = makeCNN(self.x,self.keep_prob)

        self._saver = tf.train.Saver()
        self._session = tf.InteractiveSession()

        init_op = tf.initialize_all_variables()
        self._session.run(init_op)

        self._saver.restore(self._session, "model.ckpt")


    def callback(self, image_msg):
        cv_image = self._cv_bridge.imgmsg_to_cv2(image_msg, "bgr8")
        cv_image_gray = cv2.cvtColor(cv_image, cv2.COLOR_RGB2GRAY)
        ret,cv_image_binary = cv2.threshold(cv_image_gray,128,255,cv2.THRESH_BINARY_INV)
        cv_image_28 = cv2.resize(cv_image_binary,(28,28))
        np_image = np.reshape(cv_image_28,(1,28,28,1))
        predict_num = self._session.run(self.y_conv, feed_dict={self.x:np_image,self.keep_prob:1.0})
        answer = np.argmax(predict_num,1)
        rospy.loginfo('%d' % answer)
        self._pub.publish(answer)

    def main(self):
        rospy.spin()

if __name__ == '__main__':
    rospy.init_node('rostensorflow')
    tensor = RosTensorFlow()
    tensor.main()

コード解説

import部分

tensorflow_in_ros_mnist.py
import rospy
from sensor_msgs.msg import Image
from std_msgs.msg import Int16
from cv_bridge import CvBridge
import cv2
import numpy as np
import tensorflow as tf

今回はPythonでROSのノードを立てたかったので, rospyを入れた.
また、画像を読み込むためのImageと書き出すためのInt16, ROSのメッセージファイルをOpenCVに渡すためのcv_bridge, OpenCV, Numpy, 機械学習のためのTensorflowを入れた.

CNN定義部分

tensorflow_in_ros_mnist.py
def weight_variable(shape):
  initial = tf.truncated_normal(shape, stddev=0.1)
  return tf.Variable(initial)

def bias_variable(shape):
  initial = tf.constant(0.1, shape=shape)
  return tf.Variable(initial)

def conv2d(x, W):
  return tf.nn.conv2d(x, W, strides=[1, 1, 1, 1], 
                      padding='SAME')

def max_pool_2x2(x):
  return tf.nn.max_pool(x, ksize=[1, 2, 2, 1],
                        strides=[1, 2, 2, 1], padding='SAME')


def makeCNN(x,keep_prob):
    # --- define CNN model
    W_conv1 = weight_variable([5, 5, 1, 32])
    b_conv1 = bias_variable([32])
    h_conv1 = tf.nn.relu(conv2d(x, W_conv1) + b_conv1)

    h_pool1 = max_pool_2x2(h_conv1)

    W_conv2 = weight_variable([3, 3, 32, 64])
    b_conv2 = bias_variable([64])
    h_conv2 = tf.nn.relu(conv2d(h_pool1, W_conv2) + b_conv2)

    h_pool2 = max_pool_2x2(h_conv2)

    W_fc1 = weight_variable([7 * 7 * 64, 1024])
    b_fc1 = bias_variable([1024])
    h_pool2_flat = tf.reshape(h_pool2, [-1, 7 * 7 * 64])
    h_fc1 = tf.nn.relu(tf.matmul(h_pool2_flat, W_fc1) + b_fc1)

    h_fc1_drop = tf.nn.dropout(h_fc1, keep_prob)

    W_fc2 = weight_variable([1024, 10])
    b_fc2 = bias_variable([10])

    y_conv = tf.nn.softmax(tf.matmul(h_fc1_drop, W_fc2) + b_fc2)

    return y_conv

TensorflowのTutorial(Deep MNIST for Expert)どおりの構成でCNNを定義する関数を作った.
構成としては畳み込みプーリングを2回行った後, 全結合層を入れ, 各数字に対しての確率を出すというモデルになっている.

classのinit部分前半

tensorflow_in_ros_mnist.py
class RosTensorFlow():
    def __init__(self):
        self._cv_bridge = CvBridge()
        self._sub = rospy.Subscriber('image', Image, self.callback, queue_size=1)
        self._pub = rospy.Publisher('result', Int16, queue_size=1)

前半はROSの処理.
ここでは, CvBridge関数の呼び出し, Subscriber, Publisherの定義を行います.
今回はSubscriberはImage型のメッセージを受け取り, PublisherはInt16型のメッセージを配信する.
基本的にはROSの例どおり.
subの部分の引数にcallbackが入ることで, Imageメッセージを受け取るごとにcallback関数を呼び出す.

classのinit部分後半

tensorflow_in_ros_mnist.py
        self.x = tf.placeholder(tf.float32, [None,28,28,1], name="x")
        self.keep_prob = tf.placeholder("float")
        self.y_conv = makeCNN(self.x,self.keep_prob)

        self._saver = tf.train.Saver()
        self._session = tf.InteractiveSession()

        init_op = tf.initialize_all_variables()
        self._session.run(init_op)

        self._saver.restore(self._session, "model.ckpt")

後半はTensorflowの処理.
ここではまず画像が入るplaceholderであるxとDropOut率が入るplaceholderであるkeep_probを定義する.
placeholderとはデータの入り口みたいなもので実行時にはここにデータがどんどん入る.

次に今回使うCNNをy_convとして定義する.
イメージ的にはx,keep_probというデータの入り口からCNNを通ってy_convという名の出口にデータが出てくるという道筋を定義する感じ.

道筋を定義したらsession関数で実際にデータを流す準備する.　tf.initialize_all_variables関数で一旦CNNの重みWとバイアスbを初期化.

ここに学習済みパラメータを読み出す. パラメータの読み出しにはtf.train.Saver関数を使ったあとにsaver.restoreをする必要あり.

callback部

tensorflow_in_ros_mnist.py
    def callback(self, image_msg):
        cv_image = self._cv_bridge.imgmsg_to_cv2(image_msg, "bgr8")
        cv_image_gray = cv2.cvtColor(cv_image, cv2.COLOR_RGB2GRAY)
        ret,cv_image_binary = cv2.threshold(cv_image_gray,128,255,cv2.THRESH_BINARY_INV)
        cv_image_28 = cv2.resize(cv_image_binary,(28,28))
        np_image = np.reshape(cv_image_28,(1,28,28,1))
        predict_num = self._session.run(self.y_conv, feed_dict={self.x:np_image,self.keep_prob:1.0})
        answer = np.argmax(predict_num,1)
        rospy.loginfo('%d' % answer)
        self._pub.publish(answer)

画像メッセージが入ってくるたびにここが読まれる.
cv_bridgeでメッセージを画像に変えたあと, グレースケールにして2値化, 白黒反転,サイズ調節をして, CNNに画像を投げ込んでいる.
帰ってきた推定結果predict_numの一番確率が高いものをanswerとし,　publishする

main関数については省略する.

あとがき

これでTensorflowの学習済みモデルさえあれば, 何でもROSで扱えるようになった.
FasterRCNNで物体認識とか顔認識とかもできると思うのでやってみたい.

13
7
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
Sign upLogin
13
7