Post

[논문 리뷰] SAM 2: Segment Anything in Images and Videos

Posted
By Luna Lee
30 min read

Meta FAIR arXiv 2024

📖 핵심 훑어보기 !!

  • 비디오와 이미지 도메인에서 segmentation을 수행하는 SAM2 제안, 비디오 전체에 걸쳐 target 객체의 모든 frame에 대한 segmentation mask(masklet)를 생성
  • 다양한 prompt와 과거 메모리를 활용해 비디오 전체에서 객체의 공간적 범위를 segmentation함. Memory bank 모듈을 통해 최근 frame과 prompt된 frame 정보를 저장하며, Memory Attention을 통해 저장된 정보를 결합함
  • SAM 2 모델 학습을 위해 3 단계로 이루어진 새로운 데이터 엔진을 도입하여 Segment Anything Video(SA-V) 데이터셋 구축




Introduction

기존에 발표되었던 Segment Anything(SA)에서는 이미지에서 promptable segmentation을 위한 foundation model을 도입했다. 그러나 AR/VR, robotics, 자율주행 등 많은 중요한 application에서는 이미지 level segmentation을 넘어선 시간적인 localization, 즉 비디오에 대한 segmentation이 필요하다.

이미지 segmentation에서 spatio, 즉 공간적인 범위를 결정하는 문제였다면, Video segmentation은 물체의 시공간적 (spatio-temporal) 범위를 결정하는 것을 목표로 한다. 하지만 비디오에서 물체는 움직임, 변형, 폐색, 조명 변화와 같은 다양한 요인으로 인해 모양이 크게 바뀔 수 있어 복잡한 문제이다. 또한 카메라 흔들림, 낮은 해상도 등의 이유로 이미지에 비해 품질이 떨어지며, 다수의 프레임을 효율적으로 처리하는 것도 주요 과제로 꼽힌다.


SAM2_1.png

본 논문에서는 이러한 문제를 해결하기 위해 이미지와 비디오 segmentation을 모두 처리할 수 있는 Segment Anything Model 2 (SAM 2)를 제안한다. SAM 2는 Promptable Visual Segmentation(PVS) 작업을 통해 비디오의 임의 프레임에 주어진 포인트, 박스, 혹은 마스크로 정의된 객체를 시공간적 마스크(masklet)로 예측하며, 프레임에 프롬프트를 제공하여 예측 결과를 반복적으로 개선할 수 있다. 또한 메모리 모듈을 사용하여 이전 프레임에서의 객체 정보와 상호작용을 저장하고 이를 바탕으로 비디오 전체에서 일관된 segmentation을 수행할 수 있다.

또한 학습을 위해 저자는 새로운 데이터 엔진을 도입하여 Human annotator와 상호작용하며 고난도의 데이터셋을 생성했다. 데이터 엔진은 특정 카테고리에 제한되지 않고 데이터를 생성하며, 기존의 비디오 segmentation 데이터셋보다 8.4배 더 빠르게 데이터를 수집할 수 있다. 이 방식으로 Segment Anything Video (SA-V) 데이터셋을 구축하였다.

SAM2의 내용을 이해하기 위해서는 SAM의 내용을 이해하고 있는 것이 도움이 된다. SAM에 대한 내용은 아래 링크의 리뷰를 참조하자.
SAM 논문 리뷰.






Task: promptable visual segmentation

Promptable Visual Segmentation(PVS)는 비디오의 모든 프레임에 대해, 모델에 프롬프트를 줄 수 있도록 한다. 프롬프트는 객체를 정의하거나 모델이 예측한 segmentation을 수정하기 위한 것이며, positive/negative 클릭, 경계 상자, 마스크 형태로 제공될 수 있다. PVS 작업에서는 특정 프레임에서 프롬프트를 받으면 모델이 즉시 해당 객체의 유효한 segmentation mask를 생성해야 한다.

SAM2_2.png

