🔥🔥 Pick of the Bunch: Detecting Infrared Small Targets Beyond Hit-Miss Trade-Offs via Selective Rank-Aware Attention 🔥🔥

Yimian Dai, Peiwen Pan, Yulei Qian, Yuxuan Li, Xiang Li, Jian Yang, Huan Wang

This repository is the official site for "Pick of the Bunch: Detecting Infrared Small Targets Beyond Hit-Miss Trade-Offs via Selective Rank-Aware Attention".

Abstract

Infrared small target detection faces the inherent challenge of precisely localizing dim targets amidst complex background clutter. Traditional approaches struggle to balance detection precision and false alarm rates. To break this dilemma, we propose SeRankDet, a deep network that achieves high accuracy beyond the conventional hit-miss trade-off, by following the ``Pick of the Bunch'' principle. At its core lies our Selective Rank-Aware Attention (SeRank) module, employing a non-linear Top-K selection process that preserves the most salient responses, preventing target signal dilution while maintaining constant complexity. Furthermore, we replace the static concatenation typical in U-Net structures with our Large Selective Feature Fusion (LSFF) module, a dynamic fusion strategy that empowers SeRankDet with adaptive feature integration, enhancing its ability to discriminate true targets from false alarms. The network's discernment is further refined by our Dilated Difference Convolution (DDC) module, which merges differential convolution aimed at amplifying subtle target characteristics with dilated convolution to expand the receptive field, thereby substantially improving target-background separation. Despite its lightweight architecture, the proposed SeRankDet sets new benchmarks in state-of-the-art performance across multiple public datasets. The code is available at https://github.com/GrokCV/SeRankDet.

• Abstract
• Installation
  • Step 1: Create a conda environment
  • Step 2: Install PyTorch
  • Step 3: Install OpenMMLab Codebases
• Dataset Preparation
  • File Structure
  • Datasets Link
• Training
  • Single GPU Training
  • Multi GPU Training
  • Notes
• Test
• Model Zoo and Benchmark
  • Leaderboard
  • Model Zoo
• Citation

Installation

Step 1: Create a conda environment

$ conda create --name SeRankDet python=3.8
$ source activate SeRankDet

Step 2: Install PyTorch

conda install pytorch torchvision torchaudio pytorch-cuda=10.0 -c pytorch -c nvidia

Step 3: Install OpenMMLab Codebases

# openmmlab codebases
pip install -U openmim
mim install mmcv-full==1.4.0
mim install mmdet==2.25.0
mim install mmsegmentation==0.28.0

Note: make sure you have cd to the root directory of SeRankDet

$ git clone git@github.com:GrokCV/SeRankDet.git
$ cd SeRankDet

Dataset Preparation

File Structure

|- datasets
   |- NUAA
      |-trainval
        |-images
          |-Misc_1.png
          ......
        |-masks
          |-Misc_1.png
          ......
      |-test
        |-images
          |-Misc_50.png
          ......
        |-masks
          |-Misc_50.png
          ......
   |-IRSTD1k
   |-NUDT
   |-SIRSTAUG

Please make sure that the path of your data set is consistent with the data_root in configs/_base_/datasets/dataset_name.py

Datasets Link

The datasets used in this project and the dataset split files can be downloaded from the following links:

• NoisySIRST Dataset
  • Baidu Netdisk
  • OneDrive
• SIRST Dataset
  • Baidu Netdisk
  • OneDrive
• SIRST-AUG Dataset
  • Baidu Netdisk
  • OneDrive
• NUDT-SIRST Dataset
  • Baidu Netdisk
  • OneDrive
• IRSTD1K Dataset
  • Baidu Netdisk
  • OneDrive

Training

Single GPU Training

python train.py <CONFIG_FILE>

For example:

python train.py configs/unetseries/unetseries_serankdet_512x512_500e_irstd1k.py

Multi GPU Training

nproc_per_node is the number of gpus you are using.

python -m torch.distributed.launch --nproc_per_node=[GPU_NUMS] train.py <CONFIG_FILE>

For example:

python -m torch.distributed.launch --nproc_per_node=4 train.py configs/unetseries/unetseries_serankdet_512x512_500e_irstd1k.py

Notes

• Be sure to set args.local_rank to 0 if using Multi-GPU training.

Test

python test.py <CONFIG_FILE> <SEG_CHECKPOINT_FILE>

For example:

python test.py configs/unetseries/unetseries_serankdet_512x512_500e_irstd1k.py work_dirs/unetseries_serankdet_512x512_500e_irstd1k/20221009_231431/best_mIoU.pth.tar

If you want to visualize the result, you only add --show at the end of the above command.

The default image save path is under <SEG_CHECKPOINT_FILE>. You can use --work-dir to specify the test log path, and the image save path is under this path by default. Of course, you can also use --show-dir to specify the image save path.

Model Zoo and Benchmark

Note: Both passwords for BaiduYun and OneDrive is grok.

Leaderboard

Model	SIRST		IRSTD1k		SIRSTAUG		NUDT-SIRST
	IoU	nIoU	IoU	nIoU	IoU	nIoU	IoU	nIoU
ACM	72.88	72.17	63.39	60.81	73.84	69.83	68.48	69.26
RDIAN	72.85	73.96	64.37	64.90	74.19	69.80	81.06	81.72
AGPCNet	77.13	75.19	68.81	66.18	74.71	71.49	88.71	87.48
DNANet	75.55	75.90	68.87	67.53	74.88	70.23	92.67	92.09
MTUNet	78.75	76.82	67.50	66.15	74.70	70.66	87.49	87.70
UIUNet	80.08	78.09	69.13	67.19	74.24	70.57	90.77	90.17
ABC	81.01	79.00	72.02	68.81	76.12	71.83	92.85	92.45
SeRankDet	81.27	79.66	73.66	69.11	76.49	71.98	94.28	93.69

Model Zoo

Checkpoint and Train log: BaiduCloud

Citation

If you use this toolbox or benchmark in your research, please cite this project.

@article{dai2024SeRankDet,
  title={Pick of the Bunch: Detecting Infrared Small Targets Beyond Hit-Miss Trade-Offs via Selective Rank-Aware Attention},
  journal={IEEE Transactions on Geoscience and Remote Sensing}, 
  author={Dai, Yimian and Pan, Peiwen and Qian, Yulei and Li, Yuxuan and Li, Xiang and Yang, Jian and Wang, Huan},
  year={2024},
  volume={62},
  number={},
  pages={1-15}
}