【Coral USB アクセラレータ】Coral USB アクセラレータのデモを改良し、リアルタイム顔検出を行う

はじめに

こんにちは、がんがんです。
前回、Google Coral Edge TPU USB アクセラレータのデモをいろいろと実験してみました。
前回の記事はこちらです。
gangannikki.hatenadiary.jp


今回はこのデモを改良してリアルタイムで顔検出出来るようにしていきます。

目的

  • Coral USB アクセラレータのプログラムを改良
  • リアルタイム物体検出(バウンディングボックス有り)の作成
  • リアルタイム顔検出(こちらもバウンディングボックス有り)の作成

デモ内容の確認と動機

デモのプログラムには以下のようなものがあります。

  • classify_image.py

入力画像に写っているものは何かを判別するものです。よくある画像認識モデルです。

  • object_detection.py

入力された画像から物体を検出し、矩形をつけて出力するというものです。

  • classify_capture.py

リアルタイム物体認識を実現する。検出物体に矩形はつかない。


デモの詳しい詳細はこちらをどうぞ。
gangannikki.hatenadiary.jp


検出物体にバウンディングボックスを表示させたり、リアルタイムで顔認識+性別推定を行わせたかったのでデモのみではどうにも中途半端でした。いろいろと調べてみると海外の人でバウンディングボックスを表示させている人がいたのでそちらを参考にしていきます。

実験1 検出結果にバウンディングボックスを表示させる

まずは検出物体にバウンディングボックスを表示させます。コードから先に記載しておきます。GitHubの方にも挙げてます。

import cv2 
import numpy as np
import argparse, time, re

from edgetpu.detection.engine import DetectionEngine
from PIL import Image, ImageDraw, ImageFont

from imutils.video import FPS
from imutils.video import VideoStream


def ReadLabelFile(file_path):
    with open(file_path, "r", encoding="utf-8") as f:
        lines = f.readlines()
    ret = {}
    for line in lines:
        pair = line.strip().split(maxsplit=1)
        ret[int(pair[0])] = pair[1].strip()
    return ret


def draw_image(image, results, labels):
    result_size = len(results)
    for idx, obj in enumerate(results):
        #  Prepare image for drawing
        draw = ImageDraw.Draw(image)

        #  Prepare boundary box
        box = obj.bounding_box.flatten().tolist()

        #  Draw rectangle to desired thickness
        for x in range( 0, 4 ):
            draw.rectangle(box, outline=(0, 0, 255))

        #  Annotate image with label and confidence score
        display_str = labels[obj.label_id] + ": " + str(round(obj.score*100, 2)) + "%"
        draw.text((box[0], box[1]), display_str, font=ImageFont.truetype("/usr/share/fonts/truetype/piboto/Piboto-Regular.ttf", 20))

        displayImage = np.asarray(image)
        cv2.imshow("Coral Live Object Detection", displayImage)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--model", help="Path of the detection model.", required=True)
    parser.add_argument(
        "--label", help="Path of the labels file.")
    parser.add_argument(
        "--maxobjects", type=int, default=3, help="Maximum objects")
    parser.add_argument(
        "--threshold", type=float, default=0.3, help="Minimum threshold")
    parser.add_argument(
        "--picamera", action="store_true",
        help="Use PiCamera for image capture", default=False)
    args = parser.parse_args()

    #  Prepare labels.
    labels = ReadLabelFile(args.label) if args.label else None
    #  Initialize engine
    engine = DetectionEngine(args.model)

    #  Initialize video stream
    vs = VideoStream(usePiCamera=args.picamera, resolution=(640, 480)).start()
    time.sleep(1)

    fps = FPS().start()

    while True:
        try:
            #  Read frame from video
            screenshot = vs.read()
            image = Image.fromarray(screenshot)

            #  Perform inference
            results = engine.DetectWithImage(image, threshold=args.threshold,
                                             keep_aspect_ratio=True,
                                             relative_coord=False,
                                             top_k=args.maxobjects)

            #  draw image
            draw_image(image, results, labels)

            #  closing confition
            if cv2.waitKey(5) & 0xFF == ord("q"):
                fps.stop()
                break

            fps.update()
        except KeyboardInterrupt:
            fps.stop()
            break
    
    print("Elapsed time: {}".format(str(fps.elapsed())))
    print("Approx FPS: {}".format(str(fps.fps())))

    cv2.destroyAllWindows()
    vs.stop()
    time.sleep(2)

    
if __name__ == "__main__":
    main()

classify_capture.pyからの変更点

デモのリアルタイム認識ではpicameraでの使用を想定していました。参考記事ではこれをVideoStreamを用いることでUSBカメラでの使用を可能にしています。
実際デモのリアルタイム認識にDetectionEngineを加えてみて失敗しました。。。

注意点としてはOpenCVやimutils(VIdeoStreamの元)をインストールし忘れないことです。
OpenCVの方は少し苦戦したので別にまとめてあります。
gangannikki.hatenadiary.jp


imutilsは以下のコマンドでOKです。imutilsはPython2系でも動くので、Python3で動くようにインストールしてやる必要があります。

$ pip3 install imutils

実行結果

実際に実行させてみます。モデルはmobile_ssd_v2_cocoを使用します。
(デモの続きとして、test_data直下で実行してますが、ここはお好みで大丈夫です)

~/edgetpu_demo/test_data $ 
python3 ../edgetpu/demo/object_rect.py \
--model mobilenet_ssd_v2_coco_quant_postprocess_edgetpu.tflite \
--label coco_labels.txt

実験2 リアルタイム顔検出を行う

次は顔検出を行わせます。実験1からの変更点を載せておきます。顔検出時はラベルが存在しないため、if条件を追加しました。

def draw_image(image, results, labels):
    result_size = len(results)
    for idx, obj in enumerate(results):
        #  Prepare image for drawing
        draw = ImageDraw.Draw(image)

        #  Prepare boundary box
        box = obj.bounding_box.flatten().tolist()

        #  Draw rectangle to desired thickness
        for x in range( 0, 4 ):
            draw.rectangle(box, outline=(0, 0, 255))

        #  Annotate image with label and confidence score
        #(-)display_str = labels[obj.label_id] + ": " + str(round(obj.score*100, 2)) + "%"
        #(-)draw.text((box[0], box[1]), display_str, font=ImageFont.truetype("/usr/share/fonts/truetype/piboto/Piboto-Regular.ttf", 20))
        """(+)"""
        if labels:
            display_str = labels[obj.label_id] + ": " + str(round(obj.score*100, 2)) + "%"
            draw.text((box[0], box[1]), display_str, font=ImageFont.truetype("/usr/share/fonts/truetype/piboto/Piboto-Regular.ttf", 20))
        """(+)"""

        displayImage = np.asarray(image)
        cv2.imshow("Coral Live Object Detection", displayImage)

実行させてみると、10fps前後で動作しています。おもったよりもなめらかです。

まとめ

今回はGoogle Coral Edge TPU USB アクセラレータのデモを改良して、リアルタイム顔検出を行えるようにしました。
エッジでここまで動くのでいろいろな活用がありそうだなと思ってます。
次は、この検出結果をArduinoや別のRaspberry Piとつなげていきます。