Skip to content

Commit

Permalink
update
Browse files Browse the repository at this point in the history
  • Loading branch information
@Shengwei-Peng
Shengwei-Peng committed Jun 1, 2024
1 parent 3cb0674 commit 706fdf5
Show file tree
Hide file tree
Showing 6 changed files with 100 additions and 17 deletions.
17 changes: 16 additions & 1 deletion README.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
- [Installation](#Installation)
- [Project Structure](#Project-Structure)
- [Datasets](#Datasets)
- [Reproduction](#Reproduction)
- [Usage](#Usage)
- [Acknowledgements](#Acknowledgement)
- [Arguments](#Arguments)
Expand Down Expand Up @@ -62,6 +63,7 @@ Generative-AI-Navigation-Information-for-UAV-Reconnaissance-in-Natural-Environme
├── requirements.txt
├── run_diffusion.sh
├── run_gan.sh
├── run_reproduce.sh
└── run_router.sh
```

Expand All @@ -77,6 +79,19 @@ The `training_dataset` and `testing_dataset` directories contain the datasets pr

**Note**: The images in `img/` and `label_img/` should have matching filenames (except for the file extensions) and consistent dimensions. Filenames for road data should include **RO** and filenames for river data should include **RI**.

## Reproduction
One-click execution to reproduce the best results:
```bash
bash run_reproduce.sh
```
The script performs the following steps:
- Uses `gdown` to download the best checkpoints to `./checkpoints`
- `gan/preprocess.py`: Preprocess the data for GAN.
- `gan/test.py`: Use GAN generate images.
- `diffusion/preprocess.py`: Preprocess the data for Diffusion.
- `diffusion/test.py`: Use Diffusion generate images.
- `router/router.py`: Use Router selects the final best results.

## Usage
**Warning**: Executing the scripts below requires approximately **32GB** of VRAM. If your hardware does not meet this requirement, you may need to adjust the [Arguments](#Arguments) accordingly.

Expand Down Expand Up @@ -107,7 +122,7 @@ Select the final images from both GAN and Diffusion models:
bash run_router.sh
```
The script performs the following steps:
- `router.py`: Selects the final images.
- `router.py`: Selects the final results.

## Acknowledgement
We extend our gratitude to the developers of [pix2pix](https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix) and [PITI](https://github.com/PITI-Synthesis/PITI) for generously sharing their code, which has been invaluable to our work. Additionally, we would like to thank the developers of [guided-diffusion](https://github.com/openai/guided-diffusion) for providing the pretrained model.
Expand Down
4 changes: 4 additions & 0 deletions environment.yml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,11 +22,13 @@ dependencies:
- xz=5.4.6=h5eee18b_1
- zlib=1.2.13=h5eee18b_1
- pip:
- beautifulsoup4==4.12.3
- blobfile==2.1.1
- certifi==2024.2.2
- charset-normalizer==3.3.2
- filelock==3.13.1
- fsspec==2024.2.0
- gdown==5.2.0
- huggingface-hub==0.23.2
- idna==3.7
- jinja2==3.1.3
Expand All @@ -51,10 +53,12 @@ dependencies:
- packaging==24.0
- pillow==10.3.0
- pycryptodomex==3.20.0
- pysocks==1.7.1
- pyyaml==6.0.1
- requests==2.32.3
- safetensors==0.4.3
- scipy==1.10.1
- soupsieve==2.5
- sympy==1.12
- timm==1.0.3
- torch==2.3.0+cu118
Expand Down
4 changes: 4 additions & 0 deletions requirements.txt
View file
Original file line number Diff line number Diff line change
@@ -1,8 +1,10 @@
beautifulsoup4==4.12.3
blobfile==2.1.1
certifi==2024.2.2
charset-normalizer==3.3.2
filelock==3.13.1
fsspec==2024.2.0
gdown==5.2.0
huggingface-hub==0.23.2
idna==3.7
Jinja2==3.1.3
Expand All @@ -27,10 +29,12 @@ opencv-python==4.9.0.80
packaging==24.0
pillow==10.3.0
pycryptodomex==3.20.0
PySocks==1.7.1
PyYAML==6.0.1
requests==2.32.3
safetensors==0.4.3
scipy==1.10.1
soupsieve==2.5
sympy==1.12
timm==1.0.3
torch==2.3.0+cu118
Expand Down
28 changes: 14 additions & 14 deletions run_diffusion.sh
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,31 +12,31 @@ python ./diffusion/preprocess.py \
echo "[Diffusion] Training for river data..."
export LOGDIR=./checkpoints/diffusion/RI/upsample
MODEL_FLAGS="--learn_sigma True --uncond_p 0 --image_size 256 --super_res 64 --num_res_blocks 2 --finetune_decoder True --model_path ./diffusion/ckpt/upsample.pt"
TRAIN_FLAGS="--lr 1e-5 --batch_size 4 --lr_anneal_steps 200"
TRAIN_FLAGS="--lr 1e-5 --batch_size 4 --lr_anneal_steps 20000"
DIFFUSION_FLAGS="--noise_schedule linear"
SAMPLE_FLAGS="--num_samples 2 --sample_c 1"
DATASET_FLAGS="--data_dir ./diffusion/datasets/RI_train.txt --val_data_dir ./diffusion/datasets/RI_train.txt --mode coco-edge"
python ./diffusion/train.py $MODEL_FLAGS $TRAIN_FLAGS $SAMPLE_FLAGS $DIFFUSION_FLAGS $DATASET_FLAGS
export CUDA_LAUNCH_BLOCKING=1
export LOGDIR=./checkpoints/diffusion/RI/base/stage1
MODEL_FLAGS="--learn_sigma True --uncond_p 0. --image_size 64 --finetune_decoder False"
TRAIN_FLAGS="--lr 3.5e-5 --batch_size 24 --schedule_sampler loss-second-moment --model_path ./diffusion/ckpt/base.pt --lr_anneal_steps 200"
TRAIN_FLAGS="--lr 3.5e-5 --batch_size 24 --schedule_sampler loss-second-moment --model_path ./diffusion/ckpt/base.pt --lr_anneal_steps 20000"
DIFFUSION_FLAGS=""
SAMPLE_FLAGS="--num_samples 2 --sample_c 1"
DATASET_FLAGS="--data_dir ./diffusion/datasets/RI_train.txt --val_data_dir ./diffusion/datasets/RI_train.txt --mode coco-edge"
python ./diffusion/train.py $MODEL_FLAGS $TRAIN_FLAGS $SAMPLE_FLAGS $DIFFUSION_FLAGS $DATASET_FLAGS
export CUDA_LAUNCH_BLOCKING=1
export LOGDIR=./checkpoints/diffusion/RI/base/stage1-cont/
MODEL_FLAGS="--learn_sigma True --uncond_p 0.2 --image_size 64 --finetune_decoder False --encoder_path ./checkpoints/diffusion/RI/base/stage1/checkpoints/ema_0.9999_000201.pt"
TRAIN_FLAGS="--lr 2e-5 --batch_size 24 --schedule_sampler loss-second-moment --model_path ./diffusion/ckpt/base.pt --lr_anneal_steps 200"
MODEL_FLAGS="--learn_sigma True --uncond_p 0.2 --image_size 64 --finetune_decoder False --encoder_path ./checkpoints/diffusion/RI/base/stage1/checkpoints/ema_0.9999_020000.pt"
TRAIN_FLAGS="--lr 2e-5 --batch_size 24 --schedule_sampler loss-second-moment --model_path ./diffusion/ckpt/base.pt --lr_anneal_steps 20000"
DIFFUSION_FLAGS=""
SAMPLE_FLAGS="--num_samples 2 --sample_c 1"
DATASET_FLAGS="--data_dir ./diffusion/datasets/RI_train.txt --val_data_dir ./diffusion/datasets/RI_train.txt --mode coco-edge"
python ./diffusion/train.py $MODEL_FLAGS $TRAIN_FLAGS $SAMPLE_FLAGS $DIFFUSION_FLAGS $DATASET_FLAGS
export CUDA_LAUNCH_BLOCKING=1
export LOGDIR=./checkpoints/diffusion/RI/base/stage2-decoder/
MODEL_FLAGS="--learn_sigma True --uncond_p 0.2 --image_size 64 --finetune_decoder True"
TRAIN_FLAGS="--lr 3.5e-5 --batch_size 20 --schedule_sampler loss-second-moment --model_path ./diffusion/ckpt/base.pt --encoder_path ./checkpoints/diffusion/RI/base/stage1-cont/checkpoints/ema_0.9999_000201.pt --lr_anneal_steps 200"
MODEL_FLAGS="--learn_sigma True --uncond_p 0.2 --image_size 64 --finetune_decoder True --encoder_path ./checkpoints/diffusion/RI/base/stage1-cont/checkpoints/ema_0.9999_020000.pt"
TRAIN_FLAGS="--lr 3.5e-5 --batch_size 20 --schedule_sampler loss-second-moment --model_path ./diffusion/ckpt/base.pt --lr_anneal_steps 20000"
DIFFUSION_FLAGS=""
SAMPLE_FLAGS="--num_samples 2 --sample_c 2.5"
DATASET_FLAGS="--data_dir ./diffusion/datasets/RI_train.txt --val_data_dir ./diffusion/datasets/RI_train.txt --mode coco-edge"
Expand All @@ -46,31 +46,31 @@ python ./diffusion/train.py $MODEL_FLAGS $TRAIN_FLAGS $SAMPLE_FLAGS $DIFFUSION_
echo "[Diffusion] Training for road data..."
export LOGDIR=./checkpoints/diffusion/RO/upsample
MODEL_FLAGS="--learn_sigma True --uncond_p 0 --image_size 256 --super_res 64 --num_res_blocks 2 --finetune_decoder True --model_path ./diffusion/ckpt/upsample.pt"
TRAIN_FLAGS="--lr 1e-5 --batch_size 4 --lr_anneal_steps 200"
TRAIN_FLAGS="--lr 1e-5 --batch_size 4 --lr_anneal_steps 20000"
DIFFUSION_FLAGS="--noise_schedule linear"
SAMPLE_FLAGS="--num_samples 2 --sample_c 1"
DATASET_FLAGS="--data_dir ./diffusion/datasets/RO_train.txt --val_data_dir ./diffusion/datasets/RO_train.txt --mode coco-edge"
python ./diffusion/train.py $MODEL_FLAGS $TRAIN_FLAGS $SAMPLE_FLAGS $DIFFUSION_FLAGS $DATASET_FLAGS
export CUDA_LAUNCH_BLOCKING=1
export LOGDIR=./checkpoints/diffusion/RO/base/stage1
MODEL_FLAGS="--learn_sigma True --uncond_p 0. --image_size 64 --finetune_decoder False"
TRAIN_FLAGS="--lr 3.5e-5 --batch_size 24 --schedule_sampler loss-second-moment --model_path ./diffusion/ckpt/base.pt --lr_anneal_steps 200"
TRAIN_FLAGS="--lr 3.5e-5 --batch_size 24 --schedule_sampler loss-second-moment --model_path ./diffusion/ckpt/base.pt --lr_anneal_steps 20000"
DIFFUSION_FLAGS=""
SAMPLE_FLAGS="--num_samples 2 --sample_c 1"
DATASET_FLAGS="--data_dir ./diffusion/datasets/RO_train.txt --val_data_dir ./diffusion/datasets/RO_train.txt --mode coco-edge"
python ./diffusion/train.py $MODEL_FLAGS $TRAIN_FLAGS $SAMPLE_FLAGS $DIFFUSION_FLAGS $DATASET_FLAGS
export CUDA_LAUNCH_BLOCKING=1
export LOGDIR=./checkpoints/diffusion/RO/base/stage1-cont/
MODEL_FLAGS="--learn_sigma True --uncond_p 0.2 --image_size 64 --finetune_decoder False --encoder_path ./checkpoints/diffusion/RO/base/stage1/checkpoints/ema_0.9999_000201.pt"
TRAIN_FLAGS="--lr 2e-5 --batch_size 24 --schedule_sampler loss-second-moment --model_path ./diffusion/ckpt/base.pt --lr_anneal_steps 200"
MODEL_FLAGS="--learn_sigma True --uncond_p 0.2 --image_size 64 --finetune_decoder False --encoder_path ./checkpoints/diffusion/RO/base/stage1/checkpoints/ema_0.9999_020000.pt"
TRAIN_FLAGS="--lr 2e-5 --batch_size 24 --schedule_sampler loss-second-moment --model_path ./diffusion/ckpt/base.pt --lr_anneal_steps 20000"
DIFFUSION_FLAGS=""
SAMPLE_FLAGS="--num_samples 2 --sample_c 1"
DATASET_FLAGS="--data_dir ./diffusion/datasets/RO_train.txt --val_data_dir ./diffusion/datasets/RO_train.txt --mode coco-edge"
python ./diffusion/train.py $MODEL_FLAGS $TRAIN_FLAGS $SAMPLE_FLAGS $DIFFUSION_FLAGS $DATASET_FLAGS
export CUDA_LAUNCH_BLOCKING=1
export LOGDIR=./checkpoints/diffusion/RO/base/stage2-decoder/
MODEL_FLAGS="--learn_sigma True --uncond_p 0.2 --image_size 64 --finetune_decoder True"
TRAIN_FLAGS="--lr 3.5e-5 --batch_size 20 --schedule_sampler loss-second-moment --model_path ./diffusion/ckpt/base.pt --encoder_path ./checkpoints/diffusion/RO/base/stage1-cont/checkpoints/ema_0.9999_000201.pt --lr_anneal_steps 200"
MODEL_FLAGS="--learn_sigma True --uncond_p 0.2 --image_size 64 --finetune_decoder True --encoder_path ./checkpoints/diffusion/RO/base/stage1-cont/checkpoints/ema_0.9999_020000.pt"
TRAIN_FLAGS="--lr 3.5e-5 --batch_size 20 --schedule_sampler loss-second-moment --model_path ./diffusion/ckpt/base.pt --lr_anneal_steps 20000"
DIFFUSION_FLAGS=""
SAMPLE_FLAGS="--num_samples 2 --sample_c 2.5"
DATASET_FLAGS="--data_dir ./diffusion/datasets/RO_train.txt --val_data_dir ./diffusion/datasets/RO_train.txt --mode coco-edge"
Expand All @@ -79,15 +79,15 @@ python ./diffusion/train.py $MODEL_FLAGS $TRAIN_FLAGS $SAMPLE_FLAGS $DIFFUSION_

echo "[Diffusion] Testining for river data..."
export LOGDIR=./results/diffusion/
MODEL_FLAGS="--learn_sigma True --model_path ./checkpoints/diffusion/RI/base/stage2-decoder/checkpoints/model000201.pt --sr_model_path ./checkpoints/diffusion/RI/upsample/checkpoints/model000201.pt"
MODEL_FLAGS="--learn_sigma True --model_path ./checkpoints/diffusion/RI/base/stage2-decoder/checkpoints/model020000.pt --sr_model_path ./checkpoints/diffusion/RI/upsample/checkpoints/model020000.pt"
SAMPLE_FLAGS="--num_samples 720 --sample_c 1.3 --batch_size 72"
DATASET_FLAGS="--data_dir ./diffusion/datasets/RI_train.txt --val_data_dir ./diffusion/datasets/RI_test.txt --mode coco-edge"
python ./diffusion/test.py $MODEL_FLAGS $SAMPLE_FLAGS $DATASET_FLAGS


echo "[Diffusion] Testining for road data..."
export LOGDIR=./results/diffusion/
MODEL_FLAGS="--learn_sigma True --model_path ./checkpoints/diffusion/RO/base/stage2-decoder/checkpoints/model000201.pt --sr_model_path ./checkpoints/diffusion/RO/upsample/checkpoints/model000201.pt"
MODEL_FLAGS="--learn_sigma True --model_path ./checkpoints/diffusion/RO/base/stage2-decoder/checkpoints/model020000.pt --sr_model_path ./checkpoints/diffusion/RO/upsample/checkpoints/model020000.pt"
SAMPLE_FLAGS="--num_samples 720 --sample_c 1.3 --batch_size 72"
DATASET_FLAGS="--data_dir ./diffusion/datasets/RO_train.txt --val_data_dir ./diffusion/datasets/RO_test.txt --mode coco-edge"
python ./diffusion/test.py $MODEL_FLAGS $SAMPLE_FLAGS $DATASET_FLAGS
4 changes: 2 additions & 2 deletions run_gan.sh
View file
Original file line number Diff line number Diff line change
Expand Up @@ -47,11 +47,11 @@ python ./gan/test.py \


echo "[GAN] Testining for river data..."
python ./GAN/test.py \
python ./gan/test.py \
--seed=0 \
--model=pix2pix \
--direction=AtoB \
--dataroot="./GAN/datasets/RO" \
--dataroot="./gan/datasets/RO" \
--netG=unet_256 \
--ngf=256 \
--name=gan/RO
60 changes: 60 additions & 0 deletions run_reproduce.sh
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,60 @@
echo "Downloading the best checkpoints"
gdown "https://drive.google.com/drive/folders/17B-PC-7MSnwbiohr4eGrcqg-U62wnHcE?usp=sharing" -O ./ --folder

echo "[GAN] Preprocessing data..."
python ./gan/preprocess.py \
--train_data_folder="./training_dataset" \
--test_data_folder="./testing_dataset" \
--output_folder="./gan/datasets"


echo "[GAN] Testining for river data..."
python ./gan/test.py \
--seed=0 \
--model=pix2pix \
--direction=AtoB \
--dataroot="./gan/datasets/RI" \
--netG=unet_256 \
--ngf=256 \
--name=best_gan/RI


echo "[GAN] Testining for river data..."
python ./gan/test.py \
--seed=0 \
--model=pix2pix \
--direction=AtoB \
--dataroot="./gan/datasets/RO" \
--netG=unet_256 \
--ngf=256 \
--name=best_gan/RO


echo "[Diffusion] Preprocessing data..."
python ./diffusion/preprocess.py \
--train_data_folder="./training_dataset" \
--test_data_folder="./testing_dataset" \
--output_folder="./diffusion/datasets"


echo "[Diffusion] Testining for river data..."
export LOGDIR=./results/best_diffusion/
MODEL_FLAGS="--learn_sigma True --model_path ./checkpoints/best_diffusion/RI/base/stage2-decoder/checkpoints/model020000.pt --sr_model_path ./checkpoints/best_diffusion/RI/upsample/checkpoints/model020000.pt"
SAMPLE_FLAGS="--num_samples 720 --sample_c 1.3 --batch_size 2"
DATASET_FLAGS="--data_dir ./diffusion/datasets/RI_train.txt --val_data_dir ./diffusion/datasets/RI_test.txt --mode coco-edge"
python ./diffusion/test.py $MODEL_FLAGS $SAMPLE_FLAGS $DATASET_FLAGS


echo "[Diffusion] Testining for road data..."
export LOGDIR=./results/best_diffusion/
MODEL_FLAGS="--learn_sigma True --model_path ./checkpoints/best_diffusion/RO/base/stage2-decoder/checkpoints/model020000.pt --sr_model_path ./checkpoints/best_diffusion/RO/upsample/checkpoints/model020000.pt"
SAMPLE_FLAGS="--num_samples 720 --sample_c 1.3 --batch_size 2"
DATASET_FLAGS="--data_dir ./diffusion/datasets/RO_train.txt --val_data_dir ./diffusion/datasets/RO_test.txt --mode coco-edge"
python ./diffusion/test.py $MODEL_FLAGS $SAMPLE_FLAGS $DATASET_FLAGS


echo "[Router] Selecting from the results..."
python ./router/router.py \
--diffusion_results_folder="./results/best_diffusion" \
--gan_results_folder="./results/best_gan" \
--output_folder="./best_submission"

0 comments on commit 706fdf5

Please sign in to comment.