🐶

WebカメラでYOLOv8を動かす【スニペット有】

2023/07/03に公開

はじめに

久々に YOLO を動かしてみた。しばらく見ないうちに色んなタイプのモデルが出ていてびっくりしました。自分が真面目に動かしていたのは yolov3 くらいだったので。今回は yolov8 を動かしてみました。

今回の環境は windows です。

ライブラリのインストール

pip install opencv-python ultralytics

物体検出

以下、スニペットです。ほとんど公式のコードをそのまま使っています。いくつかのモデルは実行時にダウンロードされるみたいですが、基本的には、Github にあるモデルをダウンロードしてきた方が迷わないと思います。

見てる感じ、model()からの出力は pandas かなんかのデータフレームぽいですね。

import cv2
from ultralytics import YOLO

model = YOLO('yolov8x.pt') # ここのモデルを変更すると別の認識modelになります。ただし、トラッキングは別です。

# Open the web camera stream
cap = cv2.VideoCapture(0) # サンプルからここだけ変更するだけ

while cap.isOpened():
    success, frame = cap.read()
    k = cv2.waitKey(1) # 一応キー入力で終了できるようにしておく

    if k != -1:
        break
    if success:
        results = model(frame)
        annotated_frame = results[0].plot()

        cv2.imshow("YOLOv8 Inference", annotated_frame)
cap.release()
cv2.destroyAllWindows()

トラッキング

トラッキングは、以下のようになります。本当は model の引数にて show=True にすると見るだけなら見れますが、ここではある程度コーディングしています。最適化していませんが、矩形、ID、クラス名を表示しています。

import cv2
from ultralytics import YOLO

cap = cv2.VideoCapture(0)
model = YOLO("yolov8n.pt")
while True:
    ret, frame = cap.read()
    if not ret:
        break
    results = model.track(frame,
                          conf=0.5, iou=0.5)
    result = results[0]
    # 矩形ごとに処理
    for box in result.boxes:
        # 引数がtensorなので、listに変換
        r = box.xyxy.tolist()

        cv2.rectangle(frame, (int(r[0][0]), int(r[0][1])),
                      (int(r[0][2]), int(r[0][3])), (0, 255, 0), 2)

        # idが振られないときもあるみたいなので、判定
        id = int(box.id) if box.id is not None else ""
        # クラス名を取得
        cls = result.names[int(box.cls)]
        cv2.putText(
            frame,
            f"Id {id} {cls}",
            (int(r[0][0]), int(r[0][1])),
            cv2.FONT_HERSHEY_SIMPLEX,
            1,
            (0, 0, 255),
            2,
        )

    cv2.imshow("frame", frame)
    if cv2.waitKey(1) & 0xFF == ord("q"):
        break
cap.release()
cv2.destroyAllWindows()

おわりに

個人的に懐かしむように、YOLO を動かしてみました。昔はもちっと動かすのに手間だったように感じますが、それは意味のない懐古なのかもしれません。トラッキング周りは少しドキュメントが分かりづらかったので、書いてみたので、参考になれば幸いです。

参考

Discussion