SIRST-5K (TGRS 2024)

SIRST-5K: Exploring Massive Negatives Synthesis with Self-supervised Learning for Robust Infrared Small Target Detection

Contents

• Introduction
  • curve
  • Overview
  • Visual
• Dependencies and Installation
• Dataset
• Codes Demos
• Usage
• Quantative Results
• Citations
• Acknowledgement

Introduction

curve

Overview

Visual

Dependencies and Installation

• Following DNANet
• Python == 3.8
• pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
• pip install scikit-image
• pip install tqdm
• pip install matplotlib
• pip install tensorboard==2.14.0
• pip install opencv-python==4.8.0.76

Dataset

Download the dataset download dir models[Baidu Drive][Google Drive]. Currently, the available dataset are:

• SIRST-5K: The dataset synthesized using negatives generation strategies (Fig 2).

Codes Demos

Noise Sampling

# Run Noise_Sampling.py directly
python codes/Noise_Sampling/Noise_Sampling.py

Noise displacement

# Run add_noise.py directly
python codes/Mix_Rot/add_noise.py

Negative

# Run rot_patch.py directly
python codes/Mix_Rot/rot_patch.py

# Run rot_mask.py directly
python codes/Mix_Rot/rot_mask.py

Our negative augmentation strategies can produce large amounts of challenging image data. You can download the SIRST-5K directly for training.

Usage

1. Train.

python train.py --base_size 256 --crop_size 256 --epochs 1500 --dataset [dataset-name] --split_method 50_50  --deep_supervision True --train_batch_size 16 --test_batch_size 16 --mode TXT

2. Test.

python test.py --base_size 256 --crop_size 256 --st_model [trained model path] --model_dir [model_dir] --dataset [dataset-name] --split_method 50_50    --deep_supervision True --test_batch_size 1 --mode TXT 

3. Visulize your predicts.

python visulization.py --base_size 256 --crop_size 256 --st_model [trained model path] --model_dir [model_dir] --dataset [dataset-name] --split_method 50_50    --deep_supervision True --test_batch_size 1 --mode TXT 

Quantative Results

Model	mIoU (x10(2))	Pd (x10(2))	Fa (x10(6))
Ours	92.78	98.84	2.735	[Weights]

Citation

If you find this project useful for your research, please consider citing our paper. 😃

@ARTICLE{10496142,
  author={Lu, Yahao and Lin, Yupei and Wu, Han and Xian, Xiaoyu and Shi, Yukai and Lin, Liang},
  journal={IEEE Transactions on Geoscience and Remote Sensing}, 
  title={SIRST-5K: Exploring Massive Negatives Synthesis with Self-supervised Learning for Robust Infrared Small Target Detection}, 
  year={2024},
  publisher={IEEE}
  doi={10.1109/TGRS.2024.3387125}
}

Acknowledgement

This project is build based on DNANet. We thank the authors for sharing their code.