Voice et bot modif

venv/lib/python3.12/site-packages/av/video/__init__.py

@@ -1,2 +1,2 @@
-from .frame import VideoFrame
-from .stream import VideoStream
+from .frame import VideoFrame as VideoFrame
+from .stream import VideoStream as VideoStream

venv/lib/python3.12/site-packages/av/video/codeccontext.pxd

@@ -16,18 +16,6 @@ cdef struct AVCodecPrivateData:
 
 
 cdef class VideoCodecContext(CodecContext):
-
     cdef AVCodecPrivateData _private_data
-
-    cdef VideoFormat _format
-    cdef _build_format(self)
-
-    cdef int last_w
-    cdef int last_h
     cdef readonly VideoReformatter reformatter
-
-    # For encoding.
-    cdef readonly int encoded_frame_count
-
-    # For decoding.
     cdef VideoFrame next_frame

venv/lib/python3.12/site-packages/av/video/codeccontext.py

@@ -1,41 +1,53 @@
-cimport libav as lib
-from libc.stdint cimport int64_t
-
-from av.codec.context cimport CodecContext
-from av.codec.hwaccel cimport HWAccel, HWConfig
-from av.error cimport err_check
-from av.frame cimport Frame
-from av.packet cimport Packet
-from av.utils cimport avrational_to_fraction, to_avrational
-from av.video.format cimport VideoFormat, get_pix_fmt, get_video_format
-from av.video.frame cimport VideoFrame, alloc_video_frame
-from av.video.reformatter cimport VideoReformatter
+import cython
+import cython.cimports.libav as lib
+from cython.cimports.av.codec.context import CodecContext
+from cython.cimports.av.codec.hwaccel import HWAccel
+from cython.cimports.av.error import err_check
+from cython.cimports.av.frame import Frame
+from cython.cimports.av.packet import Packet
+from cython.cimports.av.utils import avrational_to_fraction, to_avrational
+from cython.cimports.av.video.format import VideoFormat, get_pix_fmt, get_video_format
+from cython.cimports.av.video.frame import VideoFrame, alloc_video_frame
+from cython.cimports.av.video.reformatter import VideoReformatter
+from cython.cimports.libc.stdint import int64_t
 
 
-cdef lib.AVPixelFormat _get_hw_format(lib.AVCodecContext *ctx, const lib.AVPixelFormat *pix_fmts) noexcept:
+@cython.cfunc
+@cython.exceptval(check=False)
+def _get_hw_format(
+    ctx: cython.pointer[lib.AVCodecContext],
+    pix_fmts: cython.pointer[cython.const[lib.AVPixelFormat]],
+) -> lib.AVPixelFormat:
     # In the case where we requested accelerated decoding, the decoder first calls this function
     # with a list that includes both the hardware format and software formats.
     # First we try to pick the hardware format if it's in the list.
     # However, if the decoder fails to initialize the hardware, it will call this function again,
     # with only software formats in pix_fmts. We return ctx->sw_pix_fmt regardless in this case,
     # because that should be in the candidate list. If not, we are out of ideas anyways.
-    cdef AVCodecPrivateData* private_data = <AVCodecPrivateData*>ctx.opaque
-    i = 0
+    private_data: cython.pointer[AVCodecPrivateData] = cython.cast(
+        cython.pointer[AVCodecPrivateData], ctx.opaque
+    )
+    i: cython.int = 0
     while pix_fmts[i] != -1:
         if pix_fmts[i] == private_data.hardware_pix_fmt:
             return pix_fmts[i]
         i += 1
-    return ctx.sw_pix_fmt if private_data.allow_software_fallback else lib.AV_PIX_FMT_NONE
+    return (
+        ctx.sw_pix_fmt if private_data.allow_software_fallback else lib.AV_PIX_FMT_NONE
+    )
 
 
-cdef class VideoCodecContext(CodecContext):
-
-    def __cinit__(self, *args, **kwargs):
-        self.last_w = 0
-        self.last_h = 0
-
-    cdef _init(self, lib.AVCodecContext *ptr, const lib.AVCodec *codec, HWAccel hwaccel):
-        CodecContext._init(self, ptr, codec, hwaccel)  # TODO: Can this be `super`?
+@cython.final
+@cython.cclass
+class VideoCodecContext(CodecContext):
+    @cython.cfunc
+    def _init(
+        self,
+        ptr: cython.pointer[lib.AVCodecContext],
+        codec: cython.pointer[cython.const[lib.AVCodec]],
+        hwaccel: HWAccel | None,
+    ):
+        CodecContext._init(self, ptr, codec, hwaccel)
 
         if hwaccel is not None:
             try:
@@ -43,9 +55,13 @@ cdef class VideoCodecContext(CodecContext):
                 self.ptr.hw_device_ctx = lib.av_buffer_ref(self.hwaccel_ctx.ptr)
                 self.ptr.pix_fmt = self.hwaccel_ctx.config.ptr.pix_fmt
                 self.ptr.get_format = _get_hw_format
-                self._private_data.hardware_pix_fmt = self.hwaccel_ctx.config.ptr.pix_fmt
-                self._private_data.allow_software_fallback = self.hwaccel.allow_software_fallback
-                self.ptr.opaque = &self._private_data
+                self._private_data.hardware_pix_fmt = (
+                    self.hwaccel_ctx.config.ptr.pix_fmt
+                )
+                self._private_data.allow_software_fallback = (
+                    self.hwaccel.allow_software_fallback
+                )
+                self.ptr.opaque = cython.address(self._private_data)
             except NotImplementedError:
                 # Some streams may not have a hardware decoder. For example, many action
                 # cam videos have a low resolution mjpeg stream, which is usually not
@@ -57,102 +73,95 @@ cdef class VideoCodecContext(CodecContext):
                 # is_hwaccel() function on each stream's codec context.
                 self.hwaccel_ctx = None
 
-        self._build_format()
-        self.encoded_frame_count = 0
-
-    cdef _prepare_frames_for_encode(self, Frame input):
-        if not input:
+    @cython.cfunc
+    def _prepare_frames_for_encode(self, input: Frame | None) -> list:
+        if input is None or not input:
             return [None]
 
-        cdef VideoFrame vframe = input
-
-        if self._format is None:
-            raise ValueError("self._format is None, cannot encode")
-
-        # Reformat if it doesn't match.
+        vframe: VideoFrame = input
         if (
-            vframe.format.pix_fmt != self._format.pix_fmt or
-            vframe.width != self.ptr.width or
-            vframe.height != self.ptr.height
+            vframe.format.pix_fmt != self.pix_fmt
+            or vframe.width != self.ptr.width
+            or vframe.height != self.ptr.height
         ):
             if not self.reformatter:
                 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.
         if vframe.ptr.pts == lib.AV_NOPTS_VALUE:
-            vframe.ptr.pts = <int64_t>self.encoded_frame_count
-
-        self.encoded_frame_count += 1
+            vframe.ptr.pts = self.ptr.frame_num
 
         return [vframe]
 
-    cdef Frame _alloc_next_frame(self):
+    @cython.cfunc
+    def _alloc_next_frame(self) -> Frame:
         return alloc_video_frame()
 
-    cdef _setup_decoded_frame(self, Frame frame, Packet packet):
+    @cython.cfunc
+    def _setup_decoded_frame(self, frame: Frame, packet: Packet):
         CodecContext._setup_decoded_frame(self, frame, packet)
-        cdef VideoFrame vframe = frame
+        vframe: VideoFrame = frame
         vframe._init_user_attributes()
 
-    cdef _transfer_hwframe(self, Frame frame):
+    @cython.cfunc
+    def _transfer_hwframe(self, frame: Frame):
         if self.hwaccel_ctx is None:
             return frame
-
         if frame.ptr.format != self.hwaccel_ctx.config.ptr.pix_fmt:
             # If we get a software frame, that means we are in software fallback mode, and don't actually
             # need to transfer.
             return frame
 
-        cdef Frame frame_sw
-
-        frame_sw = self._alloc_next_frame()
+        if self.hwaccel_ctx.is_hw_owned:
+            cython.cast(VideoFrame, frame)._device_id = self.hwaccel_ctx.device_id
+            return frame
 
+        frame_sw: Frame = self._alloc_next_frame()
         err_check(lib.av_hwframe_transfer_data(frame_sw.ptr, frame.ptr, 0))
-
-        # TODO: Is there anything else to transfer?!
+        # TODO: Is there anything else to transfer?
         frame_sw.pts = frame.pts
-
         return frame_sw
 
-    cdef _build_format(self):
-        self._format = get_video_format(<lib.AVPixelFormat>self.ptr.pix_fmt, self.ptr.width, self.ptr.height)
-
     @property
     def format(self):
-        return self._format
+        return get_video_format(
+            cython.cast(lib.AVPixelFormat, self.ptr.pix_fmt),
+            self.ptr.width,
+            self.ptr.height,
+        )
 
     @format.setter
-    def format(self, VideoFormat format):
+    def format(self, format: VideoFormat):
         self.ptr.pix_fmt = format.pix_fmt
         self.ptr.width = format.width
         self.ptr.height = format.height
-        self._build_format()  # Kinda wasteful.
 
     @property
     def width(self):
-        if self.ptr is NULL:
+        if self.ptr is cython.NULL:
             return 0
         return self.ptr.width
 
     @width.setter
-    def width(self, unsigned int value):
+    def width(self, value: cython.uint):
         self.ptr.width = value
-        self._build_format()
 
     @property
     def height(self):
-        if self.ptr is NULL:
+        if self.ptr is cython.NULL:
             return 0
         return self.ptr.height
 
     @height.setter
-    def height(self, unsigned int value):
+    def height(self, value: cython.uint):
         self.ptr.height = value
-        self._build_format()
 
     @property
     def bits_per_coded_sample(self):
@@ -160,18 +169,17 @@ cdef class VideoCodecContext(CodecContext):
         The number of bits per sample in the codedwords. It's mandatory for this to be set for some formats to decode properly.
 
         Wraps :ffmpeg:`AVCodecContext.bits_per_coded_sample`.
-        
+
         :type: int
         """
         return self.ptr.bits_per_coded_sample
-      
+
     @bits_per_coded_sample.setter
-    def bits_per_coded_sample(self, int value):
+    def bits_per_coded_sample(self, value: cython.int):
         if self.is_encoder:
             raise ValueError("Not supported for encoders")
 
         self.ptr.bits_per_coded_sample = value
-        self._build_format()
 
     @property
     def pix_fmt(self):
@@ -180,12 +188,14 @@ cdef class VideoCodecContext(CodecContext):
 
         :type: str | None
         """
-        return getattr(self._format, "name", None)
+        desc: cython.pointer[cython.const[lib.AVPixFmtDescriptor]] = (
+            lib.av_pix_fmt_desc_get(cython.cast(lib.AVPixelFormat, self.ptr.pix_fmt))
+        )
+        return cython.cast(str, desc.name)
 
     @pix_fmt.setter
     def pix_fmt(self, value):
         self.ptr.pix_fmt = get_pix_fmt(value)
-        self._build_format()
 
     @property
     def framerate(self):
@@ -194,11 +204,11 @@ cdef class VideoCodecContext(CodecContext):
 
         :type: fractions.Fraction
         """
-        return avrational_to_fraction(&self.ptr.framerate)
+        return avrational_to_fraction(cython.address(self.ptr.framerate))
 
     @framerate.setter
     def framerate(self, value):
-        to_avrational(value, &self.ptr.framerate)
+        to_avrational(value, cython.address(self.ptr.framerate))
 
     @property
     def rate(self):
@@ -213,7 +223,7 @@ cdef class VideoCodecContext(CodecContext):
     def gop_size(self):
         """
         Sets the number of frames between keyframes. Used only for encoding.
