YOLOv8でポリープ検出した結果を動画にして保存

ポリープ検出結果を反映した内視鏡動画をつくって保存

内視鏡動画の各フレームごとにポリープ検出を実行して、その結果を反映した動画をつくりたいことは多いと思います。

ここでは、各フレームごとの画像に対してdetectionを行ったあと、どのようにして再び動画にするのかについて紹介します。

なお以下の例では内視鏡画像として、"HyperKvasir"のものを利用しています。

お持ちのデータを利用する場合はファイルの読み込みの箇所を変えるといったように適宜修正してお使いください。

また、ポリープ検出モデルの構築、内視鏡動画を各フレームごとに画像として抽出する方法については以下の記事を参考にしてください。

「YOLOv8を使って内視鏡画像からポリープを検出」

「Pythonを使って内視鏡動画からフレームごとに画像を抽出して保存」

アップロードしたファイルの確認

今回は

1. 動画ファイル(0220d11b-ab12-4b02-93ce-5d7c205c7043.avi)

2. 学習済みモデル(best.pt)

が用意できていて、それらがアップロードされた状態からスタートします。

Googleドライブから読み込む場合などは適宜画像ファイルのパスを変更してください。

ホームディレクトリ直下にアップロードした場合は下記のコマンドを実行すると、アップロードしたファイルが表示されるはずです。

!ls

内視鏡動画からフレームごとに画像として保存

ここでは、"img_動画名"というディレクトリを作ってそこに画像を保存していくことにします。

画像を保存するときの名前は「フレーム数(5桁).jpg」とします。

import os
import cv2

file_video = '0220d11b-ab12-4b02-93ce-5d7c205c7043.avi'
dir_save = 'img_' + file_video.replace('.avi' , '')

# 保存用ディレクトリを作成
os.mkdir(dir_save)

# 動画ファイルの読み込み
cap = cv2.VideoCapture(file_video)

# フレーム数を取得&確認
frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
print('総フレーム数: ', frame_count)

# フレームごとに画像を保存
for i in range(frame_count):
  # フレームを取得
  ret, frame = cap.read()

  # 画像を保存
  cv2.imwrite(f'{dir_save}/frame_{i:05}.jpg', frame)

# キャプチャを閉じる
cap.release()

画像パスの取得

import glob

# 画像パスの取得
files = glob.glob(dir_save+'/*')
# 名前順にソート
files = sorted(files)

# 並び順とファイル数を確認
print(files)
print(len(files))

YOLOv8を使うためのパッケージインストール

!pip install ultralytics

YOLOv8を使って各画像に対してポリープ検出を実行

YOLOv8を使ってポリープ検出をしていきます。以下の例では、検出に関する閾値として

1. 信頼度(conf): 0.2

2. IoU: 0.5

で行っています。
偽陽性が多い場合は信頼度の閾値を挙げて実行すると良いかもしれません。

from ultralytics import YOLO

# 学習済みモデルを読み込み
model = YOLO('best.pt')

# 各画像に対してポリープ検出して、検出結果の画像を保存
for f in files:
    results = model(f,save=True, conf=0.2, iou=0.5)

検出結果画像のパスの取得

# 検出結果画像パスの取得
files_res = glob.glob('runs/detect/predict/*')
# 上のセルで推論を実行した回数によってはpredict2など数字がついている場合あり
# 上のセルで出力される"Results saved to ..."のところを見て合わせると良い

# ソート
files_res = sorted(files_res)

# 並び順とファイル数を確認
print(files_res)
print(len(files_res))

画像群から動画をつくる

最後に画像をまとめて動画にすれば完成です。
今回はmp4動画をつくっています。

# 動画のファイル名指定
file_output = 'detect.mp4'

# 動画のfpsを指定
fps = 25

# 画像の幅、高さを取得
# .shapeの出力は高さ、幅で逆になっているので注意
img_w, img_h = cv2.imread(files[0]).shape[1], cv2.imread(files[0]).shape[0]

# 動画ファイルを作成
video = cv2.VideoWriter(file_output, cv2.VideoWriter_fourcc('M', 'P', '4', 'V'), fps, (img_w, img_h))

# 画像ファイルを順番に読み込み、動画ファイルに書き込む
for f in files_res:
    img = cv2.imread(f)
    video.write(img)

# 動画ファイルを閉じる
video.release()

こうしてつくられた動画が以下のものになります。

detect.mp4

29.8 MB

ファイルダウンロードについて

ダウンロード

検出結果の動画は、学会での発表や論文のsupplemental dataなど使える場面は多いので、ぜひ活用してみてください。

また、有料部分では上記のコードが既に書かれたnotebook(.ipynb)を配布しています。
おまけとして元動画のfpsやエンコードの調べ方も書かれているので、ご興味のある方はよろしくお願いいたします。

配布notebookの使い方についてはこちら

他にもこういった記事が読みたい！これを解説してほしい！などご意見がありましたら、アンケートフォームからご意見をいただけるとうれしいです。