The official PyTorch implementation of "SOSNet: Real-Time Small Object Segmentation via Hierarchical Decoding and Example Mining"

Improvement

Method overview

Usage

Installation

• python >= 3.6
• torch >= 1.8.1
• torchvision >= 0.9.1

Then, clone the repo and install the project with:

$ git clone https://github.com/StuLiu/sosnet
$ cd sosnet
$ pip install -e .

Configuration (click to expand)

Create a configuration file in configs. Sample configuration for segformer dataset can be found here. Then edit the fields you think if it is needed. This configuration file is needed for all of training, evaluation and prediction scripts.

Training (click to expand)

Prepare data:

• Camvid 360*480
• Vaihingen
• UAVid2020

Download pretrained module

• mit-b0 and mobilenetV3-large

To train with a single GPU:

$ python tools/train_sosnet.py --cfg configs/CONFIG_FILE.yaml --hier 1 --soem 1

Evaluation (click to expand)

Make sure to set MODEL_PATH of the configuration file to your trained model directory.

$ python tools/val.py --cfg configs/<CONFIG_FILE_NAME>.yaml

To evaluate with multi-scale and flip, change ENABLE field in MSF to true and run the same command as above.

Inference

To make an inference, edit the parameters of the config file from below.

• Change MODEL >> NAME and BACKBONE to your desired pretrained model.
• Change DATASET >> NAME to the dataset name depending on the pretrained model.
• Set TEST >> MODEL_PATH to pretrained weights of the testing model.
• Change TEST >> FILE to the file or image folder path you want to test.
• Testing results will be saved in SAVE_DIR.

## example using ade20k pretrained models
$ python tools/infer.py --cfg configs/CONFIGFILE.yaml
$ python tools/infer_single.py --img_path demo/camvid_0.png --cfg configs/segformer/camvid_mitb0.yaml

Cite/Reference

If you find SOSNet useful in your research, please consider citing:

@ARTICLE{10359121,
  author={Liu, Wang and Kang, Xudong and Duan, Puhong and Xie, Zhuojun and Wei, Xiaohui and Li, Shutao},
  journal={IEEE Transactions on Neural Networks and Learning Systems}, 
  title={SOSNet: Real-Time Small Object Segmentation via Hierarchical Decoding and Example Mining}, 
  year={2023},
  volume={},
  number={},
  pages={1-13},
  doi={10.1109/TNNLS.2023.3338732}}

Acknowledge

This project is based on the repository 'semantic-segmentation'.

• semantic-segmentation