diff --git a/RELEASE.md b/RELEASE.md index 9f274159b7..764bc93748 100644 --- a/RELEASE.md +++ b/RELEASE.md @@ -14,6 +14,13 @@ * Make Kedro instantiate datasets from `kedro_datasets` with higher priority than `kedro.extras.datasets`. `kedro_datasets` is the namespace for the new `kedro-datasets` python package. * The config loader objects now implement `UserDict` and the configuration is accessed through `conf_loader['catalog']`. * You can configure config file patterns through `settings.py` without creating a custom config loader. +* Added `VideoDataSet` to read and write video files and interact with their frames: + +| Type | Description | Location | +| ------------------------------------ | -------------------------------------------------------------------------- | ----------------------------- | +| `video.VideoDataSet` | Read and write video files from a filesystem | `kedro.extras.datasets.video` | +| `video.video_dataset.SequenceVideo` | Create a video object from an iterable sequence to use with `VideoDataSet` | `kedro.extras.datasets.video` | +| `video.video_dataset.GeneratorVideo` | Create a video object from a generator to use with `VideoDataSet` | `kedro.extras.datasets.video` | ## Bug fixes and other changes * Fixed `kedro micropkg pull` for packages on PyPI. diff --git a/kedro/extras/datasets/video/__init__.py b/kedro/extras/datasets/video/__init__.py new file mode 100644 index 0000000000..f5f7af9461 --- /dev/null +++ b/kedro/extras/datasets/video/__init__.py @@ -0,0 +1,5 @@ +"""Dataset implementation to load/save data from/to a video file.""" + +__all__ = ["VideoDataSet"] + +from kedro.extras.datasets.video.video_dataset import VideoDataSet diff --git a/kedro/extras/datasets/video/video_dataset.py b/kedro/extras/datasets/video/video_dataset.py new file mode 100644 index 0000000000..b610006e99 --- /dev/null +++ b/kedro/extras/datasets/video/video_dataset.py @@ -0,0 +1,358 @@ +"""``VideoDataSet`` loads/saves video data from an underlying +filesystem (e.g.: local, S3, GCS). It uses OpenCV VideoCapture to read +and decode videos and OpenCV VideoWriter to encode and write video. +""" +import itertools +import tempfile +from collections import abc +from copy import deepcopy +from pathlib import Path, PurePosixPath +from typing import Any, Dict, Generator, Optional, Sequence, Tuple, Union + +import cv2 +import fsspec +import numpy as np +import PIL.Image + +from kedro.io.core import AbstractDataSet, get_protocol_and_path + + +class SlicedVideo: + """A representation of slices of other video types""" + + def __init__(self, video, slice_indexes): + self.video = video + self.indexes = range(*slice_indexes.indices(len(video))) + + def __getitem__(self, index: Union[int, slice]) -> PIL.Image.Image: + if isinstance(index, slice): + return SlicedVideo(self, index) + return self.video[self.indexes[index]] + + def __len__(self) -> int: + return len(self.indexes) + + def __getattr__(self, item): + return getattr(self.video, item) + + +class AbstractVideo(abc.Sequence): + """Base class for the underlying video data""" + + _n_frames = 0 + _index = 0 # Next available frame + + @property + def fourcc(self) -> str: + """Get the codec fourcc specification""" + raise NotImplementedError() + + @property + def fps(self) -> float: + """Get the video frame rate""" + raise NotImplementedError() + + @property + def size(self) -> Tuple[int, int]: + """Get the resolution of the video""" + raise NotImplementedError() + + def __len__(self) -> int: + return self._n_frames + + def __getitem__(self, index: Union[int, slice]): + """Get a frame from the video""" + raise NotImplementedError() + + +class FileVideo(AbstractVideo): + """A video object read from a file""" + + def __init__(self, filepath: str) -> None: + self._filepath = filepath + self._cap = cv2.VideoCapture(filepath) + self._n_frames = self._get_length() + + @property + def fourcc(self) -> str: + fourcc = self._cap.get(cv2.CAP_PROP_FOURCC) + return int(fourcc).to_bytes(4, "little").decode("ascii") + + @property + def fps(self) -> float: + return self._cap.get(cv2.CAP_PROP_FPS) + + @property + def size(self) -> Tuple[int, int]: + width = int(self._cap.get(cv2.CAP_PROP_FRAME_WIDTH)) + height = int(self._cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) + return width, height + + def __getitem__(self, index: Union[int, slice]): + if isinstance(index, slice): + return SlicedVideo(self, index) + + if index < 0: + index += len(self) + if index >= len(self): + raise IndexError() + + if index != self._index: + self._cap.set(cv2.CAP_PROP_POS_FRAMES, index) + self._index = index + 1 # Next frame to decode after this + ret, frame_bgr = self._cap.read() + if not ret: + raise IndexError() + + height, width = frame_bgr.shape[:2] + return PIL.Image.frombuffer( # Convert to PIL image with RGB instead of BGR + "RGB", (width, height), frame_bgr, "raw", "BGR", 0, 0 + ) + + def _get_length(self) -> int: + # OpenCV's frame count might be an approximation depending on what + # headers are available in the video file + length = int(round(self._cap.get(cv2.CAP_PROP_FRAME_COUNT))) + if length >= 0: + return length + + # Getting the frame count with OpenCV can fail on some video files, + # counting the frames would be too slow so it is better to raise an exception. + raise ValueError( + "Failed to load video since number of frames can't be inferred" + ) + + +class SequenceVideo(AbstractVideo): + """A video object read from an indexable sequence of frames""" + + def __init__( + self, frames: Sequence[PIL.Image.Image], fps: float, fourcc: str = "mp4v" + ) -> None: + self._n_frames = len(frames) + self._frames = frames + self._fourcc = fourcc + self._size = frames[0].size + self._fps = fps + + @property + def fourcc(self) -> str: + return self._fourcc + + @property + def fps(self) -> float: + return self._fps + + @property + def size(self) -> Tuple[int, int]: + return self._size + + def __getitem__(self, index: Union[int, slice]): + if isinstance(index, slice): + return SlicedVideo(self, index) + return self._frames[index] + + +class GeneratorVideo(AbstractVideo): + """A video object with frames yielded by a generator""" + + def __init__( + self, + frames: Generator[PIL.Image.Image, None, None], + length, + fps: float, + fourcc: str = "mp4v", + ) -> None: + self._n_frames = length + first = next(frames) + self._gen = itertools.chain([first], frames) + self._fourcc = fourcc + self._size = first.size + self._fps = fps + + @property + def fourcc(self) -> str: + return self._fourcc + + @property + def fps(self) -> float: + return self._fps + + @property + def size(self) -> Tuple[int, int]: + return self._size + + def __getitem__(self, index: Union[int, slice]): + raise NotImplementedError("Underlying video is a generator") + + def __next__(self): + return next(self._gen) + + def __iter__(self): + return self + + +class VideoDataSet(AbstractDataSet[AbstractVideo, AbstractVideo]): + """``VideoDataSet`` loads / save video data from a given filepath as sequence + of PIL.Image.Image using OpenCV. + + Example adding a catalog entry with + `YAML API + `_: + + .. code-block:: yaml + >>> cars: + >>> type: video.VideoDataSet + >>> filepath: data/01_raw/cars.mp4 + >>> + >>> cars: + >>> type: video.VideoDataSet + >>> filepath: data/01_raw/cars.mp4 + >>> filepath: s3://your_bucket/data/02_intermediate/company/motorbikes.mp4 + >>> credentials: dev_s3 + >>> + + + Example using Python API: + :: + + >>> from kedro.extras.datasets.video import VideoDataSet + >>> import numpy as np + >>> + >>> video = VideoDataSet(filepath='/video/file/path.mp4').load() + >>> frame = video[0] + >>> np.sum(np.asarray(frame)) + + + Example creating a video from numpy frames using Python API: + :: + + >>> from kedro.extras.datasets.video.video_dataset import VideoDataSet, SequenceVideo + >>> import numpy as np + >>> from PIL import Image + >>> + >>> frame = np.ones((640,480,3), dtype=np.uint8) * 255 + >>> imgs = [] + >>> for i in range(255): + >>> imgs.append(Image.fromarray(frame)) + >>> frame -= 1 + >>> + >>> video = VideoDataSet("my_video.mp4") + >>> video.save(SequenceVideo(imgs, fps=25)) + + + Example creating a video from numpy frames using a generator and Python API: + :: + + >>> from kedro.extras.datasets.video.video_dataset import VideoDataSet, GeneratorVideo + >>> import numpy as np + >>> from PIL import Image + >>> + >>> def gen(): + >>> frame = np.ones((640,480,3), dtype=np.uint8) * 255 + >>> for i in range(255): + >>> yield Image.fromarray(frame) + >>> frame -= 1 + >>> + >>> video = VideoDataSet("my_video.mp4") + >>> video.save(GeneratorVideo(gen(), fps=25, length=None)) + + """ + + # pylint: disable=too-many-arguments + def __init__( + self, + filepath: str, + fourcc: Optional[str] = "mp4v", + credentials: Dict[str, Any] = None, + fs_args: Dict[str, Any] = None, + ) -> None: + """Creates a new instance of VideoDataSet to load / save video data for given filepath. + + Args: + filepath: The location of the video file to load / save data. + fourcc: The codec to use when writing video, note that depending on how opencv is + installed there might be more or less codecs avaiable. If set to None, the + fourcc from the video object will be used. + credentials: Credentials required to get access to the underlying filesystem. + E.g. for ``GCSFileSystem`` it should look like `{"token": None}`. + fs_args: Extra arguments to pass into underlying filesystem class constructor + (e.g. `{"project": "my-project"}` for ``GCSFileSystem``). + """ + # parse the path and protocol (e.g. file, http, s3, etc.) + protocol, path = get_protocol_and_path(filepath) + self._protocol = protocol + self._filepath = PurePosixPath(path) + self._fourcc = fourcc + _fs_args = deepcopy(fs_args) or {} + _credentials = deepcopy(credentials) or {} + self._storage_options = {**_credentials, **_fs_args} + self._fs = fsspec.filesystem(self._protocol, **self._storage_options) + + def _load(self) -> AbstractVideo: + """Loads data from the video file. + + Returns: + Data from the video file as a AbstractVideo object + """ + with fsspec.open( + f"filecache::{self._protocol}://{self._filepath}", + mode="rb", + **{self._protocol: self._storage_options}, + ) as fs_file: + return FileVideo(fs_file.name) + + def _save(self, data: AbstractVideo) -> None: + """Saves video data to the specified filepath.""" + if self._protocol == "file": + # Write directly to the local file destination + self._write_to_filepath(data, str(self._filepath)) + else: + # VideoWriter can't write to an open file object, instead write to a + # local tmpfile and then copy that to the destination with fsspec. + # Note that the VideoWriter fails to write to the file on Windows if + # the file is already open, thus we can't use NamedTemporaryFile. + with tempfile.TemporaryDirectory() as tmp_dir: + tmp_file = Path(tmp_dir) / self._filepath.name + self._write_to_filepath(data, str(tmp_file)) + with fsspec.open( + f"{self._protocol}://{self._filepath}", + "wb", + **self._storage_options, + ) as f_target: + with tmp_file.open("r+b") as f_tmp: + f_target.write(f_tmp.read()) + + def _write_to_filepath(self, video: AbstractVideo, filepath: str) -> None: + # TODO: This uses the codec specified in the VideoDataSet if it is not None, this is due + # to compatibility issues since e.g. h264 coded is licensed and is thus not included in + # opencv if installed from a binary distribution. Since a h264 video can be read, but not + # written, it would be error prone to use the videos fourcc code. Further, an issue is + # that the video object does not know what container format will be used since that is + # selected by the suffix in the file name of the VideoDataSet. Some combinations of codec + # and container format might not work or will have bad support. + fourcc = self._fourcc or video.fourcc + + writer = cv2.VideoWriter( + filepath, cv2.VideoWriter_fourcc(*fourcc), video.fps, video.size + ) + if not writer.isOpened(): + raise ValueError( + "Failed to open video writer with params: " + + f"fourcc={fourcc} fps={video.fps} size={video.size[0]}x{video.size[1]} " + + f"path={filepath}" + ) + try: + for frame in iter(video): + writer.write( # PIL images are RGB, opencv expects BGR + np.asarray(frame)[:, :, ::-1] + ) + finally: + writer.release() + + def _describe(self) -> Dict[str, Any]: + return dict(filepath=self._filepath, protocol=self._protocol) + + def _exists(self) -> bool: + return self._fs.exists(self._filepath) diff --git a/pyproject.toml b/pyproject.toml index 9c3e01124a..9b17440f44 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -17,6 +17,7 @@ load-plugins = [ "pylint.extensions.docparams", "pylint.extensions.no_self_use" ] +extension-pkg-whitelist = "cv2" unsafe-load-any-extension = false [tool.pylint.messages_control] disable = [ diff --git a/setup.py b/setup.py index 24e6d6e25d..c01f192e3a 100644 --- a/setup.py +++ b/setup.py @@ -80,6 +80,9 @@ def _collect_requirements(requires): } pickle_require = {"pickle.PickleDataSet": ["compress-pickle[lz4]~=2.1.0"]} pillow_require = {"pillow.ImageDataSet": ["Pillow~=9.0"]} +video_require = { + "video.VideoDataSet": ["opencv-python~=4.5.5.64"] +} plotly_require = { "plotly.PlotlyDataSet": [PANDAS, "plotly>=4.8.0, <6.0"], "plotly.JSONDataSet": ["plotly>=4.8.0, <6.0"], @@ -124,6 +127,7 @@ def _collect_requirements(requires): "pandas": _collect_requirements(pandas_require), "pickle": _collect_requirements(pickle_require), "pillow": _collect_requirements(pillow_require), + "video": _collect_requirements(video_require), "plotly": _collect_requirements(plotly_require), "redis": _collect_requirements(redis_require), "spark": _collect_requirements(spark_require), @@ -139,6 +143,7 @@ def _collect_requirements(requires): **pandas_require, **pickle_require, **pillow_require, + **video_require, **plotly_require, **spark_require, **tensorflow_required, diff --git a/test_requirements.txt b/test_requirements.txt index 52305e865a..5d10f0371d 100644 --- a/test_requirements.txt +++ b/test_requirements.txt @@ -30,6 +30,7 @@ memory_profiler>=0.50.0, <1.0 moto==1.3.7; python_version < '3.10' moto==3.0.4; python_version == '3.10' networkx~=2.4 +opencv-python~=4.5.5.64 openpyxl>=3.0.3, <4.0 pandas-gbq>=0.12.0, <1.0 pandas~=1.3 # 1.3 for read_xml/to_xml diff --git a/tests/extras/datasets/video/conftest.py b/tests/extras/datasets/video/conftest.py new file mode 100644 index 0000000000..ff084cdb5e --- /dev/null +++ b/tests/extras/datasets/video/conftest.py @@ -0,0 +1,107 @@ +from pathlib import Path + +import pytest +from PIL import Image +from utils import TEST_FPS, TEST_HEIGHT, TEST_WIDTH + +from kedro.extras.datasets.video.video_dataset import ( + FileVideo, + GeneratorVideo, + SequenceVideo, +) + + +@pytest.fixture(scope="module") +def red_frame(): + return Image.new("RGB", (TEST_WIDTH, TEST_HEIGHT), (255, 0, 0)) + + +@pytest.fixture(scope="module") +def green_frame(): + return Image.new("RGB", (TEST_WIDTH, TEST_HEIGHT), (0, 255, 0)) + + +@pytest.fixture(scope="module") +def blue_frame(): + return Image.new("RGB", (TEST_WIDTH, TEST_HEIGHT), (0, 0, 255)) + + +@pytest.fixture(scope="module") +def yellow_frame(): + return Image.new("RGB", (TEST_WIDTH, TEST_HEIGHT), (255, 255, 0)) + + +@pytest.fixture(scope="module") +def purple_frame(): + return Image.new("RGB", (TEST_WIDTH, TEST_HEIGHT), (255, 0, 255)) + + +@pytest.fixture +def color_video(red_frame, green_frame, blue_frame, yellow_frame, purple_frame): + return SequenceVideo( + [red_frame, green_frame, blue_frame, yellow_frame, purple_frame], + fps=TEST_FPS, + ) + + +@pytest.fixture +def color_video_generator( + red_frame, green_frame, blue_frame, yellow_frame, purple_frame +): + sequence = [red_frame, green_frame, blue_frame, yellow_frame, purple_frame] + + def generator(): + yield from sequence + + return GeneratorVideo( + generator(), + length=len(sequence), + fps=TEST_FPS, + ) + + +@pytest.fixture +def filepath_mp4(): + """This is a real video converted to mp4/h264 with ffmpeg command""" + return str(Path(__file__).parent / "data/video.mp4") + + +@pytest.fixture +def filepath_mkv(): + """This a a real video recoreded with an Axis network camera""" + return str(Path(__file__).parent / "data/video.mkv") + + +@pytest.fixture +def filepath_mjpeg(): + """This is a real video recorded with an Axis network camera""" + return str(Path(__file__).parent / "data/video.mjpeg") + + +@pytest.fixture +def filepath_color_mp4(): + """This is a video created with the OpenCV VideoWriter + + it contains 5 frames which each is a single color: red, green, blue, yellow, purple + """ + return str(Path(__file__).parent / "data/color_video.mp4") + + +@pytest.fixture +def mp4_object(filepath_mp4): + return FileVideo(filepath_mp4) + + +@pytest.fixture +def mkv_object(filepath_mkv): + return FileVideo(filepath_mkv) + + +@pytest.fixture +def mjpeg_object(filepath_mjpeg): + return FileVideo(filepath_mjpeg) + + +@pytest.fixture +def color_video_object(filepath_color_mp4): + return FileVideo(filepath_color_mp4) diff --git a/tests/extras/datasets/video/data/color_video.mp4 b/tests/extras/datasets/video/data/color_video.mp4 new file mode 100644 index 0000000000..01944b1b78 Binary files /dev/null and b/tests/extras/datasets/video/data/color_video.mp4 differ diff --git a/tests/extras/datasets/video/data/video.mjpeg b/tests/extras/datasets/video/data/video.mjpeg new file mode 100644 index 0000000000..cab90dda94 Binary files /dev/null and b/tests/extras/datasets/video/data/video.mjpeg differ diff --git a/tests/extras/datasets/video/data/video.mkv b/tests/extras/datasets/video/data/video.mkv new file mode 100644 index 0000000000..2710c022ff Binary files /dev/null and b/tests/extras/datasets/video/data/video.mkv differ diff --git a/tests/extras/datasets/video/data/video.mp4 b/tests/extras/datasets/video/data/video.mp4 new file mode 100644 index 0000000000..4c4b974d92 Binary files /dev/null and b/tests/extras/datasets/video/data/video.mp4 differ diff --git a/tests/extras/datasets/video/test_sliced_video.py b/tests/extras/datasets/video/test_sliced_video.py new file mode 100644 index 0000000000..e2e4975d1a --- /dev/null +++ b/tests/extras/datasets/video/test_sliced_video.py @@ -0,0 +1,56 @@ +import numpy as np +from utils import TEST_HEIGHT, TEST_WIDTH + + +class TestSlicedVideo: + def test_slice_sequence_video_first(self, color_video): + """Test slicing and then indexing a SequenceVideo""" + slice_red_green = color_video[:2] + red = np.array(slice_red_green[0]) + assert red.shape == (TEST_HEIGHT, TEST_WIDTH, 3) + assert np.all(red[:, :, 0] == 255) + assert np.all(red[:, :, 1] == 0) + assert np.all(red[:, :, 2] == 0) + + def test_slice_sequence_video_last_as_index(self, color_video): + """Test slicing and then indexing a SequenceVideo""" + slice_blue_yellow_purple = color_video[2:5] + purple = np.array(slice_blue_yellow_purple[2]) + assert purple.shape == (TEST_HEIGHT, TEST_WIDTH, 3) + assert np.all(purple[:, :, 0] == 255) + assert np.all(purple[:, :, 1] == 0) + assert np.all(purple[:, :, 2] == 255) + + def test_slice_sequence_video_last_as_end(self, color_video): + """Test slicing and then indexing a SequenceVideo""" + slice_blue_yellow_purple = color_video[2:] + purple = np.array(slice_blue_yellow_purple[-1]) + assert purple.shape == (TEST_HEIGHT, TEST_WIDTH, 3) + assert np.all(purple[:, :, 0] == 255) + assert np.all(purple[:, :, 1] == 0) + assert np.all(purple[:, :, 2] == 255) + + def test_slice_sequence_attribute(self, color_video): + """Test that attributes from the base class are reachable from sliced views""" + slice_red_green = color_video[:2] + assert slice_red_green.fps == color_video.fps + + def test_slice_sliced_video(self, color_video): + """Test slicing and then indexing a SlicedVideo""" + slice_green_blue_yellow = color_video[1:4] + slice_green_blue = slice_green_blue_yellow[:-1] + blue = np.array(slice_green_blue[1]) + assert blue.shape == (TEST_HEIGHT, TEST_WIDTH, 3) + assert np.all(blue[:, :, 0] == 0) + assert np.all(blue[:, :, 1] == 0) + assert np.all(blue[:, :, 2] == 255) + + def test_slice_file_video_first(self, mp4_object): + """Test slicing and then indexing a FileVideo""" + sliced_video = mp4_object[:2] + assert np.all(np.array(sliced_video[0]) == np.array(mp4_object[0])) + + def test_slice_file_video_last(self, mp4_object): + """Test slicing and then indexing a FileVideo""" + sliced_video = mp4_object[-2:] + assert np.all(np.array(sliced_video[-1]) == np.array(mp4_object[-1])) diff --git a/tests/extras/datasets/video/test_video_dataset.py b/tests/extras/datasets/video/test_video_dataset.py new file mode 100644 index 0000000000..81c1952513 --- /dev/null +++ b/tests/extras/datasets/video/test_video_dataset.py @@ -0,0 +1,186 @@ +import boto3 +import pytest +from moto import mock_s3 +from utils import TEST_FPS, assert_videos_equal + +from kedro.extras.datasets.video import VideoDataSet +from kedro.extras.datasets.video.video_dataset import FileVideo, SequenceVideo +from kedro.io import DataSetError + +S3_BUCKET_NAME = "test_bucket" +S3_KEY_PATH = "video" +S3_FULL_PATH = f"s3://{S3_BUCKET_NAME}/{S3_KEY_PATH}/" +AWS_CREDENTIALS = {"key": "FAKE_ACCESS_KEY", "secret": "FAKE_SECRET_KEY"} + + +@pytest.fixture +def tmp_filepath_mp4(tmp_path): + return (tmp_path / "test.mp4").as_posix() + + +@pytest.fixture +def tmp_filepath_avi(tmp_path): + return (tmp_path / "test.mjpeg").as_posix() + + +@pytest.fixture +def empty_dataset_mp4(tmp_filepath_mp4): + return VideoDataSet(filepath=tmp_filepath_mp4) + + +@pytest.fixture +def empty_dataset_avi(tmp_filepath_avi): + return VideoDataSet(filepath=tmp_filepath_avi) + + +@pytest.fixture +def mocked_s3_bucket(): + """Create a bucket for testing using moto.""" + with mock_s3(): + conn = boto3.client( + "s3", + region_name="us-east-1", + aws_access_key_id=AWS_CREDENTIALS["key"], + aws_secret_access_key=AWS_CREDENTIALS["secret"], + ) + conn.create_bucket(Bucket=S3_BUCKET_NAME) + yield conn + + +class TestVideoDataSet: + def test_load_mp4(self, filepath_mp4, mp4_object): + """Loading a mp4 dataset should create a FileVideo""" + ds = VideoDataSet(filepath_mp4) + loaded_video = ds.load() + assert_videos_equal(loaded_video, mp4_object) + + def test_save_and_load_mp4(self, empty_dataset_mp4, mp4_object): + """Test saving and reloading the data set.""" + empty_dataset_mp4.save(mp4_object) + reloaded_video = empty_dataset_mp4.load() + assert_videos_equal(mp4_object, reloaded_video) + assert reloaded_video.fourcc == empty_dataset_mp4._fourcc + + @pytest.mark.skip( + reason="Only one available codec that is typically installed when testing" + ) + def test_save_with_other_codec(self, tmp_filepath_mp4, mp4_object): + """Test saving the video with another codec than default.""" + save_fourcc = "xvid" + ds = VideoDataSet(filepath=tmp_filepath_mp4, fourcc=save_fourcc) + ds.save(mp4_object) + reloaded_video = ds.load() + assert reloaded_video.fourcc == save_fourcc + + def test_save_with_derived_codec(self, tmp_filepath_mp4, color_video): + """Test saving video by the codec specified in the video object""" + ds = VideoDataSet(filepath=tmp_filepath_mp4, fourcc=None) + ds.save(color_video) + reloaded_video = ds.load() + assert reloaded_video.fourcc == color_video.fourcc + + def test_saved_fps(self, empty_dataset_mp4, color_video): + """Verify that a saved video has the same framerate as specified in the video object""" + empty_dataset_mp4.save(color_video) + reloaded_video = empty_dataset_mp4.load() + assert reloaded_video.fps == TEST_FPS + + def test_save_sequence_video(self, color_video, empty_dataset_mp4): + """Test save (and load) a SequenceVideo object""" + empty_dataset_mp4.save(color_video) + reloaded_video = empty_dataset_mp4.load() + assert_videos_equal(color_video, reloaded_video) + + def test_save_generator_video( + self, color_video_generator, empty_dataset_mp4, color_video + ): + """Test save (and load) a GeneratorVideo object + + Since the GeneratorVideo is exhaused after saving the video to file we use + the SequenceVideo (color_video) which has the same frames to compare the + loaded video to. + """ + empty_dataset_mp4.save(color_video_generator) + reloaded_video = empty_dataset_mp4.load() + assert_videos_equal(color_video, reloaded_video) + + def test_exists(self, empty_dataset_mp4, mp4_object): + """Test `exists` method invocation for both existing and + nonexistent data set.""" + assert not empty_dataset_mp4.exists() + empty_dataset_mp4.save(mp4_object) + assert empty_dataset_mp4.exists() + + @pytest.mark.skip(reason="Can't deal with videos with missing time info") + def test_convert_video(self, empty_dataset_mp4, mjpeg_object): + """Load a file video in mjpeg format and save in mp4v""" + empty_dataset_mp4.save(mjpeg_object) + reloaded_video = empty_dataset_mp4.load() + assert_videos_equal(mjpeg_object, reloaded_video) + + def test_load_missing_file(self, empty_dataset_mp4): + """Check the error when trying to load missing file.""" + pattern = r"Failed while loading data from data set VideoDataSet\(.*\)" + with pytest.raises(DataSetError, match=pattern): + empty_dataset_mp4.load() + + def test_save_s3(self, mp4_object, mocked_s3_bucket, tmp_path): + """Test to save a VideoDataSet to S3 storage""" + video_name = "video.mp4" + + dataset = VideoDataSet( + filepath=S3_FULL_PATH + video_name, credentials=AWS_CREDENTIALS + ) + dataset.save(mp4_object) + + tmp_file = tmp_path / video_name + mocked_s3_bucket.download_file( + Bucket=S3_BUCKET_NAME, + Key=S3_KEY_PATH + "/" + video_name, + Filename=str(tmp_file), + ) + reloaded_video = FileVideo(str(tmp_file)) + assert_videos_equal(reloaded_video, mp4_object) + + @pytest.mark.xfail + @pytest.mark.parametrize( + "fourcc, suffix", + [ + ("mp4v", "mp4"), + ("mp4v", "mjpeg"), + ("mp4v", "avi"), + ("avc1", "mp4"), + ("avc1", "mjpeg"), + ("avc1", "avi"), + ("mjpg", "mp4"), + ("mjpg", "mjpeg"), + ("mjpg", "avi"), + ("xvid", "mp4"), + ("xvid", "mjpeg"), + ("xvid", "avi"), + ("x264", "mp4"), + ("x264", "mjpeg"), + ("x264", "avi"), + ("divx", "mp4"), + ("divx", "mjpeg"), + ("divx", "avi"), + ("fmp4", "mp4"), + ("fmp4", "mjpeg"), + ("fmp4", "avi"), + ], + ) + def test_video_codecs(self, fourcc, suffix, color_video): + """Test different codec and container combinations + + Some of these are expected to fail depending on what + codecs are installed on the machine. + """ + video_name = f"video.{suffix}" + video = SequenceVideo(color_video._frames, 25, fourcc) + ds = VideoDataSet(video_name, fourcc=None) + ds.save(video) + # We also need to verify that the correct codec was used + # since OpenCV silently (with a warning in the log) fall backs to + # another codec if one specified is not compatible with the container + reloaded_video = ds.load() + assert reloaded_video.fourcc == fourcc diff --git a/tests/extras/datasets/video/test_video_objects.py b/tests/extras/datasets/video/test_video_objects.py new file mode 100644 index 0000000000..66a284fa60 --- /dev/null +++ b/tests/extras/datasets/video/test_video_objects.py @@ -0,0 +1,170 @@ +import numpy as np +import pytest +from utils import ( + DEFAULT_FOURCC, + MJPEG_FOURCC, + MJPEG_FPS, + MJPEG_LEN, + MJPEG_SIZE, + MKV_FOURCC, + MKV_FPS, + MKV_LEN, + MKV_SIZE, + MP4_FOURCC, + MP4_FPS, + MP4_LEN, + MP4_SIZE, + TEST_FPS, + TEST_HEIGHT, + TEST_NUM_COLOR_FRAMES, + TEST_WIDTH, + assert_images_equal, +) + +from kedro.extras.datasets.video.video_dataset import ( + FileVideo, + GeneratorVideo, + SequenceVideo, +) + + +class TestSequenceVideo: + def test_sequence_video_indexing_first(self, color_video, red_frame): + """Test indexing a SequenceVideo""" + red = np.array(color_video[0]) + assert red.shape == (TEST_HEIGHT, TEST_WIDTH, 3) + assert np.all(red == red_frame) + + def test_sequence_video_indexing_last(self, color_video, purple_frame): + """Test indexing a SequenceVideo""" + purple = np.array(color_video[-1]) + assert purple.shape == (TEST_HEIGHT, TEST_WIDTH, 3) + assert np.all(purple == purple_frame) + + def test_sequence_video_iterable(self, color_video): + """Test iterating a SequenceVideo""" + for i, img in enumerate(map(np.array, color_video)): + assert np.all(img == np.array(color_video[i])) + assert i == TEST_NUM_COLOR_FRAMES - 1 + + def test_sequence_video_fps(self, color_video): + # Test the one set by the fixture + assert color_video.fps == TEST_FPS + + # Test creating with another fps + test_fps_new = 123 + color_video_new = SequenceVideo(color_video._frames, fps=test_fps_new) + assert color_video_new.fps == test_fps_new + + def test_sequence_video_len(self, color_video): + assert len(color_video) == TEST_NUM_COLOR_FRAMES + + def test_sequence_video_size(self, color_video): + assert color_video.size == (TEST_WIDTH, TEST_HEIGHT) + + def test_sequence_video_fourcc_default_value(self, color_video): + assert color_video.fourcc == DEFAULT_FOURCC + + def test_sequence_video_fourcc(self, color_video): + fourcc_new = "mjpg" + assert ( + DEFAULT_FOURCC != fourcc_new + ), "Test does not work if new test value is same as default" + color_video_new = SequenceVideo( + color_video._frames, fps=TEST_FPS, fourcc=fourcc_new + ) + assert color_video_new.fourcc == fourcc_new + + +class TestGeneratorVideo: + def test_generator_video_iterable(self, color_video_generator, color_video): + """Test iterating a GeneratorVideo + + The content of the mock GeneratorVideo should be the same as the SequenceVideo, + the content in the later is tested in other unit tests and can thus be trusted + """ + for i, img in enumerate(map(np.array, color_video_generator)): + assert np.all(img == np.array(color_video[i])) + assert i == TEST_NUM_COLOR_FRAMES - 1 + + def test_generator_video_fps(self, color_video_generator): + # Test the one set by the fixture + assert color_video_generator.fps == TEST_FPS + + # Test creating with another fps + test_fps_new = 123 + color_video_new = GeneratorVideo( + color_video_generator._gen, length=TEST_NUM_COLOR_FRAMES, fps=test_fps_new + ) + assert color_video_new.fps == test_fps_new + + def test_generator_video_len(self, color_video_generator): + assert len(color_video_generator) == TEST_NUM_COLOR_FRAMES + + def test_generator_video_size(self, color_video_generator): + assert color_video_generator.size == (TEST_WIDTH, TEST_HEIGHT) + + def test_generator_video_fourcc_default_value(self, color_video_generator): + assert color_video_generator.fourcc == DEFAULT_FOURCC + + def test_generator_video_fourcc(self, color_video_generator): + fourcc_new = "mjpg" + assert ( + DEFAULT_FOURCC != fourcc_new + ), "Test does not work if new test value is same as default" + color_video_new = GeneratorVideo( + color_video_generator._gen, + length=TEST_NUM_COLOR_FRAMES, + fps=TEST_FPS, + fourcc=fourcc_new, + ) + assert color_video_new.fourcc == fourcc_new + + +class TestFileVideo: + @pytest.mark.skip(reason="Can't deal with videos with missing time info") + def test_file_props_mjpeg(self, mjpeg_object): + assert mjpeg_object.fourcc == MJPEG_FOURCC + assert mjpeg_object.fps == MJPEG_FPS + assert mjpeg_object.size == MJPEG_SIZE + assert len(mjpeg_object) == MJPEG_LEN + + def test_file_props_mkv(self, mkv_object): + assert mkv_object.fourcc == MKV_FOURCC + assert mkv_object.fps == MKV_FPS + assert mkv_object.size == MKV_SIZE + assert len(mkv_object) == MKV_LEN + + def test_file_props_mp4(self, mp4_object): + assert mp4_object.fourcc == MP4_FOURCC + assert mp4_object.fps == MP4_FPS + assert mp4_object.size == MP4_SIZE + assert len(mp4_object) == MP4_LEN + + def test_file_index_first(self, color_video_object, red_frame): + assert_images_equal(color_video_object[0], red_frame) + + def test_file_index_last_by_index(self, color_video_object, purple_frame): + assert_images_equal(color_video_object[TEST_NUM_COLOR_FRAMES - 1], purple_frame) + + def test_file_index_last(self, color_video_object, purple_frame): + assert_images_equal(color_video_object[-1], purple_frame) + + def test_file_video_failed_capture(self, mocker): + """Validate good behavior on failed decode + + The best behavior in this case is not obvious, the len property of the + video object specifies more frames than is actually possible to decode. We + cannot know this in advance without spending loads of time to decode all frames + in order to count them.""" + mock_cv2 = mocker.patch("kedro.extras.datasets.video.video_dataset.cv2") + mock_cap = mock_cv2.VideoCapture.return_value = mocker.Mock() + mock_cap.get.return_value = 2 # Set the length of the video + ds = FileVideo("/a/b/c") + + mock_cap.read.return_value = True, np.zeros((1, 1)) + assert ds[0] + + mock_cap.read.return_value = False, None + with pytest.raises(IndexError): + ds[1] diff --git a/tests/extras/datasets/video/utils.py b/tests/extras/datasets/video/utils.py new file mode 100644 index 0000000000..6b675aed2f --- /dev/null +++ b/tests/extras/datasets/video/utils.py @@ -0,0 +1,49 @@ +import itertools + +import numpy as np +from PIL import ImageChops + +TEST_WIDTH = 640 # Arbitrary value for testing +TEST_HEIGHT = 480 # Arbitrary value for testing +TEST_FPS = 1 # Arbitrary value for testing + +TEST_NUM_COLOR_FRAMES = ( + 5 # This should be the same as number of frames in conftest videos +) +DEFAULT_FOURCC = "mp4v" # The expected default fourcc value + +# This is video data extracted from the video files with ffmpeg command +MKV_SIZE = (640, 360) +MKV_FPS = 50 +MKV_FOURCC = "h264" +MKV_LEN = 109 # from ffprobe + +MP4_SIZE = (640, 360) +MP4_FPS = 50 +MP4_FOURCC = "avc1" +MP4_LEN = 109 # from ffprobe + +MJPEG_SIZE = (640, 360) +MJPEG_FPS = 25 # From ffprobe, not reported by ffmpeg command +# I'm not sure that MJPE is the correct fourcc code for +# mjpeg video since I cannot find any official reference to +# that code. This is however what the openCV VideoCapture +# reports for the video, so we leave it like this for now.. +MJPEG_FOURCC = "mjpe" +MJPEG_LEN = 24 # from ffprobe + + +def assert_images_equal(image_1, image_2): + """Assert that two images are approximately equal, allow for some + compression artifacts""" + assert image_1.size == image_2.size + diff = np.asarray(ImageChops.difference(image_1, image_2)) + assert np.mean(diff) < 5 + assert np.mean(diff > 50) < 0.01 # Max 1% of pixels + + +def assert_videos_equal(video_1, video_2): + assert len(video_1) == len(video_2) + + for image_1, image_2 in itertools.zip_longest(video_1, video_2): + assert_images_equal(image_1, image_2)