[Paper Review] ICCV 2023, Segment Anything 논문리뷰

Segment Anything

ICCV 2023

• 논문 링크: https://arxiv.org/abs/2304.02643
• project 링크: https://aidemos.meta.com/segment-anything

0. Abstract

먼저, promptable segmentation task를 정의함. segmentation task는 향후 medical segmentation을 진행함에 있어서 중요함. 따라서 어떤 구조, 어떤 방식으로 진행되는지 파악하는 것이 중요함.

 

promptable segmentation task를 한 문장으로 말하면, '이미지를 주고 정답 mask를 맞혀라가 아닌, 점·박스·마스크 같은 프롬프트가 주어졌을 때 그 의도에 맞는 영역을 분할하하는 것'임. → 사용자의 의도에 따라 유연하게 정의될 수 있음.

 

다음으로, SAM(Segment Anything Model) 을 보면 다음과 같음. SAM은 이 promptable segmentation task를 수행하도록 학습된 segmentation 모델임. 결과적으로는 SAM이 data annotation을 만드는 도구로도 사용됨. 즉, 새로운 이미지에 대해 SAM이 먼저 segmentation을 suggest하고, user는 최소한의 프롬프트나 수정만 제공함. 그래서, 새로운 도메인이나 새로운 segmentation task에서도 추가 학습 없이(zero-shot) 프롬프트 설계만으로 바로 적용될 수 있음.

 

다음으로, data engine 을 보면 다음과 같음. SAM에서 사용하고 있는 data engine의 경우에는 사람과 모델이 상호작용하는 대규모 데이터 수집 시스템임. SAM을 활용해 빠르게 mask를 생성하고 이를 정제·누적하는 역할을 함. 이 data engine을 통해 만들어진 결과물이 바로 SA-1B 데이터셋이며, 여기에는 1,100만 장의 이미지와 10억 개 이상의 segmentation mask가 포함됨.

 

결과적으로 보면, promptable task가 있어야 SAM이 설계될 수 있고, SAM이 있어야 대규모 annotation이 가능하고, 그렇게 모인 SA-1B 데이터가 다시 SAM을 더 강력하게 만듦.

 

1. Introduction

For example, CLIP and ALIGN use contrastive learning to train text and image encoders that align the two modalities. Once trained, engineered text prompts enable zero-shot generalization to novel visual concepts and data distributions. → CLIP이나 ALIGN은 contrastive learning 사용.

 

While much progress has been made on vision and language encoders, computer vision includes a wide range of problems beyond this scope, and for many of these, abundant training data does not exist. → 당연하게도 task를 수행하는 데에 있어서 학습을 진행할 때 충분한 학습 데이터가 부족하다는 점은 언제나 존재함.

 

In this work, our goal is to build a foundation model for image segmentation. That is, we seek to develop a promptable model and pre-train it on a broad dataset using a task that enables powerful generalization. With this model, we aim to solve a range of downstream segmentation problems on new data distributions using prompt engineering. The success of this plan hinges on three components: task, model, and data. To develop them, we address the following questions about image segmentation:

1. What task will enable zero-shot generalization?

2. What is the corresponding model architecture?

3. What data can power this task and model?

 

Unfortunately, there is no web-scale data source for segmentation; to address this, we build a “data engine”, i.e., we iterate between using our efficient model to assist in data collection and using the newly collected data to improve the model. → 모델을 통한 데이터 수집, 새로 수집된 데이터를 사용하여 모델을 개선하는 것을 반복함.

 

1. What task will enable zero-shot generalization? 

Answer: Inspired by this line of work, we propose the promptable segmentation task, where the goal is to return a valid segmentation mask given any segmentation prompt (see Fig. 1a). The requirement of a valid output mask means that even when a prompt is ambiguous and could refer to multiple objects (for example, a point on a shirt may indicate either the shirt or the person wearing it), the output should be a reasonable mask for at least one of those objects. We use the promptable segmentation task as both a pre-training objective and to solve general downstream segmentation tasks via prompt engineering. → promptable segmentation task를 통해서 pre-training과 downstream task를 모두 한다는 거임.

 

2. What is the corresponding model architecture?

Answer: The promptable segmentation task and the goal of real-world use impose constraints on the model architecture. In particular, the model must support flexible prompts, needs to compute masks in amortized real-time to allow interactive use, and must be ambiguity-aware. → prompt는 model이 특정 task를 수행하는 데에 있어서 필요한 지시사항이나 조건을 제공하는 것임. 여기서 flexible한 prompt를 지원해야 한다는 말은, model이 다양한 종류의 prompt를 잘 처리할 수 있어야 한다는 뜻임. 또한, 실시간으로 mask를 계산할 수 있어야 하고 이를 통한 interactive 사용을 가능하게 해야함. 그리고 반드시 모호성에 대해서 인식하고 있어야 함.

 

Surprisingly, we find that a simple design satisfies all three constraints: a powerful image encoder computes an image embedding, a prompt encoder embeds prompts, and then the two information sources are combined in a lightweight mask decoder that predicts segmentation masks. We refer to this model as the Segment Anything Model, or SAM (see Fig. 1b).

 