위의 그림과 같이 초기 prompt를 받은 후, 모델은 이를 비디오 전체에 걸쳐 전파하여 target 객체의 모든 프레임에 대한 segmentation mask(masklet)를 생성한다. 이 프롬프트가 중간에 전파되지 않을 수도 있는데, 이런 경우 해당 프레임에서도 추가 프롬프트를 제공하여 segmentation을 지속적으로 개선할 수 있다(위 그림외 frame 3).

SAM 2는 이러한 PVS 작업의 데이터 수집 도구로 사용되어 SA-V 데이터셋을 구축하는 데 기여한다.






Model

SAM2_3.png

SAM 2는 SAM을 확장하여 비디오와 이미지 도메인 모두에서 작동하도록 설계된 모델이다. 모델은 각 프레임에서 point, box, mask 프롬프트를 입력으로 처리하여 비디오 전체에서 객체의 공간 범위를 segmentation한다. Mask를 개선하기 위해 프레임에 프롬프트를 반복적으로 추가할 수 있다.

하지만 SAM과 다르게, SAM 2 Decoder에서 사용하는 frame embedding 방법은 Image Encoder에서 직접 가져오지 않는다. 대신 과거 예측 memory와 prompted frame에 대해 조건화된다.

또한 prompted frame은 현재보다 미래 frame을 사용할 수도 있다.

Frame의 memory는 현재 예측을 기반으로 Memory Encoder에 의해 생성되며, 후속 frame에 사용될 수 있도록 Memory bank에 들어간다.


이제 개별 구성 요소를 살펴보자.

하지만 그 전에 “Object Pointer(객체 포인터)”와 “Spatial Memory Features”라는 개념에 대해 알고 넘어가자(그래야 아래 내용의 이해가 빠르다). SAM 2에서는 초기 frame에 대해 prompt를 받은 후, 이것을 후속 frame에 전파한다고 했는데, 논문에서는 이 방법으로 object pointer와 과거 frame에 대한 Spatial feature를 이용했다.

  • Object Pointer: 현재 프레임에 대해 Decoder에서 Mask를 생성할 때, 해당 객체의 Point를 추가적으로 예측하여 이 점을 다음 frame의 prompt로 넣어준다(일반적으로 객체의 중심점을 예측하면, 다음 frame에서도 미세한 위치는 바뀌지만 중심점의 위치에 객체가 남아있을 것이라는 생각인 것 같다). 다음 frame에 넣어주기 위해 이 pointer를 memory bank에 저장하고, 이 점을 다음 frame에 대한 prompt로 추가해주는 방식이다.
  • Spatial Memory Feature: 최근 N개의 frame feature. **메모리 뱅크에 저장된 과거 프레임의 **spatial feature map이다. 자세한 설명은 뒤에서 확인하자.



Image encoder.

  • 임의의 긴 비디오에 대해 real-time으로 처리하기 위해 streaming 방식을 사용하여, 각 비디오 frame이 입력으로 주어지면 Image embedding을 생성한다. (cf. SAM에서는 이미지 한 장에 대해 Image emdebbing을 생성하고, 이것을 바로 Decoder의 입력으로 사용함)
  • Image Encoder는 전체 interaction(전체 frame)에 대해 한 번만 실행되고, 각 frame에 대해 unconditioned tokens (feature embeddings)을 생성한다.
  • Hiera로 사전 학습된 MAE(ViT)를 이미지 Encoder로 사용한다. MAE는 계층구조로, Decoding 중에 multiscale features를 사용할 수 있다.



