Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Video handle persistence #64

Merged
merged 6 commits into from
Sep 12, 2023
Merged

Video handle persistence #64

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
a7d728d
Initial MediaVideo handle persistence
talmo Sep 12, 2023
44ba1d4
Fix using other backends
talmo Sep 12, 2023
d56e11d
Refactor and add more tests for backend reader type
talmo Sep 12, 2023
1efa15c
Test fixes
talmo Sep 12, 2023
97fb1f8
Fix frame counting with imageio backends
talmo Sep 12, 2023
27740b1
Merge branch 'main' into talmo/keep_video_handles_open
talmo Sep 12, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
143 changes: 112 additions & 31 deletions sleap_io/io/video.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,8 @@

def _get_valid_kwargs(cls, kwargs: dict) -> dict:
"""Filter a list of kwargs to the ones that are valid for a class."""
return {k: v for k, v in kwargs.items() if k in cls.__attrs_attrs__}
valid_fields = [a.name for a in attrs.fields(cls)]
return {k: v for k, v in kwargs.items() if k in valid_fields}


@attrs.define
Expand All @@ -43,18 +44,25 @@ class VideoBackend:
Attributes:
filename: Path to video file.
grayscale: Whether to force grayscale. If None, autodetect on first frame load.
keep_open: Whether to keep the video reader open between calls to read frames.
If False, will close the reader after each call. If True (the default), it
will keep the reader open and cache it for subsequent calls which may
enhance the performance of reading multiple frames.
"""

filename: str
grayscale: Optional[bool] = None
keep_open: bool = True
_cached_shape: Optional[Tuple[int, int, int, int]] = None
_open_reader: Optional[object] = None

@classmethod
def from_filename(
cls,
filename: str,
dataset: Optional[str] = None,
grayscale: Optional[bool] = None,
keep_open: bool = True,
**kwargs,
) -> VideoBackend:
"""Create a VideoBackend from a filename.
Expand All @@ -64,6 +72,10 @@ def from_filename(
dataset: Name of dataset in HDF5 file.
grayscale: Whether to force grayscale. If None, autodetect on first frame
load.
keep_open: Whether to keep the video reader open between calls to read
frames. If False, will close the reader after each call. If True (the
default), it will keep the reader open and cache it for subsequent calls
which may enhance the performance of reading multiple frames.

Returns:
VideoBackend subclass instance.
Expand All @@ -73,13 +85,17 @@ def from_filename(

if filename.endswith(MediaVideo.EXTS):
return MediaVideo(
filename, grayscale=grayscale, **_get_valid_kwargs(MediaVideo, kwargs)
filename,
grayscale=grayscale,
keep_open=keep_open,
**_get_valid_kwargs(MediaVideo, kwargs),
)
elif filename.endswith(HDF5Video.EXTS):
return HDF5Video(
filename,
dataset=dataset,
grayscale=grayscale,
keep_open=keep_open,
**_get_valid_kwargs(HDF5Video, kwargs),
)
else:
Expand Down Expand Up @@ -274,6 +290,10 @@ class MediaVideo(VideoBackend):
Attributes:
filename: Path to video file.
grayscale: Whether to force grayscale. If None, autodetect on first frame load.
keep_open: Whether to keep the video reader open between calls to read frames.
If False, will close the reader after each call. If True (the default), it
will keep the reader open and cache it for subsequent calls which may
enhance the performance of reading multiple frames.
plugin: Video plugin to use. One of "opencv", "FFMPEG", or "pyav". If `None`,
will use the first available plugin in the order listed above.
"""
Expand All @@ -297,13 +317,33 @@ def _default_plugin(self) -> str:
"No video plugins found. Install opencv-python, imageio-ffmpeg, or av."
)

@property
def reader(self) -> object:
"""Return the reader object for the video, caching if necessary."""
if self.keep_open:
if self._open_reader is None:
if self.plugin == "opencv":
self._open_reader = cv2.VideoCapture(self.filename)
elif self.plugin == "pyav" or self.plugin == "FFMPEG":
self._open_reader = iio.imopen(
self.filename, "r", plugin=self.plugin
)
return self._open_reader
else:
if self.plugin == "opencv":
return cv2.VideoCapture(self.filename)
elif self.plugin == "pyav" or self.plugin == "FFMPEG":
return iio.imopen(self.filename, "r", plugin=self.plugin)

@property
def num_frames(self) -> int:
"""Number of frames in the video."""
if self.plugin == "opencv":
return int(cv2.VideoCapture(self.filename).get(cv2.CAP_PROP_FRAME_COUNT))
return int(self.reader.get(cv2.CAP_PROP_FRAME_COUNT))
else:
return iio.improps(self.filename, plugin="pyav").shape[0]
n_frames = iio.improps(self.filename, plugin=self.plugin).shape[0]
if np.isinf(n_frames):
return self.reader.legacy_get_reader().count_frames()

def _read_frame(self, frame_idx: int) -> np.ndarray:
"""Read a single frame from the video.
Expand All @@ -319,12 +359,15 @@ def _read_frame(self, frame_idx: int) -> np.ndarray:
`get_frame` method of the `VideoBackend` class instead.
"""
if self.plugin == "opencv":
reader = cv2.VideoCapture(self.filename)
reader.set(cv2.CAP_PROP_POS_FRAMES, frame_idx)
_, img = reader.read()
else:
with iio.imopen(self.filename, "r", plugin=self.plugin) as vid:
img = vid.read(index=frame_idx)
if self.reader.get(cv2.CAP_PROP_POS_FRAMES) != frame_idx:
self.reader.set(cv2.CAP_PROP_POS_FRAMES, frame_idx)
_, img = self.reader.read()
elif self.plugin == "pyav" or self.plugin == "FFMPEG":
if self.keep_open:
img = self.reader.read(index=frame_idx)
else:
with iio.imopen(self.filename, "r", plugin=self.plugin) as reader:
img = reader.read(index=frame_idx)
return img

def _read_frames(self, frame_inds: list) -> np.ndarray:
Expand All @@ -341,7 +384,13 @@ def _read_frames(self, frame_inds: list) -> np.ndarray:
`get_frames` method of the `VideoBackend` class instead.
"""
if self.plugin == "opencv":
reader = cv2.VideoCapture(self.filename)
if self.keep_open:
if self._open_reader is None:
self._open_reader = cv2.VideoCapture(self.filename)
reader = self._open_reader
else:
reader = cv2.VideoCapture(self.filename)

reader.set(cv2.CAP_PROP_POS_FRAMES, frame_inds[0])
imgs = []
for idx in frame_inds:
Comment on lines 389 to 401
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Similar to the _read_frame method, the _read_frames method has also been updated to use the cached reader if keep_open is True. Again, make sure that the reader is properly closed when it's no longer needed.

+            if self.keep_open:
+                if self._open_reader is None:
+                    self._open_reader = cv2.VideoCapture(self.filename)
+                reader = self._open_reader
+            else:
+                reader = cv2.VideoCapture(self.filename)

Expand All @@ -352,9 +401,19 @@ def _read_frames(self, frame_inds: list) -> np.ndarray:
imgs.append(img)
imgs = np.stack(imgs, axis=0)

else:
with iio.imopen(self.filename, "r", plugin=self.plugin) as vid:
imgs = np.stack([vid.read(index=idx) for idx in frame_inds], axis=0)
elif self.plugin == "pyav" or self.plugin == "FFMPEG":
if self.keep_open:
if self._open_reader is None:
self._open_reader = iio.imopen(
self.filename, "r", plugin=self.plugin
)
reader = self._open_reader
imgs = np.stack([reader.read(index=idx) for idx in frame_inds], axis=0)
else:
with iio.imopen(self.filename, "r", plugin=self.plugin) as reader:
imgs = np.stack(
[reader.read(index=idx) for idx in frame_inds], axis=0
)
return imgs


Expand All @@ -381,6 +440,10 @@ class HDF5Video(VideoBackend):
Attributes:
filename: Path to HDF5 file (.h5, .hdf5 or .slp).
grayscale: Whether to force grayscale. If None, autodetect on first frame load.
keep_open: Whether to keep the video reader open between calls to read frames.
If False, will close the reader after each call. If True (the default), it
will keep the reader open and cache it for subsequent calls which may
enhance the performance of reading multiple frames.
dataset: Name of dataset to read from. If `None`, will try to find a rank-4
dataset by iterating through datasets in the file. If specifying an embedded
dataset, this can be the group containing a "video" dataset or the dataset
Expand Down Expand Up @@ -537,19 +600,28 @@ def _read_frame(self, frame_idx: int) -> np.ndarray:
This does not apply grayscale conversion. It is recommended to use the
`get_frame` method of the `VideoBackend` class instead.
"""
with h5py.File(self.filename, "r") as f:
ds = f[self.dataset]
if self.keep_open:
if self._open_reader is None:
self._open_reader = h5py.File(self.filename, "r")
f = self._open_reader
else:
f = h5py.File(self.filename, "r")

ds = f[self.dataset]

if self.frame_map:
frame_idx = self.frame_map[frame_idx]
if self.frame_map:
frame_idx = self.frame_map[frame_idx]

img = ds[frame_idx]
img = ds[frame_idx]

if "format" in ds.attrs:
img = self.decode_embedded(img, ds.attrs["format"])
if "format" in ds.attrs:
img = self.decode_embedded(img, ds.attrs["format"])

if self.input_format == "channels_first":
img = np.transpose(img, (2, 1, 0))

if not self.keep_open:
f.close()
return img

def _read_frames(self, frame_inds: list) -> np.ndarray:
Expand All @@ -565,20 +637,29 @@ def _read_frames(self, frame_inds: list) -> np.ndarray:
This does not apply grayscale conversion. It is recommended to use the
`get_frames` method of the `VideoBackend` class instead.
"""
with h5py.File(self.filename, "r") as f:
if self.frame_map:
frame_inds = [self.frame_map[idx] for idx in frame_inds]
if self.keep_open:
if self._open_reader is None:
self._open_reader = h5py.File(self.filename, "r")
f = self._open_reader
else:
f = h5py.File(self.filename, "r")

ds = f[self.dataset]
imgs = ds[frame_inds]
if self.frame_map:
frame_inds = [self.frame_map[idx] for idx in frame_inds]

if "format" in ds.attrs:
imgs = np.stack(
[self.decode_embedded(img, ds.attrs["format"]) for img in imgs],
axis=0,
)
ds = f[self.dataset]
imgs = ds[frame_inds]

if "format" in ds.attrs:
imgs = np.stack(
[self.decode_embedded(img, ds.attrs["format"]) for img in imgs],
axis=0,
)

if self.input_format == "channels_first":
imgs = np.transpose(imgs, (0, 3, 2, 1))

if not self.keep_open:
f.close()

return imgs
11 changes: 10 additions & 1 deletion sleap_io/model/video.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,7 @@ def from_filename(
filename: str,
dataset: Optional[str] = None,
grayscale: Optional[str] = None,
keep_open: bool = True,
**kwargs,
) -> VideoBackend:
"""Create a Video from a filename.
Expand All @@ -55,14 +56,22 @@ def from_filename(
dataset: Name of dataset in HDF5 file.
grayscale: Whether to force grayscale. If None, autodetect on first frame
load.
keep_open: Whether to keep the video reader open between calls to read
frames. If False, will close the reader after each call. If True (the
default), it will keep the reader open and cache it for subsequent calls
which may enhance the performance of reading multiple frames.

Returns:
Video instance with the appropriate backend instantiated.
"""
return cls(
filename=filename,
backend=VideoBackend.from_filename(
filename, dataset=dataset, grayscale=grayscale, **kwargs
filename,
dataset=dataset,
grayscale=grayscale,
keep_open=keep_open,
**kwargs,
),
)

Expand Down
34 changes: 28 additions & 6 deletions tests/io/test_video_backends.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
import numpy as np
from numpy.testing import assert_equal
import h5py
import pytest


def test_video_backend_from_filename(centered_pair_low_quality_path, slp_minimal_pkg):
Expand Down Expand Up @@ -55,35 +56,53 @@ def test_get_frame(centered_pair_low_quality_path):
assert_equal(backend[-3:-1], backend.get_frames(range(1097, 1099)))


def test_mediavideo(centered_pair_low_quality_path):
@pytest.mark.parametrize("keep_open", [False, True])
def test_mediavideo(centered_pair_low_quality_path, keep_open):
backend = VideoBackend.from_filename(
centered_pair_low_quality_path, plugin="FFMPEG"
centered_pair_low_quality_path, plugin="FFMPEG", keep_open=keep_open
)
assert type(backend) == MediaVideo
assert backend.filename == centered_pair_low_quality_path
talmo marked this conversation as resolved.
Show resolved Hide resolved
assert backend.shape == (1100, 384, 384, 1)
assert backend[0].shape == (384, 384, 1)
assert backend[:3].shape == (3, 384, 384, 1)
if keep_open:
assert backend._open_reader is not None
assert backend[0].shape == (384, 384, 1)
assert type(backend._open_reader).__name__ == "LegacyPlugin"

backend = VideoBackend.from_filename(centered_pair_low_quality_path, plugin="pyav")
backend = VideoBackend.from_filename(
centered_pair_low_quality_path, plugin="pyav", keep_open=keep_open
)
assert type(backend) == MediaVideo
assert backend.filename == centered_pair_low_quality_path
assert backend.shape == (1100, 384, 384, 1)
assert backend[0].shape == (384, 384, 1)
assert backend[:3].shape == (3, 384, 384, 1)
if keep_open:
assert backend._open_reader is not None
assert backend[0].shape == (384, 384, 1)
assert type(backend._open_reader).__name__ == "PyAVPlugin"

backend = VideoBackend.from_filename(
centered_pair_low_quality_path, plugin="opencv"
centered_pair_low_quality_path, plugin="opencv", keep_open=keep_open
)
assert type(backend) == MediaVideo
assert backend.filename == centered_pair_low_quality_path
assert backend.shape == (1100, 384, 384, 1)
assert backend[0].shape == (384, 384, 1)
assert backend[:3].shape == (3, 384, 384, 1)
if keep_open:
assert backend._open_reader is not None
assert backend[0].shape == (384, 384, 1)
assert type(backend._open_reader).__name__ == "VideoCapture"


def test_hdf5video_rank4(centered_pair_low_quality_path, tmp_path):
backend = VideoBackend.from_filename(centered_pair_low_quality_path)
@pytest.mark.parametrize("keep_open", [False, True])
def test_hdf5video_rank4(centered_pair_low_quality_path, tmp_path, keep_open):
backend = VideoBackend.from_filename(
centered_pair_low_quality_path, keep_open=keep_open
)
imgs = backend[:3]
assert imgs.shape == (3, 384, 384, 1)

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Note: This review was outside of the patch, so it was mapped to the patch with the greatest overlap. Original lines [59-122]

The addition of the keep_open parameter to the test_mediavideo and test_hdf5video_rank4 functions is a good improvement. It allows testing the behavior of the VideoBackend class when the video reader is kept open between calls to read frames. The assertions added to check the _open_reader attribute when keep_open is True are also correct. However, there should be an else clause to assert that _open_reader is None when keep_open is False. This will ensure that the video reader is indeed closed after each call when keep_open is False.

  if keep_open:
      assert backend._open_reader is not None
      assert backend[0].shape == (384, 384, 1)
+ else:
+     assert backend._open_reader is None

Expand All @@ -97,6 +116,9 @@ def test_hdf5video_rank4(centered_pair_low_quality_path, tmp_path):
assert backend[0].shape == (384, 384, 1)
assert backend[:].shape == (3, 384, 384, 1)
assert not backend.has_embedded_images
if keep_open:
assert backend._open_reader is not None
assert backend[0].shape == (384, 384, 1)


def test_hdf5video_embedded(slp_minimal_pkg):
Expand Down
Loading