Single-frame Infrared Small Target Detction

Supported Datasets

• SIRST
• NUDT-SIRST
• IRSTD-1k
• MWIRSTD

Supported Models

• ACM
• ACLNet
• ISNet
• UIUNet
• DNANet
• SCTransNet
• RDIAN
• EGEUNet
• EffiSegNet
• UNet Series
• ...

Train And Test

Train Code Example

CUDA_VISIBLE_DEVICES=2,3 python train.py --dataset 'SIRST' --model_name 'SCTransNet' --train 1 --test 0 --deep_supervision True --batchsize 16 --epochs 1000 --lr 0.01 --base_size 256 256 --crop_size 256 --optimizer_name 'Adam' --test_epoch 50

Test Code Example

python train.py --dataset 'SIRST' --model_name 'SCTransNet' --train 0 --test 1 --base_size 256 256 --crop_size 256 --save_pred_img  True --pth_path your_pth_path

For the SIRST and NUDT-SIRST datasets, it is recommended to employ the parameters --base_size 256 256 --crop_size 256. Conversely, for the IRSTD-1k dataset, the suggested parameters are --base_size 512 512 --crop_size 512.

References

I sincerely appreciate the following outstanding work and code !

Paper List: Awesome Infrared Small Targets

DTUM: [code] | [paper]

UIUNet: [code] | [paper]

SCTransNet: [code] | [paper]

DNANet: [code] | [paper]

MSHNet: [code] | [paper]

SIRST-5K: [code] | [paper]

MiM-ISTD: [code] | [paper]

IRSAM: [code] | [paper]