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

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
...

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:
    # やりたい処理