Tech Featured

[논문 정리] Segment Anything

Jinwoo Park

Apr 6, 2023 • 14 min read

작성자: 박진우 (GitHub)

Motivation

최근 대규모의 언어 모델들은 놀라운 Zero-shot / Few-shot Generalization 성능을 보이고 있다. 이러한 Foundation Model들은 종종 Prompt Engineering을 통해 여러가지 Task에 대한 적절한 텍스트 응답을 생성해주는 능력을 보여주기도 한다.

이 연구에서는 Image Segmentation에 대한 Foundation Model을 만드는 것을 목표로 한다. 또한 강력한 Generalization을 바탕으로 Prompt Engineering을 통해 여러 종류의 Image Segmentation 하위 문제들을 풀어내고자 한다.

이것을 가능하게 하기 위해서는 다음과 같은 세 가지 질문에서 출발할 필요가 있다.

Zero-shot Generalization을 가능하게 하는 Task는 무엇인가? (Task)
적절한 Model 구조는 무엇인가? (Model)
1의 Task와 2의 Model에 적절한 Data는 무엇인가? (Data)

Figure 1: We aim to build a foundation model for segmentation by introducing three interconnected components.

Segment Anything Task

NLP (Natural Language Process)의 경우, Foundation Model을 학습하기 위해 Next Token Prediction 과 같은 Task에 대한 Pre-training을 수행한다. 마찬가지로 Segmentation에서의 Foundation Model을 도출하기 위해 이와 유사한 효과를 기대할 수 있는 task를 정의하는 것을 목표로 한다.

저자들은 Forground / Backgroud Points, Box / Masks / Free-form Text 등을 이미지에 대한 Prompt로 정의한다. (2023년 4월 6일 기준 Free-form Text 입력은 미공개) 일명 Promptable Segmentation Task는 그 어떤 모호한 정보가 Prompt로 입력되더라도 유효한 마스크를 반환하는 것을 목표로 한다.

Training의 경우 Points, Boxes, Masks와 같은 Prompt Sequence를 시뮬레이트해서 마스크를 취득한 뒤 Ground-Truth Mask와 비교하는 방식을 사용한다(*).

* The promptable segmentation task suggests a natural pre-training algorithm that simulates a sequence of prompts (e.g., points, boxes, masks) for each training sample and compares the model’s mask predictions against the ground truth.

이렇게 얻어낸 Pre-trained Model은 Prompt Engineering을 통해 여러가지 Task에 유연하게 적용될 수 있다. 좀 더 자세히 살펴보자.

참고1, 참고2의 방식에서처럼 Training Process를 위해 Interactive Segmentation을 시뮬레이션한다.

Target Mask의 Foreground Point 또는 Bounding Box를 랜덤하게 선택한다. Point는 Target Mask에서 랜덤하게 샘플링하며, Box는 Ground Truth Mask의 Bounding Box를 사용한다. 박스의 좌표에는 박스 Sidelength의 10%의 STD로 생성한 랜덤 노이즈를 추가한다 (최대 20픽셀). - 유저가 Loose Box를 그리는 경우에 대비하기 위함.
첫 번째 Prompt로 Prediction을 한 뒤, 그 다음 Point는 Previous Mask Prediction과 Ground Truth Mask 사이의 Error Region에서 Uniformly Random으로 선택한다. 만약 Error Region이 False Negative면 Point는 Forground가 되고, False Positive면 Background가 된다. 또한 Previous Mask를 모델에 추가적인 Prompt로 입력한다. 이때 Mask는 정보량을 최대화하기 위해 {0,1}로 Thresholding되지 않은 값을 사용한다. 만약 여러개의 Mask가 반환되었을 경우에는 Predicted IoU가 가장 높은 것을 사용한다.
마스크 정보를 제공하는 것에서 이득을 얻기 위해 랜덤하게 2개 이상의 Iteration에서 Point 정보를 입력하지 않는다. (모델은 이전 마스크 정보를 Refine하게 됨)

Segment Anything Model

