323 lines
12 KiB
Python
323 lines
12 KiB
Python
|
from enum import IntEnum
|
||
|
|
|
||
|
|
import cython
|
||
|
|
import cython.cimports.libav as lib
|
||
|
|
from cython.cimports.av.error import err_check
|
||
|
|
from cython.cimports.av.video.format import VideoFormat, get_pix_fmt
|
||
|
|
from cython.cimports.av.video.frame import alloc_video_frame
|
||
|
|
|
||
|
|
|
||
|
|
class Interpolation(IntEnum):
|
||
|
|
FAST_BILINEAR: "Fast bilinear" = SWS_FAST_BILINEAR
|
||
|
|
BILINEAR: "Bilinear" = SWS_BILINEAR
|
||
|
|
BICUBIC: "2-tap cubic B-spline" = SWS_BICUBIC
|
||
|
|
X: "Experimental" = SWS_X
|
||
|
|
POINT: "Nearest neighbor / point" = SWS_POINT
|
||
|
|
AREA: "Area averaging" = SWS_AREA
|
||
|
|
BICUBLIN: "Bicubic luma / Bilinear chroma" = SWS_BICUBLIN
|
||
|
|
GAUSS: "Gaussian approximation" = SWS_GAUSS
|
||
|
|
SINC: "Unwindowed Sinc" = SWS_SINC
|
||
|
|
LANCZOS: "3-tap sinc/sinc" = SWS_LANCZOS
|
||
|
|
SPLINE: "Unwindowed natural cubic spline" = SWS_SPLINE
|
||
|
|
|
||
|
|
|
||
|
|
class Colorspace(IntEnum):
|
||
|
|
ITU709 = SWS_CS_ITU709
|
||
|
|
FCC = SWS_CS_FCC
|
||
|
|
ITU601 = SWS_CS_ITU601
|
||
|
|
ITU624 = SWS_CS_ITU624
|
||
|
|
SMPTE170M = SWS_CS_SMPTE170M
|
||
|
|
SMPTE240M = SWS_CS_SMPTE240M
|
||
|
|
DEFAULT = SWS_CS_DEFAULT
|
||
|
|
# Lowercase for b/c.
|
||
|
|
itu709 = SWS_CS_ITU709
|
||
|
|
fcc = SWS_CS_FCC
|
||
|
|
itu601 = SWS_CS_ITU601
|
||
|
|
itu624 = SWS_CS_ITU624
|
||
|
|
smpte170m = SWS_CS_SMPTE170M
|
||
|
|
smpte240m = SWS_CS_SMPTE240M
|
||
|
|
default = SWS_CS_DEFAULT
|
||
|
|
|
||
|
|
|
||
|
|
class ColorRange(IntEnum):
|
||
|
|
UNSPECIFIED: "Unspecified" = lib.AVCOL_RANGE_UNSPECIFIED
|
||
|
|
MPEG: "MPEG (limited) YUV range, 219*2^(n-8)" = lib.AVCOL_RANGE_MPEG
|
||
|
|
JPEG: "JPEG (full) YUV range, 2^n-1" = lib.AVCOL_RANGE_JPEG
|
||
|
|
NB: "Not part of ABI" = lib.AVCOL_RANGE_NB
|
||
|
|
|
||
|
|
|
||
|
|
class ColorTrc(IntEnum):
|
||
|
|
"""Transfer characteristic (gamma curve) of a video frame.
|
||
|
|
|
||
|
|
Maps to FFmpeg's ``AVColorTransferCharacteristic``.
|
||
|
|
"""
|
||
|
|
|
||
|
|
BT709: "BT.709" = lib.AVCOL_TRC_BT709
|
||
|
|
UNSPECIFIED: "Unspecified" = lib.AVCOL_TRC_UNSPECIFIED
|
||
|
|
GAMMA22: "Gamma 2.2 (BT.470M)" = lib.AVCOL_TRC_GAMMA22
|
||
|
|
GAMMA28: "Gamma 2.8 (BT.470BG)" = lib.AVCOL_TRC_GAMMA28
|
||
|
|
SMPTE170M: "SMPTE 170M" = lib.AVCOL_TRC_SMPTE170M
|
||
|
|
SMPTE240M: "SMPTE 240M" = lib.AVCOL_TRC_SMPTE240M
|
||
|
|
LINEAR: "Linear" = lib.AVCOL_TRC_LINEAR
|
||
|
|
LOG: "Logarithmic (100:1 range)" = lib.AVCOL_TRC_LOG
|
||
|
|
LOG_SQRT: "Logarithmic (100*sqrt(10):1 range)" = lib.AVCOL_TRC_LOG_SQRT
|
||
|
|
IEC61966_2_4: "IEC 61966-2-4 (sRGB)" = lib.AVCOL_TRC_IEC61966_2_4
|
||
|
|
BT1361_ECG: "BT.1361 extended colour gamut" = lib.AVCOL_TRC_BT1361_ECG
|
||
|
|
IEC61966_2_1: "IEC 61966-2-1 (sYCC)" = lib.AVCOL_TRC_IEC61966_2_1
|
||
|
|
BT2020_10: "BT.2020 10-bit" = lib.AVCOL_TRC_BT2020_10
|
||
|
|
BT2020_12: "BT.2020 12-bit" = lib.AVCOL_TRC_BT2020_12
|
||
|
|
SMPTE2084: "SMPTE 2084 (PQ, HDR10)" = lib.AVCOL_TRC_SMPTE2084
|
||
|
|
SMPTE428: "SMPTE 428-1" = lib.AVCOL_TRC_SMPTE428
|
||
|
|
ARIB_STD_B67: "ARIB STD-B67 (HLG)" = lib.AVCOL_TRC_ARIB_STD_B67
|
||
|
|
|
||
|
|
|
||
|
|
class ColorPrimaries(IntEnum):
|
||
|
|
"""Color primaries of a video frame.
|
||
|
|
|
||
|
|
Maps to FFmpeg's ``AVColorPrimaries``.
|
||
|
|
"""
|
||
|
|
|
||
|
|
BT709: "BT.709 / sRGB / sYCC" = lib.AVCOL_PRI_BT709
|
||
|
|
UNSPECIFIED: "Unspecified" = lib.AVCOL_PRI_UNSPECIFIED
|
||
|
|
BT470M: "BT.470M" = lib.AVCOL_PRI_BT470M
|
||
|
|
BT470BG: "BT.470BG / BT.601-6 625" = lib.AVCOL_PRI_BT470BG
|
||
|
|
SMPTE170M: "SMPTE 170M / BT.601-6 525" = lib.AVCOL_PRI_SMPTE170M
|
||
|
|
SMPTE240M: "SMPTE 240M" = lib.AVCOL_PRI_SMPTE240M
|
||
|
|
FILM: "Generic film (Illuminant C)" = lib.AVCOL_PRI_FILM
|
||
|
|
BT2020: "BT.2020 / BT.2100" = lib.AVCOL_PRI_BT2020
|
||
|
|
SMPTE428: "SMPTE 428-1 / XYZ" = lib.AVCOL_PRI_SMPTE428
|
||
|
|
SMPTE431: "SMPTE 431-2 (DCI-P3)" = lib.AVCOL_PRI_SMPTE431
|
||
|
|
SMPTE432: "SMPTE 432-1 (Display P3)" = lib.AVCOL_PRI_SMPTE432
|
||
|
|
EBU3213: "EBU 3213-E / JEDEC P22" = lib.AVCOL_PRI_EBU3213
|
||
|
|
|
||
|
|
|
||
|
|
@cython.cfunc
|
||
|
|
@cython.inline
|
||
|
|
def _resolve_enum_value(
|
||
|
|
value: object, enum_class: object, default: cython.int
|
||
|
|
) -> cython.int:
|
||
|
|
# Helper function to resolve enum values from different input types.
|
||
|
|
if value is None:
|
||
|
|
return default
|
||
|
|
if isinstance(value, enum_class):
|
||
|
|
return value.value
|
||
|
|
if isinstance(value, int):
|
||
|
|
return value
|
||
|
|
if isinstance(value, str):
|
||
|
|
return enum_class[value].value
|
||
|
|
raise ValueError(f"Cannot convert {value} to {enum_class.__name__}")
|
||
|
|
|
||
|
|
|
||
|
|
@cython.cfunc
|
||
|
|
@cython.inline
|
||
|
|
def _resolve_format(format: object, default: lib.AVPixelFormat) -> lib.AVPixelFormat:
|
||
|
|
if format is None:
|
||
|
|
return default
|
||
|
|
if isinstance(format, VideoFormat):
|
||
|
|
return cython.cast(VideoFormat, format).pix_fmt
|
||
|
|
return get_pix_fmt(format)
|
||
|
|
|
||
|
|
|
||
|
|
@cython.cfunc
|
||
|
|
def _set_frame_colorspace(
|
||
|
|
frame: cython.pointer(lib.AVFrame),
|
||
|
|
colorspace: cython.int,
|
||
|
|
color_range: cython.int,
|
||
|
|
):
|
||
|
|
"""Set AVFrame colorspace/range from SWS_CS_* and AVColorRange values."""
|
||
|
|
if color_range != lib.AVCOL_RANGE_UNSPECIFIED:
|
||
|
|
frame.color_range = cython.cast(lib.AVColorRange, color_range)
|
||
|
|
# Mapping from SWS_CS_* (swscale colorspace) to AVColorSpace (frame metadata).
|
||
|
|
# Note: SWS_CS_ITU601, SWS_CS_ITU624, SWS_CS_SMPTE170M, and SWS_CS_DEFAULT all have
|
||
|
|
# the same value (5), so we map 5 -> AVCOL_SPC_SMPTE170M as the most common case.
|
||
|
|
# SWS_CS_DEFAULT is handled specially by not setting frame metadata.
|
||
|
|
if colorspace == SWS_CS_ITU709:
|
||
|
|
frame.colorspace = lib.AVCOL_SPC_BT709
|
||
|
|
elif colorspace == SWS_CS_FCC:
|
||
|
|
frame.colorspace = lib.AVCOL_SPC_FCC
|
||
|
|
elif colorspace == SWS_CS_ITU601:
|
||
|
|
frame.colorspace = lib.AVCOL_SPC_SMPTE170M
|
||
|
|
elif colorspace == SWS_CS_SMPTE240M:
|
||
|
|
frame.colorspace = lib.AVCOL_SPC_SMPTE240M
|
||
|
|
|
||
|
|
|
||
|
|
@cython.final
|
||
|
|
@cython.cclass
|
||
|
|
class VideoReformatter:
|
||
|
|
"""An object for reformatting size and pixel format of :class:`.VideoFrame`.
|
||
|
|
|
||
|
|
It is most efficient to have a reformatter object for each set of parameters
|
||
|
|
you will use as calling :meth:`reformat` will reconfigure the internal object.
|
||
|
|
|
||
|
|
"""
|
||
|
|
|
||
|
|
def __dealloc__(self):
|
||
|
|
with cython.nogil:
|
||
|
|
sws_free_context(cython.address(self.ptr))
|
||
|
|
|
||
|
|
def reformat(
|
||
|
|
self,
|
||
|
|
frame: VideoFrame,
|
||
|
|
width=None,
|
||
|
|
height=None,
|
||
|
|
format=None,
|
||
|
|
src_colorspace=None,
|
||
|
|
dst_colorspace=None,
|
||
|
|
interpolation=None,
|
||
|
|
src_color_range=None,
|
||
|
|
dst_color_range=None,
|
||
|
|
dst_color_trc=None,
|
||
|
|
dst_color_primaries=None,
|
||
|
|
threads=None,
|
||
|
|
):
|
||
|
|
"""Create a new :class:`VideoFrame` with the given width/height/format/colorspace.
|
||
|
|
|
||
|
|
Returns the same frame untouched if nothing needs to be done to it.
|
||
|
|
|
||
|
|
:param int width: New width, or ``None`` for the same width.
|
||
|
|
:param int height: New height, or ``None`` for the same height.
|
||
|
|
:param format: New format, or ``None`` for the same format.
|
||
|
|
:type format: :class:`.VideoFormat` or ``str``
|
||
|
|
:param src_colorspace: Current colorspace, or ``None`` for the frame colorspace.
|
||
|
|
:type src_colorspace: :class:`Colorspace` or ``str``
|
||
|
|
:param dst_colorspace: Desired colorspace, or ``None`` for the frame colorspace.
|
||
|
|
:type dst_colorspace: :class:`Colorspace` or ``str``
|
||
|
|
:param interpolation: The interpolation method to use, or ``None`` for ``BILINEAR``.
|
||
|
|
:type interpolation: :class:`Interpolation` or ``str``
|
||
|
|
:param src_color_range: Current color range, or ``None`` for the ``UNSPECIFIED``.
|
||
|
|
:type src_color_range: :class:`ColorRange` or ``str``
|
||
|
|
:param dst_color_range: Desired color range, or ``None`` for the ``UNSPECIFIED``.
|
||
|
|
:type dst_color_range: :class:`ColorRange` or ``str``
|
||
|
|
:param dst_color_trc: Desired transfer characteristic to tag on the output frame,
|
||
|
|
or ``None`` to preserve the source frame's value. This sets frame metadata only;
|
||
|
|
it does not perform a pixel-level transfer function conversion.
|
||
|
|
:type dst_color_trc: :class:`ColorTrc` or ``int``
|
||
|
|
:param dst_color_primaries: Desired color primaries to tag on the output frame,
|
||
|
|
or ``None`` to preserve the source frame's value.
|
||
|
|
:type dst_color_primaries: :class:`ColorPrimaries` or ``int``
|
||
|
|
:param int threads: How many threads to use for scaling, or ``0`` for automatic
|
||
|
|
selection based on the number of available CPUs. Defaults to ``0`` (auto).
|
||
|
|
|
||
|
|
"""
|
||
|
|
c_dst_format = _resolve_format(format, frame.format.pix_fmt)
|
||
|
|
c_src_colorspace = _resolve_enum_value(
|
||
|
|
src_colorspace, Colorspace, frame.ptr.colorspace
|
||
|
|
)
|
||
|
|
c_dst_colorspace = _resolve_enum_value(
|
||
|
|
dst_colorspace, Colorspace, frame.ptr.colorspace
|
||
|
|
)
|
||
|
|
c_interpolation = _resolve_enum_value(
|
||
|
|
interpolation, Interpolation, SWS_BILINEAR
|
||
|
|
)
|
||
|
|
c_src_color_range = _resolve_enum_value(src_color_range, ColorRange, 0)
|
||
|
|
c_dst_color_range = _resolve_enum_value(dst_color_range, ColorRange, 0)
|
||
|
|
c_dst_color_trc = _resolve_enum_value(dst_color_trc, ColorTrc, 0)
|
||
|
|
c_dst_color_primaries = _resolve_enum_value(
|
||
|
|
dst_color_primaries, ColorPrimaries, 0
|
||
|
|
)
|
||
|
|
c_threads: cython.int = threads if threads is not None else 0
|
||
|
|
c_width: cython.int = width if width is not None else frame.ptr.width
|
||
|
|
c_height: cython.int = height if height is not None else frame.ptr.height
|
||
|
|
|
||
|
|
# Track whether user explicitly specified destination metadata
|
||
|
|
set_dst_color_trc: cython.bint = dst_color_trc is not None
|
||
|
|
set_dst_color_primaries: cython.bint = dst_color_primaries is not None
|
||
|
|
|
||
|
|
return self._reformat(
|
||
|
|
frame,
|
||
|
|
c_width,
|
||
|
|
c_height,
|
||
|
|
c_dst_format,
|
||
|
|
c_src_colorspace,
|
||
|
|
c_dst_colorspace,
|
||
|
|
c_interpolation,
|
||
|
|
c_src_color_range,
|
||
|
|
c_dst_color_range,
|
||
|
|
c_dst_color_trc,
|
||
|
|
c_dst_color_primaries,
|
||
|
|
set_dst_color_trc,
|
||
|
|
set_dst_color_primaries,
|
||
|
|
c_threads,
|
||
|
|
)
|
||
|
|
|
||
|
|
@cython.cfunc
|
||
|
|
def _reformat(
|
||
|
|
self,
|
||
|
|
frame: VideoFrame,
|
||
|
|
width: cython.int,
|
||
|
|
height: cython.int,
|
||
|
|
dst_format: lib.AVPixelFormat,
|
||
|
|
src_colorspace: cython.int,
|
||
|
|
dst_colorspace: cython.int,
|
||
|
|
interpolation: cython.int,
|
||
|
|
src_color_range: cython.int,
|
||
|
|
dst_color_range: cython.int,
|
||
|
|
dst_color_trc: cython.int,
|
||
|
|
dst_color_primaries: cython.int,
|
||
|
|
set_dst_color_trc: cython.bint,
|
||
|
|
set_dst_color_primaries: cython.bint,
|
||
|
|
threads: cython.int,
|
||
|
|
):
|
||
|
|
if frame.ptr.format < 0:
|
||
|
|
raise ValueError("Frame does not have format set.")
|
||
|
|
|
||
|
|
src_format = cython.cast(lib.AVPixelFormat, frame.ptr.format)
|
||
|
|
|
||
|
|
# Shortcut!
|
||
|
|
if (
|
||
|
|
dst_format == src_format
|
||
|
|
and width == frame.ptr.width
|
||
|
|
and height == frame.ptr.height
|
||
|
|
and dst_colorspace == src_colorspace
|
||
|
|
and src_color_range == dst_color_range
|
||
|
|
and not set_dst_color_trc
|
||
|
|
and not set_dst_color_primaries
|
||
|
|
):
|
||
|
|
return frame
|
||
|
|
|
||
|
|
if frame.ptr.hw_frames_ctx:
|
||
|
|
frame_sw = alloc_video_frame()
|
||
|
|
err_check(lib.av_hwframe_transfer_data(frame_sw.ptr, frame.ptr, 0))
|
||
|
|
frame_sw.pts = frame.pts
|
||
|
|
frame_sw._init_user_attributes()
|
||
|
|
frame = frame_sw
|
||
|
|
src_format = cython.cast(lib.AVPixelFormat, frame.ptr.format)
|
||
|
|
|
||
|
|
if self.ptr == cython.NULL:
|
||
|
|
self.ptr = sws_alloc_context()
|
||
|
|
if self.ptr == cython.NULL:
|
||
|
|
raise MemoryError("Could not allocate SwsContext")
|
||
|
|
self.ptr.threads = threads
|
||
|
|
self.ptr.flags = cython.cast(cython.uint, interpolation)
|
||
|
|
|
||
|
|
new_frame: VideoFrame = alloc_video_frame()
|
||
|
|
new_frame._copy_internal_attributes(frame)
|
||
|
|
new_frame._init(dst_format, width, height)
|
||
|
|
|
||
|
|
# Set source frame colorspace/range so sws_scale_frame can read it
|
||
|
|
frame_src_colorspace: lib.AVColorSpace = frame.ptr.colorspace
|
||
|
|
frame_src_color_range: lib.AVColorRange = frame.ptr.color_range
|
||
|
|
_set_frame_colorspace(frame.ptr, src_colorspace, src_color_range)
|
||
|
|
_set_frame_colorspace(new_frame.ptr, dst_colorspace, dst_color_range)
|
||
|
|
|
||
|
|
with cython.nogil:
|
||
|
|
ret = sws_scale_frame(self.ptr, new_frame.ptr, frame.ptr)
|
||
|
|
|
||
|
|
# Restore source frame colorspace/range to avoid side effects
|
||
|
|
frame.ptr.colorspace = frame_src_colorspace
|
||
|
|
frame.ptr.color_range = frame_src_color_range
|
||
|
|
|
||
|
|
err_check(ret)
|
||
|
|
|
||
|
|
# Set metadata-only properties on the output frame if explicitly specified
|
||
|
|
if set_dst_color_trc:
|
||
|
|
new_frame.ptr.color_trc = cython.cast(
|
||
|
|
lib.AVColorTransferCharacteristic, dst_color_trc
|
||
|
|
)
|
||
|
|
if set_dst_color_primaries:
|
||
|
|
new_frame.ptr.color_primaries = cython.cast(
|
||
|
|
lib.AVColorPrimaries, dst_color_primaries
|
||
|
|
)
|
||
|
|
|
||
|
|
return new_frame
|