Add crestereo dataset

test/datasets_utils.py

@@ -16,14 +16,15 @@
 from collections import defaultdict
 from typing import Any, Callable, Dict, Iterator, List, Optional, Sequence, Tuple, Union
 
+import numpy as np
 import PIL
 import PIL.Image
 import pytest
 import torch
 import torchvision.datasets
 import torchvision.io
 from common_utils import get_tmp_dir, disable_console_output
-
+from torchvision.transforms.functional import get_dimensions
 
 __all__ = [
     "UsageError",
@@ -561,7 +562,9 @@ def test_feature_types(self, config):
     @test_all_configs
     def test_num_examples(self, config):
         with self.create_dataset(config) as (dataset, info):
-            assert len(dataset) == info["num_examples"]
+            assert (
+                len(dataset) == info["num_examples"]
+            ), f"The number of examples {len(dataset)} does not match the expected {info['num_examples']}"
 
     @test_all_configs
     def test_transforms(self, config):
@@ -931,6 +934,48 @@ def create_random_string(length: int, *digits: str) -> str:
     return "".join(random.choice(digits) for _ in range(length))
 
 
+def shape_test_for_stereo_gt_w_mask(
+    left: PIL.Image.Image, right: PIL.Image.Image, disparity: np.ndarray, valid_mask: np.ndarray
+):
+    left_dims = get_dimensions(left)
+    right_dims = get_dimensions(right)
+    c, h, w = left_dims
+    # check that left and right are the same size
+    assert left_dims == right_dims
+    # check general shapes
+    assert c == 3
+    assert disparity.ndim == 3
+    assert disparity.shape == (1, h, w)
+    # check that valid mask is the same size as the disparity
+    _, dh, dw = disparity.shape
+    mh, mw = valid_mask.shape
+    assert dh == mh
+    assert dw == mw
+
+
+def shape_test_for_stereo_gt_no_mask(left: PIL.Image.Image, right: PIL.Image.Image, disparity: np.ndarray):
+    left_dims = get_dimensions(left)
+    right_dims = get_dimensions(right)
+    c, h, w = left_dims
+    # check that left and right are the same size
+    assert left_dims == right_dims
+    # check general shapes
+    assert c == 3
+    assert disparity.ndim == 3
+    assert disparity.shape == (1, h, w)
+
+
+def shape_test_for_stereo_no_gt(left: PIL.Image.Image, right: PIL.Image.Image, disparity: None):
+    left_dims = get_dimensions(left)
+    right_dims = get_dimensions(right)
+    c, _, _ = left_dims
+    # check that left and right are the same size
+    assert left_dims == right_dims
+    # check general shapes
+    assert c == 3
+    assert disparity is None
+
+
 def make_fake_pfm_file(h, w, file_name):
     values = list(range(3 * h * w))
     # Note: we pack everything in little endian: -1.0, and "<"