【Python】スマブラーが動画読み込み処理をジェネレータを使った関数にしてみた

ことのはじまり

とある依頼があり、10時間ぐらいのゲーム動画を機械学習することになったので
進捗状況をコンソール出力しながら動画を読み込むモジュールを作りました。

普通のプログラム

import cv2
import sys

fpath = 'data/videos/sample_video.mp4'

video = cv2.VideoCapture(fpath)
can_read, first = video.read()

if not can_read:
    sys.exit()

progress_step = 60    # 60フレームに1回、進捗を表示する
count = 0
max = int(video.get(cv2.CAP_PROP_FRAME_COUNT))

while video.isOpened():
    if(count % progress_step == 0):
        print(f'Progress rate: {count}/{max}')

    ok, frame = video.read()
    if ok:
        # やりたい処理
    else:
        break

video.release()

実行結果

Progress rate: 0/2700
Progress rate: 60/2700
Progress rate: 120/2700
Progress rate: 180/2700
...

所感

目的は達成できてるのですが、他のプログラムでも動画読み込み系の処理があり、そこでも使いまわしたいので、どこからでも呼べるような関数にすることに。

関数化Ver

関数

reader.py

def video_read(input):
    """
    Parameters
    ----------
    input: str
        Input video path

    Returns
    ------
    Iterator[np.ndarray]
        Generator of video frame
    """

    # video capture start
    video = cv2.VideoCapture(input)
    can_read, first = video.read()

    if not can_read:
        print('Cannot read video file')
        return

    max = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
    yield first

    progress_step = 60    # 60フレームに1回、進捗を表示する
    count = 0
    while video.isOpened():
        if(count % progress_step == 0):
            print(f'Progress rate: {count}/{max}')

        ok, frame = video.read()
        if ok:
            count += 1
            yield frame
        else:
            break
    video.release()

呼び出し元

import reader

it = reader.video_read('data/videos/sample_video.mp4')
try:
    first = next(it)
except StopIteration:
    sys.exit()

for image in it:
    # やりたい処理

まとめ

呼び出し元はすごくスッキリしました。
動画読み込み系の処理と、メインのやりたい処理が分離され明確になったと思います。
Pythonのジェネレータを上手に活かせた気がします。
今回の記事、スマブラー関係なくてすみません...