Memory attention.

  • Memory attention의 역할은 과거 frame feature, 새 prompt를 사용하여 현재 frame feature를 조절하는 것이다.
  • L개의 transformer block으로 구성되어 있고, 첫 번째 block은 Image encoder에서 생성된 현재 frame의 image encoding을 입력으로 받는다.
  • 각각의 block은 self-attention + cross-attention으로 구성되어 있다.
  • Self-attention 이후, cross-attention은 memory bank에 저장된 과거 frame(Spatial memory feature)과 Object pointer에 대해 수행된다.(Decoder 섹션에서 Object pointer 생성에 대해 언급됨)
    • 두 가지에 대한 Cross-attention은, Multi-head 구조에 의해 수행된다. 각 Attention Head는 “Spatial Memory Features” 또는 “Object Pointer”를 선택적으로 처리할 수 있도록 설계된다.
  • Cross-attention이 수행된 이후 MLP를 수행한다.
  • Attention 연산은 최근 효율적인 attention kernel*을 활용하여 성능을 최적화한다.

FlashAttention-2*📄
FlashAttention-2라는 논문에서 제안된 새로운 Attention Kernel을 사용했다. 이 방법은 기존 FlashAttention의 개선된 버전으로, Transformer 모델에서 Attention 계산의 효율성을 극대화하기 위해 설계되었다. FlashAttention-2는 작업 분배와 병렬 처리 방식을 개선하여 Attention 계산의 속도를 크게 향상 시켰다. 주요 개선사항은 다음과 같이 정리할 수 있다.

  1. 비행렬 곱셈 연산 감소: GPU의 행렬 곱셈 연산에 특화된 유닛을 최대한 활용하기 위해, 비행렬 곱셈 부동소수점 연산(FLOPs)의 수를 줄임
  2. 스레드 블록 간 병렬 처리 강화: 단일 Attention Head 내에서도 다양한 스레드 블록 간에 Attention 계산을 병렬화하여 GPU의 자원을 더 효율적으로 활용
  3. 스레드 블록 내 작업 분배 최적화: 각 스레드 블록 내에서 작업을 Warps 간에 최적으로 분배하여 공유 메모리 접근을 줄이고 통신 오버헤드를 최소화함





Prompt encoder and mask decoder.

  • Prompt encoder는 SAM과 동일한 구조를 사용하며, 클릭(positive/negative), bounding box, mask 등의 prompt를 처리한다.
  • Sparse prompt는 positional encoding + 학습된 embedding으로 표현되고, mask는 convolution을 통해 embedding되어 frame embedding과 합산된다.

SAM2_4.png

  • Mask decoder는 거의 SAM과 유사하지만, “양방향(two-way)” transformer block으로 구성되어 prompt와 frame embedding을 업데이트한다.
  • SAM에서와 마찬가지로, 모호한 prompt에 대해서는 다중 마스크를 예측하며, 비디오에서는 추가 prompt가 없을 경우 가장 높은 IoU를 가진 마스크를 전파한다. (eg. 모호한 prompt → point에 대한 객체가 가리키는 대상이 가방인지, 가방을 메고 있는 사람인지 모호함. SAM에서 언급됨!)

  • SAM과 달리, PVS task에서는 일부 frame에 유효한 객체가 존재하지 않을 수도 있다(occlusion*으로 인해). 이렇게 새로 발생한 출력 현상을 설명하기 위해 추가 헤드를 도입했다. 이 헤드는 관심 객체가 현재 frame에 존재하는 지 여부를 예측한다.

    Occlusion*: 객체가 frame의 일부에서 가려지거나 보이지 않는 상태

  • SAM과의 또 다른 차이점은 계층적 image encoder(memory attention을 지나쳐 바로 decoder로)의 skip-connection을 사용하여 mask decoding을 위한 고해상도 정보를 통합한다는 점이다.



Memory encoder.

  • Memory encoder는 마스크 출력을 다운샘플링하고, 이미지 인코더에서 생성된 unconditioned frame embedding과 결합(element-wise sum)하여 메모리를 생성(이렇게 생성된 메모리가 Spatial Memory Features)한다.
  • 결합된 feature는 가벼운 convolutional layer를 통해 정보를 통합한다.



Memory bank.

