敵ロボット検出を学習する(めざせ!ROBO-ONE auto出場)

第7回ROBO-ONE autoの「リモート Target を倒せ」の敵ロボット検出をやってみる。ちなみに大会自体はすでに終了している。

f:id:touch-sp:20201029220551p:plain
第7回ROBO-ONE auto リモート競技規則

ロボットの写真は2L版(127×178mm)に印刷した。
50㎝の距離での物体検出なので2Lペットボトルの大きさは必要ない。
実際のサイズ感はこんな感じ。真ん中は500mlのペットボトル。
f:id:touch-sp:20220412213519j:plain:w320

今まで

など断片的にブログの記事にしてきたが今回はその集大成。

結果

先に結果を示す。GIFに変換した動画を載せておく。
f:id:touch-sp:20201101062146g:plain
ロボット写真が赤で囲われて、緑画用紙が緑で囲われているのでうまくいっている。
アクションカメラの映像をWi-Fi(RTSP)経由でPCに送信し、そのPCで画像処理を行っているので遅延はほとんどない。
 
最初に撮影する動画の長さにもよるが、今回は学習を含めたすべての過程で1時間かかっていない。

以下に実際やった過程を示す

①ビデオを撮る

動画は640×480サイズ、fps 30でロボットを映したもの1分20秒(video_1.mp4)、緑を映したもの47秒(video_2.mp4)撮影した。
撮影はカメラを固定して、撮影範囲内で対象物を動かした。
↓サンプル
https://raw.githubusercontent.com/dmlc/web-data/master/gluoncv/tracking/Coke.mp4
最終的にロボット画像 80×30で約2400枚、緑の画像 47×30で約1400枚準備できる。

②Object Trackingを行い、結果をVOC formatで出力する

最初の物体の位置を知る方法はこちらを参照。
touch-sp.hatenablog.com

ロボットの動画:
ファイル名「video_1.mp4」
最初の物体の位置[462, 100, 84, 120]
物体の名前「target」
保存先「VOC2020」
以下のスクリプトを実行してVOC formatで出力する。

import os
import mxnet as mx
from gluoncv import model_zoo
from gluoncv.model_zoo.siamrpn.siamrpn_tracker import SiamRPNTracker
import cv2
import xml.etree.cElementTree as ET

#=========================================================
ctx = mx.gpu()

video_path = 'video_1.mp4'

target_name = 'target'

#最初のポジション
#(左上X座標、左上Y座標、横の大きさ、縦の大きさ)
gt_bbox = [462, 100, 84, 120]

out_path = 'VOC2020'
#=========================================================

annotation_dir = os.path.join('new_data', out_path, 'Annotations')
main_dir =  os.path.join('new_data', out_path, 'ImageSets/Main')
jpegimages_dir = os.path.join('new_data', out_path, 'JPEGImages')

os.makedirs(annotation_dir)
os.makedirs(main_dir)
os.makedirs(jpegimages_dir)

# mp4データを読み込む
video_frames = []
cap = cv2.VideoCapture(video_path)
while(True):
    ret, img = cap.read()
    if not ret:
        break
    video_frames.append(img)

# モデルを取得する
net = model_zoo.get_model('siamrpn_alexnet_v2_otb15', pretrained=True, root='./models', ctx=ctx)
tracker = SiamRPNTracker(net)

jpeg_filenames_list = []