By separating SAM into an image encoder and a fast prompt encoder / mask decoder, the same image embedding can be reused (and its cost amortized) with different prompts. → SAM을 image encoder와 fast prompt encode/mask decoder로 나누면, 같은 이미지 임베딩을 다양한 prompt에 재사용할 수 있고, 그 비용을 분산시킬 수 있음.

 

To make SAM ambiguity-aware, we design it to predict multiple masks for a single prompt allowing SAM to naturally handle ambiguity, such as the shirt vs. person example. → ambiguity와 관련한 개념이 조금 모호해서 정리할 필요가 있음.

Ambiguity

Ambiguity는 image에서 하나 이상의 해석이 가능한 상황을 말함. 가령, 셔츠 vs 사람 이미지를 보면 셔츠가 사람의 일부일 수도 있고 사람 전체로 보일 수도 있음. 이런 경우에 셔츠와 사람을 명확히 구분하는 것에 대해서 모호한 상황이 발생할 수 있음.

SAM은 기본적으로 이미지에서 특정 영역을 mask로 분리하는 model이라고 할 수 있음. 그래서 ambiguity를 다루기 위해서 하나의 프롬프트에 대해서 여러 개의 mask를 예측하는 방식으로 설계됨.

• 예를 들어서 셔츠 vs 사람 이라는 프롬프트에 대해서 SAM은 두 가지 가능한 해석을 모두 고려하여 각각 다른 해석을 반영하는 mask를 만들 수 있음. 하나의 mask는 shirt만을 표시하는 mask, 또 다른 mask는 사람을 포함하는 mask. 이런 식으로 다른 mask를 만들 수 있음.
• 모델이 여러 마스크를 생성하는 이유는 다음과 같음. 모호한 상황에서는 정확히 하나의 해석만을 내리는 것이 불가능하거나 혹은 부적합할 수 있음. 그래서 모델이 여러 가지의 mask를 prediction하고 그 결과로 각각의 해석을 나타낼 수 있게 함. 이 방법을 통해서 사용자는 여러 해석을 한 번에 제공받을 수 있어서 모호성을 자연스럽게 처리할 수 있게 됨. 

모호성을 자연스럽게 처리?

• 위에서 언급했다시피, SAM은 모호한 이미지에서 여러 가지 해석을 동시에 제공함. 동시에 제공함으로써 결과에 대한 확신을 높이고 flexibility를 제공한다고 할 수 있음.
• 예를 들어 셔츠 vs 사람 이미지를 처리함에 있어서, 두 가지 가능한 영역을 모두 분리하게 된다면 사용자 입장에서는 두 가지 해석을 모두 확인할 수 있음. 애초에 두 가지 해석을 모두 확인할 수 있다는 것 자체가 사용자가 더 많은 정보를 얻을 수 있고 모델이 다양한 상황을 유연하게 처리하는 방식으로 모호성을 해결할 수 있다고 봐도 무방함.

여러 mask를 예측하는 게 효과적임?

• ambiguity 문제는 보통 하나의 해석으로는 조금은 불만족스러울 때, 발생한다고 할 수 있는게, 사람과 그 이외의 배경을 구분하는 task가 있다고 했을 때, task에서 사람과 그 이외의 배경이 모호한 경우, 모델이 두 가지의 mask를 예측함으로써 각각의 해석을 처리할 수 있음.

결국에는 SAM을 쓰는 이유가 불확실한 경우에도 뭔가 좀 flexible하게 대처할 수 있다는 걸 의미함.

 

3. What data can power this task and model?

Answer: To achieve strong generalization to new data distributions, we found it necessary to train SAM on a large and diverse set of masks, beyond any segmentation dataset that already exists. Masks are not naturally abundant and thus we need an alternative strategy. Our solution is to build a “data engine”, i.e., we co-develop our model with model-in-the-loop dataset annotation (see Fig. 1c). → 당연히 기존에 존재하는 mask가 충분하지 않기 떄문에 이 논문에서는 data engine을 만듦. 모델과 함께 데이터셋 annotation을 생성하는 과정으로 co-develop한다는 것으로 이해하면 될듯함. Data engine 자체가 모델을 훈련시키면서 동시에 dataset을 생성하는 시스템. 막바지에 서술되는 model-in-the-loop는 model을 훈련하면서 data annotation을 다는 방식으로, train과 dataset 생성을 동시에 하는 거라고 이해하면 될듯.

 

