Awesome

BasicIRSTD

BasicIRSTD is a PyTorch-based open-source and easy-to-use toolbox for infrared small target detction (IRSTD). This toolbox introduces a simple pipeline to train/test your methods, and builds a benchmark to comprehensively evaluate the performance of existing methods. Our BasicIRSTD can help researchers to get access to infrared small target detction quickly, and facilitates the development of novel methods. Welcome to contribute your own methods to the benchmark.

Note: This repository will be updated on a regular basis. Please stay tuned!

Contributions

• We provide a PyTorch-based open-source and easy-to-use toolbox for IRSTD.
• We re-implement a number of existing methods on the unified datasets, and develop a benchmark for performance evaluation.
• We share the codes, models and results of existing methods to help researchers better get access to this area.

News & Updates

• Apirl 4, 2022: Pulic The BasicIRSTD ToolBox.

• July 24, 2023: Update BasicIRSTD ToolBox (stable version).

  Code: We fix bugs to ensure the stability of training results. <br style="line-height:2em"> Results: We reprodude all models and update the results in table.

• April 19, 2024: Update README.md.

  New section "Build": We add instructions for DCNv2 compiling of ISNet. <br> New section "Train on your own models": We add instructions for self-defined model usage.

• May 11, 2024: Update README.md.

  Updated section "Recources": Update the link to the pre-trained models and result files. <br> <br><br>

Requirements

• Python 3
• pytorch 1.2.0 or higher
• numpy, PIL, tqdm, shutil <br><br>

Datasets

• NUST-SIRST   [download]   [paper]
• SIRST-v1   [download]   [paper]
• SIRST-v2   [download]   [paper]
• NUDT-SIRST   [download]   [paper]
• IRSTD-1K   [download dir]   [paper]
• NUDT-SIRST-Sea   [download]   [paper]
• IRDST   [download]   [paper]

We used the SIRST-v1, NUDT-SIRST, IRSTD-1K for both training and test. Please first download our datasets via Baidu Drive (key:1113) or Google Drive, and place the 3 datasets to the folder ./datasets/. More results will be released soon!

• Our project has the following structure:

  ├──./datasets/
  │    ├── SIRST-v1
  │    │    ├── images
  │    │    │    ├── XDU0.png
  │    │    │    ├── XDU1.png
  │    │    │    ├── ...
  │    │    ├── masks
  │    │    │    ├── XDU0.png
  │    │    │    ├── XDU1.png
  │    │    │    ├── ...
  │    │    ├── img_idx
  │    │    │    ├── train_SIRST-v1.txt
  │    │    │    ├── test_SIRST-v1.txt
  │    ├── NUDT-SIRST
  │    │    ├── images
  │    │    │    ├── 000001.png
  │    │    │    ├── 000002.png
  │    │    │    ├── ...
  │    │    ├── masks
  │    │    │    ├── 000001.png
  │    │    │    ├── 000002.png
  │    │    │    ├── ...
  │    │    ├── img_idx
  │    │    │    ├── train_NUDT-SIRST.txt
  │    │    │    ├── test_NUDT-SIRST.txt
  │    ├── ...  
  

Build

Compile DCN for ISNet: <br>

1. Cd to model/ISNet/DCNv2.
2. run bash make.sh. The scripts will build DCNv2 automatically and create some folders.
3. To skip the use of DCNv2, you have to annotate ISNet in model/__init__.py.

Commands

Commands for training

• Run train.py to perform network training in single GPU and multiple GPUs. Example for training [model_name] on [dataset_name] datasets:

  $ python train.py --model_names ACM ALCNet --dataset_names SIRST-v1
  $ CUDA_VISIBLE_DEVICES=0,1 python train.py --model_names ACM ALCNet --dataset_names SIRST-v1
  

• Checkpoints and Logs will be saved to ./log/, and the ./log/ has the following structure:

  ├──./log/
  │    ├── [dataset_name]
  │    │    ├── [model_name]_eopch400.pth.tar
  

Train on your own models

• Create a folder in ./model, and put your own model in this folder.

  ├──./model/
  │    ├── xxxNet
  │    │    ├── model.py
  

• Add the model in model/__init__.py..

  from model.ACM.model_ACM import ASKCResUNet as ACM
  ...
  from model.xxxNet.model import net as xxxNet
  

• Add the model in net.py..

  if model_name == 'DNANet':
     self.model = DNANet(mode='train')
  ...
  elif model_name == 'xxxNet':
     self.model = xxxNet()
  ...
  

Commands for test

• Run test.py to perform network inference and evaluation. Example for test [model_name] on [dataset_name] datasets:

  $ python test.py --model_names ACM ALCNet --dataset_names SIRST-v1
  

• The PA/mIoU and PD/FA values of each dataset will be saved to ./test_[current time].txt<br>

• Network predictions will be saved to ./results/ that has the following structure:

  ├──./results/
  │    ├── [dataset_name]
  │    │   ├── [model_name]
  │    │   │    ├── XDU0.png
  │    │   │    ├── XDU1.png
  │    │   │    ├── ...
  │    │   │    ├── XDU20.png
  

Commands for inference only with images

• Run inference.py to inference only with images. Examples:

  $ python inference.py --model_names ACM --dataset_names SIRST-v1
  

• Network predictions will be saved to ./results/ that has the following structure:

  ├──./results/
  │    ├── [dataset_name]
  │    │   ├── [model_name]
  │    │   │    ├── XDU0.png
  │    │   │    ├── XDU1.png
  │    │   │    ├── ...
  │    │   │    ├── XDU20.png
  

Commands for evaluate on your own results

• Please first put your results on ./results/ that has the following structure:

  ├──./results/
  │    ├── [dataset_name]
  │    │   ├── [method_name]
  │    │   │    ├── XDU0.png
  │    │   │    ├── XDU1.png
  │    │   │    ├── ...
  │    │   │    ├── XDU20.png
  

• Run evaluate.py for direct eevaluation. Example for evaluate [model_name] on [dataset_name] datasets:

  $ python evaluate.py --model_names ACM --dataset_names SIRST-v1
  

• The PA/mIoU and PD/FA values of each dataset will be saved to ./eval_[current time].txt<br><br>

Commands for parameters/FLOPs calculation

• Run cal_params.py for parameters and FLOPs calculation. Examples:

  $ python cal_params.py --model_names ACM ALCNet
  

• The parameters and FLOPs of each method will be saved to ./params_[current time].txt<br><br> <be>

Benchmark

We benchmark several methods on the above datasets. mIoU, PD and FA metrics under threshold=0.5 are used for quantitative evaluation.

Note: A detailed review of existing IRSTD methods can be referred to Tianfang-Zhang/awesome-infrared-small-targets.

mIoU/PD/FA values achieved by different methods:

Recources

• We provide the result files generated by the aforementioned methods, and researchers can download the results via Baidu Drive (key:1113) and One Drive.
• The pre-trained models of the aforementioned methods can be downlaod via Baidu Drive (key:1113) and One Drive. <br><br>

Acknowledgement

We would like to thank Boyang Li, Ruojing Li, Tianhao Wu and Ting Liu for the helpful discussions and insightful suggestions regarding this repository. <br><br>

Contact

Welcome to raise issues or email to yingxinyi18@nudt.edu.cn for any question regarding our BasicIRSTD.