27054202885

Committed 06 Jun 2026 05:49AM UTC coverage: 97.114% (-0.7%) from 97.821%

Build # 27054202885

Build Type

Pull #69

github

Committed by

web-flow

Commit Message

Merge c1c4272d0 into 952a740e0

Pull Request Pull Request #69: aw shit it's mypy plugin time

Coverage Stats

376 of 403 new or added lines in 14 files covered. (93.3%)

4 existing lines in 1 file now uncovered.

1918 of 1975 relevant lines covered (97.11%)

6.78 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

96.35

/src/numpydantic/interface/video.py

"""
Interface to support treating videos like arrays using OpenCV
"""

from pathlib import Path
from typing import Any, Literal, TypeAlias

import numpy as np
from pydantic_core.core_schema import SerializationInfo

from numpydantic.interface import JsonDict, Proxy
from numpydantic.interface.interface import Interface

try:
    import cv2
    from cv2 import VideoCapture

    _CaptureUnion: TypeAlias = VideoCapture | None
except ImportError:  # pragma: no cover
    cv2 = None
    VideoCapture = None
    _CaptureUnion: TypeAlias = None

VIDEO_EXTENSIONS = (".mp4", ".avi", ".mov", ".mkv")


class VideoJsonDict(JsonDict):
    """Json-able roundtrip representation of a video file"""

    type: Literal["video"]
    file: str

    def to_array_input(self) -> "VideoProxy":
        """
        Construct a :class:`.VideoProxy`
        """
        return VideoProxy(path=Path(self.file))