-        
+
         :type: int
         """
         if self.is_decoder:
@@ -221,29 +231,31 @@ cdef class VideoCodecContext(CodecContext):
         return self.ptr.gop_size
 
     @gop_size.setter
-    def gop_size(self, int value):
+    def gop_size(self, value: cython.int):
         if self.is_decoder:
             raise RuntimeError("Cannot access 'gop_size' as a decoder")
         self.ptr.gop_size = value
 
     @property
     def sample_aspect_ratio(self):
-        return avrational_to_fraction(&self.ptr.sample_aspect_ratio)
+        return avrational_to_fraction(cython.address(self.ptr.sample_aspect_ratio))
 
     @sample_aspect_ratio.setter
     def sample_aspect_ratio(self, value):
-        to_avrational(value, &self.ptr.sample_aspect_ratio)
+        to_avrational(value, cython.address(self.ptr.sample_aspect_ratio))
 
     @property
     def display_aspect_ratio(self):
-        cdef lib.AVRational dar
-
+        dar: lib.AVRational
         lib.av_reduce(
-            &dar.num, &dar.den,
+            cython.address(dar.num),
+            cython.address(dar.den),
             self.ptr.width * self.ptr.sample_aspect_ratio.num,
-            self.ptr.height * self.ptr.sample_aspect_ratio.den, 1024*1024)
+            self.ptr.height * self.ptr.sample_aspect_ratio.den,
+            1024 * 1024,
+        )
 
-        return avrational_to_fraction(&dar)
+        return avrational_to_fraction(cython.address(dar))
 
     @property
     def has_b_frames(self):
@@ -252,6 +264,20 @@ cdef class VideoCodecContext(CodecContext):
         """
         return bool(self.ptr.has_b_frames)
 
+    @property
+    def reorder_depth(self):
+        """Raw ``has_b_frames`` value from FFmpeg (int, not bool).
+
+        After :meth:`flush_buffers`, FFmpeg may reset the internal reorder
+        heuristic. Set this to the known reorder depth *after* seeking to
+        avoid dropped hierarchical B-frames.
+        """
+        return self.ptr.has_b_frames
+
+    @reorder_depth.setter
+    def reorder_depth(self, value: cython.int):
+        self.ptr.has_b_frames = value
+
     @property
     def coded_width(self):
         """

venv/lib/python3.12/site-packages/av/video/codeccontext.pyi

@@ -19,6 +19,7 @@ class VideoCodecContext(CodecContext):
     sample_aspect_ratio: Fraction | None
     display_aspect_ratio: Fraction | None
     has_b_frames: bool
+    reorder_depth: int
     max_b_frames: int
     coded_width: int
     coded_height: int

venv/lib/python3.12/site-packages/av/video/format.pxd

@@ -2,26 +2,19 @@ cimport libav as lib
 
 
 cdef class VideoFormat:
-
     cdef lib.AVPixelFormat pix_fmt
     cdef const lib.AVPixFmtDescriptor *ptr
     cdef readonly unsigned int width, height
-
-    cdef readonly tuple components
-
     cdef _init(self, lib.AVPixelFormat pix_fmt, unsigned int width, unsigned int height)
-
     cpdef chroma_width(self, int luma_width=?)
     cpdef chroma_height(self, int luma_height=?)
 
 
 cdef class VideoFormatComponent:
-
     cdef VideoFormat format
     cdef readonly unsigned int index
     cdef const lib.AVComponentDescriptor *ptr
 
 
 cdef VideoFormat get_video_format(lib.AVPixelFormat c_format, unsigned int width, unsigned int height)
-
 cdef lib.AVPixelFormat get_pix_fmt(const char *name) except lib.AV_PIX_FMT_NONE

venv/lib/python3.12/site-packages/av/video/format.py

@@ -1,31 +1,40 @@
+import cython
+from cython import uint as cuint
 
-cdef object _cinit_bypass_sentinel = object()
+_cinit_bypass_sentinel = cython.declare(object, object())
 
-cdef VideoFormat get_video_format(lib.AVPixelFormat c_format, unsigned int width, unsigned int height):
+
+@cython.cfunc
+def get_video_format(
+    c_format: lib.AVPixelFormat, width: cuint, height: cuint
+) -> VideoFormat | None:
     if c_format == lib.AV_PIX_FMT_NONE:
         return None
 
-    cdef VideoFormat format = VideoFormat.__new__(VideoFormat, _cinit_bypass_sentinel)
+    format: VideoFormat = VideoFormat.__new__(VideoFormat, _cinit_bypass_sentinel)
     format._init(c_format, width, height)
     return format
 
-cdef lib.AVPixelFormat get_pix_fmt(const char *name) except lib.AV_PIX_FMT_NONE:
+
+@cython.cfunc
+@cython.exceptval(lib.AV_PIX_FMT_NONE, check=False)
+def get_pix_fmt(name: cython.p_const_char) -> lib.AVPixelFormat:
     """Wrapper for lib.av_get_pix_fmt with error checking."""
 
-    cdef lib.AVPixelFormat pix_fmt = lib.av_get_pix_fmt(name)
-
+    pix_fmt: lib.AVPixelFormat = lib.av_get_pix_fmt(name)
     if pix_fmt == lib.AV_PIX_FMT_NONE:
         raise ValueError("not a pixel format: %r" % name)
-
     return pix_fmt
 
 
-cdef class VideoFormat:
+@cython.final
+@cython.cclass
+class VideoFormat:
     """
 
-        >>> format = VideoFormat('rgb24')
-        >>> format.name
-        'rgb24'
+    >>> format = VideoFormat('rgb24')
+    >>> format.name
+    'rgb24'
 
     """
 
@@ -33,24 +42,20 @@ cdef class VideoFormat:
         if name is _cinit_bypass_sentinel:
             return
 
-        cdef VideoFormat other
         if isinstance(name, VideoFormat):
-            other = <VideoFormat>name
+            other: VideoFormat = cython.cast(VideoFormat, name)
             self._init(other.pix_fmt, width or other.width, height or other.height)
             return
 
-        cdef lib.AVPixelFormat pix_fmt = get_pix_fmt(name)
+        pix_fmt: lib.AVPixelFormat = get_pix_fmt(name)
         self._init(pix_fmt, width, height)
 
-    cdef _init(self, lib.AVPixelFormat pix_fmt, unsigned int width, unsigned int height):
+    @cython.cfunc
+    def _init(self, pix_fmt: lib.AVPixelFormat, width: cuint, height: cuint):
         self.pix_fmt = pix_fmt
         self.ptr = lib.av_pix_fmt_desc_get(pix_fmt)
         self.width = width
         self.height = height
-        self.components = tuple(
-            VideoFormatComponent(self, i)
-            for i in range(self.ptr.nb_components)
-        )
 
     def __repr__(self):
         if self.width or self.height:
@@ -64,54 +69,54 @@ cdef class VideoFormat:
     @property
     def name(self):
         """Canonical name of the pixel format."""
-        return <str>self.ptr.name
+        return cython.cast(str, self.ptr.name)
+
+    @property
+    def components(self):
+        return tuple(
+            VideoFormatComponent(self, i) for i in range(self.ptr.nb_components)
+        )
 
     @property
     def bits_per_pixel(self):
         return lib.av_get_bits_per_pixel(self.ptr)
 
     @property
-    def padded_bits_per_pixel(self): return lib.av_get_padded_bits_per_pixel(self.ptr)
+    def padded_bits_per_pixel(self):
+        return lib.av_get_padded_bits_per_pixel(self.ptr)
 
     @property
     def is_big_endian(self):
         """Pixel format is big-endian."""
         return bool(self.ptr.flags & lib.AV_PIX_FMT_FLAG_BE)
 
-
     @property
     def has_palette(self):
         """Pixel format has a palette in data[1], values are indexes in this palette."""
         return bool(self.ptr.flags & lib.AV_PIX_FMT_FLAG_PAL)
 
-
     @property
     def is_bit_stream(self):
         """All values of a component are bit-wise packed end to end."""
         return bool(self.ptr.flags & lib.AV_PIX_FMT_FLAG_BITSTREAM)
 
-
-    # Skipping PIX_FMT_HWACCEL
-    # """Pixel format is an HW accelerated format."""
-
     @property
     def is_planar(self):
         """At least one pixel component is not in the first data plane."""
         return bool(self.ptr.flags & lib.AV_PIX_FMT_FLAG_PLANAR)
 
-
     @property
     def is_rgb(self):
         """The pixel format contains RGB-like data (as opposed to YUV/grayscale)."""
         return bool(self.ptr.flags & lib.AV_PIX_FMT_FLAG_RGB)
-    
 
     @property
     def is_bayer(self):
         """The pixel format contains Bayer data."""
         return bool(self.ptr.flags & lib.AV_PIX_FMT_FLAG_BAYER)
 
-    cpdef chroma_width(self, int luma_width=0):
+    @cython.ccall
+    def chroma_width(self, luma_width: cython.int = 0):
         """chroma_width(luma_width=0)
 
         Width of a chroma plane relative to a luma plane.
@@ -122,7 +127,8 @@ cdef class VideoFormat:
         luma_width = luma_width or self.width
         return -((-luma_width) >> self.ptr.log2_chroma_w) if luma_width else 0
 
