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.
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.

Requirements

Python 3
pytorch 1.2.0 or higher
numpy, PIL, tqdm, shutil

Datasets

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
│    ├── ...

Build

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.

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 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

### 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 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 ``` ### Commands for inference only with images * **Run **`inference.py`** to 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 ``` ### 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 NUAA-SIRST ``` * **The PA/mIoU and PD/FA values of each dataset will be saved to** **`./eval_[current time].txt`**

### 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`**

## 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](https://github.com/Tianfang-Zhang/awesome-infrared-small-targets).** ### **mIoU/PD/FA values achieved by different methods:**

Methods	#Params	FLOPs	NUAA-SIRST			NUDT-SIRST			IRSTD-1K
Methods	#Params	FLOPs	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
ACM	0.398M	0.402G	69.440	92.015	22.707	64.855	96.720	28.587	60.326	93.266	68.494
ALCNet	0.427M	0.378G	61.047	87.072	55.978	61.131	97.249	29.093	58.088	92.929	74.453
ISNet	0.966M	30.618G	70.491	95.057	67.983	81.236	97.778	6.343	61.852	90.236	31.561
RDIAN	0.217M	3.718G	70.737	95.057	48.158	82.419	98.836	14.845	59.939	87.205	33.307
DNA-Net	4.697M	14.261G	74.815	93.536	38.279	94.192	99.259	2.436	65.735	89.562	12.336
ISTDU-Net	2.752M	7.944G	75.928	96.198	38.897	91.762	98.519	3.769	65.014	93.939	26.437
UIU-Net	50.540M	54.426G	77.531	92.395	9.330	90.517	98.836	8.342	65.690	91.246	13.475

## Recources * **We provide the result files generated by the aforementioned methods, and researchers can download the results via [Baidu Drive](https://pan.baidu.com/s/1wXxlqp9R6G5aeBCy2ftl-A?pwd=1113) (key:1113) and [One Drive](https://1drv.ms/f/s!AoFCxCGMfhW6qX3-2B54zKzPGI5b?e=rQoxAs).** * **The pre-trained models of the aforementioned methods can be downlaod via [Baidu Drive](https://pan.baidu.com/s/1mhHz0A-6H38kBBJf9hRY6Q?pwd=1113) (key:1113) and [One Drive](https://1drv.ms/f/s!AoFCxCGMfhW6qX4lS1I9V_B9JCgU?e=y3laRj).**

## Acknowledgement **We would like to thank [Boyang Li](https://github.com/YingqianWang), [Ruojing Li](https://github.com/TinaLRJ), [Tianhao Wu](https://github.com/YingqianWang) and [Ting Liu](https://github.com/LiuTing20a) for the helpful discussions and insightful suggestions regarding this repository.**

## Contact **Welcome to raise issues or email to [yingxinyi18@nudt.edu.cn](yingxinyi18@nudt.edu.cn) for any question regarding our BasicIRSTD.**

XinyiYing / BasicIRSTD

readme