A Unified SAM-Guided Self-Prompt Learning Framework for Infrared Small Target Detection (TGRS' 2025)

Yimin Fu   Jialin Lyu   Peiyuan Ma   Zhunga Liu   Michael K. Ng

Pytorch implementation for "A Unified SAM-Guided Self-Prompt Learning Framework for Infrared Small Target Detection"

Abstract: Infrared small target detection (ISTD) aims to precisely capture the location and morphology of small targets under all-weather conditions. Compared with generic objects, infrared targets in remote fields of view are smaller in size and exhibit lower signal-to-clutter ratios. This poses a significant challenge in simultaneously preserving low-level target details and understanding high-level contextual semantics, forcing a trade-off between reducing miss detection and suppressing false alarms. In addition, most existing ISTD methods are designed for specific target types under certain infrared platforms, rather than as a unified framework broadly applicable across diverse infrared sensing scenarios. To address these challenges, we propose a unified self-prompt learning framework for ISTD under the guidance of the Segment Anything Model (SAM). Specifically, the model is incorporated with SAM in the encoding stage through a consult-guide manner, adapting the general knowledge to facilitate task-specific contextual understanding. Then, shallow-layer features are employed to generate self-derived prompts, which bidirectionally interact with encoded latent representations to complement subtle low-level details. Moreover, the semantic inconsistency during resolution recovery is mitigated by integrating a mutual calibration module into skip connections, ensuring coherent spatial-semantic fusion. Extensive experiments are conducted on four public ISTD datasets, and the results demonstrate that the proposed method consistently achieves superior performance across different infrared sensing platforms and target types.

Requirements

To install the requirements, you can run the following in your environment first:

pip install -r requirements.txt

To run the code with CUDA properly, you can comment out torch and torchvision in requirement.txt, and install the appropriate version of torch and torchvision according to the instructions on PyTorch.

Or you can use uv to install dependencies:

uv sync

Datasets

For the dataset used in this paper, please download the following datasets NUDT-SIRST / IRSTD-1k / IRSTDID-SKY / NUDT-Sea and move them to ./dataset.

Or you can access all the datasets we have collected via Baidu Netdisk.

Run The Code

1. Download SAM2 Pretrained Weights

Before running the code, please download the SAM-related pretrained weights by executing:

bash sam_spl/sam2_ckpt/download_ckpts.sh

Or you can access them via Baidu Netdisk.

2. Train or Test the Model

You can train or test the model using the provided scripts. For example:

# Training example
python training.py --batch_size 12 --image_size 256 --lr 1e-2 --dataset NUDT-SIRST --save_dir ./checkpoints/NUDT-SIRST --gpu 0

# Distributed Data Parallel Training example
CUDA_VISIBLE_DEVICES="0, 1, 2, 3" torchrun --nproc_per_node=4 --nnodes=1 training.py --batch_size 12 --image_size 256 --lr 1e-2 --dataset NUDT-SIRST --save_dir ./checkpoints/NUDT-SIRST --use_ddp

# Testing example
python testing.py --dataset NUDT-SIRST --image_size 256 --weights ./checkpoints/NUDT-SIRST.pt --device cuda:0

Results

Quantative Results

Dataset	IoU (%)	F1 (%)	Pd (%)	Fa (10^-6)	Weight
NUDT-SIRST	94.63	97.24	99.47	2.55	Weight
IRSTD-1k	74.09	85.11	92.59	9.28	Weight
IRSTDID-SKY	73.40	84.66	98.72	0.97	Weight
NUDT-sea	60.93	75.72	82.33	14.92	Weight

Qualitative Results

Citation

If you find our work and this repository useful. Please consider giving a star ⭐ and citation.

@article{fu2025unified,
  title={A Unified SAM-Guided Self-Prompt Learning Framework for Infrared Small Target Detection},
  author={Fu, Yimin and Lyu, Jialin and Ma, Peiyuan and Liu, Zhunga and Ng, Michael K},
  journal={IEEE Transactions on Geoscience and Remote Sensing},
  year={2025},
  volume={63},
  pages={1-14},
  doi={10.1109/TGRS.2025.3610919},
  publisher={IEEE}
}