cilinyan / piclick

Official PyTorch implementation of PiClick: Picking the desired mask in click-based interactive segmentation.

PiClick: Picking the desired mask in click-based interactive segmentation

Visual Results

Environment

Training and evaluation environment: Python3.8.16, PyTorch 1.11.0, CentOS 7.9, CUDA 11.4, mmcv-full 1.6.2. Run the following command to install required packages.

pip3 install -r requirements.txt

You also need to configue the paths to the datasets in config.yml before training or testing.

Evaluation

Before evaluation, please download the datasets and models, and then configure the path in config.yml.

Use the following code to evaluate the base model.

# ViT-B
python scripts/evaluate_model.py NoBRS --gpu=0 \
  --checkpoint=./weights/piclick_base448.pth \
  --eval-mode=cvpr \
  --datasets=GrabCut,Berkeley,SBD,DAVIS,PascalVOC,COCO_MVal,ssTEM,BraTS,OAIZIB
# ViT-L
python scripts/evaluate_model.py NoBRS --gpu=0 \
  --checkpoint=./weights/piclick_large448.pth \
  --eval-mode=cvpr \
  --datasets=GrabCut,Berkeley,SBD,DAVIS,PascalVOC,COCO_MVal,ssTEM,BraTS,OAIZIB
# ViT-H
python scripts/evaluate_model.py NoBRS --gpu=0 \
  --checkpoint=./weights/piclick_huge448.pth \
  --eval-mode=cvpr \
  --datasets=GrabCut,Berkeley,SBD,DAVIS,PascalVOC,COCO_MVal,ssTEM,BraTS,OAIZIB

Training

Before training, please download the MAE pretrained weights (click to download: ViT-Base, ViT-Large, ViT-Huge).

Use the following code to train a base model on COCO+LVIS dataset:

# ViT-B
python -m torch.distributed.launch --nproc_per_node=8 --master_port=59566 --use_env train.py \
  models/iter_mask/piclick_base448_cocolvis_itermask.py \
  --batch-size=136 \
  --ngpus=8
# ViT-L
python -m torch.distributed.launch --nproc_per_node=8 --master_port=59516 --use_env train.py \
  models/iter_mask/piclick_large448_cocolvis_itermask.py \
  --batch-size=56 \
  --ngpus=8
# ViT-H (A100, 80G)
python -m torch.distributed.launch --nproc_per_node=8 --master_port=59516 --use_env train.py \
  models/iter_mask/piclick_huge448_cocolvis_itermask.py \
  --batch-size=96 \
  --ngpus=8 \
  --workers 48

Download

PiClick models: Google Drive, Baidu Wangpan

BraTS dataset (369 cases): Google Drive

OAI-ZIB dataset (150 cases): Google Drive

SBD dataset: Google Drive

Other datasets: RITM Github

License

The code is released under the MIT License. It is a short, permissive software license. Basically, you can do whatever you want as long as you include the original copyright and license notice in any copy of the software/source.

Note

The according annotation platform will be released upon publication.

Citation

@article{yan2024piclick,
  title={PiClick: Picking the desired mask from multiple candidates in click-based interactive segmentation},
  author={Yan, Cilin and Wang, Haochen and Liu, Jie and Jiang, Xiaolong and Hu, Yao and Tang, Xu and Kang, Guoliang and Gavves, Efstratios},
  journal={Neurocomputing},
  pages={128083},
  year={2024},
  publisher={Elsevier}
}

Acknowledgement

Our project is developed based on RITM, SimpleClick and mmdetection. We thank the authors for their great work.

piclick's Issues

I have only one GPU(A6000) and can my device train your model?

about the label

https://github.com/cilinyan/PiClick/blob/e472a436959f394480fda2308fd6f104efccc633/isegm/model/losses_despair.py#L552C14-L552C96

In the generated mask in SampleClick, 0 represents the background and 1 represents the foreground target. Here, all labels are set to 0. Why is this? Thanks!

Demo

Can this work use demo.py in SimpleClick to show the visualization?

Difference between SAM and PiClick

The framework seems to be very similar to SAM (segment everything model).

Evaluation without reference masks supervision

Hello, thanks for the interesting work. Currently the model uses reference masks to select the best candidate. How to test a model based on its own estimates of mask scores?

SBD 数据集链接没有了

semantic segmentation

我最近也在研究clicked baesd segmentation ，發現都是做instance segmentation 。
那是否有可能藉由類似的方法來達到在同一場景下，只需點擊少量的次數即把同一類（例如：人）的部分給分割出來。