for ind, frame in enumerate(video_frames):
    if ind == 0:
        tracker.init(frame, gt_bbox, ctx=ctx)
        pred_bbox = gt_bbox
    else:
        outputs = tracker.track(frame, ctx=ctx)
        pred_bbox = outputs['bbox']

    pred_bbox = list(map(int, pred_bbox))

    filename = '%06d'%(ind)

    #画像の保存
    jpeg_filename = filename + '.jpg'
    cv2.imwrite(os.path.join(jpegimages_dir, jpeg_filename), frame)

    #テキストファイルの作成
    jpeg_filenames_list.append(filename)

    #XMLファイルの保存
    xml_filename = filename + '.xml'
    
    new_root = ET.Element('annotation')
    
    new_filename = ET.SubElement(new_root, 'filename')
    new_filename.text = jpeg_filename

    Size = ET.SubElement(new_root, 'size')
    Width = ET.SubElement(Size, 'width')
    Height = ET.SubElement(Size, 'height')
    Depth = ET.SubElement(Size, 'depth')

    Width.text = str(frame.shape[1])
    Height.text = str(frame.shape[0])
    Depth.text = str(frame.shape[2])

    Object = ET.SubElement(new_root, 'object')
    
    Name = ET.SubElement(Object, 'name')
    Name.text = target_name

    Difficult = ET.SubElement(Object, 'difficult')
    Difficult.text = '0'

    Bndbox = ET.SubElement(Object, 'bndbox')
    Xmin = ET.SubElement(Bndbox, 'xmin')
    Ymin = ET.SubElement(Bndbox, 'ymin')
    Xmax = ET.SubElement(Bndbox, 'xmax')
    Ymax = ET.SubElement(Bndbox, 'ymax')

    Xmin.text = str(pred_bbox[0])
    Ymin.text = str(pred_bbox[1])
    Xmax.text = str(pred_bbox[0]+pred_bbox[2])
    Ymax.text = str(pred_bbox[1]+pred_bbox[3])

    new_tree = ET.ElementTree(new_root) 

    new_tree.write(os.path.join(annotation_dir, xml_filename))

#テキストファイルの保存
text = "\n".join(jpeg_filenames_list)
with open(os.path.join(main_dir, 'train.txt'), "w") as f:
    f.write(text)


緑の動画:
ファイル名「video_2.mp4」
最初の物体の位置[429, 182, 84, 118]
物体の名前「green」
保存先「VOC2021」
以下のスクリプトを実行してVOC formatで出力する。

import os
import mxnet as mx
from gluoncv import model_zoo
from gluoncv.model_zoo.siamrpn.siamrpn_tracker import SiamRPNTracker
import cv2
import xml.etree.cElementTree as ET

#=========================================================
ctx = mx.gpu()

video_path = 'video_2.mp4'

target_name = 'green'

#最初のポジション
#(左上X座標、左上Y座標、横の大きさ、縦の大きさ)
gt_bbox = [429, 182, 84, 118]

out_path = 'VOC2021'
#=========================================================

annotation_dir = os.path.join('new_data', out_path, 'Annotations')
main_dir =  os.path.join('new_data', out_path, 'ImageSets/Main')
jpegimages_dir = os.path.join('new_data', out_path, 'JPEGImages')

os.makedirs(annotation_dir)
os.makedirs(main_dir)
os.makedirs(jpegimages_dir)

# mp4データを読み込む
video_frames = []
cap = cv2.VideoCapture(video_path)
while(True):
    ret, img = cap.read()
    if not ret:
        break
    video_frames.append(img)

# モデルを取得する
net = model_zoo.get_model('siamrpn_alexnet_v2_otb15', pretrained=True, root='./models', ctx=ctx)
tracker = SiamRPNTracker(net)

jpeg_filenames_list = []

