View-Consistent Hierarchical 3D Segmentation Using Ultrametric Feature Fields

We take as input a set of multi-view images, along with their camera parameters. We run SAM on every image, typically yielding 50-150 masks per image. Our goal is to learn an implicit feature field that encodes a hiearchical understanding of these masks.

The SAM masks are inconsistent across views, but carry a great deal of information within each view. We use contrastive learning  [22, 11, 41] to distill this information into a 3D feature field.

We sample a pair of pixels in an image, and define it as a positive pair if both lie within a same mask, and a negative pair otherwise. We supervise the features of positive pairs to be more similar than those of negative pairs. Concretely, we follow Chen et al. [11] and minimize a binary cross entropy loss on distances between positive pairs and negative pairs. Given a positive pair $s_{p}=\{v_{p1},v_{p2}\}$ and a negative pair $s_{n}=\{v_{n1},v_{n2}\}$, our contrastive loss is

where $d$ is a distance metric, and $\tau$ is the temperature.

Minimizing the loss in Eq. 3 across many pairs $s_{p}$ and $s_{n}$ yields a feature space where segmentations can be recovered by querying a point and finding the neighborhood where pairwise feature distances are all below some threshold. As discussed in earlier sections, whether or not the resulting segmentation will be hierarchical depends on whether the distance metric $d$ is an ultrametric or a standard Euclidean metric. Our main goal is to reduce ultrametric distances, but to improve optimization, we apply loss in both ultrametric space and Euclidean space:

The ultrametric term forces the optimization to find a hierarchical decomposition of the scene, while the Euclidean term serves as regularization. Note that the Euclidean term has the advantage of directly providing a gradient to all considered feature pairs, whereas the ultrametric loss only provides a loss to the active (maximal) edge along the path between two segments.

As described in Eqs. 1 and 2, ultrametric distances are defined on a graph. At each training step, we form an approximate graph by sampling 4096 pixels within the image, connect each pixel to its 10 nearest neighbors, and use feature distances as edge weights. We use the binary partition tree algorithm [13] to efficiently compute ultrametric distances during training.

Because our input segmentations overlap within and across views, it is ambiguous whether a pair of pixels “lie within the same mask” and therefore represent a positive pair in our contrastive formulation. For example, two pixels may be within the same mask at a coarse granularity but not for a finer granularity. We address this ambiguity with a simple tree-based strategy for sampling positive and negative pairs of pixels.

For each image, we organize the segmentation masks into a hierarchical structure determined by the inclusion ratio between them. We additionally include an all-positive mask as the root of the tree. Then, for each view, we sample positive and negative pairs starting from the leaf nodes of the hierarchical tree, i.e. masks at the finest granularity. For a leaf mask $A$ with parent mask $B$, we randomly select two pixels in $A$ and designate these as a positive pair. We then randomly select one pixel in $A$ and one pixel in $\bar{A}\cap B$ and designate these as a negative pair. We then move to the parent mask and repeat. Additional details are provided in the supplementary.

We note that masks grow in size as we rise up the hierarchy, and pairs which were declared negative at finer granularity will be declared positive in courser segmentations. We mitigate this conflict by computing the contrastive loss within each level. This asks that the positive distances be smaller than the negative distances within a single level, and avoids the possibility of a pair being simultaneously positive and negative in Eq. 3. Taken together, our sampling strategy reflects an ultrametric structure and helps the segmentation hierarchy propagate across the scene.

Our approach shows that the 3D structure of a scene can resolve conflicts in segmentation cues. Conversely, can segmentation cues help resolve ambiguity in 3D structure? To explore this possibility, we add an assumption that regions belonging to the same semantic mask have smoothly changing depth. We propose a regularization that penalizes changes in curvature (i.e. the third derivative of depth) within a segment. Concretely, at each training iteration, we sample $k_{depth}$ local patches of $4\times 4$ pixels, ensuring that each sampled patch lies within one of the SAM’s finest-grained mask predictions. The depth continuity loss is defined as

where $\Delta\theta$ represents the ray angle difference between adjacent pixels, $t$ denotes a threshold, and $P_{D}$ contains the sampled pixels, with $p_{0},\ldots,p_{3}$ denoting adjacent pixels in a row or column. For training stability, we start using this loss halfway through training. Fig. 5 presents the benefits of the depth continuity loss.

