Telloドローンでプログラミング！ーディープラーニングで物体認識編ー

こんにちは。ともよです。前回はgo言語でtello使ってカメラ画像を取得する方法をご紹介しました。

Telloドローンでプログラミング！ーMacでカメラ画像を取得する方法ー - 30歳からはじめる生きがい探し

今回は、telloで撮影した画像を使って、リアルタイムに物体認識をしてみたいと思います。

ディープラーニング環境の構築

Anacondaのインストール

python本体とpythonで良く使われるライブラリをセットにしたパッケージであるAnacondaをまずインストールします。この辺が参考になります。 qiita.com

Anacondaで仮想環境作成

"droneenv"という名前のPython3.6環境を作ります。

conda create -n droneenv python=3.6 anaconda

今回は、初心者にも理解しやすいKerasを使ってディープラーニングを行います。バックエンドはTensorFlowにするので、TensorFlowのインストールも必要です。仮想環境を有効化した上でこれらをインストールします。

source activate droneenv
sudo pip install tensorflow --upgrade
pip install keras

kerasのインストールが終わったら、OpenCVをインストールします。

pip install opencv-python

今回は、SSD(Single Shot MultiBox Detector)のKeras実装を使って、物体検出を行いたいと思います。git cloneでソースコードを取ってきます。こちらのQiitaがとても参考になりました。 qiita.com 上の記事ではPython2.X系でやっていますが、今回はPython 3.6なので、何点か修正しないと動かない箇所がありました。

Telloのカメラ画像取得

Telloのカメラ画像取得といえば、Go言語で書かれたGobotに実装されているものが一番有名なようですが、今回はこの後ディープラーニングを行うことを見越して、python実装を探しました。見つかったのは、以下の2つ。

