BasicISOS

English | 简体中文

BasicISOS is a framework for Infrared Small Object Segmentation, Which supports:

Training
Resume training from checkpoint
Multi-GPUs Training (DDP)
Test and inference

It also including the official implementation of our paper Local Contrast and Global Contextual Information Make Infrared Small Object Salient Again

It also can be applied to other binary segmentation task (including: medical segmentation , salient object detection) of course.

Train your own model

Just follow custom_configs/custom_model.md, and you can easily do experiments with your own network.

Requirements

numpy >= 1.21.2 (maybe lower is ok)

opencv >= 4.5.3 (maybe lower is ok)

pytorch >= 1.9.1 (maybe >= 1.8.1 is also ok)

albumentations >= 0.5.2

etc.

📋 Note: The whole environment we use is in requirements.txt, but you don't have to use exactly the same version as we use.

Dataset Structe

If you want to train on custom datasets you should paper dataset as following structure:

|-SIRST
    |-trainval
        |-images
            |-xxx.png
        |-masks
            |-xxx.png
    |-test
        |-images
            |-xxx.png
        |-masks
            |-xxx.png

We have re-organized SIRST and IRSTD, you can just download and unzip.

Training

To train the model, run this command:

python train.py --opt ./options/ucf_train.yaml

📋 Once you run this train command it will automatically create a folder under ./experiments and save all the logs, tensorboard for losses and metrics, and the model params and training states.

Evaluation

To evaluate pretrained model, run:

python test.py --opt ./test_options/ucf_test.yaml

📋 usage example:

download and unzip datasets

download pretrained model

modified the config file ./test_options/test_demo.yaml, specifically: device (cpu are not support right now), data_root and net_path in your own case

Pre-trained Models and Results

You can download pretrained models (we also provide the whole training logs) here:

UCF for SIRST

Model name	IoU	nIoU
UCF Net	80.89	78.72

UCF for IRSTD

Model name	IoU	nIoU
UCF Net	68.92	69.26

Advanced

We have provided some networks: FPN, HRnet, SwinT, U2net, unet and so on, you can try yourself.

We will introduce how to use your own data augmentation function , how to use your own network and so on.

Contributing

📋 BasicISOS v1.0 is almost build by wcyjerry (some codes are modified from BasicSR), if you want to help us to maintain and improve it, you can email me at: wcyjerry@qq.com

Feel free to pull requests and issues.

kettogg / BasicISOS