SAM은 세 개의 컴포넌트로 구성되어 있다: Image Encoder, Flexible Prompt Encoder, Fast Mask Decoder.

Image Encoder:
고해상도의 이미지를 처리하기 위해 Masked Autoencoder로 Pre-training을 한 Vision Transformer (ViT) 기반의 구조를 사용한다(*). Image encoder의 출력은 입력 이미지 크기를 기준으로 16배 다운스케일된 Embedding이다. 참고로, Image Encoding은 Prompt처럼 실시간성이 요구되지 않기 때문에 높은 스케일로 사용하는 것이 가능하다. (프롬프트와 달리 이미지 당 최초 한번만 수행)

1024x1024의 이미지 입력을 사용했으며 좌우/위아래 중 짧은 부분에는 padding을 한다. 즉, Embedding의 크기는 1024 / 16에 해당하는 64 x 64가 된다. 또한 1 x 1 Convolution을 사용해서 256 채널의 출력을 얻으며 이어서 256채널의 3 x 3 Convolution을 적용한다. 각 Convolution 연산 뒤에는 Layer Normalization을 적용한다.

*we use an MAE [47] pre-trained Vision Transformer (ViT) [33] with minimal adaptations to process high resolution inputs, specifically a ViT-H/16 with 14×14 windowed attention and four equally-spaced global attention blocks, following [62].

Flexible Prompt Encoder:
Point, Box, Text와 같은 Sparse Prompt는 256 차원의 Embedding으로 매핑된다. Point는 Point Location과 학습 가능한 Forground / Background Embedding 중 하나와의 합으로 표현한다. Box는 Top-left Corner를 나타내는 Learned Embedding과 Bottom-right Corner를 나타내는 Learned Embedding의 합으로 표현한다. 끝으로 Free-form Text Encoder는 CLIP의 것을 차용한다.

Dense Prompt의 경우 이미지와 공간적으로 대응되는 정보다 (e.g. 마스크). 마스크는 입력 이미지보다 1/4의 크기로 입력한다. 여기에 2x2, stride 2의 Convolution을 적용하여 추가적으로 크기를 4배 더 줄인다. 마지막으로 1x1 Convolution을 사용하여 채널을 256으로 변환한다. 각 Layer는 GELU activation과 Layer Normalization을 적용한다. 마스크와 Image Embedding은 Element-wise로 더해진다. 만약에 Mask입력이 없는 경우에는 "no mask"를 나타내는 embedding을 사용한다.
Fast Mask Decoder:
Prompt Embedding과 Image Embedding을 빠르게 Output Mask에 매핑하는 역할을 한다. 두 개의 입력을 합치는 방법으로는 Trasformer Segmentation Models를 참고했으며 (참고1, 참고2), Standard Transformer Decoder를 약간 수정한 구조를 사용했다.
Decoder 적용 이전 단게에서는 Prompt Embedding에 Learned Ouput Token Embedding을 삽입한다 (Vision Transformer (ViT) 의 Class Token과 유사).
각 Decoder Layer는 다음 4가지 Step으로 동작한다.

1. Token에 대한 Self-attention
2. Token을 Query로 하여 Image Embedding에 Cross-Attention
3. Point-wise MLP로 각 Token을 업데이트
4. Image Embedding을 Query로 하여 Token에 Cross-Attention
(Cross-Attention 과정에서 Image Embedding은 64x64의 256 차원 벡터)

각 Self/Cross-attention과 MLP에는 Residual Connection, Layer Normalization, Dropout(0.1)을 적용한다. 다음 Decoder Layer에서는 이전 레이어에서 업데이트된 Image Embedding과 Token을 사용한다. (Decoder는 2층으로 구성)
또한 중요한 Geometric 정보를 잃지 않기위해 Attention Layer에 입력할 때마다 Image Embedding에 Positional Encoding을 합산하며, 마찬가지로 Attention Layer에 입력할 때마다 Updated Token에 Original Prompt Token 전체를 합산한다.

