fangyuanmao / PID

Code for PID: Physics-Informed Diffusion Model for Infrared Image Generation
MIT License
58 stars 5 forks source link

PID: Physics-Informed Diffusion Model for Infrared Image Generation

PID

Paper

The paper is under review: https://arxiv.org/abs/2407.09299

Environment

We recommend you to install the environment with environment.yaml.

conda env create --file=environment.yaml

Datasets

Download KAIST dataset from https://github.com/SoonminHwang/rgbt-ped-detection.

Download FLIR dataset from https://www.flir.com/oem/adas/adas-dataset-form/.

Checkpoint

VQGAN can be downloaded from https://github.com/CompVis/latent-diffusion.

Name Note Link
TeVNet TeVNet checkpoint for KAIST, epoch=0.95k TeVNet_KAIST.zip
TeVNet TeVNet checkpoint for FLIR, epoch=1k TeVNet_FLIR.zip
PID PID checkpoint for KAIST, k1=50, k2=5 PID_KAIST.zip
PID PID checkpoint for FLIR, k1=k2=50 PID_FLIR.zip

Evaluation

Use the shellscript to evaluate. indir is the input directory of visible RGB images, outdir is the output directory of translated infrared images, config is the chosen config in configs/latent-diffusion/config.yaml. We prepare some RGB images in dataset/KAIST for quick evaluation.

bash run_test_kaist512_vqf8.sh

Train

Dataset preparation

Prepare corresponding RGB and infrared images with same names in two directories.

Stage 1: Train TeVNet

cd TeVNet
bash shell/train.sh

Stage 2: Train PID

Use the shellscript to train. It is recommended to use our pretrained model to accelerate the train process.

bash shell/run_train_kaist512_vqf8.sh

Acknowledgements

Our code is built upon LDM and HADAR. We thank the authors for their excellent work.

Citation

If you find this work helpful in your research, please consider citing our paper:

@inproceedings{Mao2024PIDPD,
  title={PID: Physics-Informed Diffusion Model for Infrared Image Generation},
  author={Fangyuan Mao and Jilin Mei and Shun Lu and Fuyang Liu and Liang Chen and Fangzhou Zhao and Yu Hu},
  year={2024},
  url={https://doi.org/10.48550/arXiv.2407.09299}
}

If you have any question, feel free to contact maofangyuan23s@ict.ac.cn .