Skip to content

Yangr116/BoxSnake

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

13 Commits
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

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

Arch Backbone lr
sched
mask
AP
mask
AP
Download
RCNN R50-FPN 1X 31.1 config weights
RCNN R50-FPN 2X 31.6 config weights
RCNN R101-FPN 1X 31.6 config weights
RCNN R101-FPN 2X 32.1 config weights
RCNN Swin-B-FPN 1X 38.3 config weights
RCNN Swin-L-FPN 1X 38.9 config weights

mask AP is the result on validation set.

Cityscapes

Arch Backbone lr
sched
mask
AP
config Download
RCNN R50-FPN 24K iter 26.3 config weights

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}
}