for ind, frame in enumerate(video_frames):
    if ind == 0:
        tracker.init(frame, gt_bbox, ctx=ctx)
        pred_bbox = gt_bbox
    else:
        outputs = tracker.track(frame, ctx=ctx)
        pred_bbox = outputs['bbox']

    pred_bbox = list(map(int, pred_bbox))

    filename = '%06d'%(ind)

    #画像の保存
    jpeg_filename = filename + '.jpg'
    cv2.imwrite(os.path.join(jpegimages_dir, jpeg_filename), frame)

    #テキストファイルの作成
    jpeg_filenames_list.append(filename)

    #XMLファイルの保存
    xml_filename = filename + '.xml'
    
    new_root = ET.Element('annotation')
    
    new_filename = ET.SubElement(new_root, 'filename')
    new_filename.text = jpeg_filename

    Size = ET.SubElement(new_root, 'size')
    Width = ET.SubElement(Size, 'width')
    Height = ET.SubElement(Size, 'height')
    Depth = ET.SubElement(Size, 'depth')

    Width.text = str(frame.shape[1])
    Height.text = str(frame.shape[0])
    Depth.text = str(frame.shape[2])

    Object = ET.SubElement(new_root, 'object')
    
    Name = ET.SubElement(Object, 'name')
    Name.text = target_name

    Difficult = ET.SubElement(Object, 'difficult')
    Difficult.text = '0'

    Bndbox = ET.SubElement(Object, 'bndbox')
    Xmin = ET.SubElement(Bndbox, 'xmin')
    Ymin = ET.SubElement(Bndbox, 'ymin')
    Xmax = ET.SubElement(Bndbox, 'xmax')
    Ymax = ET.SubElement(Bndbox, 'ymax')

    Xmin.text = str(pred_bbox[0])
    Ymin.text = str(pred_bbox[1])
    Xmax.text = str(pred_bbox[0]+pred_bbox[2])
    Ymax.text = str(pred_bbox[1]+pred_bbox[3])

    new_tree = ET.ElementTree(new_root) 

    new_tree.write(os.path.join(annotation_dir, xml_filename))

#テキストファイルの保存
text = "\n".join(jpeg_filenames_list)
with open(os.path.join(main_dir, 'train.txt'), "w") as f:
    f.write(text)


最終的なファイル構造はこのようになる。

new_data
|   
|   
+---VOC2020
|   +---Annotations
|   |       000000.xml
|   |       000001.xml
|   |       000002.xml
|   |       .....
|   |       .....
|   |       002384.xml
|   |       002385.xml
|   |       002386.xml
|   |       
|   +---ImageSets
|   |   +---Main
|   |           train.txt
|   |           
|   +---JPEGImages
|           000000.jpg
|           000001.jpg
|           000002.jpg
|           .....
|           .....
|           002384.jpg
|           002385.jpg
|           002386.jpg
|           
+---VOC2021
    +---Annotations
    |       000000.xml
    |       000001.xml
    |       000002.xml
    |       .....
    |       .....
    |       001404.xml
    |       001405.xml
    |       001406.xml
    |       
    +---ImageSets
    |   +---Main
    |           train.txt
    |           
    +---JPEGImages
            000000.jpg
            000001.jpg
            000002.jpg
            .....
            .....
            001404.jpg
            001405.jpg
            001406.jpg

③物体検出モデルを学習させる

以下のスクリプトを実行すると「ssd_512_mobilenet1.0_2classes.params」というパラメーターが保存される。

import time
import mxnet as mx
from mxnet import gluon, autograd
from mxnet.gluon.data import DataLoader

from gluoncv import model_zoo
from gluoncv.data import VOCDetection
from gluoncv.data.transforms import presets
from gluoncv.data.batchify import Tuple, Stack
from gluoncv.loss import SSDMultiBoxLoss

ctx = [mx.gpu()]

classes = ['target', 'green']

net = model_zoo.get_model('ssd_512_mobilenet1.0_voc', pretrained=True, ctx = ctx[0], root='./models')
net.reset_class(classes)
net.hybridize()

x = mx.nd.zeros(shape=(1, 3, 512, 512),ctx=ctx[0])
with autograd.train_mode():
    _, _, anchors = net(x)

batch_size = 16
num_workers = 0
epochs = 5

width, height = 512, 512  
train_transform = presets.ssd.SSDDefaultTrainTransform(width, height, anchors.as_in_context(mx.cpu()))
batchify_fn = Tuple(Stack(), Stack(), Stack())

VOCDetection.CLASSES = classes
train_dataset = VOCDetection(root='new_data',splits=((2020, 'train'),(2021, 'train')))

train_loader = DataLoader(
    train_dataset.transform(train_transform),
    batch_size,
    shuffle=True,
    batchify_fn=batchify_fn,
    last_batch='rollover',
    num_workers=num_workers)

