Add Sintel Dataset to the dataset prototype API

test/builtin_dataset_mocks.py

@@ -10,7 +10,7 @@
 import numpy as np
 import pytest
 import torch
-from datasets_utils import create_image_folder, make_tar, make_zip
+from datasets_utils import create_image_folder, make_tar, make_zip, make_fake_flo_file
 from torch.testing import make_tensor as _make_tensor
 from torchdata.datapipes.iter import IterDataPipe
 from torchvision.prototype import datasets
@@ -490,3 +490,42 @@ def imagenet(info, root, config):
     make_tar(root, f"{devkit_root}.tar.gz", devkit_root, compression="gz")
 
     return num_samples
+
+
+@dataset_mocks.register_mock_data_fn
+def sintel(info, root, config):
+    FLOW_H, FLOW_W = 3, 4
+
+    num_images_per_scene = 3 if config["split"] == "train" else 4
+    num_scenes = 2
+
+    for split_dir in ("training", "test"):
+        for pass_name in ("clean", "final"):
+            image_root = root / split_dir / pass_name
+
+            for scene_id in range(num_scenes):
+                scene_dir = image_root / f"scene_{scene_id}"
+                create_image_folder(
+                    image_root,
+                    name=str(scene_dir),
+                    file_name_fn=lambda image_idx: f"frame_000{image_idx}.png",
+                    num_examples=num_images_per_scene,
+                )
+
+    flow_root = root / "training" / "flow"
+    for scene_id in range(num_scenes):
+        scene_dir = flow_root / f"scene_{scene_id}"
+        scene_dir.mkdir(exist_ok=True, parents=True)
+        for i in range(num_images_per_scene - 1):
+            file_name = str(scene_dir / f"frame_000{i}.flo")
+            make_fake_flo_file(h=FLOW_H, w=FLOW_W, file_name=file_name)
+
+    # with e.g. num_images_per_scene = 3, for a single scene with have 3 images
+    # which are frame_0000, frame_0001 and frame_0002
+    # They will be consecutively paired as (frame_0000, frame_0001), (frame_0001, frame_0002),
+    # that is 3 - 1 = 2 examples. Hence the formula below
+    num_passes = 2 if config["pass_name"] == "both" else 1
+    num_examples = (num_images_per_scene - 1) * num_scenes * num_passes
+
+    make_zip(root, "MPI-Sintel-complete.zip", "training", "test")
+    return num_examples

test/test_prototype_builtin_datasets.py

@@ -30,6 +30,7 @@ def dataset_parametrization(*names, decoder=to_bytes):
             "caltech256",
             "caltech101",
             "imagenet",
+            "sintel",
         )
 
     params = []

torchvision/prototype/datasets/_builtin/__init__.py

@@ -6,4 +6,5 @@
 from .mnist import MNIST, FashionMNIST, KMNIST, EMNIST, QMNIST
 from .sbd import SBD
 from .semeion import SEMEION
+from .sintel import SINTEL
 from .voc import VOC

torchvision/prototype/datasets/_builtin/sintel.py

