Beta/venv/lib/python3.12/site-packages/av/sidedata/motionvectors.py
2026-06-16 17:09:34 +00:00

136 lines
3.5 KiB
Python

import cython
from cython.cimports import libav as lib
from cython.cimports.av.sidedata.sidedata import SideData
from cython.cimports.libc.stdint import uintptr_t
_cinit_bypass_sentinel = cython.declare(object, object())
@cython.final
@cython.cclass
class MotionVectors(SideData):
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(uintptr_t, 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.final
@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})>"
)
@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