Memory bank에는 크게 두가지 형태의 정보가 저장된다.

  1. Spatial feature map: Memory bank는 FIFO 큐를 사용해 최근 N개의 frame memory와 M개의 prompted frame memory를 저장한다. 각 메모리는 Spatial feature map으로 저장된다.
    • Frame memory: 최근 N개의 Spatial Memory Feature
    • Prompted frame memory: prompted frame의 Spatial Memory Features와 관련 프롬프트 정보(예: 클릭, 경계 상자, 마스크)로 구성
  2. Object pointer: 경량 벡터(lightweight vector) 형태로 저장되어 고차원 semantic 정보를 제공하는 역할로 사용된다.

  • 메모리는 최근 N개의 frame에 시간 위치 정보(Temporal Position Information)를 포함해 단기적인 객체 이동을 모델링했다. 즉, 모델은 저장된 과거 프레임이 현재 프레임 기준 몇 번째 프레임인지를 각각 알 수 있다. 이를 통해 모델은 단기적인 객체 움직임을 추적하고, 객체가 시간적으로 이동하는 패턴을 이해할 수 있다.
    • 시간적 위치 정보를 활용하면, 객체가 여러 프레임에 걸쳐 움직이거나 변형될 때 모델이 이를 연속적으로 학습하고 예측할 수 있다.
  • Prompted frame에는 시간 정보를 포함하지 않아 다양한 시간적 범위에 대해 일반화가 가능하다.
    • 프롬프트된 프레임은 학습 시와 추론 시 시간적 위치가 달라질 수 있기 때문에, Temporal Position Information을 내장하지 않는다.
    • 학습 중에는 특정 시점의 프롬프트가 사용되지만, 추론 중에는 프롬프트가 다른 시간 범위에서 제공될 수 있기 때문에, 이러한 불일치를 방지하고 모델의 일반화 성능을 높이기 위해 시간 정보를 포함하지 않는다.



Training.

  • 모델은 이미지와 비디오 데이터에서 동시에 학습된다.
  • 8개의 프레임을 샘플링하고, 랜덤하게 최대 2개의 프레임에 prompt를 제공한다.
  • Prompt는 50% 확률로 ground-truth mask, 25% 확률로 ground-truth mask 내에서의 positive click, 나머지 25% 확률로 bounding box를 입력으로 제공한다.
  • 학습 중, 모델은 상호작용적 방식으로 ground-truth masklet을 순차적으로 예측하도록 훈련된다.






Data

비디오에서 “Segment Anything” 능력을 개발하기 위해, 대규모 비디오 segmentation 데이터셋을 수집하기 위한 데이터 엔진을 구축했다. 이 데이터 엔진은 세 가지 단계로 나뉘며, 단계별로 모델이 annotator를 지원하는 방식과 데이터 품질을 개선하는 방식에 따라 구분되었다.


1. Data engine

Phase 1. SAM per frame.

  • 작업 방식:
    • 이미지 기반 대화형 SAM을 사용하여 각 frame에서 객체를 segmentation하고, annotator가 브러시와 지우개 같은 도구로 수동 편집. 초당 6 frame의 이미지를 편집했다.
    • 모든 frame이 독립적으로 작업되며, 시간적 추적 모델은 사용하지 않음.
  • 효율성 및 데이터 수집:
    • SAM을 이용하여 하나씩 mask를 생성해야하므로 느림. 평균 37.8초/프레임 소요.
    • 16K masklets를 1.4K 비디오에서 수집.
  • 결과:
    • 매우 높은 spatial 품질의 프레임 단위 세그멘테이션 생성.
    • 이 데이터를 기반으로 SA-V val/test 세트를 구성해 평가 bias을 줄임.


