PyAV:Examples:RelayFrame

# -*- coding: utf-8 -*-

from typing import Any, Callable, Optional

from av import VideoFrame  # noqa
from av import open as av_open  # noqa
from av.container import InputContainer, OutputContainer
from av.video.stream import VideoStream
from numpy import ndarray


def media_relay(
    input_file: str,
    output_file: str,
    input_video_format: Optional[str] = None,
    output_video_format: Optional[str] = None,
    input_video_options: Optional[dict] = None,
    output_video_options: Optional[dict] = None,
    input_stream_index=0,
    output_codec_name="h264",
    output_pixel_format="yuv420p",
    frame_callback: Optional[Callable[[ndarray, int, int, Any], ndarray]] = None,
    user_data: Optional[Any] = None,
) -> None:
    input_container = av_open(
        file=input_file,
        format=input_video_format,
        mode="r",
        options=input_video_options,
    )
    output_container = av_open(
        file=output_file,
        format=output_video_format,
        mode="w",
        options=output_video_options,
    )
    assert isinstance(input_container, InputContainer)
    assert isinstance(output_container, OutputContainer)

    input_video_stream = input_container.streams.video[input_stream_index]
    assert isinstance(input_video_stream, VideoStream)

    output_video_stream = output_container.add_stream(
        codec_name=output_codec_name,
        rate=input_video_stream.codec_context.framerate,
    )
    output_video_stream.width = input_video_stream.width
    output_video_stream.height = input_video_stream.height
    output_video_stream.pix_fmt = output_pixel_format
    output_video_stream.bit_rate = input_video_stream.codec_context.bit_rate
    assert isinstance(output_video_stream.options, dict)
    assert isinstance(output_video_stream, VideoStream)

    frame_index = 0
    total_frames = input_video_stream.frames

    for frame in input_container.decode(video=input_stream_index):
        image = frame.to_ndarray(format="bgr24")  # noqa

        if frame_callback:
            result = frame_callback(image, frame_index, total_frames, user_data)
        else:
            result = image
        frame_index += 1

        new_frame = VideoFrame.from_ndarray(result, format="bgr24")
        new_frame.pts = frame.pts
        new_frame.time_base = frame.time_base
        for packet in output_video_stream.encode(new_frame):
            output_container.mux(packet)

    # Flush stream
    for packet in output_video_stream.encode():
        output_container.mux(packet)

    output_container.close()
    input_container.close()