このシリーズでは物体検出でお馴染みの「YOLOv5」を用いて、物体検出の実装を基礎から学ぶことができます。
環境構築から学習の方法、さらに活用方法までをまとめています。
Google colabを使用して簡単に物体検出のモデルを実装することができますので、ぜひ最後までご覧ください。
第7回目はYOLOv5の活用事例として、動画から座標やスコアを出力する方法を考えます。
今回の内容
・動画からフレームごとの物体検出
・座標・スコアの出力
YOLOv5の導入
Google colabを立ち上げて、早速実装していきましょう。
まずはGPUを使用できるように設定をします。
「編集」タブから「ノートブックの設定」の中の「ハードウェア アクセラレータ」を「GPU」に設定。
なお、今回紹介するコードは以下のボタンからコピーして使用していただくことも可能です。
まずはGoogle ドライブのファイルにアクセスするため、マウントします。
from google.colab import drive
drive.mount('/content/drive')
%cd ./drive/MyDrive
今回初めてYOLOv5を使用する方はクローンしましょう。
すでに前回クローンしている方は実行不要です。
!git clone https://github.com/ultralytics/yolov5
必要なライブラリをインポートします。
%cd yolov5/
!pip install -qr requirements.txt
以上で準備は終了です。
OpenCVによる動画処理
OpenCVを使用することで、動画から各フレーム(静止画)に対して処理を施すことができます。
この処理の中にYOLOv5による物体検出を加えることで、動画全体に対して連続的に座標やスコアの取得が可能となります。
動画情報取得
まずはOpenCVによる動画処理の代表例を紹介します。
処理に必要な動画の情報を取得してみましょう。
サンプルとしてこちらのフリー動画を使用します。
以下のように動画の情報を出力する関数を定義します。
import cv2
def video_info(video):
video = cv2.VideoCapture(video)
w = int(video.get(cv2.CAP_PROP_FRAME_WIDTH))
h = int(video.get(cv2.CAP_PROP_FRAME_HEIGHT))
frame_count = video.get(cv2.CAP_PROP_FRAME_COUNT)
frame_rate = video.get(cv2.CAP_PROP_FPS)
return w,h,frame_count,frame_rate
先ほどのフリー動画で実行してみましょう。
video = 'road.mp4'
video_info(video)
以下のような結果を得ます。
#出力結果
(1280, 720, 250.0, 25.0)
左から順に動画の幅、動画の高さ、フレーム数、フレームレートとなります。
フレームごとの処理を記述
次にフレームごとの処理を記述する方法を紹介します。
今回は例としてフレームごとに画像を保存することを実装してみます。
各フレームごとに画像として書き出す処理を加えることで、連続した画像を出力することができます。
import os
import cv2
def save_all_frames(video_files):
video = cv2.VideoCapture(video_files)
num = 0
while(video.isOpened()):
ret, frame = video.read()
if ret == True:
# ここに動画内での処理を書く
cv2.imwrite(os.path.basename(video_files).strip(".mp4") + "_{:0=10}".format(num) + ".jpg",frame)
num += 1
else:
break
video.release()
先ほどと同様に関数を実行しましょう。
video = 'road.mp4'
save_all_frames(video)
フレーム数分の動画から切り出した画像を出力することができました。
各フレームごとの座標とスコアをCSVファイルとして出力
先程までの内容をもとに動画から物体検出を実装し、座標やスコアを出力していきます。
今回は学習済みモデルを使用します。
自分で作成したモデルを使用したい場合は、モデルのパスを書き換えてください。
# 動画から各フレームの座標・スコアをCSVに出力(関数)
import torch
import os
import cv2
# yolov5のモデルを読み込み
model =torch.hub.load('ultralytics/yolov5', 'yolov5x', pretrained=True)
def test_all_frames(video_files):
video = cv2.VideoCapture(video_files)
num = 0
while(video.isOpened()):
ret, frame = video.read()
if ret == True:
results = model(frame)
results.pandas().xyxy[0].to_csv(os.path.basename(video_files).strip(".mp4") + "_{:0=10}".format(num) + '_result.csv')
num += 1
else:
break
video.release()
動画名を指定して関数を実行します。
# 動画から各フレームの座標・スコアをCSVに出力
video = 'road.mp4'
test_all_frames(video)
各フレームごとの座標とスコアをCSVファイルとして出力することができました。
まとめ
最後までご覧いただきありがとうございました。
今回はYOLOv5の活用事例として、動画から座標やスコアを出力する方法をまとめました。