Replace sws_scale with sws_scale_frame and enable multithreaded reformatting

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.

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.
 

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: ...

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,57 +27,16 @@ 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
     cdef _reformat(self, VideoFrame frame, int width, int height,
                    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)

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

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: ...