137 lines
3.5 KiB
Python
137 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
|