PyAV/av/sidedata/motionvectors.py at 7d48a9c3c80cf116cd9d1d6ee78209ffbac9edbd · PyAV-Org/PyAV · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
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):
    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})>"
        )

    @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