lengmo1996 / ECDM

This is the official repository of Data Generation Scheme for Thermal Modality with Edge-Guided Adversarial Conditional Diffusion Model (ACM MM'24)

🆕 [2024-07-30] upload base resource code

ECDM

This is the official repository of our work: Data Generation Scheme for Thermal Modality with Edge-Guided Adversarial Conditional Diffusion Model (ACM MM'24)

Guoqing Zhu, Honghu Pan, Qiang Wang, Chao Tian, Chao Yang, Zhenyu He

ACM MM'24 | arXiv | GitHub | Project Page

TODO list

✔️ upload base resource code

❌ update weights

✔️ update related datasets

✔️ update datasets processing scripts

✔️ update evaluation scripts

✔️ update generated thermal images from different methods

❌ support more sampler

Update

• [2024-09-07] upload generated thermal images
• [2024-07-30] upload base resource code

Contents

• User guide
  • 1. Environment configuration
    • 1.1 Docker (Recommended)
    • 1.2 Conda (Recommended)
    • 1.3 Pip (Python >= 3.10)
  • 2. Dataset Preparation
    • 2.1 Datasets
    • 2.2 Generating edge images
  • 3. Training
    • 3.1 Training first stage model
    • 3.2 Training second stage model
    • 3.3 Resuming training process
  • 4. Evaluation
    • 4.1 Generating themal images
    • 4.2 Evaluation
  • 5. Computing model Flops
  • 6. Download links
    • 6.1 Processed datasets
    • 6.2 Model weights
    • 6.3 Generated thermal images from other methods
• Acknowledgments
• Citation

User guide

1. Environment configuration

We offer guides on how to install dependencies via docker and conda.

First, clone the reposity:

git clone https://github.com/lengmo1996/ECDM.git
cd ECDM

Then, use one of the below methods to configure the environment.

1.1 Docker (Recommended)

Building image from Dockerfile

# build image from Dockerfile
docker build -t ecdm:1.0 .
# or from the mirrors
docker build -t ecdm:1.0 -f Dockerfile_mirror

or pulling from Docker Hub

docker pull lengmo1996/ecdm:1.0

Then creating container from docker image.

docker run -it --shm-size 100g --gpus all -v /path_to_dataset:/path_to_dataset -v /path_to_log:/path_to_log -v /path_to_ECDM:/path_to_ECDM --name ECDM ecdm:1.0 /bin/bash

1.2 Conda (Recommended)

conda env create -f conda.yaml
conda activate ecdm

1.3 Pip (Python >= 3.10)

pip install torch==2.1.2 torchvision==0.16.2 torchaudio==2.1.2 --index-url https://download.pytorch.org/whl/cu121 
pip install -r requirements.txt

2. Dataset Preparation

2.1 Datasets

We use LLVIP and PRW dataset in our experiments. You can download LLVIP dataset from here and PRW dataset from here.

2.2 Generating edge images

We use scripts/generate_edge_img.py to generate edge images. You should modify the ==lines 19-20== to specify your custom dataset path. Then, execute the following command to generate the edge images::

python scripts/generate_edge_img.py

This method applies a high-pass filter to the original images. Edge images can also be obtained by other edge detection methods, such as pidinet or teed.

3. Training

3.1 Training first stage model

Please modify the ==lines 66, 73, 81== in configs/ecdm_first_stage.yaml to specify your custom paths. Then, run the following command to train the first stage model:

python main.py fit -c configs/base_config.yaml -c configs/ecdm_first_stage.yaml --trainer.devices 0,1,2,3 

3.2 Training second stage model

Please modify the ==lines 74, 81, 89== in configs/ecdm_second_stage.yaml to your custom paths. Additionally, modify ==line 16== to point to the path of the first stage model weight. Then, run the following command to train the second stage model:

python main.py fit -c configs/base_config.yaml -c configs/ecdm_second_stage.yaml --trainer.devices 0,1,2,3 

If you want to generalize the model to the PRW dataset, please use the configuration of configs/ecdm_second_stage_with_PRW.yaml.

3.3 Resuming training process

If you want to resume the training process, please use the argument '--ckpt_path'. For example:

python main.py fit -c configs/base_config.yaml -c configs/ecdm_second_stage.yaml --trainer.devices 0,1,2,3 --ckpt_path 'logs/checkpoints/last.ckpt'

4. Evaluation

4.1 Generating themal images

For evaluation, we need to run the following command to generate thermal images:

python main.py test -c configs/base_config.yaml -c configs/ecdm_second_stage.yaml --trainer.devices 0,1,2,3 --ckpt_path 'logs/checkpoints/last.ckpt'

4.2 Evaluation

Please modify the ==lines 196-198== in scripts/metrics.py to your custom paths. Then, run the following command to evaluate the second stage model:

python scripts/metrics.py

Note: If version of scikit-image >= 0.16, you may trigger the following error

cannot import name 'compare_ssim' from 'skimage.measure'. 

To fix this bug, you may need to modify the lines ==23-25== in /opt/conda/lib/python3.10/site-packages/lpips/__init__.py.

def dssim(p0, p1, range=255.):
    from skimage.metrics import structural_similarity
    return (1 - structural_similarity(p0, p1, data_range=range, channel_axis=2)) / 2.

More information can be found in GitHub issue and release note of scikit-image 0.16.2 (2019-10-22).

5. Computing model Flops

You can obtain the number of parameters and FLOPs of the sampling process using the script located at scripts/compute_flops_macs_params.py

6. Download links

6.1 Processed datasets

Baidu Drive: LLVIP, PRW

6.2 Model weights

First stage weight: Second stage weight:

6.3 Generated thermal images from other methods

The generated thermal images which used for visualization (Under conditions of LLVIP edge images): Baidu Drive. The generated thermal images which used for training object detection (Under conditions of PRW edge images): Baidu Drive

Acknowledgments

This implementation is based on / inspired by: denoising-diffusion-pytorch, improved-diffusion and latent-diffusion,

Citation

If you find our work helpful for your research, please consider citing the following BibTeX entry.