Home

Awesome

BoxSnake: Polygonal Instance Segmentation with Box Supervision

Rui Yang, Lin Song, Yixiao Ge, Xiu Li

BoxSnake is an end-to-end training technique to achieve effective polygonal instance segmentation using only box annotations. It consists of two loss functions: (1) a point-based unary loss that constrains the bounding box of predicted polygons to achieve coarse-grained segmentation; and (2) a distance-aware pairwise loss that encourages the predicted polygons to fit the object boundaries.

Intro Arxiv Paper | [Video Demo]

Installation


Requirements:

CUDA kernel for MSDeformAttn and Diff rasterizer

BoxSnake also uses the deformable attention modules introduced in Deformable-DETR and the differentiable rasterizer introduced in BoundaryFormer. Please build them on your system:

bash scripts/auto_build.sh

or

cd ./modeling/layers/deform_attn
sh ./make.sh
cd ./modeling/layers/diff_ras
python setup.py build install

Example conda environment setup

conda create --name boxsnake python=3.8 -y
conda activate boxsnake

conda install pytorch==1.9.0 torchvision==0.10.0 cudatoolkit=11.1 -c pytorch -c nvidia
python -m pip install 'git+https://github.com/facebookresearch/detectron2.git'

git clone https://github.com/Yangr116/BoxSnake.git
cd BoxSnake
pip install -r requirements.txt
bash scripts/auto_build.sh

Model Zoo


COCO

ArchBackbonelr<br>schedmask <br>APmask <br>APDownload
RCNNR50-FPN1X31.1configweights
RCNNR50-FPN2X31.6configweights
RCNNR101-FPN1X31.6configweights
RCNNR101-FPN2X32.1configweights
RCNNSwin-B-FPN1X38.3configweights
RCNNSwin-L-FPN1X38.9configweights

mask AP is the result on validation set.

Cityscapes

ArchBackbonelr<br>schedmask <br>APconfigDownload
RCNNR50-FPN24K iter26.3configweights

Getting Start


We use the COCO dataset and Cityscapes dataset. Please following here to prepare them.

If you would like to use swin transformer backbone, please download swin weights from here and convert them to pkl format:

python tools/convert-pretrained-model-to-d2.py ${your_swin_pretrained.pth} ${yout_swin_pretrained.pkl}

Training

To train on COCO dataset using the R50 backbone at a 1X schedule:

# 8 gpus
python train_net.py --num-gpus 8 --config-file configs/COCO-InstanceSegmentation/BoxSnake_RCNN/boxsnake_R_50_FPN_1x.yaml

You can also run below code:

bash scripts/auto_run.sh $CONFIG  # your config

Inference

To inference on COCO validation set using trained weights:

# 8 gpus
python train_net.py --num-gpus 8 --config-file configs/COCO-InstanceSegmentation/BoxSnake_RCNN/boxsnake_R_50_FPN_1x.yaml
 --eval-only MODEL.WEIGHTS ${your/checkpoints/boxsnake_rcnn_R_50_FPN_coco_1x.pth}

Inference on a single image using trained weights:

python demo/demo.py --config-file configs/COCO-InstanceSegmentation/BoxSnake_RCNN/boxsnake_R_50_FPN_1x.yaml --input demo/demo.jpg --output ${/your/visualized/dir} --confidence-threshold 0.5 --opts MODEL.WEIGHTS ${your/checkpoints/boxsnake_rcnn_R_50_FPN_coco_1x.pth}

Others


BoxSnake is inspired by traditional levelset (including boxlevelset) and GVF methods, and you can check below links to learn them:

Some geometric knowledge may help readers to understand the BoxSnake better:

Acknowledgement

If you find BoxSnake helpful, please cite:

@misc{BoxSnake,
      title={BoxSnake: Polygonal Instance Segmentation with Box Supervision}, 
      author={Rui Yang and Lin Song and Yixiao Ge and Xiu Li},
      year={2023},
      eprint={2303.11630},
      archivePrefix={arXiv},
      primaryClass={cs.CV}
}