-    cpdef chroma_height(self, int luma_height=0):
+    @cython.ccall
+    def chroma_height(self, luma_height: cython.int = 0):
         """chroma_height(luma_height=0)
 
         Height of a chroma plane relative to a luma plane.
@@ -134,11 +140,13 @@ cdef class VideoFormat:
         return -((-luma_height) >> self.ptr.log2_chroma_h) if luma_height else 0
 
 
-cdef class VideoFormatComponent:
-    def __cinit__(self, VideoFormat format, size_t index):
+@cython.final
+@cython.cclass
+class VideoFormatComponent:
+    def __cinit__(self, format: VideoFormat, index: cython.size_t):
         self.format = format
         self.index = index
-        self.ptr = &format.ptr.comp[index]
+        self.ptr = cython.address(format.ptr.comp[index])
 
     @property
     def plane(self):
@@ -153,22 +161,25 @@ cdef class VideoFormatComponent:
     @property
     def is_alpha(self):
         """Is this component an alpha channel?"""
-        return ((self.index == 1 and self.format.ptr.nb_components == 2) or
-                (self.index == 3 and self.format.ptr.nb_components == 4))
+        return (self.index == 1 and self.format.ptr.nb_components == 2) or (
+            self.index == 3 and self.format.ptr.nb_components == 4
+        )
 
     @property
     def is_luma(self):
         """Is this component a luma channel?"""
         return self.index == 0 and (
-            self.format.ptr.nb_components == 1 or
-            self.format.ptr.nb_components == 2 or
-            not self.format.is_rgb
+            self.format.ptr.nb_components == 1
+            or self.format.ptr.nb_components == 2
+            or not self.format.is_rgb
         )
 
     @property
     def is_chroma(self):
         """Is this component a chroma channel?"""
-        return (self.index == 1 or self.index == 2) and (self.format.ptr.log2_chroma_w or self.format.ptr.log2_chroma_h)
+        return (self.index == 1 or self.index == 2) and (
+            self.format.ptr.log2_chroma_w or self.format.ptr.log2_chroma_h
+        )
 
     @property
     def width(self):
@@ -190,7 +201,7 @@ cdef class VideoFormatComponent:
 
 
 names = set()
-cdef const lib.AVPixFmtDescriptor *desc = NULL
+desc = cython.declare(cython.pointer[cython.const[lib.AVPixFmtDescriptor]], cython.NULL)
 while True:
     desc = lib.av_pix_fmt_desc_next(desc)
     if not desc:

venv/lib/python3.12/site-packages/av/video/format.pyi

@@ -6,14 +6,14 @@ class VideoFormat:
     has_palette: bool
     is_bit_stream: bool
     is_planar: bool
+    width: int
+    height: int
+    @property
+    def components(self) -> tuple[VideoFormatComponent, ...]: ...
     @property
     def is_rgb(self) -> bool: ...
     @property
     def is_bayer(self) -> bool: ...
-    width: int
-    height: int
-    components: tuple[VideoFormatComponent, ...]
-
     def __init__(self, name: str, width: int = 0, height: int = 0) -> None: ...
     def chroma_width(self, luma_width: int = 0) -> int: ...
     def chroma_height(self, luma_height: int = 0) -> int: ...

venv/lib/python3.12/site-packages/av/video/frame.pxd

@@ -6,15 +6,21 @@ from av.video.format cimport VideoFormat
 from av.video.reformatter cimport VideoReformatter
 
 
-cdef class VideoFrame(Frame):
-    # This is the buffer that is used to back everything in the AVFrame.
-    # We don't ever actually access it directly.
-    cdef uint8_t *_buffer
-    cdef object _np_buffer
+cdef class CudaContext:
+    cdef readonly int device_id
+    cdef readonly bint primary_ctx
+    cdef lib.AVBufferRef* _device_ref
+    cdef dict _frames_cache
+    cdef lib.AVBufferRef* _get_device_ref(self)
+    cdef public lib.AVBufferRef* get_frames_ctx(
+        self, lib.AVPixelFormat sw_fmt, int width, int height
+    )
 
+cdef class VideoFrame(Frame):
+    cdef CudaContext _cuda_ctx
     cdef VideoReformatter reformatter
     cdef readonly VideoFormat format
-
+    cdef readonly int _device_id
     cdef _init(self, lib.AVPixelFormat format, unsigned int width, unsigned int height)
     cdef _init_user_attributes(self)
     cpdef save(self, object filepath)

venv/lib/python3.12/site-packages/av/video/frame.py

@@ -2,14 +2,180 @@ import sys
 from enum import IntEnum
 
 import cython
+import cython.cimports.libav as lib
+from cython.cimports.av.dictionary import Dictionary
 from cython.cimports.av.error import err_check
 from cython.cimports.av.sidedata.sidedata import get_display_rotation
 from cython.cimports.av.utils import check_ndarray
 from cython.cimports.av.video.format import get_pix_fmt, get_video_format
-from cython.cimports.av.video.plane import VideoPlane
-from cython.cimports.libc.stdint import uint8_t
+from cython.cimports.av.video.plane import DLManagedTensor, VideoPlane, kCPU, kCuda
+from cython.cimports.cpython.exc import PyErr_Clear
+from cython.cimports.cpython.pycapsule import (
+    PyCapsule_GetPointer,
+    PyCapsule_IsValid,
+    PyCapsule_SetName,
+)
+from cython.cimports.cpython.ref import Py_DECREF, Py_INCREF
+from cython.cimports.libc.stdint import int64_t, uint8_t
 
-_cinit_bypass_sentinel = object()
+
+@cython.final
+@cython.cclass
+class CudaContext:
+    def __cinit__(self, device_id: cython.int = 0, primary_ctx: cython.bint = True):
+        self.device_id = device_id
+        self.primary_ctx = primary_ctx
+        self._device_ref = cython.NULL
+        self._frames_cache = {}
+
+    def __dealloc__(self):
+        ref: cython.pointer[lib.AVBufferRef]
+
+        for v in self._frames_cache.values():
+            ref = cython.cast(
+                cython.pointer[lib.AVBufferRef],
+                cython.cast(cython.size_t, v),
+            )
+            lib.av_buffer_unref(cython.address(ref))
+        self._frames_cache.clear()
+
+        ref = self._device_ref
+        if ref != cython.NULL:
+            lib.av_buffer_unref(cython.address(ref))
+            self._device_ref = cython.NULL
+
+    @cython.cfunc
+    def _get_device_ref(self) -> cython.pointer[lib.AVBufferRef]:
+        device_ref: cython.pointer[lib.AVBufferRef] = self._device_ref
+        if device_ref != cython.NULL:
+            return device_ref
+
+        device_ref = cython.NULL
+        device_bytes = f"{self.device_id}".encode()
+        c_device: cython.p_char = device_bytes
+        options: Dictionary = Dictionary(
+            {"primary_ctx": "1" if self.primary_ctx else "0"}
+        )
+        err_check(
+            lib.av_hwdevice_ctx_create(
+                cython.address(device_ref),
+                lib.AV_HWDEVICE_TYPE_CUDA,
+                c_device,
+                options.ptr,
+                0,
+            )
+        )
+        self._device_ref = device_ref
+        return device_ref
+
+    @cython.cfunc
+    def get_frames_ctx(
+        self,
+        sw_fmt: lib.AVPixelFormat,
+        width: cython.int,
+        height: cython.int,
+    ) -> cython.pointer[lib.AVBufferRef]:
+        key = (int(sw_fmt), int(width), int(height))
+        cached = self._frames_cache.get(key)
+        cached_ref: cython.pointer[lib.AVBufferRef]
+        out_ref: cython.pointer[lib.AVBufferRef]
+
+        if cached is not None:
+            cached_ref = cython.cast(
+                cython.pointer[lib.AVBufferRef],
+                cython.cast(cython.size_t, cached),
+            )
+            out_ref = lib.av_buffer_ref(cached_ref)
+            if out_ref == cython.NULL:
+                raise MemoryError("av_buffer_ref() failed")
+            return out_ref
+
+        device_ref = self._get_device_ref()
+
+        frames_ref: cython.pointer[lib.AVBufferRef] = lib.av_hwframe_ctx_alloc(
+            device_ref
+        )
+        if frames_ref == cython.NULL:
+            raise MemoryError("av_hwframe_ctx_alloc() failed")
+
+        try:
+            frames_ctx: cython.pointer[lib.AVHWFramesContext] = cython.cast(
+                cython.pointer[lib.AVHWFramesContext], frames_ref.data
+            )
+            frames_ctx.format = get_pix_fmt(b"cuda")
+            frames_ctx.sw_format = sw_fmt
+            frames_ctx.width = int(width)
+            frames_ctx.height = int(height)
+            err_check(lib.av_hwframe_ctx_init(frames_ref))
+        except Exception:
+            lib.av_buffer_unref(cython.address(frames_ref))
+            raise
+
+        out_ref = lib.av_buffer_ref(frames_ref)
+        if out_ref == cython.NULL:
+            lib.av_buffer_unref(cython.address(frames_ref))
+            raise MemoryError("av_buffer_ref() failed")
+
+        self._frames_cache[key] = cython.cast(cython.size_t, frames_ref)
+        return out_ref
+
+
+@cython.cfunc
+def _consume_dlpack(obj: object, stream: object) -> cython.pointer[DLManagedTensor]:
+    capsule: object
+    managed: cython.pointer[DLManagedTensor]
+
+    if hasattr(obj, "__dlpack__"):
+        capsule = obj.__dlpack__() if stream is None else obj.__dlpack__(stream=stream)
+    else:
+        capsule = obj
+
+    if not PyCapsule_IsValid(capsule, b"dltensor"):
+        PyErr_Clear()
+        raise TypeError(
+            "expected a DLPack capsule or an object implementing __dlpack__"
+        )
+
+    managed = cython.cast(
+        cython.pointer[DLManagedTensor],
+        PyCapsule_GetPointer(capsule, b"dltensor"),
+    )
+    if managed == cython.NULL:
+        raise ValueError("PyCapsule_GetPointer returned NULL")
+
+    if PyCapsule_SetName(capsule, b"used_dltensor") != 0:
+        raise RuntimeError("PyCapsule_SetName failed")
+
+    return managed
+
+
+@cython.cfunc
+@cython.nogil
+@cython.exceptval(check=False)
+def _dlpack_avbuffer_free(
+    opaque: cython.p_void,
+    data: cython.pointer[uint8_t],
+) -> cython.void:
+    managed: cython.pointer[DLManagedTensor] = cython.cast(
+        cython.pointer[DLManagedTensor], opaque
+    )
+    if managed != cython.NULL:
+        managed.deleter(managed)
+
+
+@cython.cfunc
+@cython.nogil
+@cython.exceptval(check=False)
+def _numpy_avbuffer_free(
+    opaque: cython.p_void,
+    data: cython.pointer[uint8_t],
+) -> cython.void:
+    if opaque != cython.NULL:
+        with cython.gil:
+            Py_DECREF(cython.cast(object, opaque))
+
+
+_cinit_bypass_sentinel = cython.declare(object, object())
 
 # `pix_fmt`s supported by Frame.to_ndarray() and Frame.from_ndarray()
 supported_np_pix_fmts = {
@@ -99,6 +265,93 @@ supported_np_pix_fmts = {
     "yuyv422",
 }
 
+# Mapping from format name to (itemsize, dtype) for formats where planes
+# are simply concatenated into shape (height, width, channels).
+_np_pix_fmt_dtypes = cython.declare(
+    dict[str, tuple[cython.uint, str]],
+    {
+        "abgr": (4, "uint8"),
+        "argb": (4, "uint8"),
+        "bayer_bggr8": (1, "uint8"),
+        "bayer_gbrg8": (1, "uint8"),
+        "bayer_grbg8": (1, "uint8"),
+        "bayer_rggb8": (1, "uint8"),
+        "bayer_bggr16le": (2, "uint16"),
+        "bayer_bggr16be": (2, "uint16"),
+        "bayer_gbrg16le": (2, "uint16"),
+        "bayer_gbrg16be": (2, "uint16"),
+        "bayer_grbg16le": (2, "uint16"),
+        "bayer_grbg16be": (2, "uint16"),
+        "bayer_rggb16le": (2, "uint16"),
+        "bayer_rggb16be": (2, "uint16"),
+        "bgr24": (3, "uint8"),
+        "bgr48be": (6, "uint16"),
+        "bgr48le": (6, "uint16"),
+        "bgr8": (1, "uint8"),
+        "bgra": (4, "uint8"),
+        "bgra64be": (8, "uint16"),
+        "bgra64le": (8, "uint16"),
+        "gbrap": (1, "uint8"),
+        "gbrap10be": (2, "uint16"),
+        "gbrap10le": (2, "uint16"),
+        "gbrap12be": (2, "uint16"),
+        "gbrap12le": (2, "uint16"),
+        "gbrap14be": (2, "uint16"),
+        "gbrap14le": (2, "uint16"),
+        "gbrap16be": (2, "uint16"),
+        "gbrap16le": (2, "uint16"),
+        "gbrapf32be": (4, "float32"),
+        "gbrapf32le": (4, "float32"),
+        "gbrp": (1, "uint8"),
+        "gbrp10be": (2, "uint16"),
+        "gbrp10le": (2, "uint16"),
+        "gbrp12be": (2, "uint16"),
+        "gbrp12le": (2, "uint16"),
+        "gbrp14be": (2, "uint16"),
+        "gbrp14le": (2, "uint16"),
+        "gbrp16be": (2, "uint16"),
+        "gbrp16le": (2, "uint16"),
+        "gbrp9be": (2, "uint16"),
+        "gbrp9le": (2, "uint16"),
+        "gbrpf32be": (4, "float32"),
+        "gbrpf32le": (4, "float32"),
+        "gray": (1, "uint8"),
+        "gray10be": (2, "uint16"),
+        "gray10le": (2, "uint16"),
+        "gray12be": (2, "uint16"),
+        "gray12le": (2, "uint16"),
+        "gray14be": (2, "uint16"),
+        "gray14le": (2, "uint16"),
+        "gray16be": (2, "uint16"),
+        "gray16le": (2, "uint16"),
+        "gray8": (1, "uint8"),
+        "gray9be": (2, "uint16"),
+        "gray9le": (2, "uint16"),
+        "grayf32be": (4, "float32"),
+        "grayf32le": (4, "float32"),
+        "rgb24": (3, "uint8"),
+        "rgb48be": (6, "uint16"),
+        "rgb48le": (6, "uint16"),
+        "rgb8": (1, "uint8"),
+        "rgba": (4, "uint8"),
+        "rgba64be": (8, "uint16"),
+        "rgba64le": (8, "uint16"),
+        "rgbaf16be": (8, "float16"),
+        "rgbaf16le": (8, "float16"),
+        "rgbaf32be": (16, "float32"),
+        "rgbaf32le": (16, "float32"),
+        "rgbf32be": (12, "float32"),
+        "rgbf32le": (12, "float32"),
+        "yuv444p": (1, "uint8"),
+        "yuv444p16be": (2, "uint16"),
+        "yuv444p16le": (2, "uint16"),
+        "yuva444p16be": (2, "uint16"),
+        "yuva444p16le": (2, "uint16"),
+        "yuvj444p": (1, "uint8"),
+        "yuyv422": (2, "uint8"),
+    },
+)
+
 
 @cython.cfunc
 def alloc_video_frame() -> VideoFrame:
@@ -122,9 +375,13 @@ class PictureType(IntEnum):
     BI = lib.AV_PICTURE_TYPE_BI  # BI type
 
 
+_is_big_endian = cython.declare(cython.bint, sys.byteorder == "big")
+
+
 @cython.cfunc
+@cython.inline
 def byteswap_array(array, big_endian: cython.bint):
-    if (sys.byteorder == "big") != big_endian:
+    if _is_big_endian != big_endian:
         return array.byteswap()
     return array
 
@@ -158,7 +415,9 @@ def copy_bytes_to_plane(
     for row in range(start_row, end_row, step):
         i_pos = row * i_stride
         if flip_horizontal:
+            i: cython.Py_ssize_t
             for i in range(0, i_stride, bytes_per_pixel):
+                j: cython.Py_ssize_t
                 for j in range(bytes_per_pixel):
                     o_buf[o_pos + i + j] = i_buf[
                         i_pos + i_stride - i - bytes_per_pixel + j
@@ -175,22 +434,39 @@ def copy_array_to_plane(array, plane: VideoPlane, bytes_per_pixel: cython.uint):
 
 
 @cython.cfunc
+@cython.inline
 def useful_array(
     plane: VideoPlane, bytes_per_pixel: cython.uint = 1, dtype: str = "uint8"
 ):
     """
