294coder / Efficient-MIF

Train your fusion model and test downstream tasks in one repo.
6 stars 1 forks source link

Effecient Multi-source Image Fusion Framework

Zihan Cao   Xiao Wu   Liang-Jian Deng
University of Electronic Science Technology of China

This a repository that allow you to efficiently train your network on multi-source image fusion task, including

We call the first two as sharpening tasks and the latter two as fusion tasks.

This framework is tested on Python 3.9.x / 3.11 and Pytorch 2.3.0 with CUDA 11.8.

NEWS✨

2024.09.27: Two papers (about Implicit Neural Representation and Diffusion Models) for pansharpening and HMIF tasks are accepted by Neurips 2024. Ready to release the code. Stay Tuned!

2024.08.31: Two papers for pansharpening and HMIF tasks are accepted by ACM MM 2024. You can find the guidances in Model Cards.

2024.07.24: Some downstream tasks: a) detection task; b) segmentation task, c) depth estimation are ready to release.

2024.06.27: We release VIS-IR and medical image fusion training scripts.

2024.04.06: Release the some pansharpening and hyperspectral and multispectral image fusion methods. You can found in model/.

Quick Start🎈

Installation Requirements

You can simply install all needed Python packages by running pip install -r requirement.txt.

Prepare Your Model

You can put in your model and register with the model name that you prefer using register_model, then import your model in model/__init__.py.

We are trying using hydra dynamic importing. So, this api may be changed.

Specify your model's train_step, val_step and patch_merge_step (optional) if you are working on Pansharpening and HMIF tasks, or, if you are working VIS-IR and medical image fusion tasks, you should specify the train_fusion_step and val_fusion_step functions.

You can find detailed example usage in the model/.

Starting Training

Here are two scripts for training your model.

  1. scripts implemented by torch-run.
  2. scripts implemented by Huggingface Accelerate.

We recommand the latter one.

We take the second choice as the example:

Sharpening Tasks

For sharpening tasks, you should modify the scripts/torch_sharpening_run.sh to feed your model name (after you register it), dataset name, and more options.

Then run it by

sh scripts/torch_sharpening_run.sh

Fusion Tasks

For fusion tasks, you should switch to scripts/accelerate_run.sh and modify it to suit your way.

Inference

After you trained your model, you need to infer model on test set.

Test Metrics

We provide the Matlab package to test the metrics on those four tasks. Please check them in Pansharpening_Hyper_SR_Matlab_Test_Package and VIS_IR_Matlab_Test_Package.

Basic Usage

cd Pansharpening_Hyper_SR_Matlab_Test_Package

%% when testing the reduced-resolution metrics on pansharpening or HMIF tasks
% Args:
% path: the saved fused image `.mat` file, find it in `visualized_img/`
% ratio: upscale ratio, e.g., 4
% full_res: we keep it to 1, not changed
% const: max value of the dataset
analysis_ref_batched_images(path, ratio, full_res, const)

%% when testing full-resolution metrics on pansharpening datasets
% Args:
% path: the saved fused image `.mat` file, find it in `visualized_img/`
% ratio: upscale ratio, e.g., 4
% sensor: the sensor name
analysis_unref_batched_images(path, ratio, sensor)

You need to define the data paths (h5 file paths). Find them in the matlab test scripts mentioned above.

Model Cards

You can find our detailed descriptions of the proposed fusion models at readmes/.

Models Tasks Guidance
LE-Mamba (ACM MM 24) Pansharpening; HMIF. README
LFormer (ACM MM 24) Pansharpening; HMIF. README
[FeINFN (Neurips 2024)]() HMIF. Ready to release the code

About Dataset

We colloct panshapening, HMIF, VIS-IR image fusion, and medical image fusion datasets from previous works. Really appreciate their great works!

The used datasets includes:

We are planning to release some of them.

Citation

If you find our works useful, please conside citing the following papers:

@inproceedings{
    cao2024novel,
    title={A novel state space model with local enhancement and state sharing for image fusion},
    author={Cao, Zihan and Wu, Xiao and Deng, Liang-Jian and Zhong, Yu},
    booktitle={ACM Multimedia 2024 (ACM MM)},
    year={2024}
}

@article{liang2024fourier,
  title={Fourier-enhanced Implicit Neural Fusion Network for Multispectral and Hyperspectral Image Fusion},
  author={Liang, Yu-Jie and Cao, Zihan and Deng, Liang-Jian and Wu, Xiao},
  journal={arXiv preprint arXiv:2404.15174},
  year={2024}

@inproceedings{
  hou2024linearlyevolved,
  title={Linearly-evolved Transformer for Pan-sharpening},
  author={Junming Hou and Zihan Cao and Naishan Zheng and Xuan Li and Xiaoyu Chen and Xinyang Liu and Xiaofeng Cong and Danfeng Hong and Man Zhou},
  booktitle={ACM Multimedia 2024},
  year={2024}
}

And some related works about image fusion may attract you:

@article{cao2024neural,
  title={Neural Shr$\backslash$" odinger Bridge Matching for Pansharpening},
  author={Cao, Zihan and Wu, Xiao and Deng, Liang-Jian},
  journal={arXiv preprint arXiv:2404.11416},
  year={2024}
}

@article{cao2024diffusion,
  title={Diffusion model with disentangled modulations for sharpening multispectral and hyperspectral images},
  author={Cao, Zihan and Cao, Shiqi and Deng, Liang-Jian and Wu, Xiao and Hou, Junming and Vivone, Gemine},
  journal={Information Fusion},
  volume={104},
  pages={102158},
  year={2024},
  publisher={Elsevier}
}

@article{zhong2024ssdiff,
  title={SSDiff: Spatial-spectral Integrated Diffusion Model for Remote Sensing Pansharpening},
  author={Zhong, Yu and Wu, Xiao and Deng, Liang-Jian and Cao, Zihan},
  journal={arXiv preprint arXiv:2404.11537},
  year={2024}
}