y4m 영상을 특정 크기의 YUV로 변환하기

#!/usr/bin/env python3
# version: 0.1
# author: Brayden Jo
import sys
import os
import subprocess
import re

DST_W = 1920
DST_H = 1088


def get_y4m_flist(path):
    files = os.listdir(path)
    y4m_list = [x for x in files if x.endswith('y4m')]
    return y4m_list


def get_frame_size(path):
    cmd = f'ffprobe {path}'
    ret = subprocess.run(cmd, capture_output=True, shell=True, encoding='utf-8')
    log = ret.stderr
    #print(log, type(log))

    # parsing
    pattern = r"\d\d+x\d\d+"
    ret = re.findall(pattern, log)
    resolution = ret[0]
    width, height = resolution.split('x')
    return int(width), int(height)


def get_prefix(path):
    fname = path.split('/')[-1]
    pattern = r"[a-zA-Z]+"
    ret = re.findall(pattern, fname)
    return ret[0]


def y4m2yuv(path, src_w, src_h, dst_w, dst_h):
    prefix = get_prefix(path)                       # aaa_11.y4m -> aaa
    output = f"{prefix}_{dst_w}x{dst_h}.yuv"        # aaa_1920x1080.yuv

    if src_w <= dst_w and src_h <= dst_h:
        cmd = f'ffmpeg -i {path} -filter_complex "[0]pad=w={dst_w}:h={dst_h}:x=0:y=0:color=black" -pixel_format yuv420p {output}'
    else:
        cmd = f'ffmpeg -i {path} -vf scale={dst_w}:{dst_h} -c:v rawvideo -pixel_format yuv420p {output}'

    subprocess.run(cmd, shell=True)


def main():
    src_path = sys.argv[1]
    y4m_list = get_y4m_flist(src_path)

    for y4m in y4m_list:
        path = os.path.join(src_path, y4m)
        src_w, src_h= get_frame_size(path)
        y4m2yuv(path, src_w, src_h, DST_W, DST_H)

    #path = "/fs1/00_YUV/vimeo-corpus-10s/autumn.y4m"
    #src_w, src_h= get_frame_size(path)
    #y4m2yuv(path, src_w, src_h, DST_W, DST_H)


if __name__ == "__main__":
    main()

단순히 영상을 resizing하는 것이 아니라 해상도를 보고 crop 하거나 0으로 패딩합니다. 예를 들어, 1920x1080 비디오를 1920x1088로 만들 때 resizing 보다는 0으로 패딩하는 것을 의미합니다.