Fix #2149 Set colorspace and color_range ...

av/codec/context.py

@@ -4,16 +4,14 @@
 from cython.cimports import libav as lib
 from cython.cimports.av.buffer import ByteSource, bytesource
 from cython.cimports.av.codec.codec import Codec, wrap_codec
-from cython.cimports.av.dictionary import _Dictionary
+from cython.cimports.av.dictionary import Dictionary
 from cython.cimports.av.error import err_check
 from cython.cimports.av.packet import Packet
 from cython.cimports.av.utils import avrational_to_fraction, to_avrational
 from cython.cimports.libc.errno import EAGAIN
 from cython.cimports.libc.stdint import uint8_t
 from cython.cimports.libc.string import memcpy
 
-from av.dictionary import Dictionary
-
 _cinit_sentinel = cython.declare(object, object())
 
 
@@ -233,7 +231,7 @@ def open(self, strict: cython.bint = True):
                 raise ValueError("CodecContext is already open.")
             return
 
-        options: _Dictionary = Dictionary()
+        options: Dictionary = Dictionary()
         options.update(self.options or {})
 
         if not self.ptr.time_base.num and self.is_encoder:

av/codec/hwaccel.py

@@ -4,12 +4,10 @@
 import cython
 import cython.cimports.libav as lib
 from cython.cimports.av.codec.codec import Codec
-from cython.cimports.av.dictionary import _Dictionary
+from cython.cimports.av.dictionary import Dictionary
 from cython.cimports.av.error import err_check
 from cython.cimports.av.video.format import get_video_format
 
-from av.dictionary import Dictionary
-
 
 class HWDeviceType(IntEnum):
     none = lib.AV_HWDEVICE_TYPE_NONE
@@ -145,7 +143,7 @@ def _initialize_hw_context(self, codec: Codec):
         if self._device:
             device_bytes = self._device.encode()
             c_device = device_bytes
