C言語版 YOLOv3 (libdarknet.so) を Python + OpenCV で使う

TL;DR

• 公式YOLOv3をビルドしてできるライブラリのlibdarknet.soをPythonから使います。
• numpy 配列を介して API を呼び出すので OpenCV と共存・協調できます。
• 大げさなフレームワークを使わないので Raspberry Pi でも軽々実装できます。

背景

Raspberry Pi でカメラを使った物体認識がしたくて、とりあえず有名どころということで YOLOv3 を使おうと思ったのですが、完全に C 言語の世界です。正直 C でカメラとか画像処理とか書く気力がなく、手っ取り早く Python + OpenCV で…… と思ったのですが、API が提供されていませんでした。いちおう公式リポジトリに Python なるフォルダがあって、そこに Python から libdarknet.so を使うサンプルがあったものの、（見ればわかるのですが）これはなんというか、本当に「呼び出してみただけ」のサンプルコードでいわゆる numpy 配列を使っていないから Python の恩恵がほとんど得られません。

一方、検索すると PyTorch などによる YOLOv3 の実装があって、これを使えば Python でばっちりYOLOが使えるものの、ちょっとした物体検出だけのために PyTorch 一式というのはいかにも襷に長し、な感じ。Raspberry Pi あたりで軽めに作りたいときは公式のlibdarknet.so程度の軽量な実装が助かります。

ポイント

• YOLO で使う画像データと OpenCV のそれとではRGBデータの格納順が違います。YOLOは Channel, Row, Column だけど、CV は Row, Column, Channel。cvimg[:,:,::-1].transpose((2,0,1)).copy() する。
• C言語とのインタフェース。YOLO のimage構造体に合わせてやる必要があります。cvimg.ctypes.data_as(POINTER(c_float)) として YOLO の image 構造体に入れるポインタを取得。
• (注) 少なくとも Raspberry pi 4 でやるときは 64-bit 版 の Raspberry Pi OS が必要です。(4 以外では未検証)

カメラで物体検出してよくある枠とラベルを描画するサンプルコード

準備

• LD_LIBRARY_PATH などを設定してリンカが libdarknet.soを見つけられるようにする。
• 実行ファイルのディレクトリに data/ を掘って、公式リポジトリに付属の coco.data、coco.names を置く。
• 同じくcfg/を掘って、公式リポジトリに付属の yolov3-tiny.cfg や yolov3.cfg を置く。
• 公式webページ に書いてある通り重み(yolov3-tiny.weightsあるいはyolov3.weights)をダウンロードする。(どちらを使うかは以下のソースコード中にハードコードしてください。)

実行

tiny かつ cfg をいじって解像度を160x160くらいに落とせば、Raspberry Pi 4 で 1 fps 以上出ました。

コード

公式リポジトリのdarknet.pyをかなり参考にして(というかパクって)います。

from ctypes import *
import cv2
import numpy as np

class BOX(Structure):
    _fields_ = [("x", c_float),
                ("y", c_float),
                ("w", c_float),
                ("h", c_float)]

class DETECTION(Structure):
    _fields_ = [("bbox", BOX),
                ("classes", c_int),
                ("prob", POINTER(c_float)),
                ("mask", POINTER(c_float)),
                ("objectness", c_float),
                ("sort_class", c_int)]

class IMAGE(Structure):
    _fields_ = [("w", c_int),
                ("h", c_int),
                ("c", c_int),
                ("data", POINTER(c_float))]

class METADATA(Structure):
    _fields_ = [("classes", c_int),
                ("names", POINTER(c_char_p))]

lib = CDLL("libdarknet.so", RTLD_GLOBAL)
lib.network_width.argtypes = [c_void_p]
lib.network_width.restype = c_int
lib.network_height.argtypes = [c_void_p]
lib.network_height.restype = c_int

predict = lib.network_predict
predict.argtypes = [c_void_p, POINTER(c_float)]
predict.restype = POINTER(c_float)

get_network_boxes = lib.get_network_boxes
get_network_boxes.argtypes = [c_void_p, c_int, c_int, c_float, c_float, POINTER(c_int), c_int, POINTER(c_int)]
get_network_boxes.restype = POINTER(DETECTION)

free_detections = lib.free_detections
free_detections.argtypes = [POINTER(DETECTION), c_int]

load_net = lib.load_network
load_net.argtypes = [c_char_p, c_char_p, c_int]
load_net.restype = c_void_p

do_nms_obj = lib.do_nms_obj
do_nms_obj.argtypes = [POINTER(DETECTION), c_int, c_int, c_float]

load_meta = lib.get_metadata
lib.get_metadata.argtypes = [c_char_p]
lib.get_metadata.restype = METADATA

predict_image = lib.network_predict_image
predict_image.argtypes = [c_void_p, IMAGE]
predict_image.restype = POINTER(c_float)

def draw_results(img, results):
    font = cv2.FONT_HERSHEY_DUPLEX
    for label, prob, (x, y, w, h) in results:
        x0, x1 = int(x - w / 2), int(x + w / 2)
        y0, y1 = int(y - h / 2), int(y + h / 2)
        cv2.rectangle(img, (x0, y0), (x1, y1), (0, 255,0), 2)
        size = cv2.getTextSize(label, font, 1 , 1)[0]
        cv2.rectangle(img, (x0, y0), (x0 + size[0] + 4, y0 - size[1] - 4), (0, 255, 0), -1)
        cv2.putText(img, label, (x0, y0 - 2), font, 1, (225, 255, 255), 1);

def detect_img(net, meta, im, thresh=.5, hier_thresh=.5, nms=.45):
    num = c_int(0)
    pnum = pointer(num)
    predict_image(net, im)
    dets = get_network_boxes(net, im.w, im.h, thresh, hier_thresh, None, 0, pnum)
    num = pnum[0]
    if (nms): do_nms_obj(dets, num, meta.classes, nms);

    res = []
    for j in range(num):
        for i in range(meta.classes):
            if dets[j].prob[i] > 0:
                b = dets[j].bbox
                res.append((meta.names[i].decode('utf-8'), dets[j].prob[i], (b.x, b.y, b.w, b.h)))
    res = sorted(res, key=lambda x: -x[1])
    free_detections(dets, num)
    return res

if __name__ == "__main__":
    cap = cv2.VideoCapture(0)
    net = load_net(b"cfg/yolov3-tiny.cfg", b"yolov3-tiny.weights", 0)
    meta = load_meta(b"data/coco.data")

    while True:
        ret, frame = cap.read()
        if ret:
             # float32, normalize, hwc > chw, contiguous
            img = (frame.astype(np.float32) / 255.0)[:,:,::-1].transpose((2,0,1)).copy()
            cimg = IMAGE(img.shape[2], img.shape[1], 3, img.ctypes.data_as(POINTER(c_float)))
            r = detect_img(net, meta, cimg)
            print(r)
            draw_results(frame, r)
            cv2.imshow('cam', frame)
            if cv2.waitKey(1) == 27: break