mbox_loss = SSDMultiBoxLoss()
ce_metric = mx.metric.Loss('CrossEntropy')
smoothl1_metric = mx.metric.Loss('SmoothL1')

trainer = gluon.Trainer(
    net.collect_params(), 'sgd',
    {'learning_rate': 0.001, 'wd': 0.0005, 'momentum': 0.9})

for epoch in range(epochs):
    ce_metric.reset()
    smoothl1_metric.reset()
    tic = time.time()
    btic = time.time()
    for i, batch in enumerate(train_loader):
        batch_size = batch[0].shape[0]
        data = gluon.utils.split_and_load(batch[0], ctx_list=ctx, batch_axis=0)
        cls_targets = gluon.utils.split_and_load(batch[1], ctx_list=ctx, batch_axis=0)
        box_targets = gluon.utils.split_and_load(batch[2], ctx_list=ctx, batch_axis=0)
        with autograd.record():
            cls_preds = []
            box_preds = []
            for x in data:
                cls_pred, box_pred, _ = net(x)
                cls_preds.append(cls_pred)
                box_preds.append(box_pred)
            sum_loss, cls_loss, box_loss = mbox_loss(
                cls_preds, box_preds, cls_targets, box_targets)
            autograd.backward(sum_loss)
        trainer.step(1)
        ce_metric.update(0, [l * batch_size for l in cls_loss])
        smoothl1_metric.update(0, [l * batch_size for l in box_loss])
        name1, loss1 = ce_metric.get()
        name2, loss2 = smoothl1_metric.get()
        if i % 20 == 0:
            print('[Epoch {}][Batch {}], Speed: {:.3f} samples/sec, {}={:.3f}, {}={:.3f}'.format(
                epoch, i, batch_size/(time.time()-btic), name1, loss1, name2, loss2))
        btic = time.time()

net.save_parameters('ssd_512_mobilenet1.0_2classes.params')

③動かしてみる。

以下のスクリプトを実行した結果が上に示した結果である。

import mxnet as mx
import gluoncv
import time
import cv2, queue, threading

class VideoCapture:

    def __init__(self, name):
        self.cap = cv2.VideoCapture(name)
        self.q = queue.Queue()
        t = threading.Thread(target=self._reader)
        t.daemon = True
        t.start()

  # read frames as soon as they are available, keeping only most recent one
    def _reader(self):
        while True:
            ret, frame = self.cap.read()
            if not ret:
                break
            if not self.q.empty():
                try:
                    self.q.get_nowait()   # discard previous (unprocessed) frame
                except queue.Empty:
                    pass
            self.q.put(frame)

    def read(self):
        return self.q.get()

ctx = mx.gpu()
# Load the model
classes = ['target', 'green']
net = gluoncv.model_zoo.get_model('ssd_512_mobilenet1.0_voc', pretrained=True, root='./models')
net.reset_class(classes)
net.load_parameters('ssd_512_mobilenet1.0_2classes.params')
net.collect_params().reset_ctx(ctx)
# Compile the model for faster speed
net.hybridize()

cap = VideoCapture('rtsp://192.72.1.1:554/liveRTSP/av4/track0')
time.sleep(1)

while True:

    frame = cap.read()
    frame = mx.nd.array(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)).astype('uint8')    
    rgb_nd, frame = gluoncv.data.transforms.presets.ssd.transform_test(frame, short=512)

    # Run frame through network
    class_IDs, scores, bounding_boxes = net(rgb_nd.as_in_context(ctx))
  
    # Display the result
    img = gluoncv.utils.viz.cv_plot_bbox(frame, bounding_boxes[0], scores[0], class_IDs[0], class_names= ['enemy', 'green_signal'], thresh=0.85)
    gluoncv.utils.viz.cv_plot_image(img)
    
    # escを押したら終了
    if cv2.waitKey(1) == 27:
        break

cv2.destroyAllWindows()

このスクリプトはアクションカメラからのRTSPを処理しているので長くなってしまった。
USB接続のWebカメラやノートPC付属のカメラであればもう少し簡単に書ける。

