Jetbot & PyTorch CNN（１）衝突回避用のデータ収集プログラム

Jetbot にはResNet18を使った衝突回避モデルがあり、それをベースにPyTorchのCNNのモデルに書き換えてみました。カメラ視野も制御に必要な下側25％にトリムして、処理の軽量化を狙います。作るプログラムは合計7つ。

　　１）衝突回避用のデータ収集プログラム
　　２）衝突回避用の学習プログラム
　　３）衝突回避用のデモプログラム
　　４）ライントレースのデータ収集プログラム
　　５）ライントレースの学習プログラム
　　６）ライントレースのデモプログラム
　　７）ライントレース＋衝突回避用のデモプログラム

今回は１）の「衝突回避用のデータ収集プログラム」です。data_collection_gamepad.ipynbを基に収集する写真のサイズを56x224変更します。

data_collection.ipynb

import traitlets
import ipywidgets.widgets as widgets
from IPython.display import display
from jetbot import Camera, bgr8_to_jpeg
import numpy as np

camera = Camera.instance(width=224, height=224)

image = widgets.Image(format='jpeg', width=224, height=224)

def display_trim(camera_image): #trim upper part # 写真サイズ変更

    image = np.copy(camera_image)
    image = image[168:224, 0:224]
    jpeg_image = bgr8_to_jpeg(image)
    return jpeg_image

#camera_link = traitlets.dlink((camera, 'value'), (image, 'value'), transform=bgr8_to_jpeg)
camera_link = traitlets.dlink((camera, 'value'), (image, 'value'), transform=display_trim)   # 写真サイズ変更

display(image)

create a few directories where we'll store all our data（変更なし）

data_collection.ipynb

import os

blocked_dir = 'dataset/blocked'
free_dir = 'dataset/free'

# we have this "try/except" statement because these next functions can throw an error if the directories exist already
try:
    os.makedirs(free_dir)
    os.makedirs(blocked_dir)
except FileExistsError:
    print('Directories not created becasue they already exist')

create and display some buttons （変更なし）

data_collection.ipynb

button_layout = widgets.Layout(width='128px', height='64px')
free_button = widgets.Button(description='add free', button_style='success', layout=button_layout)
blocked_button = widgets.Button(description='add blocked', button_style='danger', layout=button_layout)
free_count = widgets.IntText(layout=button_layout, value=len(os.listdir(free_dir)))
blocked_count = widgets.IntText(layout=button_layout, value=len(os.listdir(blocked_dir)))

display(widgets.HBox([free_count, free_button]))
display(widgets.HBox([blocked_count, blocked_button]))

uuid package in python（変更なし）

data_collection.ipynb

from uuid import uuid1

def save_snapshot(directory):
    image_path = os.path.join(directory, str(uuid1()) + '.jpg')
    with open(image_path, 'wb') as f:
        f.write(image.value)

def save_free():
    global free_dir, free_count
    save_snapshot(free_dir)
    free_count.value = len(os.listdir(free_dir))

def save_blocked():
    global blocked_dir, blocked_count
    save_snapshot(blocked_dir)
    blocked_count.value = len(os.listdir(blocked_dir))

# attach the callbacks, we use a 'lambda' function to ignore the
# parameter that the on_click event would provide to our function
# because we don't need it.
free_button.on_click(lambda x: save_free())
blocked_button.on_click(lambda x: save_blocked())

display(image)（変更なし）

data_collection.ipynb

display(image)
display(widgets.HBox([free_count, free_button]))
display(widgets.HBox([blocked_count, blocked_button]))

close the camera conneciton （変更なし）

data_collection.ipynb

camera.stop()

これで1/4サイズのデータが集められます。
意外と簡単！