Add crestereo dataset

test/datasets_utils.py

@@ -16,6 +16,7 @@
 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
@@ -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,38 @@ def create_random_string(length: int, *digits: str) -> str:
     return "".join(random.choice(digits) for _ in range(length))
 
 
+def shape_test_for_stereo_disp(
+    left: PIL.Image.Image, right: PIL.Image.Image, disparity: np.ndarray, valid_mask: np.ndarray
+):
+    left_array = np.array(left)
+    right_array = np.array(right)
+    h, w, c = left_array.shape
+    # check that left and right are the same size
+    assert left_array.shape == right_array.shape
+    # check general shapes
+    assert c == 3
+    assert len(disparity.shape) == 3
+    assert len(valid_mask.shape) == 2
+    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_none(left: PIL.Image.Image, right: PIL.Image.Image, disparity: None, valid_mask: None):
+    left_array = np.array(left)
+    right_array = np.array(right)
+    _, _, c = left_array.shape
+    # check that left and right are the same size
+    assert left_array.shape == right_array.shape
+    # check general shapes
+    assert c == 3
+    assert disparity is None
+    assert valid_mask 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 "<"