import mxnet as mx
import gluoncv

import time
import cv2

ctx = mx.gpu()
# Load the model
classes = ['target', 'green']
net = gluoncv.model_zoo.get_model('ssd_512_mobilenet1.0_voc', pretrained=True, root='./models')
net.reset_class(classes)
net.load_parameters('ssd_512_mobilenet1.0_2classes.params')
net.collect_params().reset_ctx(ctx)
# Compile the model for faster speed
net.hybridize()

# Load the webcam handler
cap = cv2.VideoCapture(0, cv2.CAP_DSHOW)
# letting the camera autofocus
time.sleep(1)

while(True):
    # Load frame from the camera
    ret, frame = cap.read()

    # Image pre-processing
    frame = mx.nd.array(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)).astype('uint8')
    rgb_nd, frame = gluoncv.data.transforms.presets.ssd.transform_test(frame, short=512, max_size=700)

    # Run frame through network
    class_IDs, scores, bounding_boxes = net(rgb_nd)

    # Display the result
    img = gluoncv.utils.viz.cv_plot_bbox(frame, bounding_boxes[0], scores[0], class_IDs[0], class_names=net.classes)
    gluoncv.utils.viz.cv_plot_image(img)

    # escを押したら終了
    if cv2.waitKey(1) == 27:
        break

cap.release()
cv2.destroyAllWindows()

動作環境

Windows10 Pro
Visual Studio 2019 communityインストール済み(←たぶんこれが必要?)
NVIDIA GeForce GTX1080
Python 3.7.9
CUDA 10.1

もしVisual Studioをインストールしていなくてスクリプトが再現できなければ、試しにインストールしてみるのが良い。
Visual Studioは2017でも良いかもしれない。
Visual Studioはcommunity版であれば無料でインストールできる。
またVisual Studioが必要だとして、Visual Studioの中のどれが必要でどれが必要でないかはよくわからない。
とりあえず自分がインストールしているのは以下の通り。
f:id:touch-sp:20201024220255p:plain

「mxnet-cu101」と「gluoncv」と「opencv-python」のみインストールした。その他は勝手についてきた。
cuDNNは別途入れていない。(こちらを参照)

pip install mxnet-cu101==1.7.0 -f https://dist.mxnet.io/python/cu101
pip install gluoncv
pip install opencv-python

certifi==2020.6.20
chardet==3.0.4
cycler==0.10.0
gluoncv==0.8.0
graphviz==0.8.4
idna==2.6
kiwisolver==1.2.0
matplotlib==3.3.2
mxnet-cu101==1.7.0
numpy==1.16.6
opencv-python==4.4.0.44
Pillow==8.0.1
portalocker==2.0.0
pyparsing==2.4.7
python-dateutil==2.8.1
pywin32==228
requests==2.18.4
scipy==1.5.3
six==1.15.0
tqdm==4.50.2
urllib3==1.22

最後に

実は同じようなことは以前にもやっっている。
touch-sp.hatenablog.com
今回は、物体検出モデルの学習までを以前より簡単にできるように工夫した。
また前回は1クラスであったが今回は2クラス。

GTX1080搭載PCで処理しているが、たとえば同じことをRaspberry Piでやるとどうなるか。同じ計算速度がでるとは思えない。軽量な物体検出モデルに変更する必要があるかもしれない。この辺りはやってみないとわからない。

f:id:touch-sp:20201102223339p:plain

ちなみに今回のように外部PC(ロボットからみて外部)で処理することが許されているので必ずしもRaspberry Piなどで処理する必要はない。
以前Wi-FiとXBeeの組み合わせで車は動かすことができた。
touch-sp.hatenablog.com
ロボットを作る技術がもしあればRobo-Oneに出場できるかも。でもそんな技術はない。

2021年1月18日追記

AutoGluonを使うとスクリプトを短くできます。
touch-sp.hatenablog.com

2022年4月9日追記

PyTorchを使って同様のことをやりました。
touch-sp.hatenablog.com


このエントリーをはてなブックマークに追加