[논문 리뷰] Segment Anything in High Quality

Ke Lei, Mingqiao Ye외 5명|arXiv (Cornell University)|2023. 06. 02.

Advanced Neural Network Applications인용 수 110

한 줄 요약

HQ-SAM은 경량의 High-Quality Output Token과 글로벌-로컬 피처 융합을 SAM에 도입하여 negligible overhead로 제로샷 분할을 가능하게 하며 44k-mask dataset에서 훈련된다.

ABSTRACT

The recent Segment Anything Model (SAM) represents a big leap in scaling up segmentation models, allowing for powerful zero-shot capabilities and flexible prompting. Despite being trained with 1.1 billion masks, SAM's mask prediction quality falls short in many cases, particularly when dealing with objects that have intricate structures. We propose HQ-SAM, equipping SAM with the ability to accurately segment any object, while maintaining SAM's original promptable design, efficiency, and zero-shot generalizability. Our careful design reuses and preserves the pre-trained model weights of SAM, while only introducing minimal additional parameters and computation. We design a learnable High-Quality Output Token, which is injected into SAM's mask decoder and is responsible for predicting the high-quality mask. Instead of only applying it on mask-decoder features, we first fuse them with early and final ViT features for improved mask details. To train our introduced learnable parameters, we compose a dataset of 44K fine-grained masks from several sources. HQ-SAM is only trained on the introduced detaset of 44k masks, which takes only 4 hours on 8 GPUs. We show the efficacy of HQ-SAM in a suite of 10 diverse segmentation datasets across different downstream tasks, where 8 out of them are evaluated in a zero-shot transfer protocol. Our code and pretrained models are at https://github.com/SysCV/SAM-HQ.

연구 동기 및 목표

• 다양한 객체에 대해 SAM의 거친 경계 너머의 마스크 품질 향상을 동기 부여한다.
• 최소한의 어댑터를 추가하면서 SAM의 제로샷 일반화 및 프롬프트 가능 디자인을 보존한다.
• 데이터 효율적인 훈련을 소형이면서도 다수의 주석이 있는 데이터셋에서 입증한다.
• 제로샷 설정에서 다양한 이미지 및 비디오 분할 벤치마크에서 robust한 성능을 보여준다.

제안 방법

• SAM의 마스크 디코더에 학습 가능한 HQ-Output Token을 주입한다.
• 초기 및 최종 ViT 인코더 계층과 마스크 특징으로부터 파생된 HQ-Features를 융합하여 더 나은 디테일을 확보한다.
• SAM을 동결한 상태에서 HQ-Output Token, 그것의 3-계층 MLP, 융합 블록만 훈련한다.
• 고품질 마스크 예측을 위한 동적 커널 생성을 위한 3-계층 MLP를 사용한다.
• HQ-Output Token 예측과 SAM의 출력 합성은 최종 마스크를 위해 원소별 덧셈으로 통합된다.
• 데이터 효율적 훈련을 가능하게 하는 6분류 HQSeg-44K를 6개 소스에서 수집한 44k-마스크 데이터셋으로 개발한다.

실험 결과

연구 질문

• RQ1 HQ-SAM이 SAM의 제로샷 성능을 해치지 않으면서 마스크 디테일과 경계 정확도를 향상시킬 수 있는가?
• RQ2 HQ-SAM 훈련 데이터가 효과적이며 최소한의 추가 매개변수로 고품질 마스크를 달성하는가?
• RQ3 글로벌-로컬 피처 융합과 HQ-Output Token이 다양한 데이터셋과 프롬프트에서 측정 가능한 이점을 제공하는가?
• RQ4 제로샷 시나리오에서 HQ-SAM이 전체 미세 조정이나 후처리 보정 접근법과 비교해 어떠한 이점을 제공하는가?

주요 결과

• HQ-SAM은 10개의 다양한 데이터셋에서 SAM보다 더 높은 품질의 마스크를 제공하면서 제로샷 기능을 유지한다.
• HQSeg-44K에서 HQ-SAM을 훈련하는 데 8대 RTX 3090 GPU에서 단 4시간이 걸리며 매개변수 오버헤드는 0.5% 미만이다.
• HQ-SAM은 경계 중심 메트릭에서 상당한 이점을 얻으며(예: 여러 미세한 데이터셋에서 mBIoU가 개선된다).
• 초기 및 최종 ViT 인코더 피처와 마스크 피처의 글로벌-로컬 융합이 SAM 피처만 사용할 때보다 분할 디테일을 향상시킨다.
• 미세 조정이나 후처리 기준선과 비교했을 때 HQ-SAM은 더 작은 매개변수 업데이트로 제로샷 성능이 우수하다.
• MobileSAM의 아주 작은 인코더를 기반으로 한 경량 HQ-SAM은 모듈러 오버헤드로 41.2 FPS를 달성하여 COCO 오픈 셋 지표를 개선한다.