Decoder를 수행한 뒤에는 2개의 Transposed Convolutional Layer를 적용하여 Image Embedding의 크기를 4배 키운다. 한편, Token을 Query로 하여 한번 더 Image Embedding에 Attention을 한뒤, Ouput Token Embedding에 3개의 MLP 를 적용하여 Channel-wise로 Upscaled Image Embedding에 적용한다.
다음은 각 컴포넌트에 대한 추가적인 정보다.

- Transformer는 Embedding Dimension으로 256을 사용.
- Transformer MLP block은 2048의 Internal Dimension 사용. (Prompt Token에만 적용)
- Cross-attention Layer에서는 64x64 Image Embedding을 사용했으며, Query / Key / Value의 채널은 계산 효율을 위해 128로 두배 줄임.
- 모든 Attention Layer는 8개의 Head를 사용.
- Ouput Image Embedding에 적용하는 Transposed Convolution은 2x2, stride 2이며 Output Channel은 64와 32 (GELU activation, Layer Normalization 사용).

Figure 14: Details of the lightweight mask decoder.

Loss:
Mask의 경우 Focal Loss와 Dice Loss를 20:1 비율로 Linear Combination하여 사용. IoU Prediction Head의 경우 IoU Prediction과 Ground-truth 사이의 Mean-square-error로 Loss를 계산. IoU Loss와 Mask Loss는 1:1로 합산.

Segment Anything Data Engine

텍스트 정보와 달리 Segmentation Mask는 인터넷에서 손쉽게 구하기 어려운 데이터다. Meta는 대규모(1.1B)의 마스크를 취득하기 위해 자체적인 Data Engine을 구현했다. Data Engine은 다음 세 가지 Stage로 구성된다.

A model-assisted manual annotation stage:
첫 번째 단계는 전통적인 interactive segmentation 과정과 유사하다. 전문 annotator들이 Segment Anything Model (이하 SAM)의 도움을 받아 웹 기반의 인터페이스를 통해 foreground / background object를 작업한다. 이때 마스크는 "stuff" 또는 "things"로 어노테이터가 자유롭게 레이블링 하도록 한다. 이 단계에서 사용하는 SAM은 공개된 segmentation dataset을 이용해 학습한다. 그리고 어노테이션 작업으로 추가 Data를 충분히 취득했을때 새로 취득한 데이터만을 사용해서 다시 학습을 수행한다. 이 과정에서 Image Encoder를 ViT-B에서부터 ViT-H까지 점진적으로 증가시키며, 총 6번의 추가 학습 과정을 수행하게 된다. 결과적으로 이 단계에서 120k 이미지로부터 4.3M개의 마스크를 취득한다.
A semi-automatic stage with a mix of automatically predicted masks and model-assisted annotation:
이 단계에서는 모델의 Segment Anything 능력을 향상시키기 위해 마스크의 종류를 다양화하는 것을 목표로 한다. 우선 모델을 이용해 Confident Mask를 잡아준 뒤에 이를 작업화면 상에 표시하면, 어노테이터들은 그 외 영역의 Object들을 작업하도록 한다. Confident 마스크를 파악하기 위해서는 첫 번째 단계에서 수집한 모든 마스크에 대한 Bounding Box Detector를 학습시킨다 ("Object" 카테고리). 이 단계를 거치며 180k 이미지로부터 5.9M 마스크를 취득한다. 첫 번째 단계에서처럼 새로 취득한 데이터를 이용해 5번의 점진적인 모델 학습을 거친다.
A fully automatic stage (without annotator input):
마지막 단계는 완전 자동화된 어노테이션 단계다.

- 모델에 32 x 32의 Regular Grid Point를 입력(prompted). 모델은 각 포인트 아래의 Object에 해당하는 마스크를 반환(Predicted).
- 모델은 Object의 Subpart, Part, Whole Object를 모두 반환.
- 모델의 IoU Prediction Module를 사용해 Confident Mask를 선택.
- 선택한 마스크 중에서 Stable한 마스크를 고름. (Probability Map을 0.5 - δ와 0.5 + δ 로 Thresholding 한 뒤에도 마스크의 형태가 유사하면 Stable하다고 판단)
- 앞서 선택한 Stable & Confident Mask들에 Non-maximal Suppression (NMS)를 적용하여 중복된 것들을 제거.
- 작은 마스크 품질을 높이기 위해 Multiple Overlapping Zoomed-in Image Crops를 이용.

