본문 바로가기
컴퓨터비전/Semantic segmentation

[논문리뷰] MaskFormer

by PIAI 2024. 3. 21.

들어가며

 

이 글은 Semantic Segmentation, Instance Segmentation, Transformer, DETR에 대한 사전 지식이 있다는 가정하에 작성되었습니다.

 

오늘 리뷰할 논문은 MaskFormer입니다. 간략하게 말하자면 이 논문은 Instance Segmentation과 Semantic Segmentaiton 사이의 차이를 메꾸려고 나온 논문입니다. 

 

 

 

위 그림을 보시면 Instance와 Semantic의 차이를 볼 수 있습니다. 왼쪽이 Semantic Segmentation으로 모든 픽셀 단위의 class를 직접적으로 예측하는 Task입니다. 반면, 오른쪽은 Instance Segmentation으로 각 클래스별 Binary Segmentation을 예측하고 결합하는 Task입니다.

 

여기서 Semantic Sementation의 단점이 있는데, 각 픽셀은 단 하나의 클래스만 예측한다는 것입니다. 이는 모든 이미지에 따른 정적인 출력값을 가져 정밀한 Segmentation이 힘들다고 합니다. 반면, Instance Segmentation은 Instance 레벨의 Region과 Segments를 가변적으로 예측하여 정밀한 Segmentation이 가능하다고 합니다.

 

그래서, 이 "기존 Instance레벨로 Region과 Sements를 예측하는 Intance Segmentation 네트워크에, Semantic Segmentation에 적용하면 더 정밀한 Semantic Segmentation이 가능하고, 하나의 네트워크로 두 가지 Task가 가능하다"가 이 논문의 골자입니다.

 

이제 어떻게 이 네트워크를 어떻게 구성하는지 자세히 살펴보도록 하겠습니다.

 

MaskFormer

 

먼저 네트워크 구성은 아래와 같습니다.

 

 

위를 보시면 Pixel-Level Module, Transformer Module, Segmentation Module로 총 3가지 모듈이 존재합니다. 

 

pixel-level module

  • backbone:  본 논문에서는 ResNet과 Swin Transformer와 같은 이미지 분류 모델을 사용하였습니다. 이 backbone을 통과하여 image features(\( \mathcal {F} \in \mathbb {R}^{C_{\mathcal {F}}  * \frac {H} {S} * \frac {W} {S}} \))을 출력합니다. 이 image features는 Transformer module로 입력되고CNN의 특성으로 local information가 전달됩니다.

  • pixel decoder: image features \( \mathcal {F} \)을 점진적으로 Upsampling하여 이미지 크기를 복원합니다(per-pixel embeddings). 후에, 이 per-pixel embeddings은 transformer module에서 출력하여 MLP을 거친 N개의 mask embeddings과 dot product를 하여, N개의 mask 예측에 사용될 것입니다. 

transformer module

이 부분은 Transformer의 Decoder의 층을 쌓아 구성되었습니다. \( N \) 개의 영역에 대한 queries를 생성하고(\( \mathcal {Q} \in \mathbb {R}^{N * C_{\mathcal {Q}}} \)), image features \( \mathcal {F} \)와 Cross-Attention으로 상호작용하여 Global information을 학습합니다. 

 

segmentation module

transformer module에서 출력된 \( N \)개의 queries는 이 모듈에서 최종적으로 Instance segmentation과 같이 classificationmask를 예측합니다. 

 

  • N class prediction: \( N \)개의 segment들의 queries에 Linear layer가 연결하고, 각 segment들의 Softmax를 취해 Classification을 합니다. \( K \) 개의 클래스들과 Background를 합친 \( K+1 \)의 클래스를 예측하는 것으로 \( \left\{p_i \in \Delta^{K+1} \right\}_{i=1}^N \)입니다. 

  • N mask embeddings: 위 Linear layer가 class 예측이였다면, 이 것은 \( N \)개의 segment들의 Mask를 예측하기 위한, 임베딩의 Linear layer입니다. 2개의 Linear layer로 구성되어 있고 \( \mathcal {E}_{mask} \in \mathbb {R}^{C_{\mathcal {E}} * N} \)의 차원입니다.

  • N mask prediction: pixel-level module에서 예측한 per-pixel embeddings \( \mathcal {E}_{pixel} \)과 mask embeddings \( \mathcal {E}_{mask} \)를 dot product로 최종 각 segment의 mask를 예측하고, Binary mask이므로 sigmoid를 취해 [0, 1]범위의 값을 가집니다. 

 

Training

위에서 \(N \) 개의 mask와 class에 대한 예측을 했습니다. 그럼 이 두가지에 대한 학습을 위한 loss함수를 구성하면 아래와 같습니다. 

 

 

앞 Term은 classification의 loss를 계산하기 위한 cross-entropy이고, 뒤의 Term은 Binary mask에 대한 loss를 구하는 것입니다. 그리고 각 Term의 아래에 \( \sigma \)는 DETR에서 제안된 각 쿼리와 GT간 최적의 Matching알고리즘은 bipartile matching입니다(자세한 설명은 생략하겠습니다).

 

추가적으로, \( \mathcal {L}_{mask} \)에는 DETR와 같이 focal, dice loss를 적용하였고 파라미터 값은 각각 20과 1로 지정하였습니다. 그리고 객체가 없는 경우의 classification 인 \( \phi \)는 0.1의 파라미터를 적용하였습니다.

 

이렇게 해서 학습과정은 기존 Instance segmentation과 동일합니다. 

 

Inference

이제 마지막 Inference 단계만 남았습니다. Instance Segmentation의 관점에서는 \( N \)의 쿼리들이 예측한 class prediction과 mask prediction의 짝들이 가장 높은 확률을 가지는 픽셀에 대응되어 Segmentation을 하면 됩니다. 출력된 Instance Segmentation의 특징은 "같은 클래스라 하더라도 다른객체로 묶는 것인데, 같은 클래스는 같은 객체로 묶으면 Semantic Segmenation 아닌가?" 하는 것이 일반적인 관점입니다.

 

 

 

하지만 위 방법보다 더 좋은 방법이 있다고 합니다. \( N * K \)의 차원을 가지는 class prediction(배경 제외)와 \( N * H * W \)의 차원을 가지는 mask prediction을 단순히 행렬곱해주니 성능이 더 잘 나왔다고 합니다. 

 

 

이렇게 해서 Semantic segmentation과 Instance segmentation을 결합할 수 있게 되었습니다.

 

 

오늘은 MaskFormer의 리뷰를 해보았습니다.

긴 글 읽어주셔서 감사합니다

 

댓글