Open-Vocabulary Panoptic Segmentation with Text-to-Image Diffusion Models
https://arxiv.org/pdf/2303.04803
0. Abstract
ODISE: Open-vocabulary DIffusion-based panptic SEgmentation을 제시한다. 이는 pre train된 text-image diffusion, discriminative model을 활용하여 open vocabulary panoptic segmentation을 수행한다. text-to-image diffusion model은 open vocabulary lanugage description에서 높은 퀄리티의 이미지를 생성하는 능력을 가지고 있다. 이는 real world에서 open concept와 높은 관련성을 가지는 internal representation space를 가짐을 의미한다. 반면에 CLIP같은 discriminative model은 open vocabulary label 속에서 이미지를 잘 구분한다. 저자들은 이 두 모델의 internal representation을 활용하여 panoptic segmentation을 실행한다. 제안한 방법은 open vocabulary panoptic segmentation과 segmentation task에서 좋은 성능을 보인다.
1. Introduction
최근 대부분의 open vocabulary recognition은 광범위한 internet scale data로 학습된 text-image discriminatvie model을 활용한다. 이러한 pre train model은 각각의 object나 pixel을 잘 구별하며, scence level의 structural understanding을 위해 최적화할 필요가 없다. 그러나 CLIP은 종종 object간의 spatial한 관계를 혼동한다. text-image discriminative model의 spatial하고 relational한 이해부족은 open vocabulary panoptic segmentation에서 한계점을 제시한다.
반면에 diffusion을 이용한 text-to-image generative model은 text를 기반으로 image를 생성하여 text와 image간의 연관성을 학습한다. diffusion모델의 cross attention은 text embedding과 visual representation 간의 관계를 학습한다. Figure 1을 보면 diffusion모델의 internal feature를 clustering했을 때 부정확하지만 semantically하게 구분되고 위치도 적합함을 알 수 있다.
저자들은 이를 활용하여 ODISE: Open-vocabulary DIffusion-based panoptic SEgmentation을 제안한다. 이는 large scale text-image diffusion과 discriminative model 모두를 활용하여 panoptic segmentation의 최신성능을 달성한다.
전반적인 파이프라인은 figure 2와 같다. pre-train한 text-image diffusion모델에 image와 caption을 입력하여 internal feature을 추출한다. 이 feature들을 이용해 mask generator가 panoptic mask를 생성한다. mask generator는 training set에서 annotate된 mask를 학습한다. 그 후 mask classification module은 각각의 예측된 mask의 diffusion feature들을 text embedding과 연관시켜 각 mask를 categorize한다. 이 classification module을 mask category label이나 image level caption으로 학습시킨다. 학습 후 text-image diffusion과 discriminative model를 사용하여 open vocabulary panoptic inference를 수행한다.
contribution은 다음과 같다.
- ODISE는 large-scale text-to-image diffusion model을 활용한 최초의 open-vocabulary segmentation task이다.
- 효과적으로 text-image diffusion과 discriminative model을 활용한 open-vocabulary panoptic segmentation 파이프라인을 구축하였다.
- 많은 open vocabulary recognition task에서 최신성능을 달성하였다.
2. Related Work
panoptic segmentation은 instance, semantic segmentation을 합친 task로, 이전의 task에서는 closed vocabulary에서 수행되었다.
open vocabulary는 대부분 instance segmentation이나 semantic segmentation을 수행하였다. 하지만 저자들은 이 둘을 모 두 수행하는 panoptic segmentation을 수행한다. 또한, classification을 위한 image discriminative task를 위한 pre train모델은 segmentation을 하기에 최적화되진 않았다. 하지만 image-to-text diffusion model을 Figure 1에 증명된 바와 같이 segmentation에 적합하다.
3. Method
3.1. Problem Definition
train category와 test category가 다르다. test category에는 train할 때 보지 못한 category도 포함된다.
train할 때 각 category에 대해 이진적 panoptic mask annotation을 제공한다. 또한 각 mask에 대한 category label이나 text caption을 제공한다.
test할 때 category label이나 text caption을 제공하지 않는다. 오직 test category name만 제공한다.(여기서 name은 category list이고 label은 그 list 안에서 각각의 객체에 해당하는 category를 지정하는 것을 말한다. open vocabulary는 unknown을 탐지하고 이를 inference하는 것인데 list를 제공한다는 건 주관식이 아니라 객관식이 되기에 이부분에 대해서는 모순이 있다는 생각이 들기도 한다. GPT의 답변에 의하면 진정한 open vocabulary와는 다르다고 할 수 있지만 test시 name을 제공하는 경우가 많다고 한다.)
3.2. Method Overview
text-to-image diffusion model에 image와 caption이 입력되고 diffusion model의 internal feature이 추출된다(Sec 3.3). 추출된 feature과 제공된 training mask annotation을 활용하여 mask generator를 학습한다. mask generator는 panoptic mask를 생성한다(Sec 3.4). training label이나 text caption을 사용하여 open vocabulary mask classfication module을 학습한다. 각 예측된 mask diffusion feature와 training category name이 있는 text encoder embedding을 사용하여 mask를 구분한다(Sec 3.5). 학습 후 text-image diffusion과 discriminative model을 활용하여 open vocabulary panoptic inference를 한다(Sec 3.6).
3.3. Text-to-Image Diffusion Model
Background
text-to-image diffusion model은 text prompt의 입력을 받으면 높은 퀄리티의 이미지를 생성할 수 있다. 이는 internet에서 크롤링된 몇백만개의 image text pair로 학습되었다. 이 text는 pre train된 text encoder를 이용해 text embedding에 encoding되어 있다. diffusion network에 입력하기 전, image는 Gaussian noise로 인해 왜곡이 된다. diffusion network는 이 왜곡된 input과 text embedding의 왜곡을 보정한다. inference동안 model은 image shape의 Gaussian noise와 text embedding을 입력으로 받아 점점 denoise하면서 realistic한 image가 된다.
Visual Representation Extraction
diffsuion기반의 text-to-image generative model은 denoising 과정을 학습하기 위해 UNet 아키텍처를 사용한다. UNet은 cross attention을 수행하기 위해 text embedding과 UNet feature 사이에 convolution block, upsampling, downsampling block, skip connection, attention block으로 구성된다. noisy input image를 올바른 방향을 잡기 위해 denoising의 각 과정에 diffusion model은 text input을 사용한다. text가 cross attention layer를 통해 model에 주입되기 때문에 visual feature와 text description이 의미적으로 연관성을 학습할 수 있게 된다. 따라서 UNet block을 통해 출력된 feature map은 panoptic segmentation을 위한 rich하고 dense한 feature를 추출할 수 있다.
저자들의 방법은 생성모델의 전체적인 단계가 아닌 visual representation을 추출하기 위해 single forward pass만 필요하다. image-text pair (x, s)를 입력하면 t단계동안 x_t는 다음과 같이 noisy image가 된다.
caption s를 pre train된 text encoder T에 encode하고 text-to-image diffusion UNet의 internal feature f를 추출한다.
이는 x의 visual representation f가 paired caption인 s에 따라 영향을 받는다는 것을 의미한다. 이는 text-to-image diffusion model이 pre training되는 동안 image-text data가 활용가능할 때 올바르게 추출될 수 있다. 그러나 paried caption이 없을 때, visual representation을 추출할 수 없다는 문제가 생긴다. 일반적으로 caption이 없을 때, empty text를 사용한다. 하지만 이는 최적의 방식이 아니며, 이 방법이 성능을 하락시킨다는 것을 저자들은 실험을 통해 증명한다. 이를 해결하기 위해 caption image data를 위한 필요성을 해결하기 위해 설계된 Implicit Captioner을 제안한다.
Implicit captioner
caption을 생성하는 기존의 captioning network을 사용하는 것 대신에, input image로부터 implicit text embedding을 생성하는 network를 학습한다. 그 후 생성한 text embedding을 diffusion model에 직접 입력한다. implicit text embedding을 활용하기 위해 CLIP의 image encoder를 사용한다. diffusion model에 입력하기 위한 implicit text embedding에 image embedding을 투영하기 위해 학습된 MLP를 활용한다. 학습동안 이 image encoder와 UNet 파라미터들은 바뀌지 않고 MLP의 파라미터만 fine tuning 한다. 최종적으로 image x의 visual representation f를 계산하는 ODISE feature extractor는 text-to-image diffusion model의 UNet과 implicit caption을 활용한다.
3.4. Mask Generator
mask genoerator에는 visual representation f가 입력되면 N개의 class와 무관한 binary mask와 N개의 각 mask에 대한 embedding feature가 출력된다. mask generator는 bounding box기반과 segmentation mask기반의 방법 모두를 활용가능하다. 본 논문에서는 dense pixel prediction을 목표로 하므로 segmentation mask기반 architecture를 사용한다. GT mask의 corresponding과 함께 pixel binary cross entropy loss를 통해 class와 무관한 binary mask를 supervise한다.
3.5. Mask Classification
예측된 각 binary mask를 category label과 매칭하기 위해 text-image discriminative model을 활용한다.
Category Label Supervision
학습동안 각 mask의 GT category label을 활용할 수 있다. 따라서 학습과정은 전통적인 closed vocabulary 학습과 유사하다. Ktrain개의 category Ctrain에 대해 각 mask embedding feature zi는 GT category yi와 매칭된다. Ctrain의 모든 category가 text encoder에 encoding한다.
그 후 Ktrain class개 중 하나가 가지는 mask embedding feature z의 probability를 통해 classification loss LC를 계산한다.
Image Caption Supervision
category label이 없다고 가정하며, 각 이미지에 대한 자연어 caption을 활용하여 예측된 mask embedding feature를 구분하기 위해 학습할 수 있다. 이를 위해 각 caption에서 명사를 추출하고 grounding category label로 사용한다. grounding loss를 활용하여 mask category label의 예측을 supervise한다. grounding loss를 계산하기 위해 mask와 caption(명사)간의 similarity를 계산한다.
최종적으로 ODISE는 Lc 또는 LG와 class agnostic binary mask loss를 함께 학습한다.
3.6. Open-Vocabulary Inference
inference동안에는 Ctest category를 사용할 수 있으며, 이는 training category와 다를 수 있다. 또한, caption이나 label이 없기 때문에 implicit caption을 얻기 위해 implicit captioner을 통과한다. UNet feature를 얻기 위해 이 caption과 image를 diffusion model에 입력한다. mask generator를 통해 binary mask를 예측한다. 각 예측된 mask를 test category와 매칭하기 위해 식 6을 이용해 probability를 계산하고 최종적으로 최대 probability를 가지는 category를 예측한다.
실험을 통해 diffusion model이 spatial하게 object instance를 하는 많은 mask를 생성하지만, classification 능력은 text-image discriminative model(CLIP)과 조합하면 향상시킬 수 있다. Figure 3과 같이 input image x가 주어지면, text-to-image discriminative model의 image encoder V를 활용하여 feature map을 encode한다. 그 후 mask에 속하는 feature와 MaskPooling을 수행한다.
diffusion 모델의 probability p와 CLIP의 probability p를 결합하여 category 예측을 얻게 된다.
4. Experiments
5. Conclusion
diffusion model을 활용한 recognition task의 첫단계를 제안한다. ODISE는 SOTA 성능을 달성하며 diffusion model이 생성뿐만 아니라 풍부한 semantic representation을 학습할 수 있음을 보여준다. 이는 diffusion모델의 internal representation을 사용하여 어떻게하면 새로운 task로 활용할 수 있을지에 대한 가능성을 제시한다.