Phase 2. SAM + SAM 2 Mask.

  • 작업 방식:
    • SAM 2를 루프에 추가했으며 SAM 2는 Mask만 prompt로 사용.
    • 첫 번째 frame에서는 SAM으로 초기 마스크 생성 후, SAM 2 Mask를 사용해 시간적 마스크 전파.
    • 잘못된 결과는 다시 수동으로 수정하고, 수정된 Mask를 이후 프레임에 대해 재전파.
  • 효율성 및 데이터 수집:
    • 평균 7.4초/프레임 소요 (Phase 1 대비 약 5.1배 빨라짐).
    • 63.5K masklets를 수집.
    • SAM 2 Mask는 Phase 1 데이터와 공개 데이터셋으로 초기 학습 후, 두 번 재학습.
  • 결과:
    • 이전 frame에 대한 메모리 없이 중간 frame에서 마스크를 완전히 새로 작성해야 하는 제약 존재.
    • 이 단계를 통해 SAM 2의 완전한 기능을 개발할 기반 마련.


Phase 3. SAM 2.

  • 작업 방식:
    • 이 단계에서는 point와 mask 등 다양한 prompt를 활용하며, 시간적 메모리를 통해 객체 mask를 예측.
    • Annotator는 필요할 때만 수정 클릭을 제공.
  • 효율성 및 데이터 수집:
    • 평균 4.5초/프레임 소요 (Phase 1 대비 약 8.4배 빠름).
    • 197.0K masklets를 수집.
    • SAM 2는 이 단계에서 다섯 번 재학습.
  • 결과:
    • 높은 품질과 효율성을 모두 달성하며 데이터 수집 속도 대폭 개선.


Quality verification.

  • Annotation의 높은 품질을 유지하기 위해 품질 검증의 단계 도입.
  • 독립된 annotator가 각 마스크를 “satisfactory” (적합) 또는 “unsatisfactory” (부적합)로 검증.
  • 부적합한 마스크는 재수정, 잘 정의되지 않은 객체를 포함하는 masklet은 삭제함.


Auto masklet generation.

  • 모델이 “모든” 것을 segment 하도록 만드려면 다양성이 중요함.
  • 인간 Annotator는 일반적으로 눈에 띄는 객체에 더 집중할 수 있으므로 자동으로 생성된 masklet으로 annotation을 보강함
  • 자동으로 masklet을 생성하기 위해, 첫 번째 frame에 대해 그리드 포인트를 prompt로 제공하여 Auto mask 생성.
  • Auto generation mask는 검증을 거쳐 “satisfactory”로 확인되면 데이터셋에 포함.
  • 실패한 mask는 Phase 3에 다시 제공되어 개선.


Analysis.

SAM2_5.png

위 표는 통제된 실험을 통해 각 데이터 엔진 단계의 annotation 프로토콜을 비교한 것이다. Phase 3는 Phase 1 대비 8.4배 빠른 속도와 유사한 품질을 달성했다.


SAM2_6.png

각 phase의 끝에서는 해당 phase의 데이터로 학습된 SAM 2의 성능을 비교할 수 있다. 해당 phase에서 생성된 추가 데이터 포함 시, SA-V val 세트와 9개 zero-shot 벤치마크에서 지속적인 성능 향상 관찰할 수 있었다.





2. SA-V datasets

SAM2_7.png

데이터 엔진으로 수집한 SA-V 데이터 세트는 642.6K masklet이 있는 50.9K 비디오로 구성되어 있다. 표 3에서는 비디오, masklet, mask 수에 걸쳐 SA-V 구성을 일반적인 VOS 데이터셋과 비교했다.

SAM2_8.png

Dataset 구성

  • 규모:
    • 50.9K 비디오, 642.6K masklets.
    • 기존 VOS 데이터셋보다 53배(Auto mask 제외 시 15배) 많은 mask 포함.
  • 비디오:
    • 50.9K 비디오는 54% 실내, 46% 실외 환경으로 구성.
    • 평균 길이는 14초이며, 일상적인 다양한 장면 포함.
  • Masklets:
    • 190.9K 수동 mask와 451.7K Auto mask.
    • 사라졌다가 다시 나타나는 객체를 포함한 42.5%의 재등장율로 기존 데이터셋과 경쟁적인 품질을 보임.

