diff --git a/imageio/plugins/pyav.py b/imageio/plugins/pyav.py index 87f973d1d..bb2a3c715 100644 --- a/imageio/plugins/pyav.py +++ b/imageio/plugins/pyav.py @@ -136,7 +136,7 @@ """ from math import ceil -from typing import Any, Dict, List, Optional, Tuple, Union +from typing import Any, Dict, List, Optional, Tuple, Union, Generator from fractions import Fraction import av @@ -412,8 +412,8 @@ def read( ffmpeg_filter = self._build_filter(filter_sequence, filter_graph) ffmpeg_filter.send(None) # init - self._seek(index, constant_framerate=constant_framerate) - desired_frame = next(self._container.decode(video=0)) + decoder = self._seek(index, constant_framerate=constant_framerate) + desired_frame = next(decoder) return self._unpack_frame(ffmpeg_filter.send(desired_frame), format=format) @@ -936,8 +936,8 @@ def metadata( metadata.update(self._video_stream.metadata) return metadata - self._seek(index, constant_framerate=constant_framerate) - desired_frame = next(self._container.decode(video=0)) + decoder = self._seek(index, constant_framerate=constant_framerate) + desired_frame = next(decoder) # useful flags defined on the frame metadata.update( @@ -953,7 +953,7 @@ def metadata( return metadata - def _seek(self, index, *, constant_framerate: bool = True) -> None: + def _seek(self, index, *, constant_framerate: bool = True) -> Generator: """Seeks to the frame at the given index.""" # this may be made faster for formats that have some kind @@ -976,6 +976,8 @@ def _seek(self, index, *, constant_framerate: bool = True) -> None: for _ in range(frames_to_yield): next(frame_generator) + return frame_generator + def close(self) -> None: """Close the Video.""" diff --git a/tests/test_pyav.py b/tests/test_pyav.py index ec0002095..a7b350824 100644 --- a/tests/test_pyav.py +++ b/tests/test_pyav.py @@ -390,3 +390,8 @@ def test_uri_reading(test_images): actual = iio.imread(uri, plugin="pyav", index=250) np.allclose(actual, expected) + + +def test_seek_last(test_images): + frame = iio.imread(test_images / "cockatoo.mp4", plugin="pyav", index=279) + assert frame.shape == (720, 1280, 3)