[IEEE TGRS] Implementation of our paper "Saliency at the Helm: Steering Infrared Small Target Detection with Learnable Kernels". Paper
Highlighting our domain-aware LLSKM, unfolding the 'Center substracts Neighbors' pattern.
- Python 3.8
- Windows10, Ubuntu18.04 or higher
- NVDIA GeForce RTX 3090
- pytorch 1.8.0 or higher
- More details from requirements.txt
We used the NUDT-SIRST, IRSTD-1K and sirst-aug for both training and test. Three datasets can be found and downloaded in: NUDT-SIRST, IRSTD-1K, SIRST-Aug.
Please first download these datasets and place the 3 datasets to the folder ./data/.
- The dataset in our project has the following structure:
├──./data/
│ ├── sirst_aug
│ │ ├── images
│ │ │ ├── 000000.png
│ │ │ ├── 000001.png
│ │ │ ├── ...
│ │ ├── img_idx
│ │ │ ├── test.txt
│ │ │ ├── train.txt
│ │ ├── masks
│ │ │ ├── 000000_mask.png
│ │ │ ├── 000001_mask.png
│ │ │ ├── ...
│ ├── ...
-
Install the environment according to
requirements.txt. -
Enter the repo, and run
train_device0.pyto perform network training:
$ python train_device0.py --model_names L2SKNet_FPN --dataset_names NUDT-SIRST IRSTD-1K SIRST-aug-
The
model_namein our code corresponds to the model name in our paper as follows:L2SKNet_FPNfor L2SKNet-FPN;L2SKNet_UNetfor L2SKNet-UNet;L2SKNet_1D_FPNfor L2SKNet-FPN*;L2SKNet_1D_UNetfor L2SKNet-UNet*.Note: The 'Recip' version is on the way out.
-
Checkpoints and Logs will be saved to
./log/, and./log/has the following structure:
├──./log/
│ ├── [dataset_name]
│ │ ├── [model_name]
│ │ │ ├── 1.pth.tar
│ │ │ ├── 2.pth.tar
│ │ │ ├── ...
│ ├── [dataset_name]_[model_name]_[time].txt
- Run
test.pyto generate file of the format .mat and .png (--test_epo 200means test with the 200th epoch model):
$ python test.py --model_names L2SKNet_FPN --dataset_names NUDT-SIRST IRSTD-1K SIRST-aug --test_epo 200- The file generated will be saved to
./result/that has the following structure:
├──./result/
│ ├── [dataset_name]
│ │ ├── img
│ │ │ ├── [model_name]
│ │ │ │ ├── 000000.png
│ │ │ │ ├── 000001.png
│ │ │ │ ├── ...
│ │ ├── mat
│ │ │ ├── [model_name]
│ │ │ │ ├── 000000.mat
│ │ │ │ ├── 000001.mat
│ │ │ │ ├── ...
- Run
cal_metrics.pyfor direct evaluation:
$ python cal_metrics.py --model_names L2SKNet_FPN --dataset_names NUDT-SIRST IRSTD-1K SIRST-aug- The file generated will be saved to
./result/that has the following structure:
├──./result/
│ ├── [dataset_name]_[model_name]_[time].txt
│ ├── [dataset_name]_[model_name].mat
- Run
t_models.pyfor parameters and FLOPs calculation:
$ python t_models.py- Run
t_time.pyfor runtimes calculation:
$ python t_time.pyWe extend our sincere gratitude to Xinyi Ying and colleagues for their outstanding toolbox, BasicIRSTD (Ver. July 24, 2023). Additionally, we would like to thank Luping Zhang for his invaluable contributions to this repository.
For any questions regarding this paper or the code, please feel free to reach out to wufengyi98@163.com.
@ARTICLE{Wu_2024_TGRS,
author = {Wu, Fengyi and Liu, Anran and Zhang, Tianfang and Zhang, Luping and Luo, Junhai and Peng, Zhenming},
title = {Saliency at the Helm: Steering Infrared Small Target Detection with Learnable Kernels},
booktitle = {IEEE Transactions on Geoscience and Remote Sensing},
year = {2024},
doi = {10.1109/TGRS.2024.3521947}
}