To segment a 2D image from our feature field, we begin by rendering our feature field to the viewpoint of the image. We then construct a graph from the feature map, denoted as $\mathcal{G}=(V,E)$, where $V$ contains the pixel features, and $E$ contains edges connecting each pixel to its 4 spatial neighbors. Edge lengths are defined as feature distances. With a given threshold $t$ as the indicator of granularity, we remove all edges longer than this threshold, resulting in a new graph denoted as $\mathcal{G}_{t}$. We then identify all connected components within this graph, and return these components as our segmentation. We use this approach to compare against 2D segmentation methods, but we note that 2D-based segmentation is not view-consistent.

To achieve 3D-consistent segmentation, we create a featurized 3D point cloud by unprojecting feature maps and depth maps rendered from the optimized implicit field. Following related NeRF works, we remove outliers and downsample the pointcloud using Open3D [64]. Then, we construct a k-nearest-neighbor (KNN) graph on the 3D point cloud, and we set each edge weight in the KNN graph to the feature distance between the connected vertices. For a threshold $t$ indicating the level of granularity, we remove all edges longer than that threshold, and the remaining connected components represent segments at the specified level of granularity. In practice, we keep the $N$ largest components as our final segmentation. Using different values of $t$ allows for segmentation at varying levels of granularity.

To propagate the 3D segmentation into a novel view, we first render a depth map of the novel view and unproject the render into 3D points. Then, for each unprojected pixel, we find its $k$ nearest neighbors to the previously estimated KNN graph and assign a segmentation label using the mode of the neighborhood.

On Blender-HS and PartNet, we train our model for $20,000$ iterations with a batch size of $4096$ and use the same optimization parameters as DFF. In contrastive learning (see Eq. 3), we set the temperature $\tau$ to $0.1$ and sample $64$ positive and negative pairs from each mask. We set the loss weight $\alpha$ of the Euclidean loss in Eq. 4 to 1. For depth continuity loss, we sample $16$ patches per mask, and begin using the depth continuity loss after $5000$ iterations. During 3D inference, we extract a point cloud from training-view depth maps, apply voxel downsampling with voxel size $2\times 10^{-3}$, and run outlier removal with distance threshold $4\times 10^{-3}$ and number threshold of $1$. We build the graph of points using $k_{graph}=16$ nearest neighbors, and we transfer point segmentation labels into a novel view using the mode of $k_{query}=5$ nearest neighbors. We retain $N=200$ graph components and set the distance threshold $d$ to be $5\times 10^{-3}$. Please refer to  Sec. 4 for definitions of the above hyperparameters.

We first organize the segmentation masks into a hierarchical structure determined by the inclusion ratio between them. One mask $A$ is designated as a child of another mask $B$ when$\frac{|A\cap B|}{|A|}>p_{in}$ and $\frac{|A\cap B|}{|A\cup B|}<p_{IoU}$. We empirically set $p_{in}=0.95$ and $p_{IoU}=0.85$. We present the hierarchical sampling algorithm we introduced in Sec. 4.1.3 in Algorithm 1.

We sample same number of positive pairs and negative pairs for training for training efficiency. Implementing a 4-1 ratio ($4\times$ more negatives than positives) instead of 1-1, the normalized covering (NC) score increases from 0.709 to 0.720. However, this slow down our training by 43%, primarily due to the time-consuming computation of ultrametric distances and the associated minimum spanning tree.

We train and perform inference on a Titan RTX GPU. Training typically takes $\sim$70 minutes, while inference takes 5 seconds per granularity for 10 views. The main expense in inference is the watershed algorithm running on 3D point clouds, which is executed once per granularity and is view-independent.

We first visualize the ground truth segmentations for the Drums scene in the BlenderHS Dataset [33] in Fig. 8. We then present qualitative results on the BlenderHS dataset [33] in Fig. 9. Our segmentations exhibit a hierarchical structure and maintain consistency across different views. We also visualize our ultrametric feature field on the Lego scene in Fig. 11, showing sharper features than DFF.

We present qualitative results on the PartNet dataset [36] in Fig. 10. Our method is able to generate hierarchical segmentation results of different objects. Leveraging the 2D masks predicted with SAM as guidance, our method proficiently segments various surfaces of sub-parts within the object, while those are not included in the PartNet ground truth annotations.

We present qualitative results on the the LLFF dataset [33] in Fig. 12. Our approach is able to generate view-consistent hierarchical segmentation results for real-world scenes.