Official implementation for IEEE Transactions on Geoscience and Remote Sensing (TGRS) paper: "Interior Attention-Aware Network for Infrared Small Target Detection". [Paper]
If you find our work useful in your research, please cite our paper
@ARTICLE{9745054,
author={Wang, Kewei and Du, Shuaiyuan and Liu, Chengxin and Cao, Zhiguo},
journal={IEEE Transactions on Geoscience and Remote Sensing},
title={Interior Attention-Aware Network for Infrared Small Target Detection},
year={2022},
doi={10.1109/TGRS.2022.3163410}
}
Packages:
- Python 3.8
- Pytorch 1.7
- opencv-python
- numpy
- tqdm
- pandas
- yaml
iaanet
├─ box_generate.py
├─ cGAN_data
│ └─ training_box_gt.csv
├─ detect.py
├─ models
│ ├─ attention.py
│ ├─ backbone.py
│ ├─ embedding.py
│ └─ transformer.py
├─ pretrained
│ ├─ iaanet.pt
│ └─ rpn.pt
├─ test.py
├─ train.py
└─ utils
├─ datasets.py
├─ general.py
└─ loss.py
cGAN_data
├─ training
│ ├─ 000000_1.png
│ ├─ 000000_2.png
│ ├─ ...
│ ├─ 009999_1.png
│ └─ 009999_2.png
├─ test_org
│ ├─ 00000.png
│ ├─ ...
│ └─ 00099.png
└─ test_gt
├─ 00000.png
├─ ...
└─ 00099.png
- Following test dirs to organize validation set:
cGAN_data
├─ ...
├─ val_org
└─ val_gt
- Use prepared bounding boxes ground truth directly
cGAN_data
├─ ...
└─ training_box_gt.csv
- Or run following command to generate bounding box ground truth from ground truth masks:
python box_generate.py --path ./cGAN_data/training/ --save_path ./cGAN_data/training_box_gt.csv --bord 4
Experiments are conducted using PyTorch with a single GeForce RTX 3090 GPU of 24 GB Memory.
Train from scratch
python train.py --batch_size 8 --epochs 10 --save_path ./outputs/demo/
Start from pretraind RPN
python train.py --rpn_pretrained ./pretrained/rpn.pt --save_path ./outputs/demo/
Run python train.py --help for more configurations
Use pretrained model for testing
python test.py --weights ./pretrained/iaanet.pt
Fast version (SG convs the proposed regions only. ):
python test.py --weights ./pretrained/iaanet.pt --fast
Run python test.py --help for more configurations
We follow MDvsFA-cGAN to calculate F-measure [Code]
| Dataset | F-measure | Precision | Recall |
|---|---|---|---|
| MDvsFA | 0.639 | 0.606 | 0.818 |
Infer a single image
python detect.py --image_path img.png --save_path ./inference/ --weights ./pretrained/iaanet.pt
Infer images in a folder
python detect.py --image_path ./folder/ --save_path ./inference/ --weights ./pretrained/iaanet.pt --folder
Fast version
python detect.py --fast
Run python detect.py --help for more configurations
MIT License
React
Typescript
Django
Laravel
Facebook
Microsoft
Google
Alibaba
D3
Tencent