RoCo Spring Competition - Optical Flow

This folder contains the dedicated scripts and configurations required to participate in the Optical Flow 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/optical_flow/

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 model:

python test.py --data.test_dataset spring-robust --model raft --ckpt_path sintel --model.corr_mode triton --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: Use the RAFT optical flow model. Check the section Adding a Custom Model to learn how to add your own custom model.
• --ckpt_path sintel: The path to the checkpoint to be used by the model. In this example, we use a custom name (sintel) that is defined in the pretrained_checkpoints property of the RAFT model. If you have a local or online checkpoint, replace sintel 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 flow 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
• sintel
• 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 model using the Spring dataset configuration:

python train.py --config models/raft/configs/raft-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., EPE).

To validate the RAFT model on the Spring validation split:

python validate.py --data.val_dataset spring-val --model raft --ckpt_path sintel --model.corr_mode triton --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 flow files, heatmaps, and EPE 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 optical flow 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_sintel/spring-robust
├── brightness
│   └── test
│       ├── 0003
│       │   ├── flow_BW_left
│       │   │   ├── flow_BW_left_0001.flo5
│       │   │   ...
│       │   ├── flow_BW_right
│       │   ├── flow_FW_left
│       │   └── flow_FW_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/optical_flow/

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

# Then generate the submission to the robustness benchmark
/path/to/subsampling_executables/flow_robust_subsampling outputs/test/raft_sintel/spring-robust/
# This will generate a file flow_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/optical_flow/models/<your_model_name>/.
2. Define your model architecture by inheriting from roco_spring_devkit.optical_flow.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.optical_flow.models.__init__.py, add one extra import:
   import roco_spring_devkit.optical_flow.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, 2, 3, height, width]. The second dimension corresponds to the number of input images, and the third to the number of image channels.

Output

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

• "flows": A 5D tensor with shape [batch_size, 1, 2, height, width] containing the estimated forwarded flow for the left camera (the "images" input). The second dimension corresponds to the number of output flows, and the third dimension to the x and y flow channels.
• "flows_b": Similar to "flows", but contains the estimated backward flow (from the second image to the first) of the "images" input.

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.optical_flow_datasets.py. For example:

   class MyFlowDataset(BaseFlowDataset):
    ...

• The easiest way is to let your new class inherit from BaseFlowDataset to let it handle the reading. However, you can also create your own full pipeline.
• If inheriting from BaseFlowDataset, your class should fill the following lists:
  • self.img_paths: Each element is a list with two elements, the paths to the first and second images.
  • self.flow_paths: (Optional) Each element is a list with one element, the path to the groundtruth flow between the first and second image.
  • self.flow_b_paths: (Optional) Each element is a list with one element, the path to the groundtruth flow between the second and first image.

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

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

This method should return an instance of MyFlowDataset.

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