SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Kunhao Liu $^{1}$  Fangneng Zhan $^{2}$  Jiahui Zhang $^{1}$  Muyu Xu $^{1}$  Yingchen Yu $^{1}$  Abdulmotaleb El Saddik $^{3,5}$  Christian Theobalt $^{2}$  Eric Xing $^{4,5}$  Shijian Lu $^{1*}$
$^{1}$ Nanyang Technological University  $^{2}$ Max Planck Institute for Informatics
<sup>3</sup>University of Ottawa <sup>4</sup>Carnegie Mellon University <sup>5</sup>MBZUAI
Open-vocabulary segmentation of 3D scenes is a fundamental function of human perception and thus a crucial objective in computer vision research. However, this task is heavily impeded by the lack of large-scale and diverse 3D open-vocabulary segmentation datasets for training robust and generalizable models. Distilling knowledge from pre-trained 2D open-vocabulary segmentation models helps but it compromises the open-vocabulary feature as the 2D models are mostly finetuned with close-vocabulary datasets. We tackle the challenges in 3D open-vocabulary segmentation by exploiting pre-trained foundation models CLIP and DINO in a weakly supervised manner. Specifically, given only the open-vocabulary text descriptions of the objects in a scene, we distill the open-vocabulary multimodal knowledge and object reasoning capability of CLIP and DINO into a neural radiance field (NeRF), which effectively lifts 2D features into view-consistent 3D segmentation. A notable aspect of our approach is that it does not require any manual segmentation annotations for either the foundation models or the distillation process. Extensive experiments show that our method even outperforms fully supervised models trained with segmentation annotations in certain scenes, suggesting that 3D open-vocabulary segmentation can be effectively learned from 2D images and text-image pairs. Code is available at https://github.com/Kunhao-Liu/3D-0VS.
Semantic segmentation of 3D scenes holds significant research value due to its broad range of applications such as robot navigation [1], object localization [2], autonomous driving [3], 3D scene editing [4], augmented/virtual reality, etc. Given the super-rich semantics in 3D scenes, a crucial aspect of this task is achieving open-vocabulary segmentation that can handle regions and objects of various semantics including those with long-tail distributions. This is a grand challenge as it necessitates a comprehensive understanding of natural language and the corresponding objects in the 3D world.
The main challenge in open-vocabulary 3D scene segmentation is the lack of large-scale and diverse 3D segmentation datasets. Existing 3D segmentation datasets like ScanNet [5] primarily focus on restricted scenes with limited object classes, making them unsuitable for training open-vocabulary models. An alternative is to distill knowledge from pre-trained 2D open-vocabulary segmentation models to 3D representations as learned with NeRF [6] or point clouds, by fitting the feature maps or segmentation probability outputs from the 2D models [4, 7]. Though this approach circumvents the need for the 3D datasets, it inherits the limitations of the 2D models which are usually finetuned with close-vocabulary datasets of limited text labels [8, 9], thereby compromising the open-vocabulary property, especially for text labels with long-tail distributions [2, 3].
Figure 1: Weakly Supervised 3D Open-vocabulary Segmentation. Given the multi-view images of a 3D scene and the open-vocabulary text descriptions, our method distills open-vocabulary multimodal knowledge from CLIP and object reasoning ability from DINO into the reconstructed NeRF, producing accurate object boundaries for the 3D scene without requiring any segmentation annotations during training.
We achieve precise and annotation-free 3D open-vocabulary segmentation by distilling knowledge from two pre-trained foundation models into NeRF in a weakly supervised manner, supervised only by the open-vocabulary text descriptions of the objects in a scene, as illustrated in Fig. 1. One foundation model is CLIP [10] which is trained with Internet-scale text-image pairs [11] capturing extensive open-vocabulary multimodal knowledge. The other is DINO [12, 13] which is trained with large-scale unlabelled images capturing superb scene layout and object boundary information. However, CLIP yields image-level features which are not suitable for pixel-level semantic segmentation. Thus certain mechanisms should be designed to extract pixel-level CLIP features without fine-tuning. Additionally, the image patches' CLIP features may have ambiguities for segmentation, which need to be regularized for accurate open-vocabulary segmentation. At the other end, DINO produces feature maps instead of explicit segmentation maps. Certain distillation techniques should be designed to extract the necessary information from DINO features to facilitate precise segmentation.
We construct a hierarchical set of image patches to extract pixel-level features from image-level CLIP features and design a 3D Selection Volume to identify the appropriate hierarchical level for each 3D point, effectively aligning CLIP features with pixel-level features without fine-tuning. In addition, we introduce a Relevancy-Distribution Alignment (RDA) loss to address CLIP feature ambiguities, aligning segmentation probability distribution with class relevancies that capture similarities between class text features and corresponding CLIP features. Moreover, we propose a novel Feature-Distribution Alignment (FDA) loss to distill object boundary information from DINO features. The FDA loss encourages close segmentation probability distributions for points with similar DINO features and distant distributions for dissimilar features. To address the training instability due to diverse distribution shapes, we further re-balance weights associated with similar and dissimilar DINO features.
Our method enables weakly supervised open-vocabulary segmentation of 3D scenes with accurate object boundaries. By distilling knowledge from CLIP without fine-tuning, our approach preserves its open-vocabulary knowledge and effectively handles text labels with long-tail distributions. A notable aspect of our approach is that it does not require any manual segmentation annotations for either the foundation models or the distillation process. Remarkably, our experiments demonstrate that our method surpasses fully supervised models trained with segmentation annotations in certain scenes, highlighting the possibility that 3D open-vocabulary segmentation can be effectively learned from large amounts of 2D images and text-image pairs.
In summary, the contributions of this work are three-fold. Firstly, we propose an innovative pipeline for weakly supervised 3D open-vocabulary segmentation by distilling knowledge from pre-trained foundation models into NeRF without requiring any annotations in training. Secondly, we introduce a Selection Volume to align image-level CLIP features with pixel-level features, supplemented by novel Relevancy-Distribution Alignment and Feature-Distribution Alignment losses that respectively resolve CLIP features' ambiguities and effectively distill DINO features for 3D scene segmentation. Lastly, extensive experiments demonstrate that our method effectively recognizes long-tail classes and produces accurate segmentation maps, even with limited input data.
Open-vocabulary Segmentation. In recent years, the field of 2D open-vocabulary segmentation has garnered significant attention, driven by the availability of extensive text-image datasets and vast computational resources. Predominant approaches [8, 14-19] typically distill knowledge from large-scale pre-trained models, such as image-text contrastive learning models [10, 20-22] and diffusion models [23]. However, the distillation process requires fine-tuning on close-vocabulary datasets, contrasting with massive datasets used for large-scale pre-trained models [11]. This leads to limited performance in recalling infrequent classes with long-tail distributions [2, 3], compromising the open-vocabulary property. OpenSeg [24] is not finetuned on a closed set of classes but is weakly supervised via image captions. However, OpenSeg has a smaller vocabulary and knowledge than CLIP as it is trained on a much smaller dataset. Our method, without fine-tuning CLIP, effectively handles such classes.
3D Scenes Segmentation. 3D scene segmentation has been a long-standing challenge in computer vision. Traditional approaches focus on point clouds or voxels with limited class variety in datasets, restricting generalizability to unseen classes [5, 25-36]. Recently, numerous point-cloud-based techniques have emerged to explore open-vocabulary 3D scene segmentation by encoding 2D open-vocabulary models' features into 3D scene points [3, 7, 37-40]. However, these methods are also mostly evaluated on datasets with restricted scenes and limited class ranges [5, 27, 28, 41], not fully exhibiting the open-vocabulary property. Moreover, point clouds have compromised geometric details, making them less suitable for precise segmentation compared to NeRF representations [6, 42]. Consequently, there has been a surge in NeRF-based 3D segmentation techniques that mainly address interactive segmentation [4, 43, 44], panoptic segmentation [45, 46], moving part segmentation [47], object part segmentation [48], object co-segmentation [49], unsupervised object segmentation [50, 51], etc. FFD [4] attempts to segment unseen text labels during training by fitting LSeg's [8] feature maps to a NeRF, but inherits LSeg's limitations, hindering generalization to long-tail distribution classes. Our method overcomes these challenges by directly using CLIP image features and distilling them into a NeRF representation [42] without fine-tuning on close-vocabulary datasets.
Foundation Models. Pre-trained foundation models [52, 53] have become a powerful paradigm in computer science due to their ability to capture general knowledge and adapt to various downstream tasks [10, 12, 13, 20, 23, 54-57]. These models are trained using various paradigms in natural language processing, such as masked language modeling [57, 58], denoising autoencoder [59], replaced token detection [60], and sentence prediction tasks[61], as well as in computer vision, including data generation [23, 62, 63], data reconstruction [64], and data contrastive learning [10, 12, 13, 20-22]. Foundation models acquire emergent capabilities for exceptional performance on downstream tasks, either in a zero-shot manner or with fine-tuning. In this work, we harness the capabilities of two prominent foundation models, CLIP [10] and DINO [12, 13]. CLIP learns associations between images and texts by mapping them to a shared space, facilitating applications in tasks like image classification, object detection, visual question-answering, and image generation [9, 10, 23, 62, 65, 66]. DINO, trained in a self-supervised manner, extracts scene layout information, particularly object boundaries, and has been successfully employed in tasks such as classification, detection, segmentation, keypoint estimation, depth estimation, and image editing [12, 13, 49, 67-69].
We propose a novel method for weakly supervised open-vocabulary segmentation of reconstructed NeRF. Given the multi-view images of a scene and the open-vocabulary text description for each class, we aim to segment the reconstructed NeRF such that every 3D point is assigned a corresponding class label.
To achieve this, we exploit the CLIP model's multimodal knowledge by mapping each 3D point to a CLIP feature representing its semantic meaning. As CLIP only generates image-level features, we extract a hierarchy of CLIP features from image patches and learn a 3D Selection Volume for pixel-level feature extraction, as described in Sec. 3.1.
Algorithm 1: Extracting pixel-level features of an image from CLIP
Input: RGB image  $I\in \mathbb{R}^{3\times H\times W}$  , number of scale  $N_{s}$  , CLIP image encoder  $E_{i}$    
Output: multi-scale pixel-level features  $F_{I}\in \mathbb{R}^{N_{s}\times D\times H\times W}$    
Initialize: the patch size of each scale patch_sizes  $\in \mathbb{R}^{N_s}$ $F_{I} =$  zeros  $(N_{s},D,H,W)$    
/\* Loop over all the scales \*/   
for scale idx, patch_size in enumerate(patch_sizes) do   
stride  $=$  patch_size/4   
count  $=$  zeros(1,1,H,W) /\* Record the patch count for each pixel \*/   
 $F_{multi\_spatial} =$  zeros(1,D,H,W)/\* Multi-spatial feature of current scale \*/   
