Diff-Mosaic

Diff-Mosaic ：Augmenting Realistic Representations in Infrared Small Target Detection via Diffusion Prior(TGRS 2024) pdf

Framework

Experiment result

Inastall

conda create -n diffmosaic python=3.9
conda activate diffmosaic
pip install -r requirements.txt

Train Pixel-Piror

Download the NUDT-SIRST dataset and SIRST dataset.

Download pretrained general_swinir_v1.ckpt to ./weight/

python train.py --config ./configs/train_pixel_prior.yaml

Train Diff-Piror

1. Download pretrained v2-1_512-ema-pruned.ckpt to ./weight/

2. Create the initial model weights.

   python scripts/make_stage2_init_weight.py \
   --cldm_config ./configs/model/pixel_prior.yaml \
   --sd_weight ./weights/v2-1_512-ema-pruned.ckpt
   --swinir_weight  ./weights/NUAA_stage1/NUAA.ckpt
   --output ./weights/stage2/nuaa_init.ckpt

3. Start training.

   python train.py --config ./configs/train_diff_prior.yaml

pre-trained model

1. Download Open clip model( [huggingface]) and place it in ./weights/
2. Download Pixel-prior model ( NUDT-SIRST , SIRST) and place in ./weights/
3. Download Diff-prior model ( NUDT-SIRST, SIRST ) and place in ./weights/

Data Preparation

1. download the NUDT-SIRST dataset and SIRST dataset
2. run mosaic.py to generate Mosaic image
3. run degrade.py to get mix image

Generate augmentation sample

python inference.py --input ./add_noise/NUDT_mosaic/  --config configs/model/diff_prior.yaml --ckpt weights/NUDT_stage2/last.ckpt --swinir_ckpt weights/NUDT_stage1/last.ckpt --steps 50 --sr_scale 1 --repeat_times 1 --color_fix_type wavelet --output results/nudt_moc/ --device cuda --use_guidance --g_scale 400 --g_t_start 200

Detection network

1. Download dataset( [google]) and place it in ./detection_model/dataset/
2. Downolad pretrained model ( [google]) and place it in ./detection_model/pretrained_model/
3. cd ./detection_model/

train

python train.py --base_size 256 --crop_size 256 --epochs 3000 --dataset enh_NUDT_aug --split_method 50_50 --model DNANet --backbone resnet_18  --deep_supervision True --train_batch_size 40 --test_batch_size 16 --mode TXT

inference

python test.py --base_size 256 --crop_size 256   --model_dir ./pretrained_model/NUDT.tar --dataset NUDT_aug --split_method 50_50 --model DNANet --backbone resnet_18  --deep_supervision True --test_batch_size 1 --mode TXT

Acknowledgement

This project is build based on DNANet and DiffBIR. We thank the authors for sharing their code.