/densefuse_pytorch

Image fusion via an autoencoder with dense-conv-block. (Pytorch)

Primary LanguagePython

Project logo

DenseFuse (Pytorch)


An pytorch implement of DenseFuse.

๐Ÿ“ Table of Contents

๐Ÿง About

This is a pytorch implement of DenseFuse proposed by this paper, H. Li, X. J. Wu, โ€œDenseFuse: A Fusion Approach to Infrared and Visible Images,โ€ IEEE Trans. Image Process., vol. 28, no. 5, pp. 2614โ€“2623, May. 2019.

The code is writted with torch 1.1.0 and pytorch-ssim.

๐ŸŽˆ Usage

Quick start

  1. Clone this repo and unpack it.
  2. Download test dataset and put test images in './images/IV_images'
  3. run 'main.py'

Training

A pretrained model is available in './train_result/model_weight.pkl'. We train it on MS-COCO 2014. There are 82783 images in total, where 1000 images (COCO_train2014_000000574951 ~ COCO_train2014_000000581921.jpg) are used for validation and the rest are for training. In the training phase, all images are resize to 256x256 and are transformed to gray pictures. Model is optimized by Adam with learning rate being 1e-4. The batch size and epoch number are 2 and 4, respectively. Loss function is MSE+lambda x SSIM, where lambda=1. The experiments were implemented with 2080ti GPU and 32GB RAM. It took about 2 hours.

If you want to re-train this net, you should download MS-COCO 2014 and run 'train.py'