HQ-SAM: Segment Anything in High Quality

[Dongwoo-Im]

github : https://github.com/SysCV/SAM-HQ

---

[Introduction]

SAM 한계

• coarse mask boundaries (ex. thin object)
  • SAM이 학습한 SA-1B는 생성된 mask로 구성되어 complex structure 퀄리티가 좋지 않음

HQ-SAM : 기존 SAM의 0.5% 정도의 추가 parameter (freeze SAM)

• learnable HQ-Output token
• global-local feature fusion

HQSeg-44K : 학습 데이터 (dataset 6개 + fine-grained mask)

---

[Related work]

high-quality segmentation

• CRF, region growing 등의 post-segmentation refinement 방식은 high-level semantic 활용이 어렵다.
• 그래서 image encoder feature와 mask decoder feature를 fusion하여 high-level semantic을 반영하고, (기존의 64x64를 256x256까지 확장)
• SAM의 output token과 GT mask 차이를 measure하는 HQ-Output Token을 추가하여 mask quality를 높인다.
• 즉, 기존 SAM의 process를 따라가면서 target(tiny object)에 대한 segmentation 성능 향상을 목적으로 함

---

[Method]

SAM

• decoder layer에서 각 attention layer에는 point embedding to token / position embedding to image 적용 (on q, k)
• 즉, image embedding과 token embedding 사이에 position-aware two-way attention 수행

HQ-SAM

• high-quality output token
  • 기존 방식 대비 적은 parameter만 추가되기 때문에 time/data efficient
  • SAM을 feeze하기에 overfitting 방지
• global-local fusion for high-quality features
  • feature 3개 사용
    • early layer : local (64x64)
    • final layer : global (64x64)
    • mask feature (256x256)
  • convolution on 3 features
• training and inference
  • sample mixed types of prompts
  • add random Gaussian noise in the boundary of GT mask

• HQ-Output Token의 attention map이 더 detail하다.

---

[Experiments]

HQ-SAM은 SAM의 boundary가 미흡하다고 tackle하기 때문에 boundary metric 추가 (B 붙은건 boundary)

• SAM (baseline)
  • DIS, ThinObject dataset에서 유난히 성능이 낮은데, 해당 dataset의 train data가 HQSeg-44K에 포함되어 있어서 그런 것으로 추정한다. (supplementary material의 table 16을 참고하면, 해당 train data를 제외하더라도 성능차가 꽤 존재한다.)
• Using SAM's mask decoder feature
  • SAM + HQ-Output Token (X Output Token) : Output Token을 HQ-Output Token으로 대체한 실험으로 추정
  • SAM + HQ-Output Token (Boundary Loss) : boundary 영역만 mask loss로 학습한 실험으로 추정
• Using Out HQ-Feature
  • 당연하지만 SAM feature보다 HQ-Feature를 사용했을 때 성능이 더 좋다.

• 흥미로운건 Deocder Mask feature를 fusion하지 않았을 때 성능 하락이 심하다는 것
  • HQ-SAM은 SAM 결과에 의존하면서 SAM이 예측하지 못하는 boundary mask를 predict하기 때문으로 추정

• Training the whole SAM : overfitting
• Finetune SAM's decoder / post-refinement : overfitting on COCO
• HQ-SAM : Finetune SAM's output token 실험과 비교하여 유의미한 성능 향상

이외에 다양한 task에서 SAM의 성능을 뛰어넘는 결과를 보여줌

• Results on the SGinW Benchmark
• Zero-Shot Open-world Segmentation
• Zero-Shot Segmentation on High-resolution BIG Dataset
• Zero-shot Instance Segmentation on COCO and LVIS
• Point-based Interactive Segmentation Comparison
• Zero-shot High-quality Video Instance Segmentation

• noise에 대해서도 훨씬 robust하다.