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/cifar.py

@@ -95,7 +95,7 @@ def _make_datapipe(
     def _generate_categories(self, root: pathlib.Path) -> List[str]:
         dp = self.resources(self.default_config)[0].to_datapipe(pathlib.Path(root) / self.name)
         dp = TarArchiveReader(dp)
-        dp = Filter(dp, path_comparator("name", self._META_FILE_NAME))
+        dp = Filter(dp, path_comparator("name", value=self._META_FILE_NAME))
         dp = Mapper(dp, self._unpickle)
         return cast(List[str], next(iter(dp))[self._CATEGORIES_KEY])
 

torchvision/prototype/datasets/_builtin/coco.py

@@ -115,7 +115,7 @@ def _make_datapipe(
         images_dp, meta_dp = resource_dps
 
         meta_dp = ZipArchiveReader(meta_dp)
-        meta_dp = Filter(meta_dp, path_comparator("name", f"instances_{config.split}{config.year}.json"))
+        meta_dp = Filter(meta_dp, path_comparator("name", value=f"instances_{config.split}{config.year}.json"))
         meta_dp = JsonParser(meta_dp)
         meta_dp = Mapper(meta_dp, getitem(1))
         meta_dp = MappingIterator(meta_dp)

torchvision/prototype/datasets/_builtin/imagenet.py

@@ -147,7 +147,7 @@ def _make_datapipe(
             dp = Mapper(dp, self._collate_train_data)
         elif config.split == "val":
             devkit_dp = TarArchiveReader(devkit_dp)
-            devkit_dp = Filter(devkit_dp, path_comparator("name", "ILSVRC2012_validation_ground_truth.txt"))
+            devkit_dp = Filter(devkit_dp, path_comparator("name", value="ILSVRC2012_validation_ground_truth.txt"))
             devkit_dp = LineReader(devkit_dp, return_path=False)
             devkit_dp = Mapper(devkit_dp, int)
             devkit_dp = Enumerator(devkit_dp, 1)
@@ -178,7 +178,7 @@ def _generate_categories(self, root: pathlib.Path) -> List[Tuple[str, ...]]:
         resources = self.resources(self.default_config)
         devkit_dp = resources[1].to_datapipe(root / self.name)
         devkit_dp = TarArchiveReader(devkit_dp)
-        devkit_dp = Filter(devkit_dp, path_comparator("name", "meta.mat"))
+        devkit_dp = Filter(devkit_dp, path_comparator("name", value="meta.mat"))
 
         meta = next(iter(devkit_dp))[1]
         synsets = read_mat(meta, squeeze_me=True)["synsets"]

torchvision/prototype/datasets/_builtin/sbd.py

@@ -156,7 +156,7 @@ def _make_datapipe(
     def _generate_categories(self, root: pathlib.Path) -> Tuple[str, ...]:
         dp = self.resources(self.default_config)[0].to_datapipe(pathlib.Path(root) / self.name)
         dp = TarArchiveReader(dp)
-        dp = Filter(dp, path_comparator("name", "category_names.m"))
+        dp = Filter(dp, path_comparator("name", value="category_names.m"))
         dp = LineReader(dp)
         dp = Mapper(dp, bytes.decode, input_col=1)
         lines = tuple(zip(*iter(dp)))[1]

torchvision/prototype/datasets/_builtin/sintel.py

@@ -0,0 +1,146 @@
+import io
+import pathlib
+import re
+from functools import partial
+from typing import Any, Callable, Dict, List, Optional, Tuple, 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, InScenePairer, path_accessor
+
+
+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[None, None], 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]] = InScenePairer(
+                pass_images_dp, scene_fn=path_accessor("parent", "name")
+            )
+            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 = InScenePairer(pass_images_dp, scene_fn=path_accessor("parent", "name"))
+            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))

torchvision/prototype/datasets/_builtin/voc.py

@@ -129,7 +129,7 @@ def _make_datapipe(
         )
 
         split_dp = Filter(split_dp, self._is_in_folder, fn_kwargs=dict(name=self._SPLIT_FOLDER[config.task]))
-        split_dp = Filter(split_dp, path_comparator("name", f"{config.split}.txt"))
+        split_dp = Filter(split_dp, path_comparator("name", value=f"{config.split}.txt"))
         split_dp = LineReader(split_dp, decode=True)
         split_dp = Shuffler(split_dp, buffer_size=INFINITE_BUFFER_SIZE)
 

torchvision/prototype/datasets/utils/_internal.py

@@ -21,6 +21,7 @@
     Optional,
     IO,
     Sized,
+    Iterable,
 )
 from typing import cast
 
