diff --git a/av/sidedata/motionvectors.pxd b/av/sidedata/motionvectors.pxd index 993c5e283..078648325 100644 --- a/av/sidedata/motionvectors.pxd +++ b/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 diff --git a/av/sidedata/motionvectors.py b/av/sidedata/motionvectors.py new file mode 100644 index 000000000..820013e43 --- /dev/null +++ b/av/sidedata/motionvectors.py @@ -0,0 +1,135 @@ +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"" + ) + + 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"" + ) + + @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 diff --git a/av/sidedata/motionvectors.pyi b/av/sidedata/motionvectors.pyi index eb514eb70..a1ab3c0fc 100644 --- a/av/sidedata/motionvectors.pyi +++ b/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]: ... diff --git a/av/sidedata/motionvectors.pyx b/av/sidedata/motionvectors.pyx deleted file mode 100644 index 8a4200232..000000000 --- a/av/sidedata/motionvectors.pyx +++ /dev/null @@ -1,104 +0,0 @@ -from collections.abc import Sequence - - -cdef object _cinit_bypass_sentinel = object() - - -# Cython doesn't let us inherit from the abstract Sequence, so we will subclass -# it later. -cdef class _MotionVectors(SideData): - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - self._vectors = {} - self._len = self.ptr.size // sizeof(lib.AVMotionVector) - - def __repr__(self): - return f"self.ptr.data:0x}" - - def __getitem__(self, int index): - - 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 __len__(self): - return self._len - - def to_ndarray(self): - 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)) - - -class MotionVectors(_MotionVectors, Sequence): - pass - - -cdef class MotionVector: - def __init__(self, sentinel, _MotionVectors parent, int index): - if sentinel is not _cinit_bypass_sentinel: - raise RuntimeError("cannot manually instantiate MotionVector") - self.parent = parent - cdef lib.AVMotionVector *base = parent.ptr.data - self.ptr = base + index - - def __repr__(self): - return f"" - - @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 diff --git a/av/video/frame.py b/av/video/frame.py index 174683727..1320afad3 100644 --- a/av/video/frame.py +++ b/av/video/frame.py @@ -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