This repository is the official implementation of BAFE-Net

• DenseSIRST
  • Model Zoo and Benchmark
• BAFE-Net
  • Installation
    • Step 1: Create a conda environment
    • Step 2: Install PyTorch
    • Step 3: Install OpenMMLab Codebases
    • Step 4: Install deepir
  • Train
  • Test

You can download our DenseSIRST dataset from Google Drive.

For both training and inference, the following dataset structure is required:

|- data
    |- SIRSTdevkit
        |-PNGImages
            |-Misc_1.png
            ......
        |-SIRST
            |-BBox
                |-Misc_1.xml
                ......
            |-BinaryMask
                |-Misc_1_pixels0.png
                |-Misc_1.xml
                ......
            |-PaletteMask
                |-Misc_1.png
                ......
            |-Point_label
                |-Misc_1_pixels0.txt
                ......
        |-SkySeg
            |-BinaryMask
                |-Misc_1_pixels0.png
                |-Misc_1.xml
                ......
            |-PaletteMask
                |-Misc_1.png
                ......
        |-Splits
            |-train_v2.txt
            |-test_v2.txt
            ......

• PNGImages is the folder for storing all images.
• SIRST and SkySeg are folders for storing annotation files.
  • SIRST corresponds to infrared small targets.
  • SkySeg corresponds to sky segmentation.

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

Checkpoint and Train log: Google Drive.

Step 1: Create a conda environment

$ conda create --name deepir python=3.9
$ conda activate deepir

Step 2: Install PyTorch

$ conda install pytorch torchvision torchaudio pytorch-cuda=11.8 -c pytorch -c nvidia

Step 3: Install OpenMMLab Codebases

pip install -U openmim
mim install mmengine
mim install "mmcv>=2.0.0"
mim install "mmdet>=3.0.0"
pip install "mmsegmentation>=1.0.0"
pip install dadaptation

Step 4: Install deepir

$ python setup.py develop

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

$ git clone git@github.com:GrokCV/BAFE-Net.git
$ cd BAFE-Net

$ CUDA_VISIBLE_DEVICES=0 python train.py <CONFIG_FILE>

For example:

$ CUDA_VISIBLE_DEVICES=0 python tools/train_det.py configs/detection/fcos_changer_seg/fcos_changer_seg_r50-caffe_fpn_gn-head_1x_densesirst.py

$ CUDA_VISIBLE_DEVICES=0 python test.py <CONFIG_FILE> <SEG_CHECKPOINT_FILE>

For example:

$ CUDA_VISIBLE_DEVICES=0 python tools/test_det.py configs/detection/fcos_changer_seg/fcos_changer_seg_r50-caffe_fpn_gn-head_1x_densesirst.py work_dirs/fcos_changer_seg_r50-caffe_fpn_gn-head_1x_densesirst/20240719_162542/best_pascal_voc_mAP_epoch_8.pth

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.

If you use our dataset or code in your research, please cite this project.

@article{xiao2024bafenet,
	title={Background Semantics Matter: Cross-Task Feature Exchange Network for Clustered Infrared Small Target Detection With Sky-Annotated Dataset}, 
	author={Mengxuan Xiao and Qun Dai and Yiming Zhu and Kehua Guo and Huan Wang and Xiangbo Shu and Jian Yang and Yimian Dai},
	year={2024},
	journal={arXiv preprint arXiv:2407.20078},
}

@article{dai2023one,
  title={One-stage cascade refinement networks for infrared small target detection},
  author={Dai, Yimian and Li, Xiang and Zhou, Fei and Qian, Yulei and Chen, Yaohong and Yang, Jian},
  journal={IEEE Transactions on Geoscience and Remote Sensing},
  volume={61},
  pages={1--17},
  year={2023},
}