[fbsync] Added InStereo2k dataset (#6347)

Summary:
* added SceneFLow variant datasets

* Changed split name to variant name

* removed trailing commented code line

* Added InStereo2k dataset

* Added Sintel Stereo dataset

* small refactor in tests

* Fixed doc formatting.

* candidate fix for FileNotFound on windows test

* Adressing comments

* Added Sintel Stereo dataset

* small refactor in tests

* Fixed doc formatting.

* candidate fix for FileNotFound on windows test

* Adressing comments

* rebased on main

* lint fix

* Added InStereo2k dataset

Reviewed By: datumbox

Differential Revision: D38824220

fbshipit-source-id: fa9ef743affa5f10a4a0a1ba130deccb6df859ca

Author: NicolasHug

docs/source/datasets.rst

@@ -113,7 +113,7 @@ Stereo Matching
     Kitti2015Stereo
     SceneFlowStereo
     SintelStereo
-
+    InStereo2k
 
 Image pairs
 ~~~~~~~~~~~

test/test_datasets.py

@@ -2863,6 +2863,11 @@ def inject_fake_data(self, tmpdir, config):
         os.makedirs(scene_flow_dir, exist_ok=True)
 
         variant_dir = scene_flow_dir / config["variant"]
+        variant_dir_prefixes = {
+            "Monkaa": 0,
+            "Driving": 2,
+            "FlyingThings3D": 2,
+        }
         os.makedirs(variant_dir, exist_ok=True)
 
         num_examples = {"FlyingThings3D": 4, "Driving": 6, "Monkaa": 5}.get(config["variant"], 0)
@@ -2880,6 +2885,12 @@ def inject_fake_data(self, tmpdir, config):
             os.makedirs(pass_dir, exist_ok=True)
             os.makedirs(disp_dir, exist_ok=True)
 
+            for i in range(variant_dir_prefixes.get(config["variant"], 0)):
+                pass_dir = pass_dir / str(i)
+                disp_dir = disp_dir / str(i)
+                os.makedirs(pass_dir, exist_ok=True)
+                os.makedirs(disp_dir, exist_ok=True)
+
             for direction in ["left", "right"]:
                 for scene_idx in range(num_examples):
                     os.makedirs(pass_dir / f"scene_{scene_idx:06d}", exist_ok=True)
@@ -2916,6 +2927,49 @@ def test_bad_input(self):
                 pass
 
 
+class InStereo2k(datasets_utils.ImageDatasetTestCase):
+    DATASET_CLASS = datasets.InStereo2k
+    FEATURE_TYPES = (PIL.Image.Image, PIL.Image.Image, (np.ndarray, type(None)))
+    ADDITIONAL_CONFIGS = datasets_utils.combinations_grid(split=("train", "test"))
+
+    @staticmethod
+    def _make_scene_folder(root: str, name: str, size: Tuple[int, int]):
+        root = pathlib.Path(root) / name
+        os.makedirs(root, exist_ok=True)
+
+        datasets_utils.create_image_file(root=root, name="left.png", size=(3, size[0], size[1]))
+        datasets_utils.create_image_file(root=root, name="right.png", size=(3, size[0], size[1]))
+        datasets_utils.create_image_file(root=root, name="left_disp.png", size=(1, size[0], size[1]))
+        datasets_utils.create_image_file(root=root, name="right_disp.png", size=(1, size[0], size[1]))
+
+    def inject_fake_data(self, tmpdir, config):
+        in_stereo_dir = pathlib.Path(tmpdir) / "InStereo2k"
+        os.makedirs(in_stereo_dir, exist_ok=True)
+
+        split_dir = pathlib.Path(in_stereo_dir) / config["split"]
+        os.makedirs(split_dir, exist_ok=True)
+
+        num_examples = {"train": 4, "test": 5}.get(config["split"], 0)
+
+        for i in range(num_examples):
+            self._make_scene_folder(split_dir, f"scene_{i:06d}", (100, 200))
+
+        return num_examples
+
+    def test_splits(self):
+        for split_name in ["train", "test"]:
+            with self.create_dataset(split=split_name) as (dataset, _):
+                for left, right, disparity in dataset:
+                    datasets_utils.shape_test_for_stereo(left, right, disparity)
+
+    def test_bad_input(self):
+        with pytest.raises(
+            ValueError, match="Unknown value 'bad' for argument split. Valid values are {'train', 'test'}."
+        ):
+            with self.create_dataset(split="bad"):
+                pass
+
+
 class SintelStereoTestCase(datasets_utils.ImageDatasetTestCase):
     DATASET_CLASS = datasets.SintelStereo
     ADDITIONAL_CONFIGS = datasets_utils.combinations_grid(pass_name=("final", "clean", "both"))

torchvision/datasets/__init__.py

@@ -1,5 +1,5 @@
 from ._optical_flow import FlyingChairs, FlyingThings3D, HD1K, KittiFlow, Sintel
-from ._stereo_matching import CarlaStereo, Kitti2012Stereo, Kitti2015Stereo, SceneFlowStereo, SintelStereo
+from ._stereo_matching import CarlaStereo, InStereo2k, Kitti2012Stereo, Kitti2015Stereo, SceneFlowStereo, SintelStereo
 from .caltech import Caltech101, Caltech256
 from .celeba import CelebA
 from .cifar import CIFAR10, CIFAR100
@@ -111,4 +111,5 @@
     "CarlaStereo",
     "SceneFlowStereo",
     "SintelStereo",
+    "InStereo2k",
 )

