Open-Vocabulary 3D Semantic Segmentation with Text-to-Image Diffusion Models
https://arxiv.org/pdf/2407.13642
0. Abstract
image-caption 쌍의 large-scale로 pre-train한 diffusion model을 사용하여 open vocabulary 3D semantic understanding을 위한 Diff2Scene를 제안한다. 이는 object, 외형, 재질, 위치, 구성요소 등을 효과적으로 구분하며, SOTA 성능을 달성하였다.
1. Introduction
open vocabulary task는 카테고리, 세부속성(모양, 색상, 재질 등), 복잡한 조합까지 text로 설명된 모든 semantic을 예측할 수 있어야 한다. 그러나 이는 쿼리의 다양성과 복잡성 때문에 challenging한 일이다. 이를 해결하기 위해 text와 image가 embedding된 CLIP같은 대형 foundation model에서 image feature를 3D로 끌어올린다. 이는 3D semantic understanding의 성과를 내지만 세밀한 class나 새로운 text query를 처리하는데 한계가 있다.
최근 개발된 text to image diffusion model은 조합된 description까지 처리하는 좋은 성능을 보인다. CLIP기반의 모델에서 한계가 보이는 반면에 이는 language로 표현된 semantic concept에도 잘 맞는 것이 증명이 되었다. 더군다나 global representation에 최적화된 CLIP과 달리 diffusion model은 dense prediction model에 필요한 local representation에서 좋은 성능을 보인다.
3D perception에서 중요한 challenge는 poin clouds와 dense label 데이터가 부족하다는 것이다. 이러한 문제를 CLIP을 이용한 zero-shot으로 많이 해결한다. 그러나 CLIP은 diffusion에 비해 fine-grained하고 localization 능력이 부족하다. 따라서 저자들은 diffusion과 mask 기반 segmentation head (Mask2Former)를 backbone으로 활용하여 mask와 semantic representation을 분리하는 방법을 제안한다. 이는 2D foundation model의 semantically rich 임베딩을 활용하여 3D branch의 geometrical 정확성 mask를 학습하는데 적합하다. 그러나 mask기반 segmentation head를 사용해 multi modal distillation을 행하는 것은 간단한 작업이 아니며, diffusion모델은 생성 목적으로 학습되어있기 때문에 추출된 frozen feature를 직접적으로 perception task에 활용하는 것은 불가능하다. 그래서 3D mask network로 mask를 예측하고 2D CLIP feature에 각 mask를 활용한다. 그러나 이를 위해서 3D label 데이터가 필요하며, 이 데이터가 부족하기 때문에 비현실적이다.
이를 위해 2D branch에서 3D branch로 지식전이를 하는 새로운 mask distillation 방법을 제안한다. 3D branch를 point clouds를 입력으로 받아 3D feature를 예측하는 방식으로 설계한다. 2D branch로부터 받은 mask embedding을 3D feature의 class probability를 예측하는데 활용한다. 이는 2D mask를 3D 기반의 pixel point correspondence로 끌어올리고 2D와 3D branch의 일관성을 학습한다.
contribution은 다음과 같다.
- 처음으로 제안하는 text image diffusion을 활용한 open vocabulary 3D semantic segmentation이다.
- 2D sementation모델로부터 knowledge distillation함으로써 3D mask 예측을 학습하는 새로운 mask distillation방법을 제안한다.
- 최신성능을 달성했다.
2. Related Work
Closed-vocabulary 3D semantic segmentation
3D point cloud는 regular한 구조가 아니기 때문에 2D task로 point cloud를 다룰수는 없다. 그래서 3D point cloud에 맞는 아키텍처들이 제안되었다. 또한 high cost, intensive labeling effot로 인해 annotation되어 있는 data가 부족하다. 이러한 제한을 해결하기 위해 zero-shot learning이 제안되었지만 여전히 어느정도의 3D point cloud annotation이 필요하다.
Open-vocabulary 3D segmentation
large-scale vision, lanugage representation learning의 진보로 open-vocabulary setting이 향상되었다. 3D semantic segmentation에서는 3D point를 2D image로 project하여 3D문제를 2D space에서 해결하는 방법이 있다. 하지만 이러한 방법은 정보손실이 일어나기에 최적의 방법이 아니다. 3D point cloud의 정보를 활용하기 위해 2D image로 생성된 caption을 이용하여 3D point cloud의 pseudo GT를 만들어 해결하는 방법이 있다. large scale vision, language model을 활용하여 3D point cloud model로부터 distilling knowledge방법을 제안하기도 했다. 이는 오직 CLIP기반의 2D open vocabulary semantic segmentation으로부터 knowledge를 distill 받아 학습한다. 그럼에도 불구하고 높은 성능을 달성하였지만 distilling의 치명적인 단점은 학습할 때 사용하지 않은 class 몇개에서 성능이 좋지 않다는 것이다. 이러한 knowledge distilling에서 영감을 받았으나 앞선 방법들과 다르게 저자들은 diffusion기반의 2D open vocabulary semantic segmentation을 teacher로 한 방법을 제안한다.
Diffusion models for scene understanding
diffusion model은 훌륭한 representation learning인 것으로 밝혀지고 있다. 그 결과 image classification, object detection 등 많은 understanding task에서 활용되며 좋은 성능을 보여주고 있다. text embedding과 연결된 internal representation을 가지는 Stable diffusion은 open vocabulary에서 좋은 성능을 보여준다.
3. Method
RGB이미지와 3D point cloud를 입력하여 각 3D point의 semantic label을 예측한다. Diff2Scene은 2개의 branch를 가진다. 2D branch는 2D open vocabulary semantic segmentation model로 설계 되었다. 이는 대량의 text-image쌍으로 pre-train된 text-to-image generative model을 활용한다. 이는 2D image가 입력되어 2D probability mask와 2D mask embedding corresponding을 예측한다. generative model 덕에 2D mask embedding은 semantically rich하다. 이는 RGB image의 주목할만한 패턴을 활용하여 2D salient mask를 생성한다. 3D branch는 point cloud와 2D mask embedding을 입력받는다. 2D mask embedding은 3D branch에서 출력되는 각각의 3D feature에 class probability를 할당하는 classifier로 사용되며, 그 결과 3D probability mask를 생성한다. 각 point의 semantic class를 예측하기 위해 해당 모델은 salient, geometric mask를 가지는 각각의 mask의 category logit을 계산하고 두 타입의 mask의 logit을 ensemble한다. 이를 통해 model은 RGB이미지로부터 salient pattern을 학습하고, point cloud로부터 geometric 정보를 학습한다.
3.2. 2D Semantic Understanding Model
3D point cloud의 GT label이 부족하기 때문에, 이를 해결하고자 vision-language foundation model을 활용하여 semantically rich 2D feature를 3D space로 transfer한다. 이를 행하는 다양한 방법이 존재하나 3D annotation 노력을 줄이기 위해 저자들은 pre train된 2D semantic segmentation 모델을 feature 추출기로 활용하여 open-vocabulary 3D segmentation을 함으로써 3D annotation GT가 필요없는 방법을 따른다.
2D segmentation model은 large-scale text-image 쌍으로 pretrain된 foundation model의 image backbone을 사용하고 semantic embedding을 예측하는 segmentation head를 가진다. 2D backbone은 contrastively pretrain이나 generatively pretrain 모두 가능하다. contrastive representation의 유명한 framework는 CLIP, ALIGN이 있고 generative pretrain은 Diffusion U-Net block이 있다.
text-image foundation모델로부터 feature representation을 추출하면, 이 feature를 segmentation head에 추가하여 각 point semantic class을 예측한다. segmentation 문제는 pixel 기반의 classification이나 mask 기반의 classification으로 정의된다. pixel 기반의 classification은 HWC(heightwidthdimension of feature embedding)이고, mask 기반의 classification은 2D feature map과 N mask queries가 입력되어 N 2D probabilistic mask와 corresponding mask embedding으로 출력된다.
저자들은 feature backbone을 localization능력이 뛰어난 diffusion을 채택하였다. 또한, mask 기반의 segmentation head를 활용한다. 이는 2D foundation model로부터 semantically-rich embedding을 활용할 수 있고 3D branch로부터 geometrically accurate mask를 학습할 수 있다.
3.3. Geometry-Aware 3D mask Model
pixel마다 semantically-rich feature를 제공지 않는 mask 기반의 segmentation head를 사용한다. 2D semantic understanding model은 image를 생성하는 stable diffusion model을 사용한다. 각 pixel feature은 고르게 추출되지 않는다. 따라서 mask기반의 2D foundation model로부터 geometry aware 3D mask model로 distill하는 mask distillation방안을 제안한다.
mask 기반의 2D foundation model은 N개의 2D probabilistic mask와 mask embedding을 예측한다. 특히 probabilistic mask는 pixel의 probability를 표현하는 probabilistic map이다. 첫번째로 pixel-point crrespondence를 계산하고 2D mask를 3D space에 올림으로써 3D probabilistic mask를 생성한다. 2D branch로부터 생성된 3D probability mask와 mask embedding의 정보를 mask distillation 방안을 제안한다. geometry-aware 3D mask를 생성하는 3D mask prediction model로 Minkowski network을 학습한다. point cloud는 memory와 계산량을 줄이기 위해 각 voxel안의 pixel을 평균화하여 voxel을 4등분한다. 3D mask prediction model은 3D feature를 생성하여 각 voxel을 represent하고 voxel안의 모든 point에 배치한다. 이는 full feature map을 만든다. semantic rich 2D mask embedding은 3D feature logit을 계산하기 위해 사용된다. 두 요소를 내적하여 sigmoid function을 취한다.
multimodal mask distillation loss를 활용하여 3D mask generator를 학습시킨다.
이는 2D, 3D branch가 일정하게 예측하도록 학습시키며 고해상도, semantically-rich feature represnetation을 학습할 수 있도록 한다.
3.4. Open-Vocabulary Inference
inference할 때는 Diff2Scene에 3D point cloud와 multiview 2D image가 입력된다. 2D semantic understanding model이 2D image를 이용하여 2D probabilistic mask와 mask embedding을 생성한다. 3D mask model은 3D point cloud과 mask embedding을 입력 받아 3D probabilistic mask를 예측한다. 3D poinc cloud의 label c를 예측하기 위해 각 2D mask를 diffusion과 discriminative model을 이용하여 label probability의 geometric mean을 계산한다. 그 후 아래 식을 이용해 label probability를 계산한다.
4. Experiments
5. Conclusion
large text-to-image diffusion model을 활용해 open vocabulary 3D semantic understanding을 행하였다. Diff2Scene은 zero-shot 3D semantic segmentation task에서 SOTA성능을 달성했으며, visual grounding task에서 뛰어난 성능을 보였다. 또한 unseen dataset, novel text query에서도 우수한 성능을 보이며 generative text-to-image foundation model을 활용한 3D semantic scene understanding task를 할 수 있는 방안을 제안한다.
제안한 model에는 여러가지 제안사항이 있다. 첫번째, 기존의 방안들보다 작은 object들에 대해서 좋은 성능을 보이지만, 여전히 rail같은 몇몇의 class에 대해서는 잘 구별하지 못한다. 둘째, 해당 모델이 비슷한 semantic meaning을 가지는 세세한 category들에 대해서는 혼돈이 있다. 예를 들어 windowsill을 window class로 구별한다.