U2Fusion: A Unified Unsupervised Image Fusion Network

This is the PyTorch implementation of U2Fusion: A Unified Unsupervised Image Fusion Network (TPAMI 2020).
The original paper can solve 3 typical image fusion tasks(multi-modal image fusion, multi-exposure image fusion and multi-focus image fusion), but in this repository, only the multi-exposure image fusion branch was implemented.

1. Environment

Python >= 3.7
PyTorch >= 1.4.0 is recommended
opencv-python = 4.5.1
matplotlib
tensorboard
pytorch_msssim

2. Dataset

The training data and testing data is from the [SICE dataset].

3. Quick Demo

Clone this repository:

git clone https://github.com/ytZhang99/U2Fusion-pytorch.git

Place the over-exposed images and under-exposed images in dataset/test_data/over and dataset/test_data/under, respectively.
Run the following command for multi-exposure fusion:
```
python main.py --test_only
```
Finally, you can find the Super-resolved and Fused results in ./test_results.

4. Training and Testing

Download the pre-trained vgg16 model from VGG16 and rename it to vgg16.pth. Place it in the same directory with vgg.py.
Place the training over-exposed images and under-exposed images in dataset/train_data/over and dataset/train_data/under, respectively.
Run the following command to train your own model:

python main.py --model mymodel.pth

Or you can fine-tune the existing model based on your own dataset:

python main.py --model model.pth

Moreover, if you want to test the model after training each epoch, run:

python main.py --model mymodel.pth --train_test

The generated model is placed in ./model/, then you can test your model with:

python main.py --test_only --model mymodel.pth

5. Citation

The following paper might be cited:

@article{xu2020u2fusion,
  title={U2Fusion: A unified unsupervised image fusion network},
  author={Xu, Han and Ma, Jiayi and Jiang, Junjun and Guo, Xiaojie and Ling, Haibin},
  journal={IEEE Transactions on Pattern Analysis and Machine Intelligence},
  year={2020},
  publisher={IEEE}
}