diff --git a/av/video/codeccontext.py b/av/video/codeccontext.py index 8ef7801bd..e62bf8fae 100644 --- a/av/video/codeccontext.py +++ b/av/video/codeccontext.py @@ -98,7 +98,11 @@ def _prepare_frames_for_encode(self, input: Frame | None) -> list: self.reformatter = VideoReformatter() vframe = self.reformatter.reformat( - vframe, self.ptr.width, self.ptr.height, self._format + vframe, + self.ptr.width, + self.ptr.height, + self._format, + threads=self.ptr.thread_count, ) # There is no pts, so create one. diff --git a/av/video/frame.py b/av/video/frame.py index 586d317b5..74be1128d 100644 --- a/av/video/frame.py +++ b/av/video/frame.py @@ -472,7 +472,7 @@ def _init(self, format: lib.AVPixelFormat, width: cython.uint, height: cython.ui self.ptr.height = height self.ptr.format = format - # We enforce aligned buffers, otherwise `sws_scale` can perform + # We enforce aligned buffers, otherwise `sws_scale_frame` can perform # poorly or even cause out-of-bounds reads and writes. if width and height: res = lib.av_frame_get_buffer(self.ptr, 16) @@ -622,7 +622,7 @@ def color_primaries(self, value): self.ptr.color_primaries = value def reformat(self, *args, **kwargs): - """reformat(width=None, height=None, format=None, src_colorspace=None, dst_colorspace=None, interpolation=None) + """reformat(width=None, height=None, format=None, src_colorspace=None, dst_colorspace=None, interpolation=None, threads=None) Create a new :class:`VideoFrame` with the given width/height/format/colorspace. diff --git a/av/video/frame.pyi b/av/video/frame.pyi index ba3c92fed..12a85182b 100644 --- a/av/video/frame.pyi +++ b/av/video/frame.pyi @@ -70,6 +70,7 @@ class VideoFrame(Frame): dst_color_range: int | str | None = None, dst_color_trc: int | ColorTrc | None = None, dst_color_primaries: int | ColorPrimaries | None = None, + threads: int | None = None, ) -> VideoFrame: ... def to_rgb(self, **kwargs: Any) -> VideoFrame: ... def save(self, filepath: str | Path) -> None: ... diff --git a/av/video/reformatter.pxd b/av/video/reformatter.pxd index f57ba7e77..e68a70105 100644 --- a/av/video/reformatter.pxd +++ b/av/video/reformatter.pxd @@ -1,14 +1,13 @@ cimport libav as lib -from libc.stdint cimport uint8_t from av.video.frame cimport VideoFrame cdef extern from "libswscale/swscale.h" nogil: cdef struct SwsContext: - pass - cdef struct SwsFilter: - pass + unsigned flags + int threads + cdef int SWS_FAST_BILINEAR cdef int SWS_BILINEAR cdef int SWS_BICUBIC @@ -28,50 +27,9 @@ cdef extern from "libswscale/swscale.h" nogil: cdef int SWS_CS_SMPTE240M cdef int SWS_CS_DEFAULT - cdef int sws_scale( - SwsContext *ctx, - const uint8_t *const *src_slice, - const int *src_stride, - int src_slice_y, - int src_slice_h, - unsigned char *const *dst_slice, - const int *dst_stride, - ) - cdef void sws_freeContext(SwsContext *ctx) - cdef SwsContext *sws_getCachedContext( - SwsContext *context, - int src_width, - int src_height, - lib.AVPixelFormat src_format, - int dst_width, - int dst_height, - lib.AVPixelFormat dst_format, - int flags, - SwsFilter *src_filter, - SwsFilter *dst_filter, - double *param, - ) - cdef const int* sws_getCoefficients(int colorspace) - cdef int sws_getColorspaceDetails( - SwsContext *context, - int **inv_table, - int *srcRange, - int **table, - int *dstRange, - int *brightness, - int *contrast, - int *saturation - ) - cdef int sws_setColorspaceDetails( - SwsContext *context, - const int inv_table[4], - int srcRange, - const int table[4], - int dstRange, - int brightness, - int contrast, - int saturation - ) + cdef SwsContext *sws_alloc_context() + cdef void sws_free_context(SwsContext **ctx) + cdef int sws_scale_frame(SwsContext *c, lib.AVFrame *dst, const lib.AVFrame *src) cdef class VideoReformatter: cdef SwsContext *ptr @@ -79,6 +37,6 @@ cdef class VideoReformatter: lib.AVPixelFormat format, int src_colorspace, int dst_colorspace, int interpolation, int src_color_range, int dst_color_range, - bint set_dst_colorspace, bint set_dst_color_range, int dst_color_trc, int dst_color_primaries, - bint set_dst_color_trc, bint set_dst_color_primaries) + bint set_dst_color_trc, bint set_dst_color_primaries, + int threads) diff --git a/av/video/reformatter.py b/av/video/reformatter.py index 9c201bbe9..a29283717 100644 --- a/av/video/reformatter.py +++ b/av/video/reformatter.py @@ -1,6 +1,7 @@ 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 from cython.cimports.av.video.frame import alloc_video_frame @@ -105,19 +106,27 @@ def _resolve_enum_value( raise ValueError(f"Cannot convert {value} to {enum_class.__name__}") -# 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. -_SWS_CS_TO_AVCOL_SPC = cython.declare( - dict, - { - SWS_CS_ITU709: lib.AVCOL_SPC_BT709, - SWS_CS_FCC: lib.AVCOL_SPC_FCC, - SWS_CS_ITU601: lib.AVCOL_SPC_SMPTE170M, - SWS_CS_SMPTE240M: lib.AVCOL_SPC_SMPTE240M, - }, -) +@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.cclass @@ -131,7 +140,7 @@ class VideoReformatter: def __dealloc__(self): with cython.nogil: - sws_freeContext(self.ptr) + sws_free_context(cython.address(self.ptr)) def reformat( self, @@ -146,6 +155,7 @@ def reformat( 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. @@ -172,6 +182,8 @@ def reformat( :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). """ @@ -193,10 +205,9 @@ def reformat( c_dst_color_primaries = _resolve_enum_value( dst_color_primaries, ColorPrimaries, 0 ) + c_threads: cython.int = threads if threads is not None else 0 # Track whether user explicitly specified destination metadata - set_dst_colorspace: cython.bint = dst_colorspace is not None - set_dst_color_range: cython.bint = dst_color_range is not None set_dst_color_trc: cython.bint = dst_color_trc is not None set_dst_color_primaries: cython.bint = dst_color_primaries is not None @@ -210,12 +221,11 @@ def reformat( c_interpolation, c_src_color_range, c_dst_color_range, - set_dst_colorspace, - set_dst_color_range, c_dst_color_trc, c_dst_color_primaries, set_dst_color_trc, set_dst_color_primaries, + c_threads, ) @cython.cfunc @@ -230,24 +240,15 @@ def _reformat( interpolation: cython.int, src_color_range: cython.int, dst_color_range: cython.int, - set_dst_colorspace: cython.bint, - set_dst_color_range: cython.bint, 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.") - # Save original values to set on the output frame (before swscale conversion) - frame_dst_colorspace = dst_colorspace - frame_dst_color_range = dst_color_range - - # The definition of color range in pixfmt.h and swscale.h is different. - src_color_range = 1 if src_color_range == ColorRange.JPEG.value else 0 - dst_color_range = 1 if dst_color_range == ColorRange.JPEG.value else 0 - src_format = cython.cast(lib.AVPixelFormat, frame.ptr.format) # Shortcut! @@ -281,81 +282,33 @@ def _reformat( ): return frame - with cython.nogil: - self.ptr = sws_getCachedContext( - self.ptr, - frame.ptr.width, - frame.ptr.height, - src_format, - width, - height, - dst_format, - interpolation, - cython.NULL, - cython.NULL, - cython.NULL, - ) - - # We want to change the colorspace/color_range transforms. - # We do that by grabbing all the current settings, changing a - # couple, and setting them all. We need a lot of state here. - inv_tbl: cython.p_int - tbl: cython.p_int - src_colorspace_range: cython.int - dst_colorspace_range: cython.int - brightness: cython.int - contrast: cython.int - saturation: cython.int - - if src_colorspace != dst_colorspace or src_color_range != dst_color_range: - with cython.nogil: - ret = sws_getColorspaceDetails( - self.ptr, - cython.address(inv_tbl), - cython.address(src_colorspace_range), - cython.address(tbl), - cython.address(dst_colorspace_range), - cython.address(brightness), - cython.address(contrast), - cython.address(saturation), - ) - err_check(ret) - - with cython.nogil: - # Grab the coefficients for the requested transforms. - # The inv_table brings us to linear, and `tbl` to the new space. - if src_colorspace != SWS_CS_DEFAULT: - inv_tbl = cython.cast( - cython.p_int, sws_getCoefficients(src_colorspace) - ) - if dst_colorspace != SWS_CS_DEFAULT: - tbl = cython.cast(cython.p_int, sws_getCoefficients(dst_colorspace)) - - ret = sws_setColorspaceDetails( - self.ptr, - inv_tbl, - src_color_range, - tbl, - dst_color_range, - brightness, - contrast, - saturation, - ) - err_check(ret) + 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 the colorspace and color_range on the output frame only if explicitly specified - if set_dst_colorspace and frame_dst_colorspace in _SWS_CS_TO_AVCOL_SPC: - new_frame.ptr.colorspace = cython.cast( - lib.AVColorSpace, _SWS_CS_TO_AVCOL_SPC[frame_dst_colorspace] - ) - if set_dst_color_range: - new_frame.ptr.color_range = cython.cast( - lib.AVColorRange, frame_dst_color_range - ) + # 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 @@ -365,15 +318,4 @@ def _reformat( lib.AVColorPrimaries, dst_color_primaries ) - with cython.nogil: - sws_scale( - self.ptr, - cython.cast("const unsigned char *const *", frame.ptr.data), - cython.cast("const int *", frame.ptr.linesize), - 0, # slice Y - frame.ptr.height, - new_frame.ptr.data, - new_frame.ptr.linesize, - ) - return new_frame diff --git a/av/video/reformatter.pyi b/av/video/reformatter.pyi index b1e51e984..c9071df49 100644 --- a/av/video/reformatter.pyi +++ b/av/video/reformatter.pyi @@ -85,4 +85,5 @@ class VideoReformatter: dst_color_range: int | str | None = None, dst_color_trc: int | ColorTrc | None = None, dst_color_primaries: int | ColorPrimaries | None = None, + threads: int | None = None, ) -> VideoFrame: ...