class VideoProxy(Proxy):
    """
    Passthrough proxy class to interact with videos as arrays
    """

    def __init__(self, path: Path | None = None, video: _CaptureUnion = None):
        if path is None and video is None:  # pragma: no cover
            raise ValueError(
                "Need to either supply a path or an opened VideoCapture object"
            )

        if path is not None:
            path = Path(path).resolve()
        self.path = path

        self._video = video  # type: Optional[VideoCapture]
        self._n_frames = None  # type: Optional[int]
        self._dtype = None  # type: Optional[np.dtype]
        self._shape = None  # type: Optional[Tuple[int, ...]]
        self._sample_frame = None  # type: Optional[np.ndarray]

    @classmethod
    def proxy_for(cls) -> type["VideoInterface"]:
        """Declare this class as a proxy for the VideoInterface"""
        return VideoInterface

    @property
    def video(self) -> VideoCapture:
        """Opened video capture object"""
        if self._video is None:
            if self.path is None:  # pragma: no cover
                raise RuntimeError(
                    "Instantiated with a VideoCapture object that has been closed, "
                    "and it cant be reopened since source path cant be gotten "
                    "from VideoCapture objects"
                )
            if not self.path.exists():
                raise FileNotFoundError(f"Video file {self.path} does not exist!")

            self._video = VideoCapture(str(self.path))
        return self._video

    def close(self) -> None:
        """Close the opened VideoCapture object"""
        if self._video is not None:
            self._video.release()
            self._video = None

    @property
    def sample_frame(self) -> np.ndarray:
        """A stored frame from the video to use when calculating shape and dtype"""
        if self._sample_frame is None:
            current_frame = int(self.video.get(cv2.CAP_PROP_POS_FRAMES))

            self.video.set(cv2.CAP_PROP_POS_FRAMES, max(0, current_frame - 1))
            status, frame = self.video.read()
            if not status:  # pragma: no cover
                raise RuntimeError("Could not read frame from video")
            self.video.set(cv2.CAP_PROP_POS_FRAMES, current_frame)
            self._sample_frame = frame
        return self._sample_frame

    @property
    def shape(self) -> tuple[int, ...]:
        """
        Shape of video like
        ``(n_frames, height, width, channels)``

        Note that this order flips the order of height and width from typical resolution
        specifications: eg. 1080p video is typically 1920x1080, but here it would be
        1080x1920. This follows opencv's ordering, which matches expectations when
        eg. an image is read and plotted with matplotlib: the first index is the
        position in the 0th dimension - the height, or "y" axis - and the second is the
        width/x.
        """
        if self._shape is None:
            self._shape = (self.n_frames, *self.sample_frame.shape)
        return self._shape

    @property
    def dtype(self) -> np.dtype:
        """Numpy dtype (from ``sample_frame`` )"""
        return self.sample_frame.dtype

    @property
    def n_frames(self) -> int:
        """
        Try to get number of frames using opencv metadata, and manually count if no
        t"""
        if self._n_frames is None:
            n_frames = self.video.get(cv2.CAP_PROP_FRAME_COUNT)
            if n_frames == 0:  # pragma: no cover
                # have to count manually for some containers with bad metadata
                # not testing for now, will wait until we encounter such a
                # video in the wild where this doesn't work.
                current_frame = self.video.get(cv2.CAP_PROP_POS_FRAMES)
                self.video.set(cv2.CAP_PROP_POS_FRAMES, 0)
                n_frames = 0
                while True:
                    status, _ = self.video.read()
                    if not status:
                        break
                    n_frames += 1
                self.video.set(cv2.CAP_PROP_POS_FRAMES, current_frame)
            self._n_frames = int(n_frames)
        return self._n_frames

    def _get_frame(self, frame: int) -> np.ndarray:
        self.video.set(cv2.CAP_PROP_POS_FRAMES, frame)
        status, frame = self.video.read()
        if not status:  # pragma: no cover
            raise ValueError(f"Could not get frame {frame}")
        return frame

    def _complete_slice(self, slice_: slice) -> slice:
        """Get a fully-built slice that can be passed to range"""
        if slice_.step is None:
            slice_ = slice(slice_.start, slice_.stop, 1)
        if slice_.stop is None:
            slice_ = slice(slice_.start, self.n_frames, slice_.step)
        if slice_.start is None:
            slice_ = slice(0, slice_.stop, slice_.step)
        return slice_

    def __array__(self) -> np.ndarray:
        """Whole video as a numpy array"""
        return self[:]

    def __getitem__(self, item: int | slice | tuple) -> np.ndarray:
        if isinstance(item, int):
            # want a single frame
            return self._get_frame(item)
        elif isinstance(item, slice):
            # slice of frames
            item = self._complete_slice(item)
            frames = [
                self._get_frame(i) for i in range(item.start, item.stop, item.step)
            ]
            return np.stack(frames)
        else:
            # slices are passed as tuples
            # first arg needs to be handled specially
            if isinstance(item[0], int):
                # single frame
                frame = self._get_frame(item[0])
                # syntax doesn't work in 3.9 but would be more explicit...
                # return frame[*item[1:]]
                return frame[item[1:]]

            elif isinstance(item[0], slice):
                frames = []
                # make a new slice since range cant take Nones, filling in missing vals
                fslice = self._complete_slice(item[0])

                for i in range(fslice.start, fslice.stop, fslice.step):
                    frames.append(self._get_frame(i))
                frame = np.stack(frames)
                # syntax doesn't work in 3.9 but would be simpler..
                # return frame[:, *item[1:]]
                # construct a new slice instead
                new_slice = (slice(None, None, None), *item[1:])
                return frame[new_slice]
            else:  # pragma: no cover
                raise ValueError(f"indices must be an int or a slice! got {item}")

    def __setitem__(self, key: int | slice, value: int | float | np.ndarray):
        raise NotImplementedError("Setting pixel values on videos is not supported!")

    def __getattr__(self, item: str):
        if item == "__name__":
            return "VideoProxy"
        return getattr(self.video, item)

    def __eq__(self, other: "VideoProxy") -> bool:
        """Check if this is a proxy to the same video file"""
        if not isinstance(other, VideoProxy):
            raise TypeError("Can only compare equality of two VideoProxies")
        return self.path == other.path

    def __len__(self) -> int:
        """Number of frames in the video"""
        return self.shape[0]


