YOLOv8(ultralytics)をGoogle Colabで動かす
2024-03-27 python yolo
YOLOのポーズ推定モデルを利用して、左右の手首の動きをトラッキングして別のムービーに書き出すデモを作ってみました。ここで使っているultraliticsというライブラリがGoogle Colab向けに見事にチューニングされていて（逆？Google Colabがultralitics向けに仕上がってる？）、本当にただ呼び出すだけです。
OpenCVの方がちょっとクセがあって、Colab上で使えるエンコーダを見つけるのに手間取りました（H264などはだめでMP4Vを指定する必要がありました）。
以下の例ではただ座標値をリストに突っ込んでいっているだけなので、カメラが固定されている必要があって厳密なトラッキングにはなっていないのですが、キーポイントの使い方のサンプルになればいいなと思って公開しています。物体そのもののトラッキングはそれに特化した関数(track)が用意されているようです。
from ultralytics import YOLO
import cv2
model = YOLO('yolov8n-pose.pt')
results = model.predict("/content/drive/MyDrive/golf.mov", stream=True) # golf.mov 452 frames
fourcc = cv2.VideoWriter_fourcc(*'MP4V')
video = cv2.VideoWriter('/content/result.mp4',fourcc,30,(1920,1080))
frame = 0
lefts, rights = [],[]
for r in results:
# r.boxes # Boxes object for bbox outputs
# r.masks # Masks object for segment masks outputs
# r.probs # Class probabilities for classification outputs
img = r.orig_img
if len(r.keypoints.xy[0]) == 17:
lefts += [r.keypoints.xy[0][9]] # left wrist
rights += [r.keypoints.xy[0][10]] # right wrist
for l in lefts[:-1]:
cv2.circle(img,(int(l[0]),int(l[1])),2,(255,0,0),2) # blue
if len(lefts) > 0:
l = lefts[-1:][0]
cv2.circle(img,(int(l[0]),int(l[1])),10,(255,0,0),3)
for r in rights[:-1]:
cv2.circle(img,(int(r[0]),int(r[1])),2,(0,0,255),2) # red
if len(rights) > 0:
r = rights[-1:][0]
cv2.circle(img,(int(r[0]),int(r[1])),10,(0,0,255),3)
video.write(img)
frame += 1
# if frame > 30: break # for debug
video.release()
出力されたムービーはこんな感じになります。ティーショットが50ヤード…ナイスショットのムービーはレアすぎて見つけきれずこんなのを恥を忍んでサンプルに(^^;
こちらはオマケで標準の書き出し機能を使って作ったムービーです。