Our data engine has three stages: assisted-manual, semi-automatic, and fully automatic. In the first stage, SAM assists annotators in annotating masks, similar to a classic interactive segmentation setup. In the second stage, SAM can automatically generate masks for a subset of objects by prompting it with likely object locations and annotators focus on annotating the remaining objects, helping increase mask diversity. In the final stage, we prompt SAM with a regular grid of foreground points, yielding on average ∼100 high-quality masks per image. → 총 3가지 단계로 나뉨, assisted-manual(보조 수동), semi-automatic(반자동), and fully automatic(완전 자동). assisted-manual 단계에서는 mask annotation을 할 때, SAM이 도와주고 기본적인 segmentation 작업을 도움. semi-automatic 단계에서는 SAM이 예상되는 객체 위치를 바탕으로 자동으로 마스크를 생성하고, 사람(annotator)은 남은 객체들만 annotation을 달면 됨. 결과적으로는 마스크의 다양성을 높이는 과정이 이루어짐. Fully automatic 단계에서는 규칙적인 그리드 포인트를 사용하여 SAM을 prompt하고, 이로 인해 매 이미지당 약 100개의 고품질 mask를 자동으로 생성하는 과정임.

SAM을 prompt하는 것

일반적으로 프롬프트(prompt)는 model에게 어떤 작업을 하라고 지시하는 입력이나 조건을 의미함. 프롬프트는 model이 무엇을 해야할지 알려주는 지시문 역할이라고 보면 될듯. Fully automatic 단계에서는 점(point)을 prompt로 사용한다고 했는데, 가령 이미지가 있고 이 이미지는 사람과 배경을 분리해야 하는 task와 관련된 이미지라고 하면, regular grid point라고 하면 이미지 전체에 걸쳐서 일정한 간격으로 배치된 점들로 이해할 수 있음. 이때 SAM에게 promt를 준다는 말은 regular grid point를 SAM에 input으로 준다는 의미인데 이 point 들은 모델에게 여기서부터 mask를 prediction하라고 지시하는 기준이 됨. 만약 점으로 prediction을 한다고 하면, grid point를 prompt로 주면 SAM은 이 점들이 배경에 해당하는지 사람에 해당하는지 예측하기 시작함. SAM은 이 점들로부터 mask를 자동으로 생성함.

2. Segment Anything Task

The promptable segmentation task, then, is to return a valid segmentation mask given any prompt. We choose this task because it leads to a natural pre-training algorithm and a general method for zero-shot transfer to downstream segmentation tasks via prompting. → it leads to a natural pre-training algorithm 이 부분이 이해가 안갔는데, 이해를 해보면, 어떠한 prompt가 주어져도 유효한 segmentation mask를 return을 하면, 되게 당연하게도 혹은 자연스럽게 pre-training 알고리즘을 이끌어낼 수 있음.

 

Our aim is to always predict a valid mask for any prompt even when the prompt is ambiguous.

 

The goal of our promptable segmentation task is to produce a broadly capable model that can adapt to many (though not all) existing and new segmentation tasks via prompt engineering. In a multi-task system, a single model performs a fixed set of tasks, e.g., joint semantic, instance, and panoptic segmentation, but the training and test tasks are the same. → 이전의 multi-task segmentation에서는 training 시점과 test 시점에서 task가 같음(-). 그러나 SAM에서는 training 시점과 test 시점에서 task가 다를 수 있음.

 

An important distinction in our work is that a model trained for promptable segmentation can perform a new, different task at inference time by acting as a component in a larger system, e.g., to perform instance segmentation, a promptable segmentation model is combined with an existing object detector

 

3. Segment Anything Model

Image encoder.

Motivated by scalability and powerful pretraining methods, we use an MAE pre-trained Vision Transformer (ViT) minimally adapted to process high resolution inputs. The image encoder runs once per image and can be applied prior to prompting the model.

 

Prompt encoder.

We consider two sets of prompts: sparse (points, boxes, text) and dense (masks). We represent points and boxes by positional encodings summed with learned embeddings for each prompt type and free-form text with an off-the-shelf text encoder from CLIP. Dense prompts (i.e., masks) are embedded using convolutions and summed element-wise with the image embedding. → points와 boxes는 위치 인코딩(positional encodings)을 사용하여 표현함. Positional encoding을 사용하여 학습된 임베딩(learned embeddings)과 합쳐져서 모델이 더 잘 이해할 수 있게 됨. 여기서 positional encoding은 각 점이나 박스가 이미지 내 어디에 위치하는지를 나타내는 정보. 텍스트(prompt)는 CLIP의 텍스트 인코더를 사용하여 처리됨. CLIP 모델은 이미지와 텍스트를 함께 학습한 모델로, 자연어로 된 지시문을 이미지와 관련된 의미 있는 정보로 변환할 수 있음.

 

Losses and training.

We supervise mask prediction with the linear combination of focal loss [65] and dice loss [73] used in [14]. → mask prediction을 focal loss와 dice loss의 선형 결합으로 supervise를 진행함. 여기서, focal loss는 불균형한 클래스에서 어려운 예제에 더 많은 가중치를 두어 학습하는 손실 함수. dice loss는 세분화 품질을 평가하는 데 사용되는 손실 함수로, 예측된 마스크와 실제 마스크의 겹침 정도를 기반으로 계산됨. 

 

 

...

 

8. Discussions

Limitations

It can miss fine structures, hallucinates small disconnected components at times, and does not produce boundaries as crisply as more computationally intensive methods that “zoom-in”, e.g..