Diff-Mosaic :Augmenting Realistic Representations in Infrared Small Target Detection via Diffusion Prior(TGRS 2024) pdf
conda create -n diffmosaic python=3.9
conda activate diffmosaic
pip install -r requirements.txt
- Download Open clip model( [huggingface]) and place it in
./weights/
- Download Pixel-prior model ( NUDT-SIRST , SIRST) and place in
./weights/
- Download Diff-prior model ( NUDT-SIRST, SIRST ) and place in
./weights/
- download the NUDT-SIRST dataset and SIRST dataset
- run
mosaic.py
to generate Mosaic image - run
degrade.py
to get mix image
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
- Download dataset( [google]) and place it in
./detection_model/dataset/
- Downolad pretrained model ( [google]) and place it in
./detection_model/pretrained_model/
- cd
./detection_model/
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
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
This project is build based on DNANet and DiffBIR. We thank the authors for sharing their code.