학습, 검증, 테스트 분할

  • SA-V val/test set:
    • 도전적인 장면을 중심으로 구성: 빠른 움직임, 복잡한 occlusion, 사라짐/재등장 패턴 포함.
    • val: 293 masklets, 155 videos.
    • test: 278 masklets, 150 videos.
  • Internal dataset:
    • 라이선스 비디오 데이터 추가: 62.9K 비디오, 69.6K masklets로 학습 데이터 확장.






Experiments

1. Zero-shot Promptable Video Segmentation

먼저 사용자 경험과 유사한 대화형 설정을 시뮬레이션하는 대화형 비디오 segmentation을 평가했다. 데이터셋은 9개의 densely annotated zero-shot 비디오 데이터셋을 사용했다. 평가는 아래 두가지 방법으로 진행했다. 평가 metric으로는 J &F accuracy를 사용했다.

  • 오프라인 평가: 비디오를 여러 번 반복하며 가장 큰 오류를 가진 frame 선택.
  • 온라인 평가: 비디오를 한 번만 순차적으로 통과하며 frame 선택.

SAM2_9.png

결과는 위 도표와 같다. SAM 2는 SAM+XMem++ 및 SAM+Cutie보다 모든 평가 설정에서 우수한 성능을 보였다. 특히 SAM 2는 3배 적은 상호작용으로 더 높은 정확도의 세그멘테이션 결과를 생성했다.



2. Semi-supervised Video Object Segmentation (VOS)

SAM2_10.png

다음으로 비디오의 첫 번째 frame에서만 click, box 또는 mask prompt를 사용하여 semi-supervised 비디오 object segmentation(VOS) 설정을 평가했다. Click prmopt를 사용할 때 첫 번째 비디오 frame에서 1, 3 또는 5번의 클릭을 대화형으로 샘플링한 다음 이러한 클릭을 기반으로 객체를 추적했다.

결과는 위 표와 같다. SAM 2는 17개 데이터셋 모두에서 XMem++와 Cutie보다 우수한 성능을 보였으며, non-interactive VOS 작업에서도 기존 방법 대비 성능이 향상되었다. 특히 SAM 2는 VOS 작업에 최적화된 기존 방법들보다 높은 정확도와 효율성을 달성했다.



3. Image Segmentation

SAM2_11.png

다음으로는 37개의 zero-shot 데이터셋에서 image segmentation을 평가했다. 1-click 및 5-click mIoU를 평가했으며, 여기에는 SA-23(기존 SAM에서 평가된 23개 데이터셋)과 새로운 비디오 데이터셋이 포함된다.

결과는 위 표와 같다. SAM 2는 SAM 대비 1-click mIoU(58.9 vs 58.1)에서 더 높은 정확도를 기록하며, 6배 더 빠른 속도를 보인다. 또한 SA-1B와 비디오 데이터로 추가 학습 시 정확도가 61.4%로 상승했으며, 비디오 벤치마크에서 큰 성능 향상 확인할 수 있다.



4. State-of-the-Art Comparison in Semi-supervised VOS

SAM2_12.png

마지막으로 기존 SOTA 방법들과 semi-supervised VOS 설정에서의 비교를 수행했다. 평가 방법으로 Hiera-B+와 Hiera-L 이미지 인코더를 사용한 SAM 2의 속도와 정확도를 비교했다(FPS 측정 및 정확도 평가). 결과적으로 SAM 2는 기존 최고 성능 대비 상당한 정확도 향상을 보여주고 있다. 뿐만 아니라 LVOS 벤치마크에서 장기적인 비디오 segmentation에서도 우수한 성능을 보인다.



AI, Paper Review
Video Segmentation Zero-shot
This post is licensed under CC BY 4.0 by the author.