RoCo Spring Competition - Stereo

This folder contains the dedicated scripts and configurations required to participate in the Stereo category of the RoCo Spring Competition.

Before proceeding, ensure you have completed the environment setup and dataset configuration outlined in the main README.md file.

Navigate to this directory in your terminal before running any scripts:

cd roco_spring_devkit/stereo/

Available Scripts

Three primary scripts are provided to manage your development workflow:

• test.py: Generates the final output predictions required for competition submission.
• train.py: Handles model training or fine-tuning pipelines.
• validate.py: Evaluates your model locally and computes standard performance metrics.

---

test.py

Use this script to run inference on the competition test sets. For example, to generate outputs using the provided baseline RAFT-Stereo model:

python test.py --data.test_dataset spring-robust --model raft_stereo --ckpt_path sceneflow --num_gpus -1 --save_viz

The resulting files will be saved to outputs/test/<model_and_checkpoint_name>.

Key Arguments Breakdown:

• --data.test_dataset spring-robust: Samples from the spring-robust test split are being used to generate the results. You can also check the section Adding a Custom Dataset to learn how to include other datasets.
• --model raft_stereo: Use the RAFT-Stereo stereo model. Check the section Adding a Custom Model to learn how to add your own custom model.
• --ckpt_path sceneflow: The path to the checkpoint to be used by the model. In this example, we use a custom name (sceneflow) that is defined in the pretrained_checkpoints property of the RAFT-Stereo model. If you have a local or online checkpoint, replace sceneflow with the path or the URL of the checkpoint you want to load.
• --num_gpus: Number of GPUs to use for inference, -1 uses all the GPUs. To select specific GPUs, set the environmental variable, for example CUDA_VISIBLE_DEVICES=0,2,4 python test.py ....
• --save_viz: (Optional) Generates image visualizations of your disparity predictions. Helpful for visual sanity checks, but can be omitted to accelerate inference and conserve disk space.

Valid Entries for --data.test_dataset

• spring
• spring-robust
• kitti-2015

To find all available CLI options, run:

python test.py -h

train.py

This script provides a standardized framework for training or fine-tuning models using the devkit's datasets.

To fine-tune the default RAFT-Stereo model using the Spring dataset configuration:

python train.py --config models/raft_stereo/configs/raft_stereo-train-spring.yaml

• Outputs: Logs, telemetry, and model checkpoints are saved automatically inside the ptlflow_logs/ folder.
• Customization: Training configurations are entirely managed via YAML files. To modify learning rates, batch sizes, or augmentations, create or adapt a configuration file. For an in-depth reference, see the PTLFlow Configuration Documentation.

validate.py

This helper script assesses your model's quality on target datasets containing ground-truth data, providing instant metric feedback (e.g., Abs).

To validate the RAFT-Stereo model on the Spring validation split:

python validate.py --data.val_dataset spring-val --model raft_stereo --ckpt_path sceneflow --write_outputs

Outputs containing validation statistics will be written to outputs/validate/<model_and_checkpoint_name>.

Key Arguments Breakdown:

• --data.val_dataset spring: Sets the validation target. You can validate on multiple datasets simultaneously by combining them with a + symbol (e.g., --data.val_dataset kitti-2015+sintel-clean).
• --write_outputs: (Optional) Exports disparity files, heatmaps, and Abs visualizations for error diagnosis.

Valid Entries for --data.val_dataset

• spring
• sintel-clean
• sintel-final
• kitti-2015

⚠️ Note: spring-robust is not a valid entry for validation because it does not provide publicly accessible validation ground truths.

---

Submission Example

This is an example of generating and submitting stereo results to the competition.

1. Run the test script as in the example of test.py above. Suppose the results are saved to

./outputs/test/raft_stereo_sceneflow/spring-robust
├── brightness
│   └── test
│       ├── 0003
│       │   ├── disp1_left
│       │   │   ├── disp1_left_0001.dsp5
│       │   │   ...
│       │   └── disp1_right
│       ...
├── clean
├── contrast
├── [other corruption names]
...

2. Use the Spring subsampling executables.

# Assuming the terminal is at /path/to/roco_spring_devkit/roco_spring_devkit/stereo/

# First, generate the submission to the main Spring benchmark using the "clean" results
/path/to/subsampling_executables/disp1_subsampling outputs/test/raft_stereo_sceneflow/spring-robust/clean/test
# This will generate a file disp1_submission.hdf5

