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!
- 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.
-
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.
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.
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.
- Python 3
- pytorch 1.2.0 or higher
- numpy, PIL, tqdm, shutil
- NUST-SIRST [download] [paper]
- NUAA-SIRST [download] [paper]
- NUDT-SIRST [download] [paper]
- IRSTD-1K [download dir] [paper]
- NUDT-SIRST-Sea [download] [paper]
- IRDST [download] [paper]
We used the NUAA-SIRST, 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/ │ ├── NUAA-SIRST │ │ ├── images │ │ │ ├── XDU0.png │ │ │ ├── XDU1.png │ │ │ ├── ... │ │ ├── masks │ │ │ ├── XDU0.png │ │ │ ├── XDU1.png │ │ │ ├── ... │ │ ├── img_idx │ │ │ ├── train_NUAA-SIRST.txt │ │ │ ├── test_NUAA-SIRST.txt │ ├── NUDT-SIRST │ │ ├── images │ │ │ ├── 000001.png │ │ │ ├── 000002.png │ │ │ ├── ... │ │ ├── masks │ │ │ ├── 000001.png │ │ │ ├── 000002.png │ │ │ ├── ... │ │ ├── img_idx │ │ │ ├── train_NUDT-SIRST.txt │ │ │ ├── test_NUDT-SIRST.txt │ ├── ...
Compile DCN for ISNet:
- Cd to
model/ISNet/DCNv2. - run
bash make.sh. The scripts will build DCNv2 automatically and create some folders. - To skip the use of DCNv2, you have to annotate ISNet in
model/__init__.py.
- Run
train.pyto 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 NUAA-SIRST $ CUDA_VISIBLE_DEVICES=0,1 python train.py --model_names ACM ALCNet --dataset_names NUAA-SIRST - Checkpoints and Logs will be saved to
./log/, and the./log/has the following structure:├──./log/ │ ├── [dataset_name] │ │ ├── [model_name]_eopch400.pth.tar
- 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() ...
-
Run
test.pyto perform network inference and evaluation. Example for test [model_name] on [dataset_name] datasets:$ python test.py --model_names ACM ALCNet --dataset_names NUAA-SIRST -
The PA/mIoU and PD/FA values of each dataset will be saved to
./test_[current time].txt -
Network predictions will be saved to
./results/that has the following structure:├──./results/ │ ├── [dataset_name] │ │ ├── [model_name] │ │ │ ├── XDU0.png │ │ │ ├── XDU1.png │ │ │ ├── ... │ │ │ ├── XDU20.png
- Run
inference.pyto inference only with images. Examples:$ python inference.py --model_names ACM --dataset_names NUAA-SIRST - Network predictions will be saved to
./results/that has the following structure:├──./results/ │ ├── [dataset_name] │ │ ├── [model_name] │ │ │ ├── XDU0.png │ │ │ ├── XDU1.png │ │ │ ├── ... │ │ │ ├── XDU20.png
- 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.pyfor direct eevaluation. Example for evaluate [model_name] on [dataset_name] datasets:$ python evaluate.py --model_names ACM --dataset_names NUAA-SIRST - The PA/mIoU and PD/FA values of each dataset will be saved to
./eval_[current time].txt
- Run
cal_params.pyfor 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
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.
|
Methods |
#Params |
FLOPs |
NUAA-SIRST |
NUDT-SIRST |
IRSTD-1K |
||||||
|
IoU |
Pd |
Fa |
IoU |
Pd |
Fa |
IoU |
Pd |
Fa |
|||
|
Top-Hat |
- |
- |
7.142 |
79.841 |
1012.003 |
20.724 |
78.408 |
166.704 |
10.062 |
75.108 |
1432.003 |
|
Max-Median |
- |
- |
1.168 |
30.196 |
55.332 |
4.201 |
58.413 |
36.888 |
7.003 |
65.213 |
59.731 |
|
RLCM |
- |
- |
21.022 |
80.612 |
199.154 |
15.139 |
66.348 |
162.996 |
14.623 |
65.658 |
17.949 |
|
WSLCM |
- |
- |
1.021 |
80.987 |
45846.164 |
0.848 |
74.574 |
52391.633 |
0.989 |
70.026 |
15027.084 |
|
TLLCM |
- |
- |
11.034 |
79.473 |
7.268 |
7.059 |
62.014 |
46.118 |
5.357 |
63.966 |
4.928 |
|
MSLCM |
- |
- |
11.557 |
78.332 |
8.374 |
6.646 |
56.827 |
25.619 |
5.346 |
59.932 |
5.410 |
|
MSPCM |
- |
- |
12.837 |
83.271 |
17.773 |
5.859 |
55.866 |
115.961 |
7.332 |
60.270 |
15.242 |
|
IPI |
- |
- |
25.674 |
85.551 |
11.470 |
17.758 |
74.486 |
41.230 |
27.923 |
81.374 |
16.183 |
|
NRAM |
- |
- |
12.164 |
74.523 |
13.852 |
6.931 |
56.403 |
19.267 |
15.249 |
70.677 |
16.926 |
|
RIPT |
- |
- |
11.048 |
79.077 |
22.612 |
29.441 |
91.850 |
344.303 |
14.106 |
77.548 |
28.310 |
|
PSTNN |
- |
- |
22.401 |
77.953 |
29.109 |
14.848 |
66.132 |
44.170 |
24.573 |
71.988 |
35.261 |
|
MSLSTIPT |
- |
- |
10.302 |
82.128 |
1131.002 |
8.341 |
47.399 |
88.102 |
11.432 |
79.027 |
1524.004 |
|
0.398M |
0.402G |
69.440 |
92.015 |
22.707 |
64.855 |
96.720 |
28.587 |
60.326 |
93.266 |
68.494 |
|
|
0.427M |
0.378G |
61.047 |
87.072 |
55.978 |
61.131 |
97.249 |
29.093 |
58.088 |
92.929 |
74.453 |
|
|
0.966M |
30.618G |
70.491 |
95.057 |
67.983 |
81.236 |
97.778 |
6.343 |
61.852 |
90.236 |
31.561 |
|
|
0.217M |
3.718G |
70.737 |
95.057 |
48.158 |
82.419 |
98.836 |
14.845 |
59.939 |
87.205 |
33.307 |
|
|
4.697M |
14.261G |
74.815 |
93.536 |
38.279 |
94.192 |
99.259 |
2.436 |
65.735 |
89.562 |
12.336 |
|
|
2.752M |
7.944G |
75.928 |
96.198 |
38.897 |
91.762 |
98.519 |
3.769 |
65.014 |
93.939 |
26.437 |
|
|
50.540M |
54.426G |
77.531 |
92.395 |
9.330 |
90.517 |
98.836 |
8.342 |
65.690 |
91.246 |
13.475 |
|
- 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.
We would like to thank Boyang Li, Ruojing Li, Tianhao Wu and Ting Liu for the helpful discussions and insightful suggestions regarding this repository.
Welcome to raise issues or email to yingxinyi18@nudt.edu.cn for any question regarding our BasicIRSTD.