Add FlyingChairs dataset for optical flow

test/datasets_utils.py

@@ -8,6 +8,7 @@
 import random
 import shutil
 import string
+import struct
 import tarfile
 import unittest
 import unittest.mock
@@ -922,3 +923,11 @@ def create_random_string(length: int, *digits: str) -> str:
         digits = "".join(itertools.chain(*digits))
 
     return "".join(random.choice(digits) for _ in range(length))
+
+
+def make_fake_flo_file(h, w, file_name):
+    """Creates a fake flow file in .flo format."""
+    values = list(range(2 * h * w))
+    content = b"PIEH" + struct.pack("i", w) + struct.pack("i", h) + struct.pack("f" * len(values), *values)
+    with open(file_name, "wb") as f:
+        f.write(content)

test/test_datasets.py

@@ -1874,11 +1874,9 @@ def _inject_pairs(self, root, num_pairs, same):
 class SintelTestCase(datasets_utils.ImageDatasetTestCase):
     DATASET_CLASS = datasets.Sintel
     ADDITIONAL_CONFIGS = datasets_utils.combinations_grid(split=("train", "test"), pass_name=("clean", "final"))
-    # We patch the flow reader, because this would otherwise force us to generate fake (but readable) .flo files,
-    # which is something we want to # avoid.
-    _FAKE_FLOW = "Fake Flow"
-    EXTRA_PATCHES = {unittest.mock.patch("torchvision.datasets.Sintel._read_flow", return_value=_FAKE_FLOW)}
-    FEATURE_TYPES = (PIL.Image.Image, PIL.Image.Image, (type(_FAKE_FLOW), type(None)))
+    FEATURE_TYPES = (PIL.Image.Image, PIL.Image.Image, (np.ndarray, type(None)))
+
+    FLOW_H, FLOW_W = 3, 4
 
     def inject_fake_data(self, tmpdir, config):
         root = pathlib.Path(tmpdir) / "Sintel"
@@ -1899,14 +1897,13 @@ def inject_fake_data(self, tmpdir, config):
                         num_examples=num_images_per_scene,
                     )
 
-        # For the ground truth flow value we just create empty files so that they're properly discovered,
-        # see comment above about EXTRA_PATCHES
         flow_root = root / "training" / "flow"
         for scene_id in range(num_scenes):
             scene_dir = flow_root / f"scene_{scene_id}"
             os.makedirs(scene_dir)
             for i in range(num_images_per_scene - 1):
-                open(str(scene_dir / f"frame_000{i}.flo"), "a").close()
+                file_name = str(scene_dir / f"frame_000{i}.flo")
+                datasets_utils.make_fake_flo_file(h=self.FLOW_H, w=self.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
@@ -1920,7 +1917,8 @@ def test_flow(self):
         with self.create_dataset(split="train") as (dataset, _):
             assert dataset._flow_list and len(dataset._flow_list) == len(dataset._image_list)
             for _, _, flow in dataset:
-                assert flow == self._FAKE_FLOW
+                assert flow.shape == (2, self.FLOW_H, self.FLOW_W)
+                np.testing.assert_allclose(flow, np.arange(flow.size).reshape(flow.shape))
 
         # Make sure flow is always None for test split
         with self.create_dataset(split="test") as (dataset, _):
@@ -2001,11 +1999,9 @@ def test_bad_input(self):
 class FlyingChairsTestCase(datasets_utils.ImageDatasetTestCase):
     DATASET_CLASS = datasets.FlyingChairs
     ADDITIONAL_CONFIGS = datasets_utils.combinations_grid(split=("train", "val"))
-    # We patch the flow reader, because this would otherwise force us to generate fake (but readable) .flo files,
-    # which is something we want to avoid.
-    _FAKE_FLOW = "Fake Flow"
-    EXTRA_PATCHES = {unittest.mock.patch("torchvision.datasets.FlyingChairs._read_flow", return_value=_FAKE_FLOW)}
-    FEATURE_TYPES = (PIL.Image.Image, PIL.Image.Image, (type(_FAKE_FLOW), type(None)))
+    FEATURE_TYPES = (PIL.Image.Image, PIL.Image.Image, (np.ndarray, type(None)))
+
+    FLOW_H, FLOW_W = 3, 4
 
     def _make_split_file(self, root, num_examples):
         # We create a fake split file here, but users are asked to download the real one from the authors website
@@ -2033,10 +2029,9 @@ def inject_fake_data(self, tmpdir, config):
             file_name_fn=lambda image_idx: f"00{image_idx}_img2.ppm",
             num_examples=num_examples_total,
         )
-        # For the ground truth flow value we just create empty files so that they're properly discovered,
-        # see comment above about EXTRA_PATCHES
         for i in range(num_examples_total):
-            open(str(root / "data" / f"00{i}_flow.flo"), "a").close()
+            file_name = str(root / "data" / f"00{i}_flow.flo")
+            datasets_utils.make_fake_flo_file(h=self.FLOW_H, w=self.FLOW_W, file_name=file_name)
 
         self._make_split_file(root, num_examples)
 
@@ -2045,10 +2040,12 @@ def inject_fake_data(self, tmpdir, config):
     @datasets_utils.test_all_configs
     def test_flow(self, config):
         # Make sure flow always exists, and make sure there are as many flow values as (pairs of) images
+        # Also make sure the flow is properly decoded
         with self.create_dataset(config=config) as (dataset, _):
             assert dataset._flow_list and len(dataset._flow_list) == len(dataset._image_list)
             for _, _, flow in dataset:
-                assert flow == self._FAKE_FLOW
+                assert flow.shape == (2, self.FLOW_H, self.FLOW_W)
+                np.testing.assert_allclose(flow, np.arange(flow.size).reshape(flow.shape))
 
 
 if __name__ == "__main__":