Make sidedata/motionvectors pure

Author: WyattBlue

av/sidedata/motionvectors.pxd

@@ -4,13 +4,11 @@ from av.frame cimport Frame
 from av.sidedata.sidedata cimport SideData
 
 
-cdef class _MotionVectors(SideData):
-
+cdef class MotionVectors(SideData):
     cdef dict _vectors
-    cdef int _len
+    cdef Py_ssize_t _len
 
 
 cdef class MotionVector:
-
-    cdef _MotionVectors parent
+    cdef MotionVectors parent
     cdef lib.AVMotionVector *ptr

av/sidedata/motionvectors.py

@@ -0,0 +1,164 @@
+from collections.abc import Sequence
+import cython
+from cython.cimports import libav as lib
+from cython.cimports.av.sidedata.sidedata import SideData
+
+_cinit_bypass_sentinel = cython.declare(object, object())
+
+
+@cython.cclass
+class MotionVectors(SideData, Sequence):
+    def __init__(self, sentinel, frame: Frame, index: cython.int):
+        SideData.__init__(self, sentinel, frame, index)
+        self._vectors = {}
+        self._len = self.ptr.size // cython.sizeof(lib.AVMotionVector)
+
+    def __repr__(self):
+        return (
+            f"<av.sidedata.MotionVectors {self.ptr.size} bytes "
+            f"of {len(self)} vectors at 0x{cython.cast(cython.uint, self.ptr.data):0x}>"
+        )
+
+    def __len__(self):
+        return self._len
+
+    def __getitem__(self, index: cython.Py_ssize_t):
+        try:
+            return self._vectors[index]
+        except KeyError:
+            pass
+
+        if index >= self._len:
+            raise IndexError(index)
+
+        vector = self._vectors[index] = MotionVector(_cinit_bypass_sentinel, self, index)
+        return vector
+
+    def __iter__(self):
+        """Iterate over all motion vectors."""
+        for i in range(self._len):
+            yield self[i]
+
+    def to_ndarray(self):
+        """
+        Convert motion vectors to a NumPy structured array.
+
+        Returns a NumPy array with fields corresponding to the AVMotionVector structure.
+        """
+        import numpy as np
+
+        return np.frombuffer(
+            self,
+            dtype=np.dtype(
+                [
+                    ("source", "int32"),
+                    ("w", "uint8"),
+                    ("h", "uint8"),
+                    ("src_x", "int16"),
+                    ("src_y", "int16"),
+                    ("dst_x", "int16"),
+                    ("dst_y", "int16"),
+                    ("flags", "uint64"),
+                    ("motion_x", "int32"),
+                    ("motion_y", "int32"),
+                    ("motion_scale", "uint16"),
+                ],
+                align=True,
+            ),
+        )
+
+@cython.cclass
+class MotionVector:
+    """
+    Represents a single motion vector from video frame data.
+
+    Motion vectors describe the motion of a block of pixels between frames.
+    """
+
+    def __init__(self, sentinel, parent: MotionVectors, index: cython.int):
+        if sentinel is not _cinit_bypass_sentinel:
+            raise RuntimeError("cannot manually instantiate MotionVector")
+        self.parent = parent
+        base: cython.pointer[lib.AVMotionVector] = cython.cast(
+            cython.pointer[lib.AVMotionVector], parent.ptr.data
+        )
+        self.ptr = base + index
+
+    def __repr__(self):
+        return (
+            f"<av.sidedata.MotionVector {self.w}x{self.h} "
+            f"from ({self.src_x},{self.src_y}) to ({self.dst_x},{self.dst_y})>"
+        )
+
+    def __eq__(self, other):
+        """Compare two motion vectors for equality."""
+        if not isinstance(other, MotionVector):
+            return NotImplemented
+        return (
+            self.source == other.source
+            and self.w == other.w
+            and self.h == other.h
+            and self.src_x == other.src_x
+            and self.src_y == other.src_y
+            and self.dst_x == other.dst_x
+            and self.dst_y == other.dst_y
+            and self.motion_x == other.motion_x
+            and self.motion_y == other.motion_y
+            and self.motion_scale == other.motion_scale
+        )
+
+    def __hash__(self):
+        return hash(
+            (
+                self.source,
+                self.w,
+                self.h,
+                self.src_x,
+                self.src_y,
+                self.dst_x,
+                self.dst_y,
+                self.motion_x,
+                self.motion_y,
+                self.motion_scale,
+            )
+        )
+
+    @property
+    def source(self):
+        return self.ptr.source
+
+    @property
+    def w(self):
+        return self.ptr.w
+
+    @property
+    def h(self):
+        return self.ptr.h
+
+    @property
+    def src_x(self):
+        return self.ptr.src_x
+
+    @property
+    def src_y(self):
+        return self.ptr.src_y
+
+    @property
+    def dst_x(self):
+        return self.ptr.dst_x
+
+    @property
+    def dst_y(self):
+        return self.ptr.dst_y
+
+    @property
+    def motion_x(self):
+        return self.ptr.motion_x
+
+    @property
+    def motion_y(self):
+        return self.ptr.motion_y
+
+    @property
+    def motion_scale(self):
+        return self.ptr.motion_scale

av/sidedata/motionvectors.pyi

@@ -6,11 +6,9 @@ from .sidedata import SideData
 
 class MotionVectors(SideData, Sequence[MotionVector]):
     @overload
-    def __getitem__(self, index: int): ...
+    def __getitem__(self, index: int) -> MotionVector: ...
     @overload
-    def __getitem__(self, index: slice): ...
-    @overload
-    def __getitem__(self, index: int | slice): ...
+    def __getitem__(self, index: slice) -> list[MotionVector]: ...
     def __len__(self) -> int: ...
     def to_ndarray(self) -> np.ndarray[Any, Any]: ...
 

av/sidedata/motionvectors.pyx

@@ -158,7 +158,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