@@ -0,0 +1,168 @@
+import io
+import pathlib
+import re
+from functools import partial
+from typing import Any, Callable, Dict, List, Optional, Tuple, Iterator, Iterable, TypeVar, BinaryIO
+
+import torch
+from torchdata.datapipes.iter import (
+    IterDataPipe,
+    Demultiplexer,
+    Mapper,
+    Shuffler,
+    Filter,
+    IterKeyZipper,
+    ZipArchiveReader,
+)
+from torchvision.prototype.datasets.utils import (
+    Dataset,
+    DatasetConfig,
+    DatasetInfo,
+    HttpResource,
+    OnlineResource,
+    DatasetType,
+)
+from torchvision.prototype.datasets.utils._internal import INFINITE_BUFFER_SIZE, read_flo
+
+T = TypeVar("T")
+
+try:
+    from itertools import pairwise  # type: ignore[attr-defined]
+except ImportError:
+    from itertools import tee
+
+    def pairwise(iterable: Iterable[T]) -> Iterable[Tuple[T, T]]:
+        a, b = tee(iterable)
+        next(b, None)
+        return zip(a, b)
+
+
+class InSceneGrouper(IterDataPipe[Tuple[Tuple[str, T], Tuple[str, T]]]):
+    def __init__(self, datapipe: IterDataPipe[Tuple[str, T]]) -> None:
+        self.datapipe = datapipe
+
+    def __iter__(self) -> Iterator[Tuple[Tuple[str, Any], Tuple[str, Any]]]:
+        for item1, item2 in pairwise(sorted(self.datapipe)):
+            if pathlib.Path(item1[0]).parent != pathlib.Path(item2[0]).parent:
+                continue
+
+            yield item1, item2
+
+
+class SINTEL(Dataset):
+
+    _FILE_NAME_PATTERN = re.compile(r"(frame|image)_(?P<idx>\d+)[.](flo|png)")
+
+    def _make_info(self) -> DatasetInfo:
+        return DatasetInfo(
+            "sintel",
+            type=DatasetType.IMAGE,
+            homepage="http://sintel.is.tue.mpg.de/",
+            valid_options=dict(
+                split=("train", "test"),
+                pass_name=("clean", "final", "both"),
+            ),
+        )
+
+    def resources(self, config: DatasetConfig) -> List[OnlineResource]:
+        archive = HttpResource(
+            "http://files.is.tue.mpg.de/sintel/MPI-Sintel-complete.zip",
+            sha256="bdc80abbe6ae13f96f6aa02e04d98a251c017c025408066a00204cd2c7104c5f",
+        )
+        return [archive]
+
+    def _filter_split(self, data: Tuple[str, Any], *, split: str) -> bool:
+        path = pathlib.Path(data[0])
+        # The dataset contains has the folder "training", while allowed options for `split` are
+        # "train" and "test", we don't check for equality here ("train" != "training") and instead
+        # check if split is in the folder name
+        return split in path.parents[2].name
+
+    def _filter_pass_name_and_flow(self, data: Tuple[str, Any], *, pass_name: str) -> bool:
+        path = pathlib.Path(data[0])
+        if pass_name == "both":
+            matched = path.parents[1].name in ["clean", "final", "flow"]
+        else:
+            matched = path.parents[1].name in [pass_name, "flow"]
+        return matched
+
+    def _classify_archive(self, data: Tuple[str, Any], *, pass_name: str) -> Optional[int]:
+        path = pathlib.Path(data[0])
+        suffix = path.suffix
+        if suffix == ".flo":
+            return 0
+        elif suffix == ".png":
+            return 1
+        else:
+            return None
+
+    def _flows_key(self, data: Tuple[str, Any]) -> Tuple[str, int]:
+        path = pathlib.Path(data[0])
+        category = path.parent.name
+        idx = int(self._FILE_NAME_PATTERN.match(path.name).group("idx"))  # type: ignore[union-attr]
+        return category, idx
+
+    def _add_fake_flow_data(self, data: Tuple[str, Any]) -> Tuple[tuple, Tuple[str, Any]]:
+        return ((None, None), data)
+
+    def _images_key(self, data: Tuple[Tuple[str, Any], Tuple[str, Any]]) -> Tuple[str, int]:
+        return self._flows_key(data[0])
+
+    def _collate_and_decode_sample(
+        self,
+        data: Tuple[Tuple[Optional[str], Optional[BinaryIO]], Tuple[Tuple[str, BinaryIO], Tuple[str, BinaryIO]]],
+        *,
+        decoder: Optional[Callable[[BinaryIO], torch.Tensor]],
+    ) -> Dict[str, Any]:
+        flow_data, images_data = data
+        flow_path, flow_buffer = flow_data
+        image1_data, image2_data = images_data
+        image1_path, image1_buffer = image1_data
+        image2_path, image2_buffer = image2_data
+
+        return dict(
+            image1=decoder(image1_buffer) if decoder else image1_buffer,
+            image1_path=image1_path,
+            image2=decoder(image2_buffer) if decoder else image2_buffer,
+            image2_path=image2_path,
+            flow=read_flo(flow_buffer) if flow_buffer else None,
+            flow_path=flow_path,
+            scene=pathlib.Path(image1_path).parent.name,
+        )
+
+    def _make_datapipe(
+        self,
+        resource_dps: List[IterDataPipe],
+        *,
+        config: DatasetConfig,
+        decoder: Optional[Callable[[io.IOBase], torch.Tensor]],
+    ) -> IterDataPipe[Dict[str, Any]]:
+        dp = resource_dps[0]
+        archive_dp = ZipArchiveReader(dp)
+
+        curr_split = Filter(archive_dp, self._filter_split, fn_kwargs=dict(split=config.split))
+        filtered_curr_split = Filter(
+            curr_split, self._filter_pass_name_and_flow, fn_kwargs=dict(pass_name=config.pass_name)
+        )
+        if config.split == "train":
+            flo_dp, pass_images_dp = Demultiplexer(
+                filtered_curr_split,
+                2,
+                partial(self._classify_archive, pass_name=config.pass_name),
+                drop_none=True,
+                buffer_size=INFINITE_BUFFER_SIZE,
+            )
+            flo_dp = Shuffler(flo_dp, buffer_size=INFINITE_BUFFER_SIZE)
+            pass_images_dp: IterDataPipe[Tuple[str, Any], Tuple[str, Any]] = InSceneGrouper(pass_images_dp)
+            zipped_dp = IterKeyZipper(
+                flo_dp,
+                pass_images_dp,
+                key_fn=self._flows_key,
+                ref_key_fn=self._images_key,
+            )
+        else:
+            pass_images_dp = Shuffler(filtered_curr_split, buffer_size=INFINITE_BUFFER_SIZE)
+            pass_images_dp = InSceneGrouper(pass_images_dp)
+            zipped_dp = Mapper(pass_images_dp, self._add_fake_flow_data)
+
+        return Mapper(zipped_dp, self._collate_and_decode_sample, fn_kwargs=dict(decoder=decoder))