위와 같은 방법을 적용해서 11M 이미지로부터 1.1B의 고품질 마스크를 추출해냈다.

Segment Anything Dataset

SA-1B 데이터셋은 11M의 다양하고, 높은 화질의 이미지를 포함하며, 또한 Data Engine으로부터 취득한 1.1B개의 고품질 마스크를 포함한다.

Image:
평균 3300×4950 pixels 이미지를 취득. 이미지 크기가 너무 커지는 것을 고려하여 Shortest Side를 기준으로 1500 pixel로 다운샘플. 사람, 차량 번호판 등은 블러 처리.
Masks:
Data Engine을 통해 1.1B의 마스크를 만들었으며, 이 중 99.1%는 완전 자동으로 취득. 즉, 자동 마스킹 기능의 성능이 마스크 품질에 있어 매우 중요. 이 부분을 확인하기 위해 전문가 어노테이터와 생성된 마스크 사이의 품질을 측정하였고, 생성된 마스크의 품질이 충분히 좋다고 결론. 이에 SA-1B에는 자동으로 취득한 마스크만 포함 (!).
Mask Quality:
마스크의 품질을 평가하기 위해 500개의 이미지를 랜덤으로 샘플링 (약 50k의 마스크 포함). 전문 어노테이터들은 샘플링한 데이터를 수정하는 작업을 수행. 이 방법으로 자동 생성 마스크와 이에 대한 Ground Truth Mask 취득했으며, 이 둘의 IoU를 측정. 결과적으로 85~91%의 IoU를 기록. 또한 실험을 통해 Human Rating으로 여러 데이터셋에 대해 생성된 마스크가 충분히 좋은 품질을 가지고 있음을 확인 (아래 Experimental Results의 Fig9-(b) 참고) .

Experimental Results

Zero-shot transfer 실험들을 살펴보도록 한다.

Zero-Shot Single Point Valid Mask Evaluation

Single Foreground Point로 Object에 대한 Segmenting을 수행하는 것이다. 다양한 이미지 분포에서 실험을 진행하기 위해 23개의 데이터셋을 사용했다. 결과적으로 23개 데이터셋 중 16개 데이터셋에서 RITM을 상회하는 성적을 보였다.

Zero-Shot Edge Detection

BSDS500 데이터셋을 이용해 Edge Detection을 수행해봤다. Edge는 Unthresholded Mask Probability Map에 Sobel Filtering을 적용해서 도출했다. Edge Detection을 학습하지 않았음에도 정성적으로 매우 그럴듯한 Edge Map을 보여준다. 또한 SAM은 HED 같은 딥러닝 기반의 Edge Detector 보다도 더 좋은 정량적 성과를 보인다.

Zero-Shot Object Proposals

Object Detection에서 매우 중요한 Task인 Object Proposal Generation의 벤치마킹 결과다. ViTDet-H가 전반적으로 가장 좋은 결과를 보였지만, 일부 Metric에서는 SAM이 약간은 앞서는 모습을 보이기도 한다.

Zero-Shot Instance Segmentation

COCO와 LVIS 데이터셋에서 ViTDet과 SAM의 성능을 비교해봤다. 정량평가로는 ViTDet에 다소 뒤쳐지지만, 흥미롭게도 Human Rating에서는 ViTDet보다 더 좋은 결과를 얻었다.

Zero-Shot Text-to-Mask

Free-fromed Text로부터 Segmentation을 수행하는 Task다. 이 실험의 경우 SAM의 능력치를 실험하기 위한 일종의 PoC (Proof of Concept)이다. Text 입력만으로는 정확한 출력이 나오지 않을 때는, PhraseClick 처럼 클릭 정보와 함께 보조적인 정보로 활용될 수 있을 것이다.