DVR: Multi-Modal Deformable Image Registration of the Brain and Abdomen with Reinforcement Learning

This repository contains the files necessary to execute the baseline, DVR, and ablation models involved in the project. Refer to the following sections for instructions regarding how to navigate the different files and implement intended functions.

Note: For each slurm file, you may need to modify the commands based on your cluster or virtual machine configurations. Each main codebase directory contains a 'requirements.txt' file you may use as reference in installing the necessary packages needed to run each model.

Pre-processing

1. Firstly, download the following datasets from the following links, requesting the necessary permissions as stipulated. Store each dataset in the 'datasets' folder.

• CHAOS: https://chaos.grand-challenge.org/Download/
• L2R: https://learn2reg.grand-challenge.org/Datasets/
• MRXFDG*: https://pmc.ncbi.nlm.nih.gov/articles/PMC8446124/

2. Go to the 'pre-processing' folder and run the files in the 'slurm' folder for each respective dataset.

NOTE*: If running from an offline cluster, the skull stripping functions in mrxfdg_pp.py may not work. Alternatively, skull stripping may be done before running the file using SynthStrip: https://surfer.nmr.mgh.harvard.edu/docs/synthstrip/

Baseline Models

Each baseline model (SyN, VoxelMorph, NestedMorph, Dino-Reg, SPAC) contains its own folder in the 'baseline' directory with imported code from their original repositories. The final recorded results of these baseline models are in the 'final_results' folder. Firstly, apply data augmentation by running aug.slurm. Then, run the respective files in the 'slurm' folder to train and test them for each of the datasets:

• To get initial metrics for datasets: test-initial.slurm
• SyN: syn.slurm
• VoxelMorph: vm-train.slurm & vm-test.slurm
• NestedMorph: nm-train.slurm & nm-test.slurm
• Dino-Reg: dino-reg-chaos.slurm, dino-reg-l2r.slurm, & dino-reg-mrxfdg.slurm
• SPAC: First run spac-data.slurm & spac-data-test.slurm to fix data format for this model then run spac-l2r.slurm, spac-mrxfdg.slurm, & spac-test.slurm for training and testing

Main DVR Models

1. Download the DINOv2 pretrained model using the code below and store it in dvr/pretrained-models:

    wget https://dl.fbaipublicfiles.com/dinov2/dinov2_vitl14/dinov2_vitl14_reg4_pretrain.pth

2. Go to the main 'slurm' directory to run the main DVR model in 'dvr2' with curriculum learning with the following files: train-l2r.slurm, train-mrxfdg.slurm, test-l2r.slurm, & test-mrxfdg.slurm.

Ablation Models

To run the ablation studies, run the files in the main 'slurm' directory as follows:

1. Oneshot registration with Dino features: oneshot-l2r.slurm & oneshot-mrxfdg.slurm
2. Vlearn registration with raw images (no dino features): raw-img-l2r.slurm & raw-img-mrxfdg.slurm
3. Vlearn registration with Dino intermediate layers: int-layer-l2r.slurm & int-layer-mrxfdg.slurm

For the these, use the main DVR models' slurm files as normal, but modify the import statements in main2.py accordingly (ex. 'from dvr2.vlearn_rob import train_vlearn' -> 'from dvr-spac-hybrid.vlearn_rob import train_vlearn'). 4. Registration with SAC Planner network: 'dvr-spac-hybrid' folder 5. Simplified Vlearn Registration without curriculum learning: 'dvr-vlearn' folder

6. Various reward systems: any of the varying reward systems in 'reward.py' can be executed by changing the 'reward_system' parameter in line 941 of 'vlearn_rob.py'

Main DVR & Ablation Results

The results are in the 'dvr-experiments' folder. The main results of the DVR model for each dataset are in csv files in the 'Models" subfolder while that of the ablation studies are in 'Ablation'.