HQ-SAM

 

 

 

 

Abstract

• SAM의 본래 디자인과 model weight를 그대로 사용하는 모델인 HQ-SAM
• High-Quality Output Token을 도입해서 SAM의 mask decoder가 high quality mask를 generate하도록 함
•  

 

Motivation

기존의 SAM이 가지고 있는 문제는 두 가지가 있다:

1. Coarse mask boundaries
2. Incorrect prediction

 

이러한 문제는 Fig. 1에서 더 잘 확인할 수 있다. 

 

하지만 이를 해결하기 위한 방법으로 SAM decoder를 fine-tuning할 경우 zero-shot capability가 크게 떨어지는 문제가 있다. 따라서 HQ-SAM은 SAM architecture를 그대로 reuse한다. 

 

또한 이를 위해서 extremely fine-grained image mask annotation이 포함된 HQSeg-44K dataset을 구성하였다.

 

Figure 2.

zero-shot capability의 preservation은 Fig. 2에서 확인된다. 

 

 

 

Method

Figure 3.

 

High-Quality Output Token

Fig. 3를 보면, SAM의 original design에서는 output token이 MLP와 함께 mask decoding에 사용된다. 이는 learnable한 것으로, ViT의 class token이나 DETR의 object query와 비슷한 역할을 한다. 

 

SAM의 coarse mask를 바로 input으로 넣지 않고, HQ-Output Token$\in \mathbb R^{1\times256}$이 SAM의 output token $\in \mathbb R^{4\times256}$와 prompt token $\in \mathbb R^{N_{\text{prompt}}\times256}$에 concatenate된다. 

 

즉, 기존의 output token에 1d짜리 HQ-Output Token을 새로 추가하여 high resolution에 대응하도록 model의 capacity를 높인 것으로 볼 수 있다. 

 

이 HQ-Output Token은 decoder 안에서 cross attention, self attention 등이 수행된다. 

 

즉, 결과적으로 fine-tuning 하는 것은 3-layer MLP와 HQ-Output Token 뿐이고, 나머지 parameter는 freeze된다. 이런 방식으로 SAM의 zero-shot segmentation capacity를 forgetting하지 않은 채 fine-tuning할 수 있는 것이다. 

 

Global-Local Fusion for High-Quality Features

아주 fine하게 segmentation하기 위해서, SAM의 mask decoder의 feature를 그냥 사용하지 않고, SAM의 여러 stage에서 feature를 추출한 뒤 fuse해서 HQ-Feature를 만들어 사용했다. 

 

총 세 종류의 feature가 사용되었는데, 이는 다음과 같다: 

1. SAM의 ViT encoder의 early layer local feature(first global attention block) $64\times 64$
   -> general image edge/boundary detail 포착
2. SAM ViT encoder의 final layer global feature $64\times 64$
3. SAM의 mask encoder의 mask feature $256\times256$

이를 fuse하기 위해서 1과 2의 feature를 transposed convolution으로 upsampling하고, 이를 element-wise로 더했다. 

 

Table 3.

 

이 방법은 Table 3의 ablation study를 통해서 효과가 있었음을 입증한다. 

 

 

 

Results

quantatitive한 result는 생략하고 qualitative한 figure만 수록한다. 

 

Figure 6.

 

Figure 7.

---

References

[1] Ke, L., Ye, M., Danelljan, M., Tai, Y. W., Tang, C. K., & Yu, F. (2024). Segment anything in high quality. Advances in Neural Information Processing Systems, 36.

Footnotes