python 从视频中提取帧范围

Posted 2021-05-08

#!/usr/bin/python3.6
import argparse
import collections
import json
import subprocess
import sys
from subprocess import Popen
from typing import List

from math import ceil

Rational = collections.namedtuple('Rational', ['num', 'den'])

def video_stream(ffp):
    return next(filter(lambda x: x['codec_type'] == 'video', ffp['streams']))


def fps(ffp) -> Rational:
    name = ffp['format']['format_name']
    vstream = video_stream(ffp)
    fps: str = vstream['avg_frame_rate'] if name != "mxf" else vstream['r_frame_rate']
    if fps is None:
        return None
    return parse_rational(fps)


def parse_rational(r: str) -> Rational:
    (num, den) = r.replace(":", "/").split("/")
    return Rational(int(num), int(den))


# for [1..255] return 256
# for 256 return 256
# for 0 return 0
def align16up(x: int) -> int:
    return x + 15 & (-15 - 1)


def ffmpeg_cmd(ffp, downscale: float = 2.0, startframe: int = 0, endframe_inclusive: int = -1) -> List[str]:
    vstream = video_stream(ffp)
    bps = int(vstream.get('bits_per_raw_sample', '8'))
    sarstr = vstream.get('sample_aspect_ratio', '1/1')
    sar = parse_rational(sarstr)

    duration: int = int(vstream['nb_frames'])
    assert duration > 0, "video has zero frames"
    if endframe_inclusive > -1:
        endframe_inclusive = min(duration - 1, endframe_inclusive)
    duration = min(duration, endframe_inclusive - startframe + 1)
    ss = 0
    if startframe > 0:
        _fps = fps(ffp)
        ss = startframe * _fps.den / _fps.num

    pix_fmt = "bgr48le" if bps > 8 else "bgr24"  # cv2.imread reads images in BGR order
    width = int(vstream['width'])
    height = int(vstream['height'])

    interlaced = False  # TODO: vstream['..']
    # fix non-square pixel videos
    width = ceil(width * sar.num / sar.den)

    width16 = align16up(width)
    height16 = align16up(height)

    if downscale != 1:
        width16 = align16up(int(width / downscale))
        height16 = align16up(int(height / downscale))

    videopath = ffp['format']['filename']

    cmd: List[str] = ['ffmpeg']
    if ss > 0:
        cmd.extend(['-ss', str(ss)])
    # if format/codec combination is seekable -ss should come BEFORE -i, otherwise slow seek will happen
    cmd.extend(['-i', videopath])

    if endframe_inclusive > -1:
        cmd.extend(['-vframes', str(duration)])

    if interlaced:
        yadif = "yadif,"
    else:
        yadif = ""
    if width != width16 or height != height16:
        cmd.extend(['-vf', "%sscale=%d:%d" % (yadif, width16, height16)])
    # cmd.extend(['-pix_fmt', pix_fmt, '-f', 'rawvideo', '-'])
    cmd.extend(['-pix_fmt', pix_fmt, '-f', 'image2', '-start_number', str(startframe), '%03d.png'])

    return cmd


parser = argparse.ArgumentParser(description='extract frames')
parser.add_argument('--input_path', type=str, required=True)
parser.add_argument('--start_frame', type=int, default=0, help='end frame (inclusive)')
parser.add_argument('--end_frame', type=int, default=-1, help='end frame (inclusive)')

cfg = parser.parse_args()

process: Popen = subprocess.Popen(
    ['ffprobe', '-v', 'quiet', '-i', cfg.input_path, '-show_streams', '-show_format', '-print_format', 'json'],
    stdout=subprocess.PIPE)
ffpjson = process.communicate()[0]
ffp = json.loads(ffpjson)

print(' '.join(ffmpeg_cmd(ffp, 1.0, cfg.start_frame, cfg.end_frame)))