Fixes2

conftest.py

@@ -0,0 +1,24 @@
+import pytest
+
+
+def pytest_addoption(parser):
+    parser.addoption(
+        "--runslow", action="store_true", default=False, help="run slow tests"
+    )
+
+
+def pytest_configure(config):
+    config.addinivalue_line(
+        "markers",
+        "slow: marks tests as slow (deselect with '-m \"not slow\"' or run only with --runslow)",
+    )
+
+
+def pytest_collection_modifyitems(config, items):
+    if config.getoption("--runslow"):
+        # --runslow flag was passed; don't skip slow tests
+        return
+    skip_slow = pytest.mark.skip(reason="need --runslow flag to run")
+    for item in items:
+        if "slow" in item.keywords:
+            item.add_marker(skip_slow)

pyproject.toml

@@ -44,6 +44,7 @@ dependencies = [
     "tabulate<0.11",
     "tensorboard<2.21",
     "timm<1.1",
+    "triton<3.8",
 ]
 
 [project.urls]

requirements.txt

@@ -11,4 +11,5 @@ requests<2.34
 scipy<1.18
 tabulate<0.11
 tensorboard<2.21
-timm<1.1
+timm<1.1
+triton<3.8

roco_spring_devkit/common/data/Spring_val.txt

@@ -1 +1 @@
-0027
+0022

roco_spring_devkit/common/data/scene_flow_datamodule.py

@@ -49,6 +49,7 @@ def __init__(
         spring_root_dir: Optional[str] = None,
         robust_spring_root_dir: Optional[str] = None,
         dataset_config_path: str = "../../datasets.yaml",
+        disparity2_in_frame1: bool = True,
     ):
         super().__init__()
         self.predict_dataset = predict_dataset
@@ -67,6 +68,9 @@ def __init__(
         self.spring_root_dir = spring_root_dir
         self.robust_spring_root_dir = robust_spring_root_dir
         self.dataset_config_path = dataset_config_path
+        # All the datasets return the second-frame disparity ('disparities'[1]) on the pixel grid of the first frame,
+        # which is the convention of the KITTI/Spring scene flow benchmarks and of the models in this repository.
+        self.disparity2_in_frame1 = disparity2_in_frame1
 
         self.predict_dataset_parsed = None
         self.test_dataset_parsed = None
@@ -321,7 +325,10 @@ def _get_kitti_dataset(self, is_train: bool, *args: str) -> Dataset:
             # These transforms are based on RAFT: https://github.com/princeton-vl/RAFT
             transform = ft.Compose(
                 [
-                    ft.ToTensor(device=device, fp16=self.train_transform_fp16),
+                    ft.ToTensor(
+                        device=device,
+                        fp16=self.train_transform_fp16,
+                    ),
                     ft.RandomScaleAndCrop(
                         (cy, cx), (-0.2, 0.4), (-0.2, 0.2), sparse=True
                     ),
@@ -342,6 +349,7 @@ def _get_kitti_dataset(self, is_train: bool, *args: str) -> Dataset:
             get_flow=True,
             get_disparity=True,
             get_intrinsics=True,
+            disparity2_in_frame1=self.disparity2_in_frame1,
         )
         return dataset
 
@@ -378,7 +386,10 @@ def _get_sintel_dataset(self, is_train: bool, *args: str) -> Dataset:
             # These transforms are based on RAFT: https://github.com/princeton-vl/RAFT
             transform = ft.Compose(
                 [
-                    ft.ToTensor(device=device, fp16=self.train_transform_fp16),
+                    ft.ToTensor(
+                        device=device,
+                        fp16=self.train_transform_fp16,
+                    ),
                     ft.RandomScaleAndCrop((cy, cx), (-0.2, 0.6), (-0.2, 0.2)),
                     ft.ColorJitter(0.4, 0.4, 0.4, 0.5 / 3.14, 0.2),
                     ft.GaussianNoise(0.02),
@@ -399,6 +410,7 @@ def _get_sintel_dataset(self, is_train: bool, *args: str) -> Dataset:
             get_flow=True,
             get_disparity=True,
             get_intrinsics=True,
+            disparity2_in_frame1=self.disparity2_in_frame1,
         )
         return dataset
 
@@ -450,7 +462,10 @@ def _get_spring_dataset(self, is_train: bool, *args: str) -> Dataset:
             # Transforms copied from SEA-RAFT
             transform = ft.Compose(
                 [
-                    ft.ToTensor(device=device, fp16=self.train_transform_fp16),
+                    ft.ToTensor(
+                        device=device,
+                        fp16=self.train_transform_fp16,
+                    ),
                     ft.RandomScaleAndCrop((cy, cx), (0.0, 0.2), (-0.2, 0.2)),
                     ft.ColorJitter(0.4, 0.4, 0.4, 0.5 / 3.14, 0.2),
                     ft.GaussianNoise(0.02),
@@ -475,6 +490,7 @@ def _get_spring_dataset(self, is_train: bool, *args: str) -> Dataset:
             get_flow=True,
             get_disparity=True,
             get_intrinsics=True,
+            disparity2_in_frame1=self.disparity2_in_frame1,
             robust_mode=robust_mode,
             robust_root_dir=self.robust_spring_root_dir,
         )
@@ -487,13 +503,22 @@ def _get_things_dataset(self, is_train: bool, *args: str) -> Dataset:
         pass_names = ["clean", "final"]
         split = "trainval"
         sintel_transform = False
+        add_reverse = False
+        get_right_disparity = False
+        get_right_flow = False
         for v in args:
             if v in ["clean", "final"]:
                 pass_names = [v]
             elif v in ["train", "val", "trainval"]:
                 split = v
             elif v == "sinteltransform":
                 sintel_transform = True
+            elif v == "rev":
+                add_reverse = True
+            elif v == "rdisp":
+                get_right_disparity = True
+            elif v == "rflow":
+                get_right_flow = True
             else:
                 raise ValueError(f"Invalid arg: {v}")
 
@@ -521,7 +546,10 @@ def _get_things_dataset(self, is_train: bool, *args: str) -> Dataset:
                 major_scale = (-0.4, 0.8)
             transform = ft.Compose(
                 [
-                    ft.ToTensor(device=device, fp16=self.train_transform_fp16),
+                    ft.ToTensor(
+                        device=device,
+                        fp16=self.train_transform_fp16,
+                    ),
                     ft.RandomScaleAndCrop((cy, cx), major_scale, (-0.2, 0.2)),
                     ft.ColorJitter(0.4, 0.4, 0.4, 0.5 / 3.14, 0.2),
                     ft.GaussianNoise(0.02),
@@ -542,5 +570,9 @@ def _get_things_dataset(self, is_train: bool, *args: str) -> Dataset:
             get_flow=True,
             get_disparity=True,
             get_intrinsics=True,
+            add_reverse=add_reverse,
+            get_right_disparity=get_right_disparity,
+            get_right_flow=get_right_flow,
+            disparity2_in_frame1=self.disparity2_in_frame1,
         )
         return dataset