TelloPy (https://github.com/hanyazou/TelloPy)
pytello (https://bitbucket.org/PingguSoft/pytello/src)

pytelloの方は動画を録画して保存する仕様でリアルタイムな映像取得までは実装されていませんでしたので、TelloPyを使うことにしました。

TelloPyのインストール

基本的にはTelloPyのサイトのREADME通りにインストールすれば良いです。ただし、私はPyAVのインストールでエラーが出たので、PyAVはcondaコマンドでインストールしました。

pip install tellopy
conda install av -c conda-forge
pip install image

TelloPyの実行

まずは、Telloの電源をONにした上で、Wi-fiでTelloとPCを接続します。Wifiアクセスポイントの中に"TELLO-XXXXX"という名前のものが現れますので、接続しましょう。接続した状態で以下のコマンドを実行してみてください。コマンドを実行すると、Telloのカメラ映像とエッジ抽出結果が表示されます。

python -m tellopy.examples.video_effect

MacBook Proで実行しましたが、やや映像の遅延が見られます。こちらに書いてあるように、4枚に1枚しか処理しないようにvideo_effect.pyを修正して実行すると、だいたいリアルタイムに表示されるようになりました。

opencv av stream · Issue #5 · hanyazou/TelloPy · GitHub

Telloのリアルタイム動画取得&画像認識

TelloPyのサイトにある、video_effect.pyをダウンロードしてきて編集しました。単純にSSDのコードと、video_effect.pyを組み合わせただけですが、今のところうまく動いています。ディープラーニング部分の処理負荷が大きいため、2秒に1回しか物体認識させないようにしています。

import cv2
import keras
from keras.applications.imagenet_utils import preprocess_input
from keras.backend.tensorflow_backend import set_session
from keras.models import Model
from keras.preprocessing import image
import matplotlib.pyplot as plt
import numpy as np
from scipy.misc import imread
import tensorflow as tf
from ssd import SSD300
from ssd_utils import BBoxUtility
from PIL import Image

import sys
import traceback
import tellopy
import av
import cv2.cv2 as cv2  # for avoidance of pylint error
import numpy

plt.rcParams['figure.figsize'] = (8, 8)
plt.rcParams['image.interpolation'] = 'nearest'

np.set_printoptions(suppress=True)

config = tf.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = 0.45
set_session(tf.Session(config=config))

voc_classes = ['Aeroplane', 'Bicycle', 'Bird', 'Boat', 'Bottle',
               'Bus', 'Car', 'Cat', 'Chair', 'Cow', 'Diningtable',
               'Dog', 'Horse','Motorbike', 'Person', 'Pottedplant',
               'Sheep', 'Sofa', 'Train', 'Tvmonitor']
NUM_CLASSES = len(voc_classes) + 1

input_shape=(300, 300, 3)
model = SSD300(input_shape, num_classes=NUM_CLASSES)
model.load_weights('weights_SSD300.hdf5', by_name=True)
bbox_util = BBoxUtility(NUM_CLASSES)


def getSSDImage(frame):

    img2 = image.img_to_array(frame.resize((300, 300)))
    img = np.asarray(frame)

    inputs = []
    inputs.append(img2.copy())
    inputs = preprocess_input(np.array(inputs))
    preds = model.predict(inputs, batch_size=1, verbose=1)
    results = bbox_util.detection_out(preds)

    # Parse the outputs.
    det_label = results[0][:, 0]
    det_conf = results[0][:, 1]
    det_xmin = results[0][:, 2]
    det_ymin = results[0][:, 3]
    det_xmax = results[0][:, 4]
    det_ymax = results[0][:, 5]

    # Get detections with confidence higher than 0.6.
    top_indices = [i for i, conf in enumerate(det_conf) if conf >= 0.6]

    top_conf = det_conf[top_indices]
    top_label_indices = det_label[top_indices].tolist()
    top_xmin = det_xmin[top_indices]
    top_ymin = det_ymin[top_indices]
    top_xmax = det_xmax[top_indices]
    top_ymax = det_ymax[top_indices]

    colors = plt.cm.hsv(np.linspace(0, 1, 21)).tolist()

    for i in range(top_conf.shape[0]):
        xmin = int(round(top_xmin[i] * img.shape[1]))
        ymin = int(round(top_ymin[i] * img.shape[0]))
        xmax = int(round(top_xmax[i] * img.shape[1]))
        ymax = int(round(top_ymax[i] * img.shape[0]))
        score = top_conf[i]
        label = int(top_label_indices[i])
        label_name = voc_classes[label - 1]
        display_txt = '{:0.2f}, {}'.format(score, label_name)

        color = colors[label]

        cv2.rectangle(img, (xmin, ymin), (xmax, ymax), (int(colors[label][0]*255), int(colors[label][1]*255), int(colors[label][2]*255)), 2)
        cv2.rectangle(img, (xmin, ymin-15), (xmin+100, ymin+5), (int(colors[label][0]*255), int(colors[label][1]*255), int(colors[label][2]*255)),-1)
        cv2.putText(img, display_txt, (xmin, ymin), cv2.FONT_HERSHEY_PLAIN, 1, (255, 255, 255), 1, cv2.LINE_AA)
    imgcv = img[:, :, ::-1].copy()
    return imgcv

def main():
    drone = tellopy.Tello()

    try:
        drone.connect()
        drone.wait_for_connection(60.0)
        drone.set_loglevel(drone.LOG_INFO)
        drone.set_exposure(0)

        container = av.open(drone.get_video_stream())
        frame_count = 0
        while True:
            for frame in container.decode(video=0):
                frame_count = frame_count + 1
                if (frame_count > 300) and (frame_count%50 == 0):
                    imgpil = frame.to_image()
                    image = getSSDImage(imgpil)
                    cv2.imshow('Original', image)
                    cv2.waitKey(1)

    except Exception as ex:
        exc_type, exc_value, exc_traceback = sys.exc_info()
        traceback.print_exception(exc_type, exc_value, exc_traceback)
        print(ex)
    finally:
        drone.quit()
        cv2.destroyAllWindows()

if __name__ == '__main__':
    main()

まとめ

トイドローンTelloで撮った画像を使って、SSDでリアルタイムに物体認識してみました。今の所、0.5FPSで動いています。（色々無駄な処理が入っているように思うのでもう少し高速化できそう...)。 Telloではコマンドから機体操縦もできるので、物体認識の結果によって、飛行経路を変えるようなアプリケーションも作れそうですね。

今回使用したドローン↓