SimpleClick: Interactive Image Segmentation with Simple Vision Transformers

Qin Liu; Zhenlin Xu; Gedas Bertasius; Marc Niethammer

Source-linked AI summary

SimpleClick: Interactive Image Segmentation with Simple Vision Transformers

Qin Liu, Zhenlin Xu, Gedas Bertasius, Marc Niethammer

arXiv:2210.11006v3Oct 20, 2022cs.CV

TL;DR

Interactive segmentation has largely relied on hierarchical backbones, leaving the plain ViT underexplored for this task. SimpleClick adapts a plain ViT with click and mask inputs, a simple feature pyramid, and a lightweight decoder, achieving strong benchmark performance and practical computational efficiency.

Problem
The central gap is whether a plain, non-hierarchical ViT can serve as an effective backbone for click-based interactive segmentation.
Method
SimpleClick adapts a plain ViT with a simple feature pyramid, lightweight MLP decoder, click encoding, previous-mask input, and end-to-end fine-tuning from MAE weights.
Results
4.15 NoC@90 on SBD improves 21.8% over the previous best result, while the method achieves state-of-the-art performance across reported benchmarks.
Takeaways & Limitations
SimpleClick shows strong generalizability on medical images and is positioned as a practical annotation tool through a tiny model and computational analysis.
Takeaways & Limitations
The ViT-H model is much larger than existing models, and challenging thin, elongated, or cluttered objects remain failure cases.

Abstract

from arXiv · show

Click-based interactive image segmentation aims at extracting objects with a limited user clicking. A hierarchical backbone is the de-facto architecture for current methods. Recently, the plain, non-hierarchical Vision Transformer (ViT) has emerged as a competitive backbone for dense prediction tasks. This design allows the original ViT to be a foundation model that can be finetuned for downstream tasks without redesigning a hierarchical backbone for pretraining. Although this design is simple and has been proven effective, it has not yet been explored for interactive image segmentation. To fill this gap, we propose SimpleClick, the first interactive segmentation method that leverages a plain backbone. Based on the plain backbone, we introduce a symmetric patch embedding layer that encodes clicks into the backbone with minor modifications to the backbone itself. With the plain backbone pretrained as a masked autoencoder (MAE), SimpleClick achieves state-of-the-art performance. Remarkably, our method achieves 4.15 NoC@90 on SBD, improving 21.8% over the previous best result. Extensive evaluation on medical images demonstrates the generalizability of our method. We further develop an extremely tiny ViT backbone for SimpleClick and provide a detailed computational analysis, highlighting its suitability as a practical annotation tool.

1. Introduction

SimpleClick addresses the underexplored use of a plain ViT backbone for click-based interactive segmentation. It reports state-of-the-art results across natural-image benchmarks, generalization to medical images, and computational efficiency for annotation.

Motivation: Interactive segmentation seeks high-quality pixel-level annotations with limited user interaction, especially clicks.Such annotations support applications including video understanding, self-driving, and medical imaging.
Motivation: Hierarchical backbones dominate current interactive segmentation, but plain ViTs can capture global information through self-attention without hierarchical downsampling or FPN-like aggregation.This design also supports transfer from readily available pretrained ViT models such as MAE.
Results: SimpleClick achieves 4.15 NoC@90 on SBD, improving 21.8% over the previous best method without complex FPN-like design or local refinement.The method was evaluated on 10 public benchmarks spanning natural and medical images.
Contributions: SimpleClick is presented as the first plain-backbone method for interactive image segmentation, with state-of-the-art natural-image performance and strong medical-image generalizability.The authors also emphasize computational efficiency for practical annotation tools.

2. Related Work

Prior interactive segmentation research commonly uses ConvNet or hierarchical ViT backbones and click-based protocols, while non-interactive segmentation has established plain ViTs as competitive alternatives. SimpleClick extends this plain-ViT direction to interactive segmentation.

Interactive Image Segmentation: Interactive segmentation methods have evolved from pixel-graph approaches toward ConvNet-based systems supporting boxes, polygons, clicks, and scribbles.Click-based approaches are especially common because of their simplicity and established training and evaluation protocols.
Vision Transformers for Non-Interactive Segmentation: ViT-based methods have achieved competitive segmentation performance, while the original ViT maintains single-scale feature maps without a hierarchy.Swin Transformer, SegFormer, and HRViT instead introduce hierarchical or multi-scale representations.
Vision Transformers for Non-Interactive Segmentation: SimpleClick explores a plain ViT as the backbone for interactive segmentation, extending recent evidence from semantic segmentation and object detection.The supplied related-work passages do not provide parameter values for Table 1.

3. Method

SimpleClick adapts a plain ViT for interactive segmentation with a single-scale backbone, a simple multi-scale feature pyramid, and a lightweight MLP decoder. Clicks and previous masks guide the backbone, which is fine-tuned end to end from MAE-pretrained weights.

Architecture: SimpleClick uses a plain ViT backbone that maintains single-scale feature maps and processes patch embeddings through Transformer blocks.The implementation includes ViT-B, ViT-L, and ViT-H backbones.
Architecture: The segmentation head uses only MLP layers and accounts for up to 1% of model parameters.It produces a 1/4-scale probability map before upsampling to the original resolution.
Training and Inference: SimpleClick uses readily available MAE-pretrained ImageNet-1K ViT weights and fine-tunes the model end to end for interactive segmentation.The supplied passages describe the MAE pretraining and fine-tuning setup but do not provide further training-result metrics.
Architecture: A simple feature pyramid generates multi-scale representations from the backbone’s last feature map using four parallel convolution or deconvolution layers.The lightweight MLP head transforms, upsamples, and combines these feature maps into a segmentation probability map.
Training and Inference: User clicks are encoded as a two-channel disk map and combined with the previous segmentation as input to support iterative refinement and clicking from existing masks.Training uses random and iterative click simulation strategies.

4. Experiments

SimpleClick is evaluated across natural and medical-image benchmarks using click-efficiency, segmentation-quality, generalization, computational, and ablation analyses. It achieves state-of-the-art results while retaining practical efficiency, and the ablations identify the importance of backbone finetuning and multi-scale features.

Evaluation setup: 10 public datasets, including seven natural-image and three medical-image datasets, are used for evaluation.The experiments cover both in-domain benchmark comparisons and out-of-domain medical-image testing.
Comparison with Previous Results: SimpleClick achieves state-of-the-art performance on all five reported natural-image benchmarks.Table 2 compares GrabCut, Berkeley, SBD, DAVIS, and Pascal VOC results with prior methods.
Comparison with Previous Results: 4.15 NoC@90 on the SBD validation set improves the previous best score by 21.8%.The ViT-H model is trained on SBD, and the SBD validation set contains 6671 instances.
Out-of-Domain Evaluation on Medical Images: The models generalize well to three medical datasets without finetuning, with stronger generalization from the larger COCO+LVIS training set.The medical evaluation includes ssTEM, BraTS, and OAIZIB; larger models and training sets also yield more prominent gains in convergence analysis.
Towards Practical Annotation Tool: The extremely tiny SimpleClick backbone outperforms FocalClick models without pretraining, while ViT-H runs at 132ms per click with 3.22G memory consumption.The computational comparison measures parameters, FLOPs, GPU memory, and seconds per click.
Ablation Study: Freezing the ViT backbone substantially worsens performance, multi-scale features matter, and the last backbone feature map suffices for an effective feature pyramid.The parallel feature pyramid does not surpass the simple feature pyramid using only the last feature map.

5. Limitations and Remarks

The authors identify model scale and challenging object configurations as limitations, while positioning SimpleClick as a baseline for future ViT-based interactive segmentation.

The best-performing ViT-H model is much larger than existing models, raising concerns about fairness in comparisons.The authors address this concern by developing a tiny model for fair comparison.
SimpleClick may fail on objects with very thin, elongated shapes or cluttered occlusions.The authors leave improvements for these scenarios to future work.
The authors hope SimpleClick will serve as a strong baseline for future high-performing interactive segmentation methods based on ViTs and large-scale pretraining.

