Pytorch implementation of Tokencut:
Self-supervised Transformers for Unsupervised Object Discovery using Normalized Cut
Yangtao Wang, Xi Shen, Shell Xu Hu, Yuan Yuan, James L. Crowley, Dominique Vaufreydaz
[Project page]
[ Github (Video Segmentation) ]
[Paper]
If our project is helpful for your research, please consider citing :
@inproceedings{wang2022tokencut,
title={Self-supervised Transformers for Unsupervised Object Discovery using Normalized Cut},
author={Wang, Yangtao and Shen, Xi and Hu, Shell Xu and Yuan, Yuan and Crowley, James L. and Vaufreydaz, Dominique},
booktitle={Conference on Computer Vision and Pattern Recognition}
year={2022}
}
09/06/2022 Extension work of TokeCut Video Segmentation is realised!
03/10/2022
Creating a 480p Demo using Gradio. Try out the Web Demo:
Internet image results:
02/26/2022
Integrated into Huggingface Spaces 🤗 using Gradio. Try out the Web Demo:
02/26/2022 A simple TokenCut Colab Demo is available.
02/21/2022 Initial commit: Code of TokenCut is released, including evaluation of unsupervised object discovery, unsupervised saliency object detection, weakly supervised object locolization.
This code was implemented with Python 3.7, PyTorch 1.7.1 and CUDA 11.2. Please refer to the official installation. If CUDA 10.2 has been properly installed :
pip install torch==1.7.1 torchvision==0.8.2
In order to install the additionnal dependencies, please launch the following command:
pip install -r requirements.txt
We provide quick download commands in DOWNLOAD_DATA.md for VOC2007, VOC2012, COCO, CUB, ImageNet, ECSSD, DUTS and DUT-OMRON as well as DINO checkpoints.
We provide TokenCut visualization for bounding box prediction and attention map. Using all for all visualization results.
python main_tokencut.py --image_path examples/VOC07_000036.jpg --visualize pred
python main_tokencut.py --image_path examples/VOC07_000036.jpg --visualize attn
python main_tokencut.py --image_path examples/VOC07_000036.jpg --visualize all
We provide TokenCut segmentation results as follows:
cd unsupervised_saliency_detection
python get_saliency.py --sigma-spatial 16 --sigma-luma 16 --sigma-chroma 8 --vit-arch small --patch-size 16 --img-path ../examples/VOC07_000036.jpg --out-dir ./output
Following are the different steps to reproduce the results of TokenCut presented in the paper.
In order to apply TokenCut and compute corloc results (VOC07 68.8, VOC12 72.1), please launch:
python main_tokencut.py --dataset VOC07 --set trainval
python main_tokencut.py --dataset VOC12 --set trainval
If you want to extract Dino features, which corresponds to the KEY features in DINO:
mkdir features
python main_lost.py --dataset VOC07 --set trainval --save-feat-dir features/VOC2007
Results are provided given the 2014 annotations following previous works. The following command line allows you to get results on the subset of 20k images of the COCO dataset (corloc 58.8), following previous litterature. To be noted that the 20k images are a subset of the train set.
python main_tokencut.py --dataset COCO20k --set train
We have tested the method on different setups of the VIT model, corloc results are presented in the following table (more can be found in the paper).
| arch | pre-training | dataset | ||
|---|---|---|---|---|
| VOC07 | VOC12 | COCO20k | ||
| ViT-S/16 | DINO | 68.8 | 72.1 | 58.8 |
| ViT-S/8 | DINO | 67.3 | 71.6 | 60.7 |
| ViT-B/16 | DINO | 68.8 | 72.4 | 59.0 |
Previous results on the dataset VOC07 can be obtained by launching:
python main_tokencut.py --dataset VOC07 --set trainval #VIT-S/16
python main_tokencut.py --dataset VOC07 --set trainval --patch_size 8 #VIT-S/8
python main_tokencut.py --dataset VOC07 --set trainval --arch vit_base #VIT-B/16
To evaluate on ECSSD, DUTS, DUT_OMRON dataset:
python get_saliency.py --out-dir ECSSD --sigma-spatial 16 --sigma-luma 16 --sigma-chroma 8 --nb-vis 1 --vit-arch small --patch-size 16 --dataset ECSSD
python get_saliency.py --out-dir DUTS --sigma-spatial 16 --sigma-luma 16 --sigma-chroma 8 --nb-vis 1 --vit-arch small --patch-size 16 --dataset DUTS
python get_saliency.py --out-dir DUT --sigma-spatial 16 --sigma-luma 16 --sigma-chroma 8 --nb-vis 1 --vit-arch small --patch-size 16 --dataset DUT
This should give:
| Method | ECSSD | DUTS | DUT-OMRON | ||||||
|---|---|---|---|---|---|---|---|---|---|
| maxF | IoU | Acc | maxF | IoU | Acc | maxF | IoU | Acc | |
| TokenCut | 80.3 | 71.2 | 91.8 | 67.2 | 57.6 | 90.3 | 60.0 | 53.3 | 88.0 |
| TokenCut + BS | 87.4 | 77.2 | 93.4 | 75.5 | 62,4 | 91.4 | 69.7 | 61.8 | 89.7 |
To finetune ViT-S/16 on CUB on a single node with 4 gpus for 1000 epochs run:
python -m torch.distributed.launch --nproc_per_node=4 main.py --data_path /path/to/data --batch_size_per_gpu 256 --dataset cub --weight_decay 0.005 --pretrained_weights ./dino_deitsmall16_pretrain.pth --epoch 1000 --output_dir ./path/to/checkpoin --lr 2e-4 --warmup-epochs 50 --num_labels 200 --num_workers 16 --n_last_blocks 1 --avgpool_patchtokens true --arch vit_small --patch_size 16
To evaluate a fine-tuned ViT-S/16 on CUB val with a single GPU run:
python eval.py --pretrained_weights ./path/to/checkpoint --dataset cub --data_path ./path/to/data --batch_size_per_gpu 1 --no_center_crop
This should give:
Top1_cls: 79.12, top5_cls94.80, gt_loc: 0.914, top1_loc:0.723
To Evaluate ViT-S/16 finetuned on ImageNet val with a single GPU run:
python eval.py --pretrained_weights /path/to/checkpoint --classifier_weights /path/to/linear_weights--dataset imagenet --data_path ./path/to/data --batch_size_per_gpu 1 --num_labels 1000 --batch_size_per_gpu 1 --no_center_crop --input_size 256 --tau 0.2 --patch_size 16 --arch vit_small
TokenCut code is built on top of LOST, DINO, Segswap, and Bilateral_Sovlver. We would like to sincerely thanks those authors for their great works.