[논문 리뷰] SimpleClick: Interactive Image Segmentation with Simple Vision Transformers

Qin Liu, Zhenlin Xu외 2명|arXiv (Cornell University)|2022. 10. 20.

Visual Attention and Saliency Detection인용 수 9

한 줄 요약

SimpleClick는 계층적 아키텍처를 회피하는 단순한 비전 트랜스포머(ViT) 기반 백본을 사용하는 최초의 상호작용 이미지 세그멘테이션 방법을 제안한다. 마스크된 오토에ncoder로 사전학습한 ViT와 대칭 패치 임베딩 레이어를 활용하여 클릭을 직접 인코딩함으로써, SBD에서 4.15 NoC@90의 최고 성능을 달성하며 이전 방법들보다 21.8% 향상되었고, 실세계 어노테이션 툴에 적합한 계산 효율성을 유지한다.

ABSTRACT

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.

연구 동기 및 목표

• 계층적 아키텍처의 지배를 뛰어넘어 단순한 비계층적 비전 트랜스포머를 상호작용 이미지 세그멘테이션의 백본으로 사용할 수 있는지 탐색한다.
• 특수 작업용 백본을 재설계하지 않고도 일반 목적의 마스크된 오토에ncoder로 사전학습한 ViT에서 상호작용 세그멘테이션으로의 전이 학습을 가능하게 한다.
• 의료 영상 및 자원이 제한된 환경에서의 실용적 어노테이션 요구사항을 충족하는 계산 효율적인 모델을 개발한다.
• 단순한 피처 피라미드가 단일 스케일의 마지막 피처 맵에서만 구성되어도 고품질의 세그멘테이션을 위해 충분한지 입증한다.
• 기초 모델을 활용한 상호작용 세그멘테이션 분야의 미래 연구를 위한 강력한 베이스라인을 수립한다.

제안 방법

• 계층적 설계와 전형적으로 사용되는 피처 피라미드 네트워크(FPN)를 피하기 위해 단일 스케일 피처 맵을 사용하는 단순한 비전 트랜스포머 백본을 사용한다.
• 사용자 클릭을 백본에 직접 인코딩할 수 있도록 최소한의 아키텍처 수정으로 가능한 대칭 패치 임베딩 레이어를 도입한다.
• 강력한 피처 표현 전이를 가능하게 하기 위해 대규모 데이터셋에서 ViT 백본을 마스크된 오토에ncoder(MAE)로 사전학습한다.
• 백본의 마지막 피처 맵만 사용하여 단순한 피처 피라미드를 구성함으로써 일반적인 피처 학습과 세그멘테이션 전용 모듈을 분리한다.
• 계산 비용을 줄이기 위해 피처 피라미드를 세그멘테이션 마스크로 변환하기 위해 경량의 MLP 디코더를 적용한다.
• 백본을 포함한 전체 모델을 엔드 투 엔드로 파인튜닝하여 성능을 최대화한다.

실험 결과

연구 질문

• RQ1 계층적 아키텍처를 수정하지 않고도 단순한 비전 트랜스포머 백본이 상호작용 이미지 세그멘테이션에서 계층적 백본을 능가할 수 있는가?
• RQ2 마스크된 오토에코딩을 통한 단순 ViT의 사전학습이 상호작용 세그멘테이션 작업으로 효과적으로 전이되는가?
• RQ3 단일 스케일 피처 맵에서 구성된 단순한 피처 피라미드가 다중 스케일 또는 FPN 유사 설계를 능가하거나 이를 충족할 수 있는가?
• RQ4 제안된 방법은 실세계 응용 분야에서 실용적인 어노테이션 툴로 사용하기에 충분히 계산 효율적인가?
• RQ5 특히 의료 영상 데이터셋에서의 도메인 간 일반화 능력은 어떠한가?

주요 결과

• SimpleClick는 SBD 벤치마크에서 4.15 NoC@90의 성능을 달성하여 이전 최고 성능 방법보다 21.8% 향상되었다.
• 모든 이전의 계층적 백본 기반 방법보다 자연 이미지 및 의료 영상 벤치마크에서 모두 뛰어난 성능을 보이며 강력한 일반화 능력을 입증했다.
• 파인튜닝 중 ViT 백본을 동결하면 성능에 심각한 하락이 발생하여, 최적의 결과를 얻기 위해 백본의 파인튜닝이 필수적임을 시사한다.
• 피처 피라미드의 다중 스케일 성질이 핵심적임을 입증하였으며, 이를 제거하면 비록 마지막 피처 맵만 사용하더라도 성능에 심각한 하락이 발생한다.
• 가장 작은 모델 버전(ViT-xT)은 클릭당 17ms의 추론 속도를 기록하고 GPU 메모리 소비량이 0.17GB 뿐이어서 실시간 어노테이션에 적합하다.
• 가장 큰 ViT-H 버전도 추론 시간이 132ms이며 GPU 메모리 사용량이 3.22GB에 불과하여 상호작용 도구의 실용적 한계 내에 유지된다.