/\* Loop over all the patches \*/   
for x idx in range((H - patch_size)/stride + 1) do   
start_x=xidx  $\times$  stride   
for y idx in range((W - patch_size)/stride + 1) do   
start_y=yidx  $\times$  stride   
/\* Get image patch's coordinates with randomness \*/ (left, upper, right, lower)  $=$  (max(start_y - randint(0, stride), 0), max(start_x - randint(0, stride), 0), min(start_y + patch_size + randint(0, stride), W), min(start_x + patch_size + randint(0, stride), H)) /\* Get image patch's CLIP feature \*/  $F_{patch} = E_i(I.$  crop(left,upper,right,lower))   
 $F_{multi\_spatial}[::,upper:lower,left:right] += F_{patch}$  count[:.,upper : lower, left : right]  $+ = 1$    
end   
end   
 $F_{multi\_spatial} / =$  count   
 $F_{I}[scale\_idx] = F_{multi\_spatial}$
However, the image patches' CLIP features may have ambiguities for segmentation, showing inaccurate absolute relevancy values which lead to misclassification. Take an image patch capturing an apple lying on a lawn as an example. The corresponding CLIP feature contains both apple and lawn information. If the apple is relatively small, the patch could be classified as lawn because lawn dominates the patch's CLIP features. Then the class apple would be ignored. To address this issue, we introduce a Relevancy-Distribution Alignment (RDA) loss, aligning the segmentation probability distribution with each class's normalized relevancy map, as described in Sec. 3.2. For precise object boundaries, we align the segmentation probability distribution with the images' DINO features. Previous segmentation approaches utilizing DINO features have focused on unsupervised segmentation [49, 68], lacking semantic meaning for segmented parts. In the open-vocabulary context, assigning accurate text labels to segmented regions requires aligning DINO feature clusters with correct text labels. To overcome this challenge, we propose a Feature-Distribution Alignment (FDA) loss, segmenting the scene based on DINO features' distribution and assigning appropriate text labels, as described in Sec. 3.3.
3.1 Distilling Pixel-level CLIP Features with a 3D Selection Volume
We propose a method based on multi-scale and multi-spatial strategies to adapt CLIP's image-level features for pixel-level segmentation, motivated by the observation that a pixel's semantic meaning should remain invariant to its surrounding pixels. The multi-scale component extracts features from patches of varying sizes around each pixel, and the multi-spatial component extracts from patches in which each pixel is at different positions. We average the multi-spatial features in each scale, attributing the primary direction of these features to the pixel's semantic meaning. We utilize a sliding-window algorithm for multi-spatial feature extraction. To prevent checkerboard patterns in the features and potential segmentation artifacts, we introduce randomness in the window size. The algorithm's pseudo-code is in Alg. 1.
After extracting pixel-level features from CLIP, we now have the multi-view RGB images and their corresponding multi-scale pixel-level features. Each ray  $\mathbf{r}$  is assigned its multi-scale feature  $F(\mathbf{r})\in \mathbb{R}^{N_s\times D}$  for supervision. Rather than simply averaging each ray's multi-scale features across scales, we introduce a 3D Selection Volume  $S$  to determine the most suitable scale indicative of the object size within a patch. Each 3D point  $\mathbf{x}\in \mathbb{R}^3$  yields a selection vector  $S_{\mathbf{x}}\in \mathbb{R}^{N_s}$  from  $S$ . Following [4, 45], we introduce an additional branch to render the CLIP feature. We can then render the RGB value, the CLIP feature, and the selection vector of each ray  $\mathbf{r}$  using volume rendering [6]:
$$
\hat {C} (\mathbf {r}) = \sum_ {i} T _ {i} \alpha_ {i} C _ {i} \in \mathbb {R} ^ {3}, \quad \hat {F} (\mathbf {r}) = \sum_ {i} T _ {i} \alpha_ {i} F _ {i} \in \mathbb {R} ^ {D}, \tag {1}
$$
Figure 2: Mitigating CLIP features' ambiguities with normalized relevancy maps. For original relevancy maps  $r_a, r_b$  of classes  $a$  and  $b$ , we note a higher relevancy for class  $b$  in Region 2 than in other image regions. Despite this, the ambiguities of CLIP features lead to Region 2's classification as  $a$  due to the higher absolute relevancy of  $a$  in Region 2, even as  $a$  is located in Region 1. To rectify this, we normalize each class's relevancy maps to a fixed range. These normalized relevancy maps,  $\bar{r_a}$  and  $\bar{r_b}$ , reduce such ambiguities, facilitating accurate region-class assignments.
$$
S (\mathbf {r}) = \operatorname {S o f t m a x} \left(\sum_ {i} T _ {i} \alpha_ {i} S _ {i}\right) \in [ 0, 1 ] ^ {N _ {s}}, \tag {2}
$$
where  $C_i, F_i, S_i$  are the color, feature, and selection vector of each sampled point along the ray,  $T_i = \Pi_{j=0}^{i-1}(1 - \alpha_i)$  is the accumulated transmittance and  $\alpha_i = 1 - \exp(-\delta_i \sigma_i)$  is the opacity of the point. We apply a Softmax function to the selection vector of each ray such that the sum of the probability of each scale is equal to 1.
For a set of rays  $\mathcal{R}$  in each training batch, the supervision loss can then be formulated as the combination of the L2 distance between rendered and ground truth RGB values and the cosine similarities  $\cos\langle, \rangle$  between the rendered features and the selected multi-scale CLIP features:
$$
\mathcal {L} _ {\text {s u p e r v i s i o n}} = \sum_ {\mathbf {r} \in \mathcal {R}} \left(\left\| \hat {C} (\mathbf {r}) - C (\mathbf {r}) \right\| _ {2} ^ {2} - \cos \langle \hat {F} (\mathbf {r}), S (\mathbf {r}) F (\mathbf {r}) \rangle\right). \tag {3}
$$
Given a set of text descriptions  $\{[\mathrm{CLASS}]_i\}_{i = 1}^C$  of  $C$  classes and the CLIP text encoder  $E_{t}$ , we can get the classes' text features  $T = E_{t}([\mathrm{CLASS}])\in \mathbb{R}^{C\times D}$ . Then we can get the segmentation logits  $z(\mathbf{r})$  of the ray  $\mathbf{r}$  by computing the cosine similarities between the rendered CLIP feature and the classes' text features:
$$
z (\mathbf {r}) = \cos \langle T, \hat {F} (\mathbf {r}) \rangle \in \mathbb {R} ^ {C}. \tag {4}
$$
We can then get the class label of the ray  $l(\mathbf{r}) = \operatorname{argmax}(z(\mathbf{r}))$ .
3.2 Relevancy-Distribution Alignment for Ambiguity Mitigation
To mitigate the ambiguities of the CLIP features, we propose to align the segmentation probability distribution with the spatially normalized relevancy maps of each class, enabling our method to identify specific image regions described by each class text, as illustrated in Fig. 2. The segmentation probability of each ray  $P(\mathbf{r})$  can be derived from the segmentation logits with a Softmax function:
$$
P (\mathbf {r}) = \operatorname {S o f t m a x} (z (\mathbf {r})) \in [ 0, 1 ] ^ {C}. \tag {5}
$$
The relevancy of a given class is determined by the similarity between the class's text feature and the selected feature from the hierarchy of image patches' CLIP features. Given an image  $I$ , we can get its multi-scale pixel-level CLIP feature  $F_{I} \in \mathbb{R}^{N_{s} \times D \times H \times W}$  using Alg. 1 and selection vector  $S_{I} \in \mathbb{R}^{N_{s} \times H \times W}$  using Eq. (2). And then we can get the image's relevancy map  $R_{I} \in \mathbb{R}^{C \times H \times W}$  as:
$$
R _ {I _ {h w}} = S _ {I _ {h w}} \cos \langle T, F _ {I _ {h w}} \rangle , \tag {6}
$$
where where  $h, w$  denotes the index in the  $H$  and  $W$  channel. We normalize each class's relevancy independently within an input view to [0, 1] to mitigate the ambiguities of CLIP features, making our method discern image regions described by each class text:
$$
\bar {R} _ {I} = \left(R _ {I} - \min  \left(R _ {I}\right)\right) / \left(\max  \left(R _ {I}\right) - \min  \left(R _ {I}\right)\right) \in [ 0, 1 ] ^ {C \times H \times W}, \tag {7}
$$
Figure 3: Difference between similar and distant distributions. Distributions having large divergence from the target distribution exhibit significantly diverse shapes, increasing the training instability (left). Conversely, distributions displaying low divergence with the target distribution consistently demonstrate a similar shape (right).
where  $\min()$  and  $\max()$  are the functions getting the lowest and highest values across the spatial dimensions (i.e.  $H$  and  $W$ ). We apply a Softmax function to  $\bar{R}_I$  to make it a probability vector. Then we can assign each ray  $\mathbf{r}$  its normalized relevancy with all the classes  $\bar{R}(\mathbf{r}) \in [0,1]^C$ . We employ the Jensen-Shannon (JS) divergence to measure the discrepancy between the normalized relevancy  $\bar{R}(\mathbf{r})$  and the segmentation probability distribution  $P(\mathbf{r})$  of each ray, formulating the Relevancy-Distribution Alignment (RDA) loss:
$$
\mathcal {L} _ {R D A} = \sum_ {\mathbf {r} \in \mathcal {R}} \sum_ {c \in C} \left(P (\mathbf {r}) _ {c} \log \left(\frac {P (\mathbf {r}) _ {c}}{M _ {P \bar {R}} (\mathbf {r}) _ {c}}\right) + \bar {R} (\mathbf {r}) _ {c} \log \left(\frac {\bar {R} (\mathbf {r}) _ {c}}{M _ {P \bar {R}} (\mathbf {r}) _ {c}}\right)\right) / 2, \tag {8}
$$
where  $M_{P\bar{R}}(\mathbf{r}) = (P(\mathbf{r}) + \bar{R} (\mathbf{r})) / 2$  is the average of the two distributions, and the subscript  $c$  denotes the probability of the  $c$ th class. By aligning the normalized relevancies and the segmentation probability distributions, our method can effectively identify the specific region corresponding to the text description of each class.
3.3 Feature-Distribution Alignment for Precise Object Boundary Segmentation
To ensure the segmentation exhibits precise object boundaries, we align the segmentation probability distribution with the images' DINO features, which have been shown to capture superb scene layouts and object boundary information [12, 13]. Following [49, 68], we extract the scene layout information with a DINO feature correlation tensor. Given a patch of size  $H_{p} \times W_{p}$ , we can get the correlation tensor  $Corr\_F \in \mathbb{R}^{H_{p}W_{p} \times H_{p}W_{p}}$  as:
$$
C o r r _ {-} F _ {h w i j} = \cos \left\langle f _ {h w}, f _ {i j} \right\rangle , \tag {9}
$$
whose entries represent the cosine similarity between the DINO features  $f$  at spatial positions  $(h, w)$  and  $(i, j)$  of the patch. In order to construct the correlation tensor for the segmentation probability distribution, we propose utilizing the JS divergence to assess the similarity between segmentation probabilities at two distinct spatial positions. The choice of JS divergence offers several advantages, including its symmetric nature and a bounded range of [0, 1], which contribute to improved numerical stability. However, since we only care about the class label of each point, i.e. the entry with the highest probability, we use a low temperature  $\tau < 1$  to get a sharper version of the segmentation probability distribution  $\hat{P}$  to let the model focus on the entry with the largest probability:
$$
\acute {P} = \text {S o f t m a x} (z / \tau) \in [ 0, 1 ] ^ {C}. \tag {10}
$$
The distribution correlation tensor  $Corr\_D \in \mathbb{R}^{H_pW_p \times H_pW_p}$  can thus be computed with:
$$
C o r r _ {-} D _ {h w i j} = \sum_ {c \in C} \left(\hat {P} _ {h w c} \log \left(\frac {\hat {P} _ {h w c}}{M _ {\hat {P} \hat {P} c}}\right) + \hat {P} _ {i j c} \log \left(\frac {\hat {P} _ {i j c}}{M _ {\hat {P} \hat {P} c}}\right)\right) / 2, \tag {11}
$$
where  $\dot{P}_{hwc},\dot{P}_{ijc}$  are the segmentation probabilities of the  $c$ th class at spatial locations  $(h,w)$  and  $(i,j)$  of the patch,  $M_{\hat{P}\hat{P}} = (\dot{P}_{hw} + \dot{P}_{ij}) / 2$  is the average of the two distributions. Thus the correlation loss [68] can be expressed as: