Structure-Aware Riemannian Flow Matching (SA-RFM)

📌 ICML 2026 Poster
Official implementation of Structure-Aware Riemannian Flow Matching for Hyperspectral and Multispectral Image Registration and Fusion.

---

📖 Overview

SA-RFM is a flow matching-based framework designed for the joint registration and fusion of hyperspectral images and multispectral images.

The method aims to improve spatial alignment and spectral reconstruction by introducing a structure-aware flow matching strategy.

---

🧩 Framework

Figure 2. The proposed SA-RFM framework. (a) Architecture of the fusion modules and velocity network, where both the encoders and decoder are constructed using ConvBlocks. (b) The training scheme involving data sampling at various time steps. (c) Illustration of the proposed Riemannian metric formulated in Eq. (7), which assigns higher transport costs across texture directions and lower costs along them.

---

🚀 Getting Started

1. Clone the Repository

Clone this repository:

git clone https://github.com/ZhangQuan-hub/SA-RFM.git
cd SA-RFM

2. Download Datasets

The benchmark hyperspectral datasets (Chikusei, Botswana and Pavia University) used in this project are available at:

🔗 ZhangQuan-hub/HyperspectralDatasets - v2 Release

For the Xiongan dataset, please also refer to the original published paper:

🔗 雄安新区马蹄湾村航空高光谱遥感影像分类数据集

3. Run Training

The main training script is:

1main_SAMB.py

You can start the training process directly by running (Note that you should first download the training datasets):

python 1main_SAMB.py

---

🛠️ Acknowledgements

The implementation was developed with the assistance of Claude 4.5 Sonnet and includes corresponding comments for better readability.

---

📝 Citation

If you find this work or code useful for your research, please consider citing our paper (not yet):