-    Return the useful part of the VideoPlane as a single dimensional array.
+    Return the useful part of the VideoPlane as a strided array.
 
-    We are simply discarding any padding which was added for alignment.
+    We are simply creating a view that discards any padding which was added for
+    alignment.
     """
     import numpy as np
 
-    total_line_size: cython.size_t = abs(plane.line_size)
-    useful_line_size: cython.size_t = plane.width * bytes_per_pixel
-    arr = np.frombuffer(plane, np.uint8)
-    if total_line_size != useful_line_size:
-        arr = arr.reshape(-1, total_line_size)[:, 0:useful_line_size].reshape(-1)
-    return arr.view(np.dtype(dtype))
+    dtype_obj = np.dtype(dtype)
+    line_size = plane.frame.ptr.linesize[plane.index]
+    total_line_size = abs(line_size)
+    itemsize = dtype_obj.itemsize
+    channels = bytes_per_pixel // itemsize
+
+    if channels == 1:
+        shape = (plane.height, plane.width)
+        strides = (total_line_size, itemsize)
+    else:
+        shape = (plane.height, plane.width, channels)
+        strides = (total_line_size, bytes_per_pixel, itemsize)
+
+    if line_size < 0:
+        offset = (plane.height - 1) * total_line_size
+        strides = (-total_line_size, *strides[1:])
+        return np.ndarray(
+            shape, dtype=dtype_obj, buffer=plane, offset=offset, strides=strides
+        )
+
+    return np.ndarray(shape, dtype=dtype_obj, buffer=plane, strides=strides)
 
 
 @cython.cfunc
@@ -199,6 +475,7 @@ def check_ndarray_shape(array: object, ok: cython.bint):
         raise ValueError(f"Unexpected numpy array shape `{array.shape}`")
 
 
+@cython.final
 @cython.cclass
 class VideoFrame(Frame):
     def __cinit__(self, width=0, height=0, format="yuv420p"):
@@ -217,13 +494,10 @@ class VideoFrame(Frame):
             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_image_alloc(
-                    self.ptr.data, self.ptr.linesize, width, height, format, 16
-                )
-                self._buffer = self.ptr.data[0]
+                res = lib.av_frame_get_buffer(self.ptr, 16)
 
         if res:
             err_check(res)
@@ -239,11 +513,7 @@ class VideoFrame(Frame):
         )
 
     def __dealloc__(self):
-        # The `self._buffer` member is only set if *we* allocated the buffer in `_init`,
-        # as opposed to a buffer allocated by a decoder.
-        lib.av_freep(cython.address(self._buffer))
-        # Let go of the reference from the numpy buffers if we made one
-        self._np_buffer = None
+        lib.av_frame_unref(self.ptr)
 
     def __repr__(self):
         return (
@@ -259,17 +529,28 @@ class VideoFrame(Frame):
         # We need to detect which planes actually exist, but also constrain ourselves to
         # the maximum plane count (as determined only by VideoFrames so far), in case
         # the library implementation does not set the last plane to NULL.
+        fmt = self.format
+        if self.ptr.hw_frames_ctx:
+            frames_ctx: cython.pointer[lib.AVHWFramesContext] = cython.cast(
+                cython.pointer[lib.AVHWFramesContext], self.ptr.hw_frames_ctx.data
+            )
+            fmt = get_video_format(
+                frames_ctx.sw_format, self.ptr.width, self.ptr.height
+            )
+
         max_plane_count: cython.int = 0
-        for i in range(self.format.ptr.nb_components):
-            count = self.format.ptr.comp[i].plane + 1
+        for i in range(fmt.ptr.nb_components):
+            count = fmt.ptr.comp[i].plane + 1
             if max_plane_count < count:
                 max_plane_count = count
-        if self.format.name == "pal8":
+        if fmt.name == "pal8":
             max_plane_count = 2
 
         plane_count: cython.int = 0
         while plane_count < max_plane_count and self.ptr.extended_data[plane_count]:
             plane_count += 1
+        if plane_count == 1:
+            return (VideoPlane(self, 0),)
         return tuple([VideoPlane(self, i) for i in range(plane_count)])
 
     @property
@@ -338,8 +619,34 @@ class VideoFrame(Frame):
     def color_range(self, value):
         self.ptr.color_range = value
 
+    @property
+    def color_trc(self):
+        """Transfer characteristic of frame.
+
+        Wraps :ffmpeg:`AVFrame.color_trc`.
+
+        """
+        return self.ptr.color_trc
+
+    @color_trc.setter
+    def color_trc(self, value):
+        self.ptr.color_trc = value
+
+    @property
+    def color_primaries(self):
+        """Color primaries of frame.
+
+        Wraps :ffmpeg:`AVFrame.color_primaries`.
+
+        """
+        return self.ptr.color_primaries
+
+    @color_primaries.setter
+    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.
 
@@ -405,17 +712,24 @@ class VideoFrame(Frame):
         plane: VideoPlane = self.reformat(format="rgb24", **kwargs).planes[0]
 
         i_buf: cython.const[uint8_t][:] = plane
-        i_pos: cython.size_t = 0
-        i_stride: cython.size_t = plane.line_size
+        line_size: cython.int = plane.line_size
+        i_stride: cython.size_t = abs(line_size)
 
         o_pos: cython.size_t = 0
         o_stride: cython.size_t = plane.width * 3
         o_size: cython.size_t = plane.height * o_stride
         o_buf: bytearray = bytearray(o_size)
 
+        # For bottom-up frames (negative line_size) the buffer protocol exposes
+        # rows from the lowest address, so the top display row is at the far end.
+        i_pos: cython.size_t = (plane.height - 1) * i_stride if line_size < 0 else 0
+
         while o_pos < o_size:
             o_buf[o_pos : o_pos + o_stride] = i_buf[i_pos : i_pos + o_stride]
-            i_pos += i_stride
+            if line_size < 0:
+                i_pos -= i_stride
+            else:
+                i_pos += i_stride
             o_pos += o_stride
 
         return Image.frombytes(
@@ -444,150 +758,67 @@ class VideoFrame(Frame):
         .. note:: For ``gbrp`` formats, channels are flipped to RGB order.
 
         """
-        frame: VideoFrame = self.reformat(**kwargs)
+        if self.ptr.hw_frames_ctx and "format" not in kwargs:
+            frames_ctx: cython.pointer[lib.AVHWFramesContext] = cython.cast(
+                cython.pointer[lib.AVHWFramesContext], self.ptr.hw_frames_ctx.data
+            )
+            kwargs = dict(kwargs)
+            kwargs["format"] = get_video_format(
+                frames_ctx.sw_format, self.ptr.width, self.ptr.height
+            ).name
+
+        frame: VideoFrame = self.reformat(**kwargs) if len(kwargs) > 0 else self
+        if frame.ptr.hw_frames_ctx:
+            raise ValueError("Cannot convert a hardware frame to numpy directly.")
 
         import numpy as np
 
         # check size
-        if frame.format.name in {
-            "yuv420p",
-            "yuvj420p",
-            "yuyv422",
-            "yuv422p10le",
-            "yuv422p",
-        }:
-            assert frame.width % 2 == 0, (
-                "the width has to be even for this pixel format"
-            )
-            assert frame.height % 2 == 0, (
-                "the height has to be even for this pixel format"
-            )
-
+        format_name = frame.format.name
+        planes: tuple[VideoPlane, ...] = frame.planes
         # cases planes are simply concatenated in shape (height, width, channels)
-        itemsize, dtype = {
-            "abgr": (4, "uint8"),
-            "argb": (4, "uint8"),
-            "bayer_bggr8": (1, "uint8"),
-            "bayer_gbrg8": (1, "uint8"),
-            "bayer_grbg8": (1, "uint8"),
-            "bayer_rggb8": (1, "uint8"),
-            "bayer_bggr16le": (2, "uint16"),
-            "bayer_bggr16be": (2, "uint16"),
-            "bayer_gbrg16le": (2, "uint16"),
-            "bayer_gbrg16be": (2, "uint16"),
-            "bayer_grbg16le": (2, "uint16"),
-            "bayer_grbg16be": (2, "uint16"),
-            "bayer_rggb16le": (2, "uint16"),
-            "bayer_rggb16be": (2, "uint16"),
-            "bgr24": (3, "uint8"),
-            "bgr48be": (6, "uint16"),
-            "bgr48le": (6, "uint16"),
-            "bgr8": (1, "uint8"),
-            "bgra": (4, "uint8"),
-            "bgra64be": (8, "uint16"),
-            "bgra64le": (8, "uint16"),
-            "gbrap": (1, "uint8"),
-            "gbrap10be": (2, "uint16"),
-            "gbrap10le": (2, "uint16"),
-            "gbrap12be": (2, "uint16"),
-            "gbrap12le": (2, "uint16"),
-            "gbrap14be": (2, "uint16"),
-            "gbrap14le": (2, "uint16"),
-            "gbrap16be": (2, "uint16"),
-            "gbrap16le": (2, "uint16"),
-            "gbrapf32be": (4, "float32"),
-            "gbrapf32le": (4, "float32"),
-            "gbrp": (1, "uint8"),
-            "gbrp10be": (2, "uint16"),
-            "gbrp10le": (2, "uint16"),
-            "gbrp12be": (2, "uint16"),
-            "gbrp12le": (2, "uint16"),
-            "gbrp14be": (2, "uint16"),
-            "gbrp14le": (2, "uint16"),
-            "gbrp16be": (2, "uint16"),
-            "gbrp16le": (2, "uint16"),
-            "gbrp9be": (2, "uint16"),
-            "gbrp9le": (2, "uint16"),
-            "gbrpf32be": (4, "float32"),
-            "gbrpf32le": (4, "float32"),
-            "gray": (1, "uint8"),
-            "gray10be": (2, "uint16"),
-            "gray10le": (2, "uint16"),
-            "gray12be": (2, "uint16"),
-            "gray12le": (2, "uint16"),
-            "gray14be": (2, "uint16"),
-            "gray14le": (2, "uint16"),
-            "gray16be": (2, "uint16"),
-            "gray16le": (2, "uint16"),
-            "gray8": (1, "uint8"),
-            "gray9be": (2, "uint16"),
-            "gray9le": (2, "uint16"),
-            "grayf32be": (4, "float32"),
-            "grayf32le": (4, "float32"),
-            "rgb24": (3, "uint8"),
-            "rgb48be": (6, "uint16"),
-            "rgb48le": (6, "uint16"),
-            "rgb8": (1, "uint8"),
-            "rgba": (4, "uint8"),
-            "rgba64be": (8, "uint16"),
-            "rgba64le": (8, "uint16"),
-            "rgbaf16be": (8, "float16"),
-            "rgbaf16le": (8, "float16"),
-            "rgbaf32be": (16, "float32"),
-            "rgbaf32le": (16, "float32"),
-            "rgbf32be": (12, "float32"),
-            "rgbf32le": (12, "float32"),
-            "yuv444p": (1, "uint8"),
-            "yuv444p16be": (2, "uint16"),
-            "yuv444p16le": (2, "uint16"),
-            "yuva444p16be": (2, "uint16"),
-            "yuva444p16le": (2, "uint16"),
-            "yuvj444p": (1, "uint8"),
-            "yuyv422": (2, "uint8"),
-        }.get(frame.format.name, (None, None))
-        if itemsize is not None:
-            layers = [
-                useful_array(plan, itemsize, dtype).reshape(
-                    frame.height, frame.width, -1
-                )
-                for plan in frame.planes
-            ]
-            if len(layers) == 1:  # shortcut, avoid memory copy
-                array = layers[0]
+        if format_name in _np_pix_fmt_dtypes:
+            if format_name == "yuyv422":
+                assert frame.ptr.width % 2 == 0, "width has to be even for yuyv422"
+                assert frame.ptr.height % 2 == 0, "height has to be even for yuyv422"
+            itemsize: cython.uint
+            itemsize, dtype = _np_pix_fmt_dtypes[format_name]
+            num_planes: cython.size_t = len(planes)
+            if num_planes == 1:  # shortcut, avoid memory copy
+                array = useful_array(planes[0], itemsize, dtype)
             else:  # general case
