🧩 LDMReg

A Hierarchical Feature-guided Latent Diffusion Model for Multimodal Image Registration

Official PyTorch implementation.

---

📖 Table of Contents

• 🛠️ Environment
• 🧪 Testing
• 🚂 Train
  • 1. Data Download
  • 2. Unregistered Data Generation
  • 3. Data Format
  • 4. Start Training
• 📝 Citation

---

🛠️ Environment

We test the code on PyTorch 1.12.1 + CUDA 11.3.

1. Create a new conda environment

conda create -n LDMReg python=3.9
conda activate LDMReg

2. Install dependencies (Please execute the code line by line.)

pip install torch==1.12.1+cu113 torchvision==0.13.1+cu113 torchaudio==0.12.1 --extra-index-url [https://download.pytorch.org/whl/cu113](https://download.pytorch.org/whl/cu113)

pip install -r requirements.txt

---

🧪 Testing

You can directly test our model to generate registered images using the following code.

💡 Note1: The pre-training weights should be saved in the ./check/ folder.

💡 Note2: Additionally, please uncomment 'img_name':img_name on line 169 in ./data/load_pre_made_dataset.py

🔗 Link: Download Checkpoint Here

# Visible and infrared rigid image registration
python test-ir_rigid.py

# Visible and near-infrared rigid image registration
python test-nir_rigid.py

# Visible and depth rigid image registration
python test-rgbd_rigid.py

# Visible and infrared non-rigid image registration
python test-ir_nonrigid.py

# PET and MRI non-rigid image registration
python test-med_rigid.py

You can find their corresponding configuration file paths in ./config/.

⚙️ Other Test: If you want to test corrupted data, you can run:

# Corrupted visible and infrared rigid image registration
python test-ir_rigid_Corrupted.py

# Corrupted visible and near-infrared rigid image registration
python test-nir_rigid_Corrupted.py

# Corrupted visible and depth rigid image registration
python test-rgbd_rigid_Corrupted.py

---

🚂 Train

1. Data Download

The datasets are available for download at the following links:

Task	Dataset Name	Link
VI-IR	M3FD	TarDAL
	RoadScene	RoadScene
	LLVIP	LLVIP
	FLIR_aligned	FLIR_aligned
VI-NIR	RGB-NIR Scene	EPFL Dataset
	MCubeS	Multimodal Material Segmentation
MED	Harvard	Harvard Medical
RGBD	DIODE	DIODE Dataset

2. Unregistered Data Generation

This subsection provides tools to generate training and testing datasets with various types of image deformations, including homography and elastic.

📌 Supported Deformation Types

Deformation Type	Description	Formula	Effect
Homography	Single-plane projective transformation	$x'=\frac{Hx}{h_3^Tx}$ (H is a 3×3 homography matrix)	Rigid projective deformation
Elastic	Local elastic distortions	$x'=x+\alpha d(x)$ (d(x) is a random smooth displacement field)	Subpixel-level local non-rigid deformation

🔮 Generate Unaliged Datasets

Run the following commands to generate datasets under homography or elastic deformations:

cd Data_generate
export PYTHONPATH=$PYTHONPATH:$(pwd)/..

# Homography
python generate_warp_dataset.py --image_root ./data/ --save_root ./hom/ --image_size 512 --transform_type hom

# Elastic
python generate_warp_dataset.py --image_root ./data/ --save_root ./elastic/ --image_size 512 --transform_type elastic

⚠️ Warning: Adjust the parameters in the scripts to control the strength and scale of the deformation according to your needs.

🌌 Generate Corrupted Datasets

cd Data_generate

python generate_corruption_dataset.py --image_root ./data/ --save_root ./hom/ --image_size 512 --transform_type hom

⚠️ Warning: Adjust the parameters in the scripts to control the strength of the corrupted degraed according to your needs.

3. Data Format

The data should be organized in the following directory structure:

train/
├── VI-IR/                                      
│   ├── ir/                                      
│   │   ├── 0000.png                            
│   ├── ir_warp/                                      
│   │   ├── 0000.png                                                            
│   ├── vi/                                      
│   │   ├── 0000.png                                                            
│   └── vi_warp/                                      
│       └── 0000.png                            
│
├── VI-NIR/                                        
│   ├── nir/
│   │   ├── 0000.png  
│   ... 
│
└── MED/                                        
    ├── mri/
    │   ├── 0000.png  
    ... 

📝 Configuration Requirement: It is worth noting that the corresponding dataset paths must be specified in the train_2D.json file prior to training:

"dataroot_elstic": "",
"dataroot_ir": "",
"dataroot_NIR": "",
"dataroot_RGBD": "",
"dataroot_Med": ""

4. Start Training

You can use the following code to train the MOFM model for different fusion tasks:

python -m torch.distributed.launch --nproc_per_node=5 train.py \
  --sample_selected Unipc \
  --model_selected NAFNet_rigid \
  --feat_network VGG_cross_align \
  --batch_size 10 \
  --reg_task VI-IR_Hom \
  --img_size 256 \
  --Corrupt_type Corrupt

🔄 Parameter Guide:

• nproc_per_node: Represents the number of GPUs. (Note: If you change this, the device number of gpu_ids in ./config/train.json needs to be updated. For example, if nproc_per_node=2, then gpu_ids=[0,1].)
• sample_selected: Choose different samplers for training. Options: DDIM, ddp-solver, ddp-solver++, Deis, Unipc, Heun, and Euler. (Default: Unipc)
• model_selected: Select the parameters for the rigid and non-rigid networks. Options: NAFNet_nonrigid and NAFNet_rigid.
• feat_network: Choose the feature-guided network. (Default: VGG_cross_align)
• batch_size: The batch size for training.
• reg_task: Registration task selection. Options: VI-IR_Elastic, MED_Elastic, VI-IR_Hom, VI-NIR_Hom, RGBD_Hom.
• img_size: Resolution of the training images.
• Corrupt_type: Select whether to train normally or with added corruption/interference. Options: No_Corrupt and Corrupt.

💡 Pro Tip: If you want to train three fusion tasks at once, you can directly run:

python sample_all.py

Citation

@article{
    author    = {},
    title     = {LDMReg: A Hierarchical Feature-guided Latent Diffusion Model for Multimodal Image Registration},
    booktitle = {-},
    year      = {-},
    pages     = {-},
    doi       = {-},
}