Quickly decode video frame by frame

decode_video.py

# !pip install pyav
import av

PATH_VIDEO = "..."

container = av.open(PATH_VIDEO)
video_stream = container.streams.video[0]
video_decoder = container.decode(video_stream)

for frame in video_decoder:
    frame = process_frame(frame)
    
def process_frame(
    f:av.VideoFrame,
    height:int = 224,
    width:int=224,
    frame_format:str="rgb24",
) -> np.ndarray:
    """
    Reformat an av.VideoFrame to a np.ndarray 
    """
    import numpy as np
    f = f.reformat(
        format=frame_format,
        interpolation=av.video.reformatter.Interpolation.FAST_BILINEAR,
        width=width,
        height=height,
    )
    #f = f.to_image()  # gets a PIL.Image.Image
    f = np.transpose(f.to_ndarray(), axes=(2,0,1))  # H,W,C -> C,H,W
    return f