class VideoInterface(Interface):
    """
    OpenCV interface to treat videos as arrays.
    """

    name = "video"
    input_types = (str, Path, VideoCapture, VideoProxy)
    return_type = VideoProxy
    json_model = VideoJsonDict

    @classmethod
    def enabled(cls) -> bool:
        """Check if opencv-python is available in the environment"""
        return cv2 is not None

    @classmethod
    def check(cls, array: Any) -> bool:
        """
        Check if array is a string or Path with a supported video extension,
        or an opened VideoCapture object
        """
        if (VideoCapture is not None and isinstance(array, VideoCapture)) or isinstance(
            array, VideoProxy
        ):
            return True

        if isinstance(array, dict):
            array = array.get("file", "")

        if isinstance(array, str):
            try:
                array = Path(array)
            except TypeError:  # pragma: no cover
                # fine, just not a video
                return False

        return isinstance(array, Path) and array.suffix.lower() in VIDEO_EXTENSIONS

    def before_validation(self, array: Any) -> VideoProxy:
        """Get a :class:`.VideoProxy` object for this video"""
        if isinstance(array, VideoCapture):
            proxy = VideoProxy(video=array)
        elif isinstance(array, VideoProxy):
            proxy = array
        else:
            proxy = VideoProxy(path=array)
        return proxy

    @classmethod
    def to_json(
        cls, array: VideoProxy, info: SerializationInfo
    ) -> list | VideoJsonDict:
        """Return a json-representation of a video"""
        if info.round_trip:
            return VideoJsonDict(type=cls.name, file=str(array.path))
        else:
            return np.array(array).tolist()