-                array = np.concatenate(layers, axis=2)
-            array = byteswap_array(array, frame.format.name.endswith("be"))
-            if array.shape[2] == 1:  # skip last channel for gray images
-                return array.squeeze(2)
-            if frame.format.name.startswith("gbr"):  # gbr -> rgb
-                buffer = array[:, :, 0].copy()
-                array[:, :, 0] = array[:, :, 2]
-                array[:, :, 2] = array[:, :, 1]
-                array[:, :, 1] = buffer
-            if not channel_last and frame.format.name in {"yuv444p", "yuvj444p"}:
+                array = np.empty(
+                    (frame.ptr.height, frame.ptr.width, num_planes), dtype=dtype
+                )
+                if format_name.startswith("gbr"):
+                    plane_indices = (2, 0, 1, *range(3, num_planes))
+                else:
+                    plane_indices = range(num_planes)
+                for i, p_idx in enumerate(plane_indices):
+                    array[:, :, i] = useful_array(planes[p_idx], itemsize, dtype)
+            array = byteswap_array(array, format_name.endswith("be"))
+            if not channel_last and format_name in {"yuv444p", "yuvj444p"}:
                 array = np.moveaxis(array, 2, 0)
             return array
 
         # special cases
-        if frame.format.name in {"yuv420p", "yuvj420p", "yuv422p"}:
+        if format_name in {"yuv420p", "yuvj420p", "yuv422p"}:
+            assert frame.ptr.width % 2 == 0, "width has to be even for this format"
+            assert frame.ptr.height % 2 == 0, "height has to be even for this format"
             return np.hstack(
                 [
-                    useful_array(frame.planes[0]),
-                    useful_array(frame.planes[1]),
-                    useful_array(frame.planes[2]),
+                    useful_array(planes[0]).reshape(-1),
+                    useful_array(planes[1]).reshape(-1),
+                    useful_array(planes[2]).reshape(-1),
                 ]
-            ).reshape(-1, frame.width)
-        if frame.format.name == "yuv422p10le":
+            ).reshape(-1, frame.ptr.width)
+        if format_name == "yuv422p10le":
+            assert frame.ptr.width % 2 == 0, "width has to be even for this format"
+            assert frame.ptr.height % 2 == 0, "height has to be even for this format"
             # Read planes as uint16 at their original width
-            y = useful_array(frame.planes[0], 2, "uint16").reshape(
-                frame.height, frame.width
-            )
-            u = useful_array(frame.planes[1], 2, "uint16").reshape(
-                frame.height, frame.width // 2
-            )
-            v = useful_array(frame.planes[2], 2, "uint16").reshape(
-                frame.height, frame.width // 2
-            )
+            y = useful_array(planes[0], 2, "uint16")
+            u = useful_array(planes[1], 2, "uint16")
+            v = useful_array(planes[2], 2, "uint16")
 
             # Double the width of U and V by repeating each value
             u_full = np.repeat(u, 2, axis=1)
@@ -595,25 +826,25 @@ class VideoFrame(Frame):
             if channel_last:
                 return np.stack([y, u_full, v_full], axis=2)
             return np.stack([y, u_full, v_full], axis=0)
-        if frame.format.name == "pal8":
-            image = useful_array(frame.planes[0]).reshape(frame.height, frame.width)
+        if format_name == "pal8":
+            image = useful_array(planes[0])
             palette = (
-                np.frombuffer(frame.planes[1], "i4")
+                np.frombuffer(planes[1], "i4")
                 .astype(">i4")
                 .reshape(-1, 1)
                 .view(np.uint8)
             )
             return image, palette
-        if frame.format.name == "nv12":
+        if format_name == "nv12":
             return np.hstack(
                 [
-                    useful_array(frame.planes[0]),
-                    useful_array(frame.planes[1], 2),
+                    useful_array(planes[0]).reshape(-1),
+                    useful_array(planes[1], 2).reshape(-1),
                 ]
-            ).reshape(-1, frame.width)
+            ).reshape(-1, frame.ptr.width)
 
         raise ValueError(
-            f"Conversion to numpy array with format `{frame.format.name}` is not yet supported"
+            f"Conversion to numpy array with format `{format_name}` is not yet supported"
         )
 
     def set_image(self, img):
@@ -857,10 +1088,6 @@ class VideoFrame(Frame):
         return frame
 
     def _image_fill_pointers_numpy(self, buffer, width, height, linesizes, format):
-        c_format: lib.AVPixelFormat
-        c_ptr: cython.pointer[uint8_t]
-        c_data: cython.size_t
-
         # If you want to use the numpy notation, then you need to include the following lines at the top of the file:
         #      cimport numpy as cnp
         #      cnp.import_array()
@@ -877,29 +1104,42 @@ class VideoFrame(Frame):
         # Using buffer.ctypes.data helps avoid any kind of usage of the c-api from
         # numpy, which avoid the need to add numpy as a compile time dependency.
 
-        c_data = buffer.ctypes.data
-        c_ptr = cython.cast(cython.pointer[uint8_t], c_data)
-        c_format = get_pix_fmt(format)
-        lib.av_freep(cython.address(self._buffer))
+        c_data: cython.Py_ssize_t = buffer.ctypes.data
+        c_ptr: cython.pointer[uint8_t] = cython.cast(cython.pointer[uint8_t], c_data)
+        c_format: lib.AVPixelFormat = get_pix_fmt(format)
+        lib.av_frame_unref(self.ptr)
 
         # Hold on to a reference for the numpy buffer so that it doesn't get accidentally garbage collected
-        self._np_buffer = buffer
         self.ptr.format = c_format
         self.ptr.width = width
         self.ptr.height = height
         for i, linesize in enumerate(linesizes):
             self.ptr.linesize[i] = linesize
 
-        res = lib.av_image_fill_pointers(
-            self.ptr.data,
-            cython.cast(lib.AVPixelFormat, self.ptr.format),
-            self.ptr.height,
-            c_ptr,
-            self.ptr.linesize,
+        required = err_check(
+            lib.av_image_fill_pointers(
+                self.ptr.data,
+                cython.cast(lib.AVPixelFormat, self.ptr.format),
+                self.ptr.height,
+                c_ptr,
+                self.ptr.linesize,
+            )
         )
 
-        if res:
-            err_check(res)
+        py_buf = cython.cast(object, buffer)
+        Py_INCREF(py_buf)
+
+        self.ptr.buf[0] = lib.av_buffer_create(
+            c_ptr,
+            required,
+            _numpy_avbuffer_free,
+            cython.cast(cython.p_void, py_buf),
+            0,
+        )
+        if self.ptr.buf[0] == cython.NULL:
+            Py_DECREF(py_buf)
+            raise MemoryError("av_buffer_create failed")
+
         self._init_user_attributes()
 
     @staticmethod
@@ -1177,3 +1417,215 @@ class VideoFrame(Frame):
         else:
             raise NotImplementedError(f"Format '{format}' is not supported.")
         return frame
+
+    @staticmethod
+    def from_dlpack(
+        planes,
+        format: str = "nv12",
+        width: int = 0,
+        height: int = 0,
+        stream=None,
+        device_id: int | None = None,
+        primary_ctx: bool = True,
+        cuda_context=None,
+    ):
+        if not isinstance(planes, (tuple, list)):
+            planes = (planes,)
+
+        if len(planes) != 2:
+            raise ValueError(
+                "from_dlpack currently supports 2-plane formats only (nv12/p010le/p016le)"
+            )
+
+        sw_fmt: lib.AVPixelFormat = get_pix_fmt(format)
+        nv12 = get_pix_fmt(b"nv12")
+        p010le = get_pix_fmt(b"p010le")
+        p016le = get_pix_fmt(b"p016le")
+
+        if sw_fmt not in (nv12, p010le, p016le):
+            raise NotImplementedError("from_dlpack supports nv12, p010le, p016le only")
+
+        expected_bits = 8 if sw_fmt == nv12 else 16
+        itemsize = 1 if expected_bits == 8 else 2
+
+        m0: cython.pointer[DLManagedTensor] = cython.NULL
+        m1: cython.pointer[DLManagedTensor] = cython.NULL
+        frame: VideoFrame = None
+
+        try:
+            m0 = _consume_dlpack(planes[0], stream)
+            m1 = _consume_dlpack(planes[1], stream)
+
+            dev_type0 = m0.dl_tensor.device_type
+            dev_type1 = m1.dl_tensor.device_type
+            if dev_type0 != dev_type1:
+                raise ValueError("plane tensors must have the same device_type")
+            if dev_type0 not in (kCuda, kCPU):
+                raise NotImplementedError(
+                    "only CPU and CUDA DLPack tensors are supported"
+                )
+
+            dev0 = m0.dl_tensor.device_id
+            dev1 = m1.dl_tensor.device_id
+            if dev0 != dev1:
+                raise ValueError("plane tensors must be on the same CUDA device")
+            if dev_type0 == kCuda:
+                if device_id is None:
+                    device_id = dev0
+                elif device_id != dev0:
+                    raise ValueError(
+                        "device_id does not match the DLPack tensor device_id"
+                    )
+            else:
+                if device_id not in (None, 0):
+                    raise ValueError("device_id must be 0 for CPU tensors")
+                device_id = 0
+            if dev_type0 == kCPU and (dev0 != 0 or dev1 != 0):
+                raise ValueError("CPU DLPack tensors must have device_id == 0")
+
+            if (
+                m0.dl_tensor.dtype.code != 1
+                or m0.dl_tensor.dtype.bits != expected_bits
+                or m0.dl_tensor.dtype.lanes != 1
+            ):
+                raise TypeError("unexpected dtype for plane 0")
+            if (
+                m1.dl_tensor.dtype.code != 1
+                or m1.dl_tensor.dtype.bits != expected_bits
+                or m1.dl_tensor.dtype.lanes != 1
+            ):
+                raise TypeError("unexpected dtype for plane 1")
+
+            if m0.dl_tensor.ndim != 2:
+                raise ValueError("plane 0 must be 2D (H, W)")
+
+            y_h = cython.cast(int64_t, m0.dl_tensor.shape[0])
+            y_w = cython.cast(int64_t, m0.dl_tensor.shape[1])
+
+            if width == 0 and height == 0:
+                width = cython.cast(int, y_w)
+                height = cython.cast(int, y_h)
+            elif width == 0 or height == 0:
+                raise ValueError("either specify both width/height or neither")
+            else:
+                if y_w != width or y_h != height:
+                    raise ValueError("plane 0 shape does not match width/height")
+
+            if width % 2 or height % 2:
+                raise ValueError("width/height must be even for nv12/p010le/p016le")
+
+            if m0.dl_tensor.strides != cython.NULL:
+                if m0.dl_tensor.strides[1] != 1:
+                    raise ValueError("plane 0 must be contiguous in the last dimension")
+                y_pitch_elems = cython.cast(int64_t, m0.dl_tensor.strides[0])
+            else:
+                y_pitch_elems = cython.cast(int64_t, width)
+
+            y_linesize = cython.cast(int, y_pitch_elems * itemsize)
+            y_size = cython.cast(int, y_linesize * height)
+
+            uv_ndim = m1.dl_tensor.ndim
+            uv_h_expected = height // 2
+
+            if uv_ndim == 2:
+                uv_h = cython.cast(int, m1.dl_tensor.shape[0])
+                uv_w = cython.cast(int, m1.dl_tensor.shape[1])
+                if uv_h != uv_h_expected or uv_w != width:
+                    raise ValueError("plane 1 must have shape (H/2, W) for 2D UV")
+                if m1.dl_tensor.strides != cython.NULL:
+                    if m1.dl_tensor.strides[1] != 1:
+                        raise ValueError(
+                            "plane 1 must be contiguous in the last dimension"
+                        )
+                    uv_pitch_elems = cython.cast(int64_t, m1.dl_tensor.strides[0])
+                else:
+                    uv_pitch_elems = cython.cast(int64_t, uv_w)
+            elif uv_ndim == 3:
+                uv_h = cython.cast(int, m1.dl_tensor.shape[0])
+                uv_w2 = cython.cast(int, m1.dl_tensor.shape[1])
+                uv_c = cython.cast(int, m1.dl_tensor.shape[2])
+                if uv_h != uv_h_expected or uv_w2 != (width // 2) or uv_c != 2:
+                    raise ValueError("plane 1 must have shape (H/2, W/2, 2) for 3D UV")
+                if m1.dl_tensor.strides != cython.NULL:
+                    if m1.dl_tensor.strides[2] != 1 or m1.dl_tensor.strides[1] != 2:
+                        raise ValueError(
+                            "unexpected UV plane strides for (H/2, W/2, 2)"
+                        )
+                    uv_pitch_elems = cython.cast(int64_t, m1.dl_tensor.strides[0])
+                else:
+                    uv_pitch_elems = cython.cast(int64_t, width)
+            else:
+                raise ValueError("plane 1 must be 2D or 3D")
+
+            uv_linesize = cython.cast(int, uv_pitch_elems * itemsize)
+            uv_size = cython.cast(int, uv_linesize * (height // 2))
+
+            frame = alloc_video_frame()
+            frame.ptr.width = width
+            frame.ptr.height = height
+            if dev_type0 == kCuda:
+                ctx: CudaContext
+                frames_ref: cython.pointer[lib.AVBufferRef]
+                if cuda_context is None:
+                    ctx = CudaContext(device_id=device_id, primary_ctx=primary_ctx)
+                else:
+                    if not isinstance(cuda_context, CudaContext):
+                        raise TypeError("cuda_context must be a CudaContext")
+                    if int(cuda_context.device_id) != int(device_id):
+                        raise ValueError(
+                            "cuda_context.device_id does not match the DLPack tensor device_id"
+                        )
+                    if bool(cuda_context.primary_ctx) != bool(primary_ctx):
+                        raise ValueError(
+                            "cuda_context.primary_ctx does not match primary_ctx"
+                        )
+                    ctx = cython.cast(CudaContext, cuda_context)
+
+                frames_ref = ctx.get_frames_ctx(sw_fmt, width, height)
+                frame.ptr.format = get_pix_fmt(b"cuda")
+                frame.ptr.hw_frames_ctx = frames_ref
+                frame._device_id = device_id
+                frame._cuda_ctx = ctx
+            else:
+                frame.ptr.format = sw_fmt
+
+            y_ptr = cython.cast(
+                cython.pointer[uint8_t], m0.dl_tensor.data
+            ) + cython.cast(cython.size_t, m0.dl_tensor.byte_offset)
+            uv_ptr = cython.cast(
+                cython.pointer[uint8_t], m1.dl_tensor.data
+            ) + cython.cast(cython.size_t, m1.dl_tensor.byte_offset)
+
+            frame.ptr.buf[0] = lib.av_buffer_create(
+                y_ptr, y_size, _dlpack_avbuffer_free, cython.cast(cython.p_void, m0), 0
+            )
+            if frame.ptr.buf[0] == cython.NULL:
+                raise MemoryError("av_buffer_create failed for plane 0")
+            frame.ptr.data[0] = y_ptr
+            frame.ptr.linesize[0] = y_linesize
+            m0 = cython.NULL
+
+            frame.ptr.buf[1] = lib.av_buffer_create(
+                uv_ptr,
+                uv_size,
+                _dlpack_avbuffer_free,
+                cython.cast(cython.p_void, m1),
+                0,
+            )
+            if frame.ptr.buf[1] == cython.NULL:
+                raise MemoryError("av_buffer_create failed for plane 1")
+            frame.ptr.data[1] = uv_ptr
+            frame.ptr.linesize[1] = uv_linesize
+            m1 = cython.NULL
+
+            frame._init_user_attributes()
+            return frame
+
+        except Exception:
+            if frame is not None:
+                lib.av_frame_unref(frame.ptr)
+            if m0 != cython.NULL:
+                m0.deleter(m0)
+            if m1 != cython.NULL:
+                m1.deleter(m1)
+            raise

venv/lib/python3.12/site-packages/av/video/frame.pyi

@@ -1,6 +1,6 @@
 from enum import IntEnum
 from pathlib import Path
-from typing import Any, ClassVar, Union
+from typing import Any, Union
 
 import numpy as np
 
@@ -8,6 +8,7 @@ from av.frame import Frame
 
 from .format import VideoFormat
 from .plane import VideoPlane
+from .reformatter import ColorPrimaries, ColorTrc
 
 _SupportedNDarray = Union[
     np.ndarray[Any, np.dtype[np.uint8]],
@@ -28,12 +29,21 @@ class PictureType(IntEnum):
     SP = 6
     BI = 7
 
+class CudaContext:
+    @property
+    def device_id(self) -> int: ...
+    @property
+    def primary_ctx(self) -> bool: ...
+    def __init__(self, device_id: int = 0, primary_ctx: bool = True) -> None: ...
+
 class VideoFrame(Frame):
     format: VideoFormat
     planes: tuple[VideoPlane, ...]
     pict_type: int
     colorspace: int
     color_range: int
+    color_trc: int
+    color_primaries: int
 
     @property
     def time(self) -> float: ...
@@ -58,6 +68,9 @@ class VideoFrame(Frame):
         interpolation: int | str | None = None,
         src_color_range: int | str | None = None,
         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: ...
@@ -84,3 +97,14 @@ class VideoFrame(Frame):
         flip_horizontal: bool = False,
         flip_vertical: bool = False,
     ) -> VideoFrame: ...
+    @staticmethod
+    def from_dlpack(
+        planes: object | tuple[object, ...],
+        format: str = "nv12",
+        width: int = 0,
+        height: int = 0,
+        stream: int | None = None,
+        device_id: int | None = None,
+        primary_ctx: bool = True,
+        cuda_context: CudaContext | None = None,
+    ) -> VideoFrame: ...

venv/lib/python3.12/site-packages/av/video/plane.pxd

@@ -1,8 +1,38 @@
+from libc.stdint cimport int64_t, uint8_t, uint16_t, uint64_t
+
 from av.plane cimport Plane
 from av.video.format cimport VideoFormatComponent
 
 
 cdef class VideoPlane(Plane):
-
     cdef readonly size_t buffer_size
     cdef readonly unsigned int width, height
+
+
+cdef enum DeviceType:
+    kCPU = 1
+    kCuda = 2
+
+cdef struct DLDataType:
+    uint8_t code
+    uint8_t bits
+    uint16_t lanes
+
+cdef struct DLTensor:
+    void* data
+    int device_type
+    int device_id
+    int ndim
+    DLDataType dtype
+    int64_t* shape
+    int64_t* strides
+    uint64_t byte_offset
+
+cdef struct DLManagedTensor
+
+ctypedef void (*DLManagedTensorDeleter)(DLManagedTensor*) noexcept nogil
+
+cdef struct DLManagedTensor:
+    DLTensor dl_tensor
+    void* manager_ctx
+    DLManagedTensorDeleter deleter

venv/lib/python3.12/site-packages/av/video/plane.py

@@ -0,0 +1,305 @@
+import cython
+import cython.cimports.libav as lib
+from cython.cimports.av.error import err_check
+from cython.cimports.av.video.format import (
+    VideoFormatComponent,
+    get_pix_fmt,
+    get_video_format,
+)
+from cython.cimports.av.video.frame import VideoFrame
+from cython.cimports.cpython import PyBUF_WRITABLE, PyBuffer_FillInfo
+from cython.cimports.cpython.buffer import Py_buffer
+from cython.cimports.cpython.pycapsule import (
+    PyCapsule_GetPointer,
+    PyCapsule_IsValid,
+    PyCapsule_New,
+)
+from cython.cimports.libc.stdlib import free, malloc
+
+
+@cython.final
+@cython.cclass
+class VideoPlane(Plane):
+    def __cinit__(self, frame: VideoFrame, index: cython.int):
+        # The palette plane has no associated component or linesize; set fields manually
+        fmt = frame.format
+        if frame.ptr.hw_frames_ctx:
+            frames_ctx: cython.pointer[lib.AVHWFramesContext] = cython.cast(
+                cython.pointer[lib.AVHWFramesContext], frame.ptr.hw_frames_ctx.data
+            )
+            fmt = get_video_format(
+                frames_ctx.sw_format, frame.ptr.width, frame.ptr.height
+            )
+
+        if index == 1 and fmt.name == "pal8":
+            self.width = 256
+            self.height = 1
+            self.buffer_size = 256 * 4
+            return
+
+        for i in range(fmt.ptr.nb_components):
+            if fmt.ptr.comp[i].plane == index:
+                component = VideoFormatComponent(fmt, i)
+                self.width = component.width
+                self.height = component.height
+                break
+        else:
+            raise RuntimeError(f"could not find plane {index} of {fmt!r}")
+
+        # Sometimes, linesize is negative (and that is meaningful). We are only
+        # insisting that the buffer size be based on the extent of linesize, and
+        # ignore it's direction.
+        self.buffer_size = abs(self.frame.ptr.linesize[self.index]) * self.height
+
+    @cython.cfunc
+    def _buffer_size(self) -> cython.size_t:
+        return self.buffer_size
+
+    @property
+    def line_size(self):
+        """
+        Bytes per horizontal line in this plane.
+
+        :type: int
+        """
+        return self.frame.ptr.linesize[self.index]
+
+    @cython.cfunc
+    def _buffer_writable(self) -> cython.bint:
+        if self.frame.ptr.hw_frames_ctx:
+            return False
+        return True
+
+    def __getbuffer__(self, view: cython.pointer[Py_buffer], flags: cython.int):
+        if self.frame.ptr.hw_frames_ctx:
+            raise TypeError(
+                "Hardware frame planes do not support the Python buffer protocol. "
+                "Use DLPack (__dlpack__) or download to a software frame."
+            )
+        if flags & PyBUF_WRITABLE and not self._buffer_writable():
+            raise ValueError("buffer is not writable")
+
+        ptr: cython.p_void = self._buffer_ptr()
+        line_size: cython.int = self.frame.ptr.linesize[self.index]
+        if line_size < 0:
+            height: cython.int = self.height
+            ptr = cython.cast(cython.p_char, ptr) + (height - 1) * line_size
+
+        PyBuffer_FillInfo(view, self, ptr, self._buffer_size(), 0, flags)
+
+    def __dlpack_device__(self):
+        if self.frame.ptr.hw_frames_ctx:
+            if cython.cast(lib.AVPixelFormat, self.frame.ptr.format) != get_pix_fmt(
+                b"cuda"
+            ):
+                raise NotImplementedError(
+                    "DLPack export is only implemented for CUDA hw frames"
+                )
+            return (kCuda, self.frame._device_id)
+        return (kCPU, 0)
+
+    def __dlpack__(self, *, stream: int | None = None):
+        if self.frame.ptr.buf[0] == cython.NULL:
+            raise TypeError(
+                "DLPack export requires a refcounted AVFrame (frame.buf[0] is NULL)"
+            )
+
+        sw_fmt: lib.AVPixelFormat
+        device_type: cython.int
+        device_id: cython.int
+
+        if self.frame.ptr.hw_frames_ctx:
+            if cython.cast(lib.AVPixelFormat, self.frame.ptr.format) != get_pix_fmt(
+                b"cuda"
+            ):
+                raise NotImplementedError(
+                    "DLPack export is only implemented for CUDA hw frames"
+                )
+            if stream is not None:
+                raise NotImplementedError(
+                    "CUDA stream synchronization is not supported. "
+                    "Pass stream=None and synchronize before calling __dlpack__."
+                )
+
+            frames_ctx: cython.pointer[lib.AVHWFramesContext] = cython.cast(
+                cython.pointer[lib.AVHWFramesContext], self.frame.ptr.hw_frames_ctx.data
+            )
+            sw_fmt = frames_ctx.sw_format
+            device_type = kCuda
+            device_id = self.frame._device_id
+        else:
+            sw_fmt = cython.cast(lib.AVPixelFormat, self.frame.ptr.format)
+            device_type = kCPU
+            device_id = 0
+
+        line_size = self.line_size
+        if line_size < 0:
+            raise NotImplementedError(
+                "negative linesize is not supported for DLPack export"
+            )
+
+        nv12 = get_pix_fmt(b"nv12")
+        p010le = get_pix_fmt(b"p010le")
+        p016le = get_pix_fmt(b"p016le")
+
+        ndim, bits, itemsize = cython.declare(cython.int)
+        s0, s1, s2 = cython.declare(int64_t)
+        st0, st1, st2 = cython.declare(int64_t)
+
+        if sw_fmt == nv12:
+            itemsize = 1
+            bits = 8
+            if self.index == 0:
+                ndim = 2
+                s0 = self.frame.ptr.height
+                s1 = self.frame.ptr.width
+                st0 = line_size
+                st1 = 1
+            elif self.index == 1:
+                ndim = 3
+                s0 = self.frame.ptr.height // 2
+                s1 = self.frame.ptr.width // 2
+                s2 = 2
+                st0 = line_size
+                st1 = 2
+                st2 = 1
+            else:
+                raise ValueError("invalid plane index for NV12")
+        elif sw_fmt == p010le or sw_fmt == p016le:
+            itemsize = 2
+            bits = 16
+            if line_size % itemsize:
+                raise ValueError("linesize is not aligned to dtype")
+            if self.index == 0:
+                ndim = 2
+                s0 = self.frame.ptr.height
+                s1 = self.frame.ptr.width
+                st0 = line_size // itemsize
+                st1 = 1
+            elif self.index == 1:
+                ndim = 3
+                s0 = self.frame.ptr.height // 2
+                s1 = self.frame.ptr.width // 2
+                s2 = 2
+                st0 = line_size // itemsize
+                st1 = 2
+                st2 = 1
+            else:
+                raise ValueError("invalid plane index for P010/P016")
+        else:
+            raise NotImplementedError("unsupported sw_format for DLPack export")
+
+        frame_ref: cython.pointer[lib.AVFrame] = lib.av_frame_alloc()
+        if frame_ref == cython.NULL:
+            raise MemoryError("av_frame_alloc() failed")
+        err_check(lib.av_frame_ref(frame_ref, self.frame.ptr))
+
+        shape = cython.cast(
+            cython.pointer[int64_t], malloc(ndim * cython.sizeof(int64_t))
+        )
+        strides = cython.cast(
+            cython.pointer[int64_t], malloc(ndim * cython.sizeof(int64_t))
+        )
+        if shape == cython.NULL or strides == cython.NULL:
+            if shape != cython.NULL:
+                free(shape)
+            if strides != cython.NULL:
+                free(strides)
+            lib.av_frame_free(cython.address(frame_ref))
+            raise MemoryError("malloc() failed")
+
+        if ndim == 2:
+            shape[0] = s0
+            shape[1] = s1
+            strides[0] = st0
+            strides[1] = st1
+        else:
+            shape[0] = s0
+            shape[1] = s1
+            shape[2] = s2
+            strides[0] = st0
+            strides[1] = st1
+            strides[2] = st2
+
+        ctx = cython.cast(
+            cython.pointer[cython.p_void], malloc(3 * cython.sizeof(cython.p_void))
+        )
+        if ctx == cython.NULL:
+            free(shape)
+            free(strides)
+            lib.av_frame_free(cython.address(frame_ref))
+            raise MemoryError("malloc() failed")
+
+        ctx[0] = cython.cast(cython.p_void, frame_ref)
+        ctx[1] = cython.cast(cython.p_void, shape)
+        ctx[2] = cython.cast(cython.p_void, strides)
+
+        managed = cython.cast(
+            cython.pointer[DLManagedTensor], malloc(cython.sizeof(DLManagedTensor))
+        )
+        if managed == cython.NULL:
+            free(ctx)
+            free(shape)
+            free(strides)
+            lib.av_frame_free(cython.address(frame_ref))
+            raise MemoryError("malloc() failed")
+
+        managed.dl_tensor.data = cython.cast(cython.p_void, frame_ref.data[self.index])
+        managed.dl_tensor.device_type = device_type
+        managed.dl_tensor.device_id = device_id
+        managed.dl_tensor.ndim = ndim
+        managed.dl_tensor.dtype = DLDataType(code=1, bits=bits, lanes=1)
+        managed.dl_tensor.shape = shape
+        managed.dl_tensor.strides = strides
+        managed.dl_tensor.byte_offset = 0
+        managed.manager_ctx = cython.cast(cython.p_void, ctx)
+        managed.deleter = _dlpack_managed_tensor_deleter
+
+        try:
+            capsule = PyCapsule_New(
+                cython.cast(cython.p_void, managed),
+                b"dltensor",
+                _dlpack_capsule_destructor,
+            )
+        except Exception:
+            _dlpack_managed_tensor_deleter(managed)
+            raise
+
+        return capsule
+
+
+@cython.cfunc
+@cython.nogil
+@cython.exceptval(check=False)
+def _dlpack_managed_tensor_deleter(
+    managed: cython.pointer[DLManagedTensor],
+) -> cython.void:
+    if managed == cython.NULL:
+        return
+    ctx = cython.cast(cython.pointer[cython.p_void], managed.manager_ctx)
+    if ctx != cython.NULL:
+        frame_ref = cython.cast(cython.pointer[lib.AVFrame], ctx[0])
+        shape = cython.cast(cython.pointer[int64_t], ctx[1])
+        strides = cython.cast(cython.pointer[int64_t], ctx[2])
+
+        if frame_ref != cython.NULL:
+            lib.av_frame_free(cython.address(frame_ref))
+        if shape != cython.NULL:
+            free(shape)
+        if strides != cython.NULL:
+            free(strides)
+        free(ctx)
+
+    free(managed)
+
+
+@cython.cfunc
+@cython.exceptval(check=False)
+def _dlpack_capsule_destructor(capsule: object) -> cython.void:
+    if PyCapsule_IsValid(capsule, b"dltensor"):
+        managed = cython.cast(
+            cython.pointer[DLManagedTensor],
+            PyCapsule_GetPointer(capsule, b"dltensor"),
+        )
+        if managed != cython.NULL:
+            managed.deleter(managed)

venv/lib/python3.12/site-packages/av/video/plane.pyi

@@ -1,3 +1,5 @@
+from types import CapsuleType
+
 from av.plane import Plane
 
 from .frame import VideoFrame
@@ -9,3 +11,5 @@ class VideoPlane(Plane):
     buffer_size: int
 
     def __init__(self, frame: VideoFrame, index: int) -> None: ...
+    def __dlpack_device__(self) -> tuple[int, int]: ...
+    def __dlpack__(self, *, stream: int | None = None) -> CapsuleType: ...

venv/lib/python3.12/site-packages/av/video/plane.pyx

@@ -1,37 +0,0 @@
-from av.video.frame cimport VideoFrame
-
-
-cdef class VideoPlane(Plane):
-    def __cinit__(self, VideoFrame frame, int index):
-        # The palette plane has no associated component or linesize; set fields manually
-        if frame.format.name == "pal8" and index == 1:
-            self.width = 256
-            self.height = 1
-            self.buffer_size = 256 * 4
-            return
-
-        for i in range(frame.format.ptr.nb_components):
-            if frame.format.ptr.comp[i].plane == index:
-                component = frame.format.components[i]
-                self.width = component.width
-                self.height = component.height
-                break
-        else:
-            raise RuntimeError(f"could not find plane {index} of {frame.format!r}")
-
-        # Sometimes, linesize is negative (and that is meaningful). We are only
-        # insisting that the buffer size be based on the extent of linesize, and
-        # ignore it's direction.
-        self.buffer_size = abs(self.frame.ptr.linesize[self.index]) * self.height
-
-    cdef size_t _buffer_size(self):
-        return self.buffer_size
-
-    @property
-    def line_size(self):
-        """
-        Bytes per horizontal line in this plane.
-
-        :type: int
-        """
-        return self.frame.ptr.linesize[self.index]

venv/lib/python3.12/site-packages/av/video/reformatter.pxd

@@ -3,11 +3,40 @@ cimport libav as lib
 from av.video.frame cimport VideoFrame
 
 
+cdef extern from "libswscale/swscale.h" nogil:
+    cdef struct SwsContext:
+        unsigned flags
+        int threads
+
+    cdef int SWS_FAST_BILINEAR
+    cdef int SWS_BILINEAR
+    cdef int SWS_BICUBIC
+    cdef int SWS_X
+    cdef int SWS_POINT
+    cdef int SWS_AREA
+    cdef int SWS_BICUBLIN
+    cdef int SWS_GAUSS
+    cdef int SWS_SINC
+    cdef int SWS_LANCZOS
+    cdef int SWS_SPLINE
+    cdef int SWS_CS_ITU709
+    cdef int SWS_CS_FCC
+    cdef int SWS_CS_ITU601
+    cdef int SWS_CS_ITU624
+    cdef int SWS_CS_SMPTE170M
+    cdef int SWS_CS_SMPTE240M
+    cdef int SWS_CS_DEFAULT
+
+    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 lib.SwsContext *ptr
-
+    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)
+                   int src_color_range, int dst_color_range,
+                   int dst_color_trc, int dst_color_primaries,
+                   bint set_dst_color_trc, bint set_dst_color_primaries,
+                   int threads)

venv/lib/python3.12/site-packages/av/video/reformatter.py

@@ -0,0 +1,322 @@
+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

venv/lib/python3.12/site-packages/av/video/reformatter.pyi

@@ -38,6 +38,39 @@ class ColorRange(IntEnum):
     JPEG = 2
     NB = 3
 
+class ColorTrc(IntEnum):
+    BT709 = cast(int, ...)
+    UNSPECIFIED = cast(int, ...)
+    GAMMA22 = cast(int, ...)
+    GAMMA28 = cast(int, ...)
+    SMPTE170M = cast(int, ...)
+    SMPTE240M = cast(int, ...)
+    LINEAR = cast(int, ...)
+    LOG = cast(int, ...)
+    LOG_SQRT = cast(int, ...)
+    IEC61966_2_4 = cast(int, ...)
+    BT1361_ECG = cast(int, ...)
+    IEC61966_2_1 = cast(int, ...)
+    BT2020_10 = cast(int, ...)
+    BT2020_12 = cast(int, ...)
+    SMPTE2084 = cast(int, ...)
+    SMPTE428 = cast(int, ...)
+    ARIB_STD_B67 = cast(int, ...)
+
+class ColorPrimaries(IntEnum):
+    BT709 = cast(int, ...)
+    UNSPECIFIED = cast(int, ...)
+    BT470M = cast(int, ...)
+    BT470BG = cast(int, ...)
+    SMPTE170M = cast(int, ...)
+    SMPTE240M = cast(int, ...)
+    FILM = cast(int, ...)
+    BT2020 = cast(int, ...)
+    SMPTE428 = cast(int, ...)
+    SMPTE431 = cast(int, ...)
+    SMPTE432 = cast(int, ...)
+    EBU3213 = cast(int, ...)
+
 class VideoReformatter:
     def reformat(
         self,
@@ -50,4 +83,7 @@ class VideoReformatter:
         interpolation: int | str | None = None,
         src_color_range: int | str | None = None,
         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: ...

venv/lib/python3.12/site-packages/av/video/reformatter.pyx

@@ -1,222 +0,0 @@
-cimport libav as lib
-from libc.stdint cimport uint8_t
-
-from av.error cimport err_check
-from av.video.format cimport VideoFormat
-from av.video.frame cimport alloc_video_frame
-
-from enum import IntEnum
-
-
-class Interpolation(IntEnum):
-    FAST_BILINEAR: "Fast bilinear" = lib.SWS_FAST_BILINEAR
-    BILINEAR: "Bilinear" = lib.SWS_BILINEAR
-    BICUBIC: "Bicubic" = lib.SWS_BICUBIC
-    X: "Experimental" = lib.SWS_X
-    POINT: "Nearest neighbor / point" = lib.SWS_POINT
-    AREA: "Area averaging" = lib.SWS_AREA
-    BICUBLIN: "Luma bicubic / chroma bilinear" = lib.SWS_BICUBLIN
-    GAUSS: "Gaussian" = lib.SWS_GAUSS
-    SINC: "Sinc" = lib.SWS_SINC
-    LANCZOS: "Bicubic spline" = lib.SWS_LANCZOS
-
-
-class Colorspace(IntEnum):
-    ITU709 = lib.SWS_CS_ITU709
-    FCC = lib.SWS_CS_FCC
-    ITU601 = lib.SWS_CS_ITU601
-    ITU624 = lib.SWS_CS_ITU624
-    SMPTE170M = lib.SWS_CS_SMPTE170M
-    SMPTE240M = lib.SWS_CS_SMPTE240M
-    DEFAULT = lib.SWS_CS_DEFAULT
-    # Lowercase for b/c.
-    itu709 = lib.SWS_CS_ITU709
-    fcc = lib.SWS_CS_FCC
-    itu601 = lib.SWS_CS_ITU601
-    itu624 = lib.SWS_CS_ITU624
-    smpte170m = lib.SWS_CS_SMPTE170M
-    smpte240m = lib.SWS_CS_SMPTE240M
-    default = lib.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
-
-
-def _resolve_enum_value(value, enum_class, default):
-    # 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__}")
-
-
-cdef 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 nogil:
-            lib.sws_freeContext(self.ptr)
-
-    def reformat(self, VideoFrame frame not None, width=None, height=None,
-                 format=None, src_colorspace=None, dst_colorspace=None,
-                 interpolation=None, src_color_range=None,
-                 dst_color_range=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:`color range` or ``str``
-        :param dst_color_range: Desired color range, or ``None`` for the ``UNSPECIFIED``.
-        :type  dst_color_range: :class:`color range` or ``str``
-
-        """
-
-        cdef VideoFormat video_format = VideoFormat(format if format is not None else frame.format)
-
-        cdef int c_src_colorspace = _resolve_enum_value(src_colorspace, Colorspace, frame.colorspace)
-        cdef int c_dst_colorspace = _resolve_enum_value(dst_colorspace, Colorspace, frame.colorspace)
-        cdef int c_interpolation = _resolve_enum_value(interpolation, Interpolation, int(Interpolation.BILINEAR))
-        cdef int c_src_color_range = _resolve_enum_value(src_color_range, ColorRange, 0)
-        cdef int c_dst_color_range = _resolve_enum_value(dst_color_range, ColorRange, 0)
-
-        return self._reformat(
-            frame,
-            width or frame.ptr.width,
-            height or frame.ptr.height,
-            video_format.pix_fmt,
-            c_src_colorspace,
-            c_dst_colorspace,
-            c_interpolation,
-            c_src_color_range,
-            c_dst_color_range,
-        )
-
-    cdef _reformat(self, VideoFrame frame, int width, int height,
-                   lib.AVPixelFormat dst_format, int src_colorspace,
-                   int dst_colorspace, int interpolation,
-                   int src_color_range, int dst_color_range):
-
-        if frame.ptr.format < 0:
-            raise ValueError("Frame does not have format set.")
-
-        # 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
-
-        cdef lib.AVPixelFormat src_format = <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
-        ):
-            return frame
-
-        with nogil:
-            self.ptr = lib.sws_getCachedContext(
-                self.ptr,
-                frame.ptr.width,
-                frame.ptr.height,
-                src_format,
-                width,
-                height,
-                dst_format,
-                interpolation,
-                NULL,
-                NULL,
-                NULL
-            )
-
-        # We want to change the colorspace/color_range transforms.
-        # We do that by grabbing all of the current settings, changing a
-        # couple, and setting them all. We need a lot of state here.
-        cdef const int *inv_tbl
-        cdef const int *tbl
-        cdef int src_colorspace_range, dst_colorspace_range
-        cdef int brightness, contrast, saturation
-        cdef int ret
-
-        if src_colorspace != dst_colorspace or src_color_range != dst_color_range:
-            with nogil:
-                # Casts for const-ness, because Cython isn't expressive enough.
-                ret = lib.sws_getColorspaceDetails(
-                    self.ptr,
-                    <int**>&inv_tbl,
-                    &src_colorspace_range,
-                    <int**>&tbl,
-                    &dst_colorspace_range,
-                    &brightness,
-                    &contrast,
-                    &saturation
-                )
-
-            err_check(ret)
-
-            with nogil:
-                # Grab the coefficients for the requested transforms.
-                # The inv_table brings us to linear, and `tbl` to the new space.
-                if src_colorspace != lib.SWS_CS_DEFAULT:
-                    inv_tbl = lib.sws_getCoefficients(src_colorspace)
-                if dst_colorspace != lib.SWS_CS_DEFAULT:
-                    tbl = lib.sws_getCoefficients(dst_colorspace)
-
-                # Apply!
-                ret = lib.sws_setColorspaceDetails(
-                    self.ptr,
-                    inv_tbl,
-                    src_color_range,
-                    tbl,
-                    dst_color_range,
-                    brightness,
-                    contrast,
-                    saturation
-                )
-
-            err_check(ret)
-
-        # Create a new VideoFrame.
-        cdef VideoFrame new_frame = alloc_video_frame()
-        new_frame._copy_internal_attributes(frame)
-        new_frame._init(dst_format, width, height)
-
-        # Finally, scale the image.
-        with nogil:
-            lib.sws_scale(
-                self.ptr,
-                # Cast for const-ness, because Cython isn't expressive enough.
-                <const uint8_t**>frame.ptr.data,
-                frame.ptr.linesize,
-                0,  # slice Y
-                frame.ptr.height,
-                new_frame.ptr.data,
-                new_frame.ptr.linesize,
-            )
-
-        return new_frame

venv/lib/python3.12/site-packages/av/video/stream.pxd

@@ -1,3 +1,5 @@
+from libc.stdint cimport int32_t, uint8_t
+
 from av.packet cimport Packet
 from av.stream cimport Stream
 
@@ -5,5 +7,12 @@ from .frame cimport VideoFrame
 
 
 cdef class VideoStream(Stream):
+    # Display matrix (9 int32, native-endian) written as AV_PKT_DATA_DISPLAYMATRIX
+    # coded side data at mux time, applied only when _has_display_matrix is set.
+    cdef int32_t _display_matrix[9]
+    cdef uint8_t _has_display_matrix
+
+    cdef _apply_display_matrix(self)
+
     cpdef encode(self, VideoFrame frame=?)
     cpdef decode(self, Packet packet=?)

venv/lib/python3.12/site-packages/av/video/stream.py

@@ -1,13 +1,19 @@
 import cython
 from cython.cimports import libav as lib
 from cython.cimports.av.packet import Packet
-from cython.cimports.av.utils import avrational_to_fraction, to_avrational
+from cython.cimports.av.stream import Stream
+from cython.cimports.av.utils import avrational_to_fraction
 from cython.cimports.av.video.frame import VideoFrame
+from cython.cimports.libc.stdint import int32_t
+from cython.cimports.libc.string import memcpy
 
 
+@cython.final
 @cython.cclass
 class VideoStream(Stream):
     def __repr__(self):
+        if self.codec_context is None:
+            return f"<av.VideoStream #{self.index} video/<nocodec> at 0x{id(self):x}>"
         return (
             f"<av.VideoStream #{self.index} {self.name}, "
             f"{self.format.name if self.format else None} {self.codec_context.width}x"
@@ -19,7 +25,10 @@ class VideoStream(Stream):
             raise AttributeError(
                 f"'{type(self).__name__}' object has no attribute '{name}'"
             )
-
+        if self.codec_context is None:
+            raise AttributeError(
+                f"'{type(self).__name__}' object has no attribute '{name}'"
+            )
         return getattr(self.codec_context, name)
 
     @cython.ccall
@@ -50,6 +59,62 @@ class VideoStream(Stream):
         """
         return self.codec_context.decode(packet)
 
+    @cython.cfunc
+    def _finalize_for_output(self):
+        Stream._finalize_for_output(self)
+        # avcodec_parameters_from_context() overwrites codecpar.coded_side_data,
+        # so inject the display matrix after it, before avformat_write_header().
+        if self.codec_context is not None and self._has_display_matrix:
+            self._apply_display_matrix()
+
+    @cython.cfunc
+    def _apply_display_matrix(self):
+        n: cython.int = 9 * cython.sizeof(int32_t)
+        sd: cython.pointer[lib.AVPacketSideData] = lib.av_packet_side_data_new(
+            cython.address(self.ptr.codecpar.coded_side_data),
+            cython.address(self.ptr.codecpar.nb_coded_side_data),
+            lib.AV_PKT_DATA_DISPLAYMATRIX,
+            n,
+            0,
+        )
+        if sd == cython.NULL:
+            raise MemoryError("could not allocate display matrix side data")
+
+        memcpy(sd.data, self._display_matrix, n)
+
+    def set_display_matrix(self, matrix):
+        """Set the display matrix written to the container as coded side data.
+
+        ``matrix`` is a sequence of 9 integers in FFmpeg's ``AV_PKT_DATA_DISPLAYMATRIX``
+        layout, or ``None`` to clear. Must be called before the first frame is
+        encoded. See :meth:`set_display_rotation` for a higher-level helper.
+        """
+        if matrix is None:
+            self._has_display_matrix = False
+            return
+
+        vals = [int(v) for v in matrix]
+        if len(vals) != 9:
+            raise ValueError("display matrix must have exactly 9 elements")
+        i: cython.int
+        for i in range(9):
+            self._display_matrix[i] = vals[i]
+        self._has_display_matrix = True
+
+    def set_display_rotation(self, degrees, hflip=False, vflip=False):
+        """Set the container display matrix from a rotation and optional flips.
+
+        ``degrees`` is a counter-clockwise rotation (matching the value read back
+        from :attr:`VideoFrame.rotation`); ``hflip`` / ``vflip`` mirror after it.
+        Together these express all eight EXIF orientations. Must be called before
+        the first frame is encoded.
+        """
+        # av_display_rotation_set() takes a clockwise angle; negate so our public
+        # `degrees` is counter-clockwise, matching VideoFrame.rotation on read.
+        lib.av_display_rotation_set(self._display_matrix, -float(degrees))
+        lib.av_display_matrix_flip(self._display_matrix, bool(hflip), bool(vflip))
+        self._has_display_matrix = True
+
     @property
     def average_rate(self):
         """

venv/lib/python3.12/site-packages/av/video/stream.pyi

@@ -1,3 +1,4 @@
+from collections.abc import Sequence
 from fractions import Fraction
 from typing import Iterator, Literal
 
@@ -20,6 +21,10 @@ class VideoStream(Stream):
     def encode(self, frame: VideoFrame | None = None) -> list[Packet]: ...
     def encode_lazy(self, frame: VideoFrame | None = None) -> Iterator[Packet]: ...
     def decode(self, packet: Packet | None = None) -> list[VideoFrame]: ...
+    def set_display_matrix(self, matrix: Sequence[int] | None) -> None: ...
+    def set_display_rotation(
+        self, degrees: float, hflip: bool = ..., vflip: bool = ...
+    ) -> None: ...
 
     # from codec context
     format: VideoFormat