@@ -33,6 +34,19 @@
 from torchdata.datapipes.iter import IoPathFileLister, IoPathFileLoader
 from torchdata.datapipes.utils import StreamWrapper
 
+K = TypeVar("K")
+D = TypeVar("D")
+
+try:
+    from itertools import pairwise  # type: ignore[attr-defined]
+except ImportError:
+    from itertools import tee
+
+    def pairwise(iterable: Iterable[D]) -> Iterable[Tuple[D, D]]:
+        a, b = tee(iterable)
+        next(b, None)
+        return zip(a, b)
+
 
 __all__ = [
     "INFINITE_BUFFER_SIZE",
@@ -48,11 +62,9 @@
     "Decompressor",
     "fromfile",
     "read_flo",
+    "InScenePairer",
 ]
 
-K = TypeVar("K")
-D = TypeVar("D")
-
 # pseudo-infinite until a true infinite buffer is supported by all datapipes
 INFINITE_BUFFER_SIZE = 1_000_000_000
 
@@ -117,17 +129,22 @@ def getitem(*items: Any) -> Callable[[Any], Any]:
     return functools.partial(_getitem_closure, items=items)
 
 
-def _path_attribute_accessor(path: pathlib.Path, *, name: str) -> D:
-    return cast(D, getattr(path, name))
+def _path_attribute_accessor(path: pathlib.Path, *, attrs: Sequence[str]) -> D:
+    obj: Any = path
+    for attr in attrs:
+        obj = getattr(obj, attr)
+    return cast(D, obj)
 
 
 def _path_accessor_closure(data: Tuple[str, Any], *, getter: Callable[[pathlib.Path], D]) -> D:
     return getter(pathlib.Path(data[0]))
 
 
-def path_accessor(getter: Union[str, Callable[[pathlib.Path], D]]) -> Callable[[Tuple[str, Any]], D]:
-    if isinstance(getter, str):
-        getter = functools.partial(_path_attribute_accessor, name=getter)
+def path_accessor(*attrs: Union[str, Callable[[pathlib.Path], D]]) -> Callable[[Tuple[str, Any]], D]:
+    if not callable(attrs[0]):
+        getter = cast(Callable[[pathlib.Path], D], functools.partial(_path_attribute_accessor, attrs=attrs))
+    else:
+        getter = attrs[0]
 
     return functools.partial(_path_accessor_closure, getter=getter)
 
@@ -136,8 +153,8 @@ def _path_comparator_closure(data: Tuple[str, Any], *, accessor: Callable[[Tuple
     return accessor(data) == value
 
 
-def path_comparator(getter: Union[str, Callable[[pathlib.Path], D]], value: D) -> Callable[[Tuple[str, Any]], bool]:
-    return functools.partial(_path_comparator_closure, accessor=path_accessor(getter), value=value)
+def path_comparator(*attrs: Union[str, Callable[[pathlib.Path], D]], value: D) -> Callable[[Tuple[str, Any]], bool]:
+    return functools.partial(_path_comparator_closure, accessor=path_accessor(*attrs), value=value)
 
 
 class CompressionType(enum.Enum):
@@ -321,3 +338,26 @@ def read_flo(file: BinaryIO) -> torch.Tensor:
     width, height = fromfile(file, dtype=torch.int32, byte_order="little", count=2)
     flow = fromfile(file, dtype=torch.float32, byte_order="little", count=height * width * 2)
     return flow.reshape((height, width, 2)).permute((2, 0, 1))
+
+
+class InScenePairer(IterDataPipe[Tuple[Tuple[str, BinaryIO], Tuple[str, BinaryIO]]]):
+    def __init__(
+        self, datapipe: IterDataPipe[Tuple[str, BinaryIO]], *, scene_fn: Callable[[Tuple[str, BinaryIO]], K]
+    ) -> None:
+        self.datapipe = datapipe
+        self.scene_fn = scene_fn
+
+    def __iter__(self) -> Iterator[Tuple[Tuple[str, BinaryIO], Tuple[str, BinaryIO]]]:
+        prev_bytes = b""
+        for (path1, stream1), (path2, stream2) in pairwise(sorted(self.datapipe)):
+            if self.scene_fn(path1) != self.scene_fn(path2):
+                prev_bytes = b""
+                continue
+
+            buffer1 = io.BytesIO(prev_bytes) if prev_bytes else stream1
+            prev_bytes = stream2.read()
+            if prev_bytes == b"":
+                print()
+            buffer2 = io.BytesIO(prev_bytes)
+
+            yield (path1, buffer1), (path2, buffer2)