torchvision/datasets/_stereo_matching.py

@@ -438,13 +438,19 @@ def __init__(
 
         root = root / variant
 
+        prefix_directories = {
+            "Monkaa": Path("*"),
+            "FlyingThings3D": Path("*") / "*" / "*",
+            "Driving": Path("*") / "*" / "*",
+        }
+
         for p in passes:
-            left_image_pattern = str(root / p / "*" / "left" / "*.png")
-            right_image_pattern = str(root / p / "*" / "right" / "*.png")
+            left_image_pattern = str(root / p / prefix_directories[variant] / "left" / "*.png")
+            right_image_pattern = str(root / p / prefix_directories[variant] / "right" / "*.png")
             self._images += self._scan_pairs(left_image_pattern, right_image_pattern)
 
-            left_disparity_pattern = str(root / "disparity" / "*" / "left" / "*.pfm")
-            right_disparity_pattern = str(root / "disparity" / "*" / "right" / "*.pfm")
+            left_disparity_pattern = str(root / "disparity" / prefix_directories[variant] / "left" / "*.pfm")
+            right_disparity_pattern = str(root / "disparity" / prefix_directories[variant] / "right" / "*.pfm")
             self._disparities += self._scan_pairs(left_disparity_pattern, right_disparity_pattern)
 
     def _read_disparity(self, file_path: str) -> Tuple:
@@ -588,3 +594,72 @@ def __getitem__(self, index: int) -> Tuple:
             the valid_mask is a numpy array of shape (H, W).
         """
         return super().__getitem__(index)
+
+
+class InStereo2k(StereoMatchingDataset):
+    """`InStereo2k <https://github.com/YuhuaXu/StereoDataset>`_ dataset.
+
+    The dataset is expected to have the following structre: ::
+
+        root
+            InStereo2k
+                train
+                    scene1
+                        left.png
+                        right.png
+                        left_disp.png
+                        right_disp.png
+                        ...
+                    scene2
+                    ...
+                test
+                    scene1
+                        left.png
+                        right.png
+                        left_disp.png
+                        right_disp.png
+                        ...
+                    scene2
+                    ...
+
+    Args:
+        root (string): Root directory where InStereo2k is located.
+        split (string): Either "train" or "test".
+        transforms (callable, optional): A function/transform that takes in a sample and returns a transformed version.
+    """
+
+    def __init__(self, root: str, split: str = "train", transforms: Optional[Callable] = None):
+        super().__init__(root, transforms)
+
+        root = Path(root) / "InStereo2k" / split
+
+        verify_str_arg(split, "split", valid_values=("train", "test"))
+
+        left_img_pattern = str(root / "*" / "left.png")
+        right_img_pattern = str(root / "*" / "right.png")
+        self._images = self._scan_pairs(left_img_pattern, right_img_pattern)
+
+        left_disparity_pattern = str(root / "*" / "left_disp.png")
+        right_disparity_pattern = str(root / "*" / "right_disp.png")
+        self._disparities = self._scan_pairs(left_disparity_pattern, right_disparity_pattern)
+
+    def _read_disparity(self, file_path: str) -> Tuple:
+        disparity_map = np.asarray(Image.open(file_path), dtype=np.float32)
+        # unsqueeze disparity to (C, H, W)
+        disparity_map = disparity_map[None, :, :] / 1024.0
+        valid_mask = None
+        return disparity_map, valid_mask
+
+    def __getitem__(self, index: int) -> Tuple:
+        """Return example at given index.
+
+        Args:
+            index(int): The index of the example to retrieve
+
+        Returns:
+            tuple: A 3-tuple with ``(img_left, img_right, disparity)``.
+            The disparity is a numpy array of shape (1, H, W) and the images are PIL images.
+            If a ``valid_mask`` is generated within the ``transforms`` parameter,
+            a 4-tuple with ``(img_left, img_right, disparity, valid_mask)`` is returned.
+        """
+        return super().__getitem__(index)