Skip to content

YonghaoXu/Txt2Img-MHN

Repository files navigation

Txt2Img-MHN: Remote Sensing Image Generation from Text Using Modern Hopfield Networks


This is the official PyTorch implementation of the paper Txt2Img-MHN: Remote Sensing Image Generation from Text Using Modern Hopfield Networks.

Table of content

  1. Preparation
  2. Training VQVAE and VQGAN
  3. Training Txt2Img-MHN
  4. Image Generation
  5. Inception Score and FID Score
  6. CLIP Score
  7. Zero-Shot Classification
  8. Paper
  9. Acknowledgement
  10. License

Preparation

  • Install required packages: pip install -r requirements.txt
  • Install Taming Transformers:
    • Download the repo
    • Run pip install -e .
    • Copy the files in the folder taming-transformers-master from this repo to the downloaded taming-transformers-master folder
  • Download the remote sensing text-image dataset RSICD used in this repo
  • Extract separate .txt files for the text descriptions of each image in RSICD: python data_preparation.py
  • Data folder structure:
├── RSICD/
│   ├── airport_1.jpg   
│   ├── airport_2.jpg  
│   ├── ...  
│   ├── viaduct_420.jpg  
│   ├── airport_1.txt   
│   ├── airport_2.txt   
│   ├── ...  
│   ├── viaduct_420.txt   

Training VQVAE and VQGAN

  • Train VQVAE:
$ cd Txt2Img-MHN-main
$ CUDA_VISIBLE_DEVICES=0 python train_vqvae.py --data_dir /Path/To/RSICD/
  • Train VQGAN:
$ cd taming-transformers-master
$ CUDA_VISIBLE_DEVICES=0 python main.py --base configs/custom_vqgan.yaml -t True --gpus 0,

Training Txt2Img-MHN

  • Train Txt2Img-MHN with the pretrained VQVAE:
$ cd Txt2Img-MHN-main
$ CUDA_VISIBLE_DEVICES=0 python train_txt2img_mhn.py --vae_type 0 --data_dir /Path/To/RSICD/ --vqvae_path /Path/To/vae.pth --batch_size 8
  • Train Txt2Img-MHN with the pretrained VQGAN:
$ cd Txt2Img-MHN-main
$ CUDA_VISIBLE_DEVICES=0 python train_txt2img_mhn.py --vae_type 1 --data_dir /Path/To/RSICD/ --vqgan_model_path /Path/To/last.ckpt --vqgan_config_path /Path/To/project.yaml --batch_size 8

Note: Training with multiple GPUs is supported. Simply specify the GPU ids with CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7,...

  • Use tensorboard to monitor the training process:
$ cd Txt2Img-MHN-main
$ tensorboard --logdir ./ --samples_per_plugin images=100

Image Generation

  • Txt2Img-MHN (VQVAE):
$ cd Txt2Img-MHN-main
$ CUDA_VISIBLE_DEVICES=0 python gen_im.py --vae_type 0 --data_dir /Path/To/RSICD/ --vqvae_path /Path/To/vae.pth --mhn_vqvae_path /Path/To/mhn_vqvae.pth --num_gen_per_image 10
  • Txt2Img-MHN (VQGAN):
$ cd Txt2Img-MHN-main
$ CUDA_VISIBLE_DEVICES=0 python gen_im.py --vae_type 1 --data_dir /Path/To/RSICD/ --vqgan_model_path /Path/To/last.ckpt --vqgan_config_path /Path/To/project.yaml  --mhn_vqgan_path /Path/To/mhn_vqgan.pth --num_gen_per_image 10

Alternatively, you can download our pretrained models for a quick look.

Inception Score and FID Score

  • Data preparation: Before training the Inception model, prepare a new data folder with the structure below:
├── RSICD_cls/
│   ├── airport/     
|   |   ├── airport_1.jpg   
|   |   ├── airport_2.jpg   
|   |   ├── ... 
│   ├── bareland/     
|   |   ├── bareland_1.jpg   
|   |   ├── bareland_2.jpg   
|   |   ├── ... 
│   ├── ...  
│   ├── viaduct/     
|   |   ├── viaduct_1.jpg   
|   |   ├── viaduct_2.jpg   
|   |   ├── ...   
  • Pretrain the Inception model:
$ cd Txt2Img-MHN-main/is_fid_score
$ CUDA_VISIBLE_DEVICES=0 python pretrain_inception.py --root_dir /Path/To/RSICD_cls/
  • Calculate the Inception score and FID score:
$ cd Txt2Img-MHN-main/is_fid_score
$ CUDA_VISIBLE_DEVICES=0 python is_fid_score.py --gen_dir /Path/To/GenImgFolder/ --data_dir /Path/To/RSICD/

CLIP Score

$ cd Txt2Img-MHN-main
$ CUDA_VISIBLE_DEVICES=0 python clip_score.py --gen_dir /Path/To/GenImgFolder/ --data_dir /Path/To/RSICD/

Zero-Shot Classification

$ cd Txt2Img-MHN-main/zero_shot_classification
$ CUDA_VISIBLE_DEVICES=0 python zero_shot_evaluation.py --gen_dir /Path/To/GenImgFolder/ --root_dir /Path/To/RSICD/

Paper

Txt2Img-MHN: Remote Sensing Image Generation from Text Using Modern Hopfield Networks

Please cite the following paper if you find it useful for your research:

@article{txt2img_mhn,
  title={Txt2Img-MHN: Remote Sensing Image Generation from Text Using Modern Hopfield Networks},
  author={Xu, Yonghao and Yu, Weikang and Ghamisi, Pedram and Kopp, Michael and Hochreiter, Sepp},
  journal={IEEE Trans. Image Process.}, 
  doi={10.1109/TIP.2023.3323799},
  year={2023}
}

Acknowledgement

DALLE-pytorch

taming-transformers

metrics

CLIP-rsicd

This research has been conducted at the Institute of Advanced Research in Artificial Intelligence (IARAI).

License

This repo is distributed under MIT License. The code can be used for academic purposes only.

About

[IEEE TIP 2023] Txt2Img-MHN: Remote Sensing Image Generation from Text Using Modern Hopfield Networks

Topics

Resources

License

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published

Languages