Skip to content

cong-yang/BlumNet

Repository files navigation

Introduction

Blumnet is an skeleton detector based on DETR and DeformableDETR. It treats the skeleton as graph, detects skeleton curves, endpoints and junctions, and reconstructs graphs.

Citation

@article{Zhang2022BlumNet,
title={BlumNet: Graph Component Detection for Object Skeleton Extraction},
author={Zhang, Yulu and Sang, Liang and Grzegorzek, Marcin and See, John and Yang, Cong},
booktitle={ACM International Conference on Multimedia},
pages={5527-5536},
year={2022}
}

License

This project is released under the Apache 2.0 license.

Usage

1. Install requirements

  • Linux, CUDA>=10.2, GCC>=5.4, Python>=3.8

We recommend you to use Anaconda to create a conda environment:

conda create -n blumnet python=3.8
conda activate blumnet
  • PyTorch>=1.8, torchvision>=0.9 (following instructions here)

For example, you could install pytorch and torchvision as following:

conda install pytorch=1.9.0 torchvision=0.10.0 cudatoolkit=10.2 -c pytorch
  • Other requirements
pip install -r requirements.txt

2. Compile CUDA operators

This is for deformable transformer.

cd ./models/ops
sh ./make.sh
# unit test (should see all checking is True)
python test.py

3. Download backbone weights

Download the weights for backbones:

bash detection/download_backbones.sh

tree models/pretrained/ # check files

Make sure you get these files for pretrained backbone weights

models/pretrained/
├── swin_base_patch4_window12_384_22k.pth
└── swin_small_patch4_window7_224.pth
└── vgg16_caffe-292e1171.pth

4. Prepare datasets

Please download sk1491 from GoogleDrive and organize them as following:

code_root/
└── datasets/
    └── sk1491/
        ├── test/
        ├── train/
            ├── im
            ├── gt
            └── train_pair.lst

We also provide the other datasets' links: sk506, SymPASCAL, WH-SYMMAX. Note that these datasets' format is a bit different from their official version, the png/jpg files are extracted from original .mat files. As the groundtruth of formal datasets do not distinguish skeletons between multiple objects, we add some stitched images to sk1491, so that we could get more than one skeleton graphs per image. Based on this augmentation, it's easy to try Blumnet on multi-skeletons situation. We name the augmentated dataset sk1491_mix.

For China mainland researchers, we provide downloading link of datasets on Baidu YunPan, its passward is p7fw.

5. Training and Eval

cd $BLUMNET
# start training
bash detection/train_gcd.sh

# eval F1-score
bash detection/test_gcd.sh

# show graph component of testset
bash detection/visualise_gcd.sh

6. Results and trained models.

Backbone sk506 sk1491 WH-SYMMAX SymPASCAL
VGG16 0.750 0.792 0.873 0.515
SwinBase 0.752 0.826 0.877 0.521

Test demo

step 1 download the blumnet(swin_small) and save it as bellow:

code_root/
└── exps/
    └── demo
        └── checkpoint.pth

step 2 run test on demo image

cd $BLUMNET
python test.py

The visualised image is generated in datasets/demo_rst/. As bellow, demo1 and demo2 show 2 skeleton graph components inferred from Blumnet, demo3 shows all the branches collected from curves and points.

Demo1 Demo2 Demo3

Reference

  1. Adalsn
  2. DETR
  3. DeformableDETR
  4. SwinTransformer
  5. HistogramSegmentation

About

No description, website, or topics provided.

Resources

License

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published