# Then generate the submission to the robustness benchmark
/path/to/subsampling_executables/disp1_robust_subsampling outputs/test/raft_stereo_sceneflow/spring-robust/
# This will generate a file disp1_robustness.hdf5

3. Submit the two respective .hdf5 files to the Spring benchmark website.

---

Adding a Custom Model

To introduce your own model design into the competition pipeline:

1. Create a new directory under roco_spring_devkit/stereo/models/<your_model_name>/.
2. Define your model architecture by inheriting from roco_spring_devkit.stereo.models.base_model.base_model BaseModel class.
   • Be sure to initialize the BaseModel by calling

   super().__init__(
    output_stride=<maximum stride of the model>,
    loss_fn=<None or check the loss_fn instructions below>,
    **kwargs
   )

3. Register your new architecture configuration inside your model initialization script.
   • Import the register_model annotation:

   from roco_spring_devkit.common.utils.registry import register_model

- Annotate your class to register it. For example, if your model is called `MyCustomModel`:

   @register_model
   class MyCustomModel(BaseModel):
    ...

4. Import your custom model in __init__.py to make it visible to the other scripts.

   # In file roco_spring_devkit.stereo.models.__init__.py, add one extra import:
   import roco_spring_devkit.stereo.models.<my_custom_model_folder_name>.<my_custom_model_file_name>

Model I/O

To interact with the existing scripts, your model should support the following I/O:

Input

The model's forward function must receive a dictionary as one of the first arguments. The input images are provided as the dictionary keys:

• "images": A 5D tensor with shape [batch_size, 1, 3, height, width] containing the left image. The second dimension corresponds to the number of input images, and the third to the number of image channels.
• "images_right": A 5D tensor with shape [batch_size, 1, 3, height, width] containing the right image.

Output

Your model must output a dictionary with at least the following keys:

• "disparities": A 5D tensor with shape [batch_size, 1, 1, height, width] containing the estimated disparity from the left to the right image.
• "disparities_right": Similar to "disparities", but contains the estimated disparity from the right to the left image.

loss_fn

If you want to use the train.py script to train your model, you need to provide a function or a callable class to loss_fn during the initialization of your custom model. To work, it must support the following I/O.

Input

It must receive at least two dictionaries as the first two inputs.

• The first dictionary contains the outputs of the model.
• The second dictionary contains the inputs of the model, such as the groundtruth information.

Output

The output must be either a single tensor containing the loss value or a dictionary. If the output is a dictionary, it must contain a key named loss whose value is a tensor containing the loss value.

Additional details about adding a model

For an explicit, step-by-step code walkthrough, please consult the PTLFlow Custom Model Guide.

Adding a Custom Dataset

If you want to train your models on additional external datasets, follow these steps:

1. Define a class for your new dataset in roco_spring_devkit.common.data.scene_flow_datasets.py. For example:

   class MyStereoDataset(BaseSceneFlowDataset):
    ...

• The easiest way is to let your new class inherit from BaseSceneFlowDataset to let it handle the reading. However, you can also create your own full pipeline.
• If inheriting from BaseSceneFlowDataset, your class should fill the following lists:
  • self.img_paths: Each element is a list with one element, the path to the left image.
  • self.img_r_paths: Each element is a list with one element, the path to the right image.
  • self.disp_paths: (Optional) Each element is a list with one element, the path to the groundtruth disparity between the left and right image.
  • self.disp_r_paths: (Optional) Each element is a list with one element, the path to the groundtruth disparity between the right and left image.

2. Add a new entry to datasets.yaml.
   • For example, add my_stereo: /path/to/my_stereo.
3. Define a new loader in roco_spring_devkit.common.data.stereo_datamodule.py.
   • Add one input argument to StereoDataModule. The argument must be called my_stereo_dataset_root_dir, where my_stereo is the same name defined in datasets.yaml.
   • Add the new argument as a class variable: self.my_stereo_dataset_root_dir = my_stereo_dataset_root_dir.
   • Create a new class method called

   def _get_my_stereo_dataset(self, is_train: bool, *args: str):
    ...

This method should return an instance of MyStereoDataset.

You can find more layout instructions and boilerplate code requirements in the PTLFlow Custom Dataset Guide.