-        c_options: _Dictionary = Dictionary(self.options)
+        c_options: Dictionary = Dictionary(self.options)
 
         err_check(
             lib.av_hwdevice_ctx_create(

av/container/core.pxd

@@ -3,7 +3,7 @@ cimport libav as lib
 from av.codec.hwaccel cimport HWAccel
 from av.container.pyio cimport PyIOFile
 from av.container.streams cimport StreamContainer
-from av.dictionary cimport _Dictionary
+from av.dictionary cimport Dictionary
 from av.format cimport ContainerFormat
 from av.stream cimport Stream
 

av/container/core.py

@@ -21,7 +21,6 @@
 from cython.cimports.libc.stdint import int64_t
 from cython.operator import dereference
 
-from av.dictionary import Dictionary
 from av.logging import Capture as LogCapture
 
 _cinit_sentinel = cython.declare(object, object())
@@ -331,7 +330,7 @@ def __cinit__(
             self.ptr.flags |= lib.AVFMT_FLAG_CUSTOM_IO
 
         ifmt: cython.pointer[lib.AVInputFormat]
-        c_options: _Dictionary
+        c_options: Dictionary
         if not self.writeable:
             ifmt = self.format.iptr if self.format else cython.NULL
             c_options = Dictionary(self.options, self.container_options)

av/container/input.py

@@ -1,16 +1,14 @@
 import cython
 from cython.cimports.av.codec.context import CodecContext, wrap_codec_context
 from cython.cimports.av.container.streams import StreamContainer
-from cython.cimports.av.dictionary import _Dictionary
+from cython.cimports.av.dictionary import Dictionary
 from cython.cimports.av.error import err_check
 from cython.cimports.av.packet import Packet
 from cython.cimports.av.stream import Stream, wrap_stream
 from cython.cimports.av.utils import avdict_to_dict
 from cython.cimports.libc.stdint import int64_t
 from cython.cimports.libc.stdlib import free, malloc
 
-from av.dictionary import Dictionary
-
 
 @cython.cfunc
 def close_input(self: InputContainer):
@@ -34,8 +32,8 @@ def __cinit__(self, *args, **kwargs):
         # If we have either the global `options`, or a `stream_options`, prepare
         # a mashup of those options for each stream.
         c_options: cython.pointer[cython.pointer[lib.AVDictionary]] = cython.NULL
-        base_dict: _Dictionary
-        stream_dict: _Dictionary
+        base_dict: Dictionary
+        stream_dict: Dictionary
         if self.options or self.stream_options:
             base_dict = Dictionary(self.options)
             c_options = cython.cast(

av/container/output.py

@@ -6,14 +6,12 @@
 from cython.cimports.av.codec.codec import Codec
 from cython.cimports.av.codec.context import CodecContext, wrap_codec_context
 from cython.cimports.av.container.streams import StreamContainer
-from cython.cimports.av.dictionary import _Dictionary
+from cython.cimports.av.dictionary import Dictionary
 from cython.cimports.av.error import err_check
 from cython.cimports.av.packet import Packet
 from cython.cimports.av.stream import Stream, wrap_stream
 from cython.cimports.av.utils import dict_to_avdict, to_avrational
 
-from av.dictionary import Dictionary
-
 
 @cython.cfunc
 def close_output(self: OutputContainer):
@@ -384,8 +382,8 @@ def start_encoding(self):
             errors=self.metadata_errors,
         )
 
-        all_options: _Dictionary = Dictionary(self.options, self.container_options)
-        options: _Dictionary = all_options.copy()
+        all_options: Dictionary = Dictionary(self.options, self.container_options)
+        options: Dictionary = all_options.copy()
         self.err_check(lib.avformat_write_header(self.ptr, cython.address(options.ptr)))
 
         # Track option usage...

av/dictionary.pxd

@@ -1,9 +1,8 @@
 cimport libav as lib
 
 
-cdef class _Dictionary:
+cdef class Dictionary:
     cdef lib.AVDictionary *ptr
-    cpdef _Dictionary copy(self)
+    cpdef Dictionary copy(self)
 
-
-cdef _Dictionary wrap_dictionary(lib.AVDictionary *input_)
+cdef Dictionary wrap_dictionary(lib.AVDictionary *input_)

av/dictionary.py

@@ -1,11 +1,9 @@
-from collections.abc import MutableMapping
-
 import cython
 from cython.cimports.av.error import err_check
 
 
 @cython.cclass
-class _Dictionary:
+class Dictionary:
     def __cinit__(self, *args, **kwargs):
         for arg in args:
             self.update(arg)
@@ -17,9 +15,8 @@ def __dealloc__(self):
             lib.av_dict_free(cython.address(self.ptr))
 
     def __getitem__(self, key: cython.str):
-        element = cython.declare(
-            cython.pointer[lib.AVDictionaryEntry],
-            lib.av_dict_get(self.ptr, key, cython.NULL, 0),
+        element: cython.pointer[lib.AVDictionaryEntry] = lib.av_dict_get(
+            self.ptr, key, cython.NULL, 0
         )
         if element == cython.NULL:
             raise KeyError(key)
@@ -35,7 +32,7 @@ def __len__(self):
         return err_check(lib.av_dict_count(self.ptr))
 
     def __iter__(self):
-        element = cython.declare(cython.pointer[lib.AVDictionaryEntry], cython.NULL)
+        element: cython.pointer[lib.AVDictionaryEntry] = cython.NULL
         while True:
             element = lib.av_dict_get(self.ptr, "", element, lib.AV_DICT_IGNORE_SUFFIX)
             if element == cython.NULL:
@@ -46,17 +43,32 @@ def __repr__(self):
         return f"av.Dictionary({dict(self)!r})"
 
     def copy(self):
-        other = cython.declare(_Dictionary, Dictionary())
+        other: Dictionary = Dictionary()
         lib.av_dict_copy(cython.address(other.ptr), self.ptr, 0)
         return other
 
+    def pop(self, key: str):
+        value = self[key]
+        del self[key]
+        return value
 
-class Dictionary(_Dictionary, MutableMapping):
-    pass
+    def update(self, other=(), /, **kwds):
+        if isinstance(other, Dictionary):
+            lib.av_dict_copy(
+                cython.address(self.ptr), cython.cast(Dictionary, other).ptr, 0
+            )
+        elif hasattr(other, "keys"):
+            for key in other.keys():
+                self[key] = other[key]
+        else:
+            for key, value in other:
+                self[key] = value
+        for key, value in kwds.items():
+            self[key] = value
 
 
 @cython.cfunc
-def wrap_dictionary(input_: cython.pointer[lib.AVDictionary]) -> _Dictionary:
-    output = cython.declare(_Dictionary, Dictionary())
+def wrap_dictionary(input_: cython.pointer[lib.AVDictionary]) -> Dictionary:
+    output: Dictionary = Dictionary()
     output.ptr = input_
     return output

av/dictionary.pyi

@@ -1,10 +1,17 @@
-from collections.abc import MutableMapping
-from typing import Iterator
+from typing import Iterable, Iterator, Mapping
 
-class Dictionary(MutableMapping[str, str]):
+class Dictionary:
     def __getitem__(self, key: str) -> str: ...
     def __setitem__(self, key: str, value: str) -> None: ...
     def __delitem__(self, key: str) -> None: ...
     def __len__(self) -> int: ...
     def __iter__(self) -> Iterator[str]: ...
     def __repr__(self) -> str: ...
+    def copy(self) -> Dictionary: ...
+    def pop(self, key: str) -> str: ...
+    def update(
+        self,
+        other: Mapping[str, str] | Iterable[tuple[str, str]] = (),
+        /,
+        **kwds: str,
+    ) -> None: ...

av/filter/context.py

@@ -3,15 +3,13 @@
 import cython
 import cython.cimports.libav as lib
 from cython.cimports.av.audio.frame import alloc_audio_frame
-from cython.cimports.av.dictionary import _Dictionary
+from cython.cimports.av.dictionary import Dictionary
 from cython.cimports.av.error import err_check
 from cython.cimports.av.filter.link import alloc_filter_pads
 from cython.cimports.av.frame import Frame
 from cython.cimports.av.utils import avrational_to_fraction
 from cython.cimports.av.video.frame import alloc_video_frame
 
-from av.dictionary import Dictionary
-
 _cinit_sentinel = cython.declare(object, object())
 
 
@@ -72,7 +70,7 @@ def init(self, args=None, **kwargs):
         if args and kwargs:
             raise ValueError("cannot init from args and kwargs")
 
-        dict_: _Dictionary = None
+        dict_: Dictionary = None
         c_args: cython.p_char = cython.NULL
         if args or not kwargs:
             if args:

av/sidedata/sidedata.pxd

@@ -1,18 +1,16 @@
 cimport libav as lib
 
 from av.buffer cimport Buffer
-from av.dictionary cimport _Dictionary, wrap_dictionary
+from av.dictionary cimport Dictionary, wrap_dictionary
 from av.frame cimport Frame
 
 
 cdef class SideData(Buffer):
     cdef Frame frame
     cdef lib.AVFrameSideData *ptr
-    cdef _Dictionary metadata
-
+    cdef Dictionary metadata
 
 cdef SideData wrap_side_data(Frame frame, int index)
-
 cdef int get_display_rotation(Frame frame)
 
 cdef class _SideDataContainer:

av/video/reformatter.pxd

@@ -10,4 +10,5 @@ cdef class VideoReformatter:
     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)
+                   int src_color_range, int dst_color_range,
+                   bint set_dst_colorspace, bint set_dst_color_range)

av/video/reformatter.py

@@ -58,6 +58,21 @@ def _resolve_enum_value(value, enum_class, default):
     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,
+    {
+        lib.SWS_CS_ITU709: lib.AVCOL_SPC_BT709,
+        lib.SWS_CS_FCC: lib.AVCOL_SPC_FCC,
+        lib.SWS_CS_ITU601: lib.AVCOL_SPC_SMPTE170M,
+        lib.SWS_CS_SMPTE240M: lib.AVCOL_SPC_SMPTE240M,
+    },
+)
+
+
 @cython.cclass
 class VideoReformatter:
     """An object for reformatting size and pixel format of :class:`.VideoFrame`.
@@ -123,6 +138,10 @@ def reformat(
             dst_color_range, ColorRange, 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
+
         return self._reformat(
             frame,
             width or frame.ptr.width,
@@ -133,6 +152,8 @@ def reformat(
             c_interpolation,
             c_src_color_range,
             c_dst_color_range,
+            set_dst_colorspace,
+            set_dst_color_range,
         )
 
     @cython.cfunc
@@ -147,10 +168,16 @@ 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,
     ):
         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
@@ -231,6 +258,16 @@ def _reformat(
         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
+            )
+
         with cython.nogil:
             lib.sws_scale(
                 self.ptr,