6. Conclusions

SimpleClick is a plain-backbone interactive segmentation method that uses a general-purpose ViT and MAE-pretrained weights. It achieves state-of-the-art performance on natural images, generalizes to medical images, and includes a tiny model and computational analysis for practical annotation.

SimpleClick is the first plain-backbone method for interactive image segmentation.It leverages a general-purpose ViT backbone that can benefit from readily available pretrained ViT models.
SimpleClick achieves state-of-the-art performance on natural images and strong generalizability on medical images.These results use MAE-pretrained weights.
A tiny SimpleClick model and detailed computational analysis highlight its suitability as a practical annotation tool.

A. Datasets

The supplementary dataset setup describes training with SBD or combined COCO+LVIS datasets and places these choices in the context of earlier interactive segmentation training data.

SimpleClick models are trained using either SBD or the combined COCO+LVIS datasets.
Before RITM, most deep learning-based interactive segmentation models used SBD or Pascal VOC.
SBD and Pascal VOC cover 20 categories of general objects, including persons, transportation vehicles, animals, and indoor objects.

B.1. Architectures

The architecture uses plain ViT feature processing with a multi-resolution segmentation head, symmetric click and image patch embeddings, and high-resolution attention settings.

Architecture parameters: The default input size is 448 × 448; ViT-B and ViT-L use 16 × 16 patches, while ViT-H uses 14 × 14 patches.The smaller ViT-H patch size produces a higher-resolution representation.
Feature pyramid: SimpleClick builds four feature-map resolutions from the ViT backbone using convolutional and transposed-convolution layers.The resolutions are {1/32, 1/16, 1/8, 1/4}.
Segmentation head: Feature maps are projected to common dimensions, concatenated, and converted into a one-channel binary segmentation with an MLP and sigmoid function.The four maps are upsampled to 1/4 resolution before concatenation.
Click encoding: Clicks and previous segmentation are encoded as a three-channel disk map and processed by a patch embedding layer matching the RGB image branch.Positive and negative clicks are encoded separately, with disks of radius 5 pixels.
High-resolution attention: High-resolution finetuning interpolates positional encodings and uses non-overlapping window attention with several global blocks.

C. Additional Comparison Results

SimpleClick performs better than other models across four additional datasets, though the gains are less compelling than on larger benchmarks. Human evaluation also shows generalization to unseen medical objects without finetuning.

Additional comparison results: SimpleClick models perform better than other models on GrabCut, Berkeley, DAVIS, and COCO.These convergence results are measured using mean IoU given k clicks (mIoU@k).
Additional comparison results: The additional-dataset results are less compelling than those on SBD or Pascal VOC, likely because these datasets contain fewer evaluation images.GrabCut has 50 instances, compared with 6671 evaluation instances in SBD.
Medical-image generalization: Human-in-the-loop evaluation on ssTEM, OAIZIB, and BraTS shows generalization to unseen medical objects without finetuning.The qualitative results are obtained with human-provided clicks.

D. Human Evaluation on Medical Images

Human evaluation tests SimpleClick on three medical-image datasets with clicks supplied by a person. Easy objects may require one click, while challenging objects can require more than ten clicks for high-quality segmentation despite no medical-data finetuning.

Evaluation setup: Human-in-the-loop evaluation provides all clicks for qualitative testing on ssTEM, OAIZIB, and BraTS.The evaluation supplements the paper’s out-of-domain medical-image experiments.
Observed interaction needs: Simple objects such as cell nuclei in ssTEM may achieve good segmentation with as little as one click.The result is reported from qualitative human evaluation.
Observed interaction needs: Challenging objects such as knee cartilage and brain tumors may require more than ten clicks for high-quality segmentation.These examples come from OAIZIB and BraTS, respectively.
Generalization: The observed medical-image performance is promising because the models are not finetuned on label-scarce medical imaging datasets.The attached videos demonstrate the evaluation process.