1	"""
2	Interface to support treating videos like arrays using OpenCV
3	"""
4
5	from pathlib import Path	7✔
6	from typing import Any, Literal, TypeAlias	7✔
7
8	import numpy as np	7✔
9	from pydantic_core.core_schema import SerializationInfo	7✔
10
11	from numpydantic.interface import JsonDict, Proxy	7✔
12	from numpydantic.interface.interface import Interface	7✔
13
14	try:	7✔
15	import cv2	7✔
16	from cv2 import VideoCapture	7✔
17
18	_CaptureUnion: TypeAlias = VideoCapture \| None	7✔
19	except ImportError: # pragma: no cover
20	cv2 = None
21	VideoCapture = None
22	_CaptureUnion: TypeAlias = None
23
24	VIDEO_EXTENSIONS = (".mp4", ".avi", ".mov", ".mkv")	7✔
25
26
27	class VideoJsonDict(JsonDict):	7✔
28	"""Json-able roundtrip representation of a video file"""
29
30	type: Literal["video"]	7✔
31	file: str	7✔
32
33	def to_array_input(self) -> "VideoProxy":	7✔
34	"""
35	Construct a :class:`.VideoProxy`
36	"""
37	return VideoProxy(path=Path(self.file))	×
38
39
40	class VideoProxy(Proxy):	7✔
41	"""
42	Passthrough proxy class to interact with videos as arrays
43	"""
44
45	def __init__(self, path: Path \| None = None, video: _CaptureUnion = None):	7✔
46	if path is None and video is None: # pragma: no cover
47	raise ValueError(
48	"Need to either supply a path or an opened VideoCapture object"
49	)
50
51	if path is not None:	7✔
52	path = Path(path).resolve()	7✔
53	self.path = path	7✔
54
55	self._video = video # type: Optional[VideoCapture]	7✔
56	self._n_frames = None # type: Optional[int]	7✔
57	self._dtype = None # type: Optional[np.dtype]	7✔
58	self._shape = None # type: Optional[Tuple[int, ...]]	7✔
59	self._sample_frame = None # type: Optional[np.ndarray]	7✔
60
61	@classmethod	7✔
62	def proxy_for(cls) -> type["VideoInterface"]:	7✔
63	"""Declare this class as a proxy for the VideoInterface"""
64	return VideoInterface	7✔
65
66	@property	7✔
67	def video(self) -> VideoCapture:	7✔
68	"""Opened video capture object"""
69	if self._video is None:	7✔
70	if self.path is None: # pragma: no cover
71	raise RuntimeError(
72	"Instantiated with a VideoCapture object that has been closed, "
73	"and it cant be reopened since source path cant be gotten "
74	"from VideoCapture objects"
75	)
76	if not self.path.exists():	7✔
77	raise FileNotFoundError(f"Video file {self.path} does not exist!")	7✔
78
79	self._video = VideoCapture(str(self.path))	7✔
80	return self._video	7✔
81
82	def close(self) -> None:	7✔
83	"""Close the opened VideoCapture object"""
84	if self._video is not None:	7✔
85	self._video.release()	7✔
86	self._video = None	7✔
87
88	@property	7✔
89	def sample_frame(self) -> np.ndarray:	7✔
90	"""A stored frame from the video to use when calculating shape and dtype"""
91	if self._sample_frame is None:	7✔
92	current_frame = int(self.video.get(cv2.CAP_PROP_POS_FRAMES))	7✔
93
94	self.video.set(cv2.CAP_PROP_POS_FRAMES, max(0, current_frame - 1))	7✔
95	status, frame = self.video.read()	7✔
96	if not status: # pragma: no cover
97	raise RuntimeError("Could not read frame from video")
98	self.video.set(cv2.CAP_PROP_POS_FRAMES, current_frame)	7✔
99	self._sample_frame = frame	7✔
100	return self._sample_frame	7✔
101
102	@property	7✔
103	def shape(self) -> tuple[int, ...]:	7✔
104	"""
105	Shape of video like
106	``(n_frames, height, width, channels)``
107
108	Note that this order flips the order of height and width from typical resolution
109	specifications: eg. 1080p video is typically 1920x1080, but here it would be
110	1080x1920. This follows opencv's ordering, which matches expectations when
111	eg. an image is read and plotted with matplotlib: the first index is the
112	position in the 0th dimension - the height, or "y" axis - and the second is the
113	width/x.
114	"""
115	if self._shape is None:	7✔
116	self._shape = (self.n_frames, *self.sample_frame.shape)	7✔
117	return self._shape	7✔
118
119	@property	7✔
120	def dtype(self) -> np.dtype:	7✔
121	"""Numpy dtype (from ``sample_frame`` )"""
122	return self.sample_frame.dtype	7✔
123
124	@property	7✔
125	def n_frames(self) -> int:	7✔
126	"""
127	Try to get number of frames using opencv metadata, and manually count if no
128	t"""
129	if self._n_frames is None:	7✔
130	n_frames = self.video.get(cv2.CAP_PROP_FRAME_COUNT)	7✔
131	if n_frames == 0: # pragma: no cover
132	# have to count manually for some containers with bad metadata
133	# not testing for now, will wait until we encounter such a
134	# video in the wild where this doesn't work.
135	current_frame = self.video.get(cv2.CAP_PROP_POS_FRAMES)
136	self.video.set(cv2.CAP_PROP_POS_FRAMES, 0)
137	n_frames = 0
138	while True:
139	status, _ = self.video.read()
140	if not status:
141	break
142	n_frames += 1
143	self.video.set(cv2.CAP_PROP_POS_FRAMES, current_frame)
144	self._n_frames = int(n_frames)	7✔
145	return self._n_frames	7✔
146
147	def _get_frame(self, frame: int) -> np.ndarray:	7✔
148	self.video.set(cv2.CAP_PROP_POS_FRAMES, frame)	7✔
149	status, frame = self.video.read()	7✔
150	if not status: # pragma: no cover
151	raise ValueError(f"Could not get frame {frame}")
152	return frame	7✔
153
154	def _complete_slice(self, slice_: slice) -> slice:	7✔
155	"""Get a fully-built slice that can be passed to range"""
156	if slice_.step is None:	7✔
157	slice_ = slice(slice_.start, slice_.stop, 1)	7✔
158	if slice_.stop is None:	7✔
159	slice_ = slice(slice_.start, self.n_frames, slice_.step)	7✔
160	if slice_.start is None:	7✔
161	slice_ = slice(0, slice_.stop, slice_.step)	7✔
162	return slice_	7✔
163
164	def __array__(self) -> np.ndarray:	7✔
165	"""Whole video as a numpy array"""
UNCOV 166	return self[:]	×
167
168	def __getitem__(self, item: int \| slice \| tuple) -> np.ndarray:	7✔
169	if isinstance(item, int):	7✔
170	# want a single frame
171	return self._get_frame(item)	7✔
172	elif isinstance(item, slice):	7✔
173	# slice of frames
174	item = self._complete_slice(item)	7✔
175	frames = [	7✔
176	self._get_frame(i) for i in range(item.start, item.stop, item.step)
177	]
178	return np.stack(frames)	7✔
179	else:
180	# slices are passed as tuples
181	# first arg needs to be handled specially
182	if isinstance(item[0], int):	7✔
183	# single frame
184	frame = self._get_frame(item[0])	7✔
185	# syntax doesn't work in 3.9 but would be more explicit...
186	# return frame[*item[1:]]
187	return frame[item[1:]]	7✔
188
189	elif isinstance(item[0], slice):	7✔
190	frames = []	7✔
191	# make a new slice since range cant take Nones, filling in missing vals
192	fslice = self._complete_slice(item[0])	7✔
193
194	for i in range(fslice.start, fslice.stop, fslice.step):	7✔
195	frames.append(self._get_frame(i))	7✔
196	frame = np.stack(frames)	7✔
197	# syntax doesn't work in 3.9 but would be simpler..
198	# return frame[:, *item[1:]]
199	# construct a new slice instead
200	new_slice = (slice(None, None, None), *item[1:])	7✔
201	return frame[new_slice]	7✔
202	else: # pragma: no cover
203	raise ValueError(f"indices must be an int or a slice! got {item}")
204
205	def __setitem__(self, key: int \| slice, value: int \| float \| np.ndarray):	7✔
206	raise NotImplementedError("Setting pixel values on videos is not supported!")	7✔
207
208	def __getattr__(self, item: str):	7✔
209	if item == "__name__":	7✔
210	return "VideoProxy"	7✔
211	return getattr(self.video, item)	7✔
212
213	def __eq__(self, other: "VideoProxy") -> bool:	7✔
214	"""Check if this is a proxy to the same video file"""
215	if not isinstance(other, VideoProxy):	7✔
216	raise TypeError("Can only compare equality of two VideoProxies")	7✔
217	return self.path == other.path	7✔
218
219	def __len__(self) -> int:	7✔
220	"""Number of frames in the video"""
221	return self.shape[0]	7✔
222
223
224	class VideoInterface(Interface):	7✔
225	"""
226	OpenCV interface to treat videos as arrays.
227	"""
228
229	name = "video"	7✔
230	input_types = (str, Path, VideoCapture, VideoProxy)	7✔
231	return_type = VideoProxy	7✔
232	json_model = VideoJsonDict	7✔
233
234	@classmethod	7✔
235	def enabled(cls) -> bool:	7✔
236	"""Check if opencv-python is available in the environment"""
237	return cv2 is not None	7✔
238
239	@classmethod	7✔
240	def check(cls, array: Any) -> bool:	7✔
241	"""
242	Check if array is a string or Path with a supported video extension,
243	or an opened VideoCapture object
244	"""
245	if (VideoCapture is not None and isinstance(array, VideoCapture)) or isinstance(	7✔
246	array, VideoProxy
247	):
248	return True	7✔
249
250	if isinstance(array, dict):	7✔
251	array = array.get("file", "")	7✔
252
253	if isinstance(array, str):	7✔
254	try:	7✔
255	array = Path(array)	7✔
256	except TypeError: # pragma: no cover
257	# fine, just not a video
258	return False
259
260	return isinstance(array, Path) and array.suffix.lower() in VIDEO_EXTENSIONS	7✔
261
262	def before_validation(self, array: Any) -> VideoProxy:	7✔
263	"""Get a :class:`.VideoProxy` object for this video"""
264	if isinstance(array, VideoCapture):	7✔
265	proxy = VideoProxy(video=array)	7✔
266	elif isinstance(array, VideoProxy):	7✔
267	proxy = array	7✔
268	else:
269	proxy = VideoProxy(path=array)	7✔
270	return proxy	7✔
271
272	@classmethod	7✔
273	def to_json(	7✔
274	cls, array: VideoProxy, info: SerializationInfo
275	) -> list \| VideoJsonDict:
276	"""Return a json-representation of a video"""
UNCOV 277	if info.round_trip:	×
UNCOV 278	return VideoJsonDict(type=cls.name, file=str(array.path))	×
279	else:
UNCOV 280	return np.array(array).tolist()	×

p2p-ld / numpydantic / 27054202885

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous