docker
python3
YOLOv3

Docker上でPython3からYOLO v3 を動かしてみた

はじめに

YOLOは C言語で作られた機械学習のフレームワークdarknet上で動く、物体検出のネットワークです。今回はYOLO v3をDocker上で動かしてみました。やったことは次の通りです。

  • darknetをCPU用にビルド
  • yolov3のモデルと、yolov3-tinyのモデルの両方を試す
  • Python2用のサンプルを拡張し、Python3から使えるようにする(GitHub)
    • ついでに枠とラベルを書き込むようにする

実験にあたっては、こちらの記事を参考にさせていただきました。

環境

  • ホストOS: Mac OS X 10.12.6 (Sierra)
  • Docker: Docker for Mac Version 17.06.0-ce-mac18 (18433)
  • コンテナのベース: Ubuntu 18.04 LTS

Dockerfile

次のモジュールやファイルをインストールしています。

  • python3
  • ツール群(git, vim, wget)
  • darknet
  • yolo v3モデル, yolo v3 tinyモデル
  • python3用のサンプル
# Ubuntu 18.04 and yolo v3
#  darknet: https://github.com/pjreddie/darknet

FROM ubuntu:18.04
MAINTAINER mganeko

#  
# -- if you are using docker behind proxy, please set ENV --
#

#ENV http_proxy "http://proxy.yourdomain.com:8080/"
#ENV https_proxy "http://proxy.yourdomain.com:8080/"

ENV DEBIAN_FRONTEND nonineractive

# -- build step --
RUN apt update && apt upgrade -y

RUN apt install python3 -y
RUN apt install python3-pip -y
RUN apt install python3-dev -y
RUN python3 -V
RUN pip3 -V
RUN pip3 install --upgrade pip
RUN pip -V

RUN apt install git -y
RUN apt install vim -y
RUN apt install wget -y

#
# ------ yolo v3 ---
#
RUN mkdir /root/work 
WORKDIR /root/work/
RUN git clone https://github.com/pjreddie/darknet.git
WORKDIR /root/work/darknet
RUN make
RUN wget https://pjreddie.com/media/files/yolov3.weights
RUN wget https://pjreddie.com/media/files/yolov3-tiny.weights

RUN ln -s /root/work/darknet/libdarknet.so /usr/lib/libdarknet.so

#-- copy ---
WORKDIR /root/work/
RUN git clone https://github.com/mganeko/python3_yolov3.git
RUN cp /root/work/python3_yolov3/darknet-tiny-label.py /root/work/darknet/python/


# --- for running --

WORKDIR /root/work/darknet/

ビルド手順

ターミナルから docker build を実行(イメージ名は適宜付けてください)

$ docker build -t mganeko/yolov3 -f Dockerfile .

起動手順

(1) ホスト側で作業ディレクトリを用意

$ mkdir temp

(2) 作業ディレクトリをマウントして、ターミナルから docker run でコンテナを起動

$ docker run -it -v `pwd`/temp:/root/work/temp mganeko/yolov3 bash

実行結果

コンテナ内のbashで実行します。

darknet (C言語サンプル)

yolo v3 通常モデル

$ ./darknet detect cfg/yolov3.cfg yolov3.weights data/dog.jpg
$ cp predictions.png ../temp

10秒程度かかりました。
data/dog.jpg: Predicted in 10.687456 seconds.
dog: 99%
truck: 92%
bicycle: 99%

dog_predictions.png

yolo v3 tiny モデル

$ ./darknet detect cfg/yolov3-tiny.cfg yolov3-tiny.weights data/dog.jpg

こちらは1秒弱です。だいぶ確信度が下がりましたが、速いです。

data/dog.jpg: Predicted in 0.968052 seconds.
dog: 57%
truck: 56%
car: 62%
bicycle: 59%

dog_tiny_predictions.png

トラックと自動車で2重に検出されています。

Python3 サンプル

darknetに付属のサンプルはPython2用の物だったので、次のように改造しました。(GitHub)

  • Python 3 用に変更
  • yolov3-tinyモデルを使う
  • 検出枠とラベルを合成した画像を出力(ファイル名は detect_result.png で固定)

コンテナ内のbashで次のように実行します。結果のファイルを、ホストからマウントしたディレクトリにコピーしています。

$ python3 python/darknet-tiny-label.py
$ cp detect_result.png ../temp

実行結果はこちらです。同じく1秒程度かかります。

result_label2.png

Python3での実装

Python3 で検出枠とラベルを合成するために、いくつかの処理を加えました。

Cのヘッダーファイルを参考に、使いたい関数の引数の型を定義します。

C関数の型定義
# --- add for save image, frame and label --
save_image = lib.save_image
save_image.argtypes = [IMAGE, c_char_p]
draw_box_width = lib.draw_box_width
draw_box_width.argtypes = [IMAGE, c_int, c_int, c_int, c_int, c_int, c_float, c_float, c_float]

load_alphabet = lib.load_alphabet
load_alphabet.restype = POINTER(POINTER(IMAGE))

get_label = lib.get_label
get_label.restype = IMAGE
get_label.argtypes = [POINTER(POINTER(IMAGE)), c_char_p, c_int]
draw_label = lib.draw_label
draw_label.argtypes = [IMAGE, c_int, c_int, IMAGE, POINTER(c_float)]

枠はdraw_box_width()で描画します。取得された領域の座標が、左上でなくて中心点の座標であることに、最初は気が付きませんでした。

枠を描画
  b = dets[j].bbox
  # -- rect -- (image, left, top, right, bottom, line_width, R, G, B)
  draw_box_width(im, int(b.x - b.w/2), int(b.y - b.h/2), int(b.x + b.w/2), int(b.y + b.h/2), int(4), 0.8, 0.8, 0.8)

ラベルの描画は2段階です。さらにその前にフォント情報(?)らしきものを読み込んでおく必要があります。

ラベルを描画
  alp = load_alphabet()

  #-- label --
  name = meta.names[i]
  label = get_label(alp, name, 18)  # ラベルをイメージとして取得
  draw_label(im, int(b.y - b.h/2), int(b.x - b.w/2), label, farray); #ラベルを描画
  free_image(label)

終わりに

darknet / YOLO v3 を動かすところまでは割とすんなり行ったのですが、Pythonから呼び出して、枠とラベルを合成するところに苦労しました。なんとかやり方は分かったので、他の応用ができそうです。
また、CPUではどうしても時間がかかるので、GPU環境でも試してみたいところです。