SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Ruihang Chu<sup>1</sup> Enze Xie<sup>2</sup> Shentong Mo<sup>3</sup>  
Zhenguo Li<sup>2</sup> Matthias Nießner<sup>4</sup> Chi-Wing Fu<sup>1</sup> Jiaya Jia<sup>1,5</sup>  
<sup>1</sup>The Chinese University of Hong Kong <sup>2</sup>Huawei Noah's Ark Lab  
<sup>3</sup>MBZUAI <sup>4</sup>Technical University of Munich <sup>5</sup>SmartMore  
https://ruihangchu.com/diffcomplete.html
We introduce a new diffusion-based approach for shape completion on 3D range scans. Compared with prior deterministic and probabilistic methods, we strike a balance between realism, multi-modality, and high fidelity. We propose DiffComplete by casting shape completion as a generative task conditioned on the incomplete shape. Our key designs are two-fold. First, we devise a hierarchical feature aggregation mechanism to inject conditional features in a spatially-consistent manner. So, we can capture both local details and broader contexts of the conditional inputs to control the shape completion. Second, we propose an occupancy-aware fusion strategy in our model to enable the completion of multiple partial shapes and introduce higher flexibility on the input conditions. DiffComplete sets a new SOTA performance (e.g.,  $40\%$  decrease on  $l_{1}$  error) on two large-scale 3D shape completion benchmarks. Our completed shapes not only have a realistic outlook compared with the deterministic methods but also exhibit high similarity to the ground truths compared with the probabilistic alternatives. Further, DiffComplete has strong generalizability on objects of entirely unseen classes for both synthetic and real data, eliminating the need for model re-training in various applications.
Figure 1: Our method is able to (a) produce realistic completed shapes from partial scans, (b) incorporate multiple incomplete scans (denoted by the plus symbol) to improve the completion accuracy, and (c) directly generalize to work on real objects of unseen classes without finetuning.
The advent of affordable range sensors, such as the Microsoft Kinect and Intel RealSense, has spurred remarkable progress in 3D reconstruction [1], facilitating applications in content creation, mixed reality, robot navigation, and more. Despite the improved reconstruction quality [2-5], typical scanning sessions cannot cover every aspect of the scene objects. Hence, the reconstructed 3D models often have incomplete geometry, thereby limiting their practical usage in downstream applications.
To fully unlock the potential of 3D reconstruction for supporting assorted applications, it is essential to address the challenges of completing incomplete shapes.
Effectively, a shape completion should produce shapes that are realistic, probabilistic, and high-fidelity. First, the produced shapes should look plausible without visual artifacts. Second, considering the under-determined nature of completion, it is desirable for the model to generate diverse, multimodal outputs when filling in the missing regions, to improve the likelihood of obtaining a superior shape and enable creative use. Third, while multiple outputs encourage the model's generative ability, maintaining effective control over the completion results is crucial to ensure coherent reconstructions that closely resemble the ground-truth shapes.
Current approaches to 3D shape completion fall into deterministic and probabilistic paradigms. The former class excels at aligning predictions with ground truths, owing to their determined mapping functions and full supervision. However, this can expose the models to a higher risk of over-fitting, leading to undesirable artifacts, such as rugged surfaces, especially on unseen objects. On the other hand, probabilistic approaches formulate the shape completion as a conditional generation task to produce plausible results, paired with techniques like autoregressive model [6], adversarial training [7-9], or diffusion models [10-13]. These approaches mainly focus on cases that prioritize completion diversity, for instance, filling in large missing regions cropped out by 3D boxes, where the algorithm has more freedom to explore various plausible shapes. However, when completing shapes obtained from range scans/reconstructions, whose geometries are often contaminated with noise, having varying degrees of incompleteness (see row 1 in Fig. 1), a promising goal is to recover the ground-truth scanned objects as accurately as possible. Therefore, relying solely on previous probabilistic approaches may compromise the completion accuracy and leads to low-fidelity results.
In this work, we focus on completing shapes acquired from range scans, aiming to produce realistic and high-fidelity completed shapes, while considering also the probabilistic uncertainty. We approach this task as a conditional generation with the proposed diffusion model named DiffComplete. To address the accuracy challenge typically in probabilistic models, we introduce a key design that incorporates range scan features in a hierarchical and spatially-consistent manner to facilitate effective controls. Additionally, we propose a novel occupancy-aware fusion strategy that allows taking multiple partial shapes as input for more precise and controllable completions.
Specifically, we formulate the diffusion process on a truncated distance field in a volumetric space. This representation allows us to encode both complete and incomplete shapes into multi-resolution structured feature volumes, thereby enabling their interactions at various network levels. At each level, features are aggregated based on voxel-to-voxel correspondence, enabling precise correlation of the difference in partialness between the shapes. By leveraging multi-level aggregation, the completion process can be controlled by both local details and broader contexts of the conditional inputs. It effectively improves the completion accuracy and network generalizability, as illustrated in Sec. 4.5.
To condition the completion on multiple partial shapes, our occupancy-aware fusion strategy adopts the occupancy mask as a weight to adaptively fuse all observed geometry features from various incomplete shapes, and then employ the fused features to control the completion process. For robustness, such an operation is performed in a multi-scale feature space. With our novel design, switching between single and multiple conditions can be efficiently achieved by finetuning an MLP layer. As Fig. 1(b) and 6 show, using multiple conditions for completion progressively refines the final shapes towards the ground truths.
DiffComplete sets a new state-of-the-art performance on two large-scale shape completion benchmarks [14, 15] in terms of completion accuracy and visual quality. Compared to prior deterministic and probabilistic approaches, we not only generate multimodal plausible shapes with minimal artifacts but also present high shape fidelity relative to the ground truths. Further, DiffComplete can directly generalize to work on unseen object categories. Without specific zero-shot designs, it surpasses all existing approaches for both synthetic and real-world data, allowing us to eliminate the need for model re-training in various applications. To sum up, our contributions are listed as follows:
- We introduce a diffusion model to complete 3D shapes on range scans, enabling realistic, probabilistic, and high-fidelity 3D geometry.
- We enhance the completion accuracy with two designs. For a single condition, we propose a hierarchical feature aggregation strategy to control the outputs. For multiple conditions, we introduce an occupancy-aware fusion strategy to incorporate more shape details.
- We show SOTA performance on shape completion on both novel instances and entirely unseen object categories, along with an in-depth analysis on a range of model characteristics.
RGB-D reconstruction. Traditional methods rely on geometric approaches for 3D reconstruction [16-20]. A pioneering method [21] proposes a volumetric fusion strategy to integrate multiple range images into a unified 3D model on truncated signed distance fields (TSDF), forming the basis for many modern techniques like KinectFusion [22, 23], VoxelHashing [24], and BundleFusion [2]. Recent learning-based approaches further improve the reconstruction quality with fewer artifacts [25-28, 3], yet the intrinsic occlusions and measurement noise of 3D scans constrain the completeness of 3D reconstructions, making them still less refined than manually-created assets.
3D shape completion is a common post-processing step to fill in missing parts in the reconstructed shapes. Classic methods mainly handle small holes and geometry primitives via Laplacian hole filling [29-31] or Poisson surface reconstruction [32, 33]. Another line exploits the structural regularities of 3D shapes, such as the symmetries, to predict the unobserved data [34-38].
The availability of large 3D data has sparked retrieval-based methods [39-42] and learning-based fitting methods [43, 44, 14, 45-49]. The former retrieves the shapes from a database that best match the incomplete inputs, whereas the latter minimizes the difference between the network-predicted shapes and ground truths. 3D-EPN [14], for instance, proposes a 3D encoder-decoder architecture to predict the complete shape from partial volumetric data. Scan2Mesh [50] converts range scans into 3D meshes with a direct optimization on the mesh surface. PatchComplete [15] further leverages local structural priors for completing shapes of unseen categories.
Generative methods, e.g., GANs [51, 9, 52, 7, 8] and AutoEncoders [53], offer an alternative to shape completion. While some generative models can generate diverse global shapes given a partial input, they potentially allow for a high generation freedom and overlook the completion accuracy. IMLE [54], for instance, adopts an Implicit Maximum Likelihood Estimation technique to specially enhance structural variance among the generated shapes. Distinctively, we formulate a diffusion model for shape completion. Our method mainly prioritizes fidelity relative to ground truths while preserving output diversity. Compared to both generative and fitting-based paradigms, our method also effectively reduces surface artifacts, producing more realistic and natural 3D shapes. In addition, we show superior generalization ability on completing objects of novel classes over SOTAs.
Diffusion models for 3D generation. Diffusion models [55-62] have shown superior performance in various generation tasks, outperforming GANs' sample quality while preserving the likelihood evaluation property of VAEs. When adopted in 3D domain, a range of works [12, 63-65] focus on point cloud generation. For more complex surface generation, some works [66, 13, 10, 67, 68, 11] adopt latent diffusion models to learn implicit neural representations and form final shapes via a decoder. For conditional shape completion, both DiffRF [10] and Diffusion-SDF [69] adopt a masked diffusion strategy to fill in missing regions cropped out by 3D boxes. However, their training processes do not leverage a paired incomplete-to-complete ground truth, which may prevent them from accurately learning the completion rules. Contrarily, our method explicitly uses the scan pairs for conditional training, yielding outputs closely resembling the actual objects scanned. Also, we apply the diffusion process in explicit volumetric TSDFs, preserving more geometrical structures during the completing process.
To prepare the training data, we generate an incomplete 3D scan from depth frames using volumetric fusion [21] and represent the scan as a truncated signed distance field (TSDF) in a volumetric grid. Yet, to accurately represent a ground-truth shape, we choose to use a volumetric truncated unsigned distance field (TUDF) instead. This is because retrieving the sign bit from arbitrary 3D CAD models (e.g., some are not closed) is non-trivial. By using TUDF, we can robustly capture the geometric features of different objects without being limited by the topology.
Given such a volume pair, i.e., the incomplete scan  $c$  and complete 3D shape  $x_0$ , we formulate the shape completion as a generation task that produces  $x_0$  conditioned on  $c$ . We employ the probabilistic diffusion model as our generative model. In particular, it contains (i) a forward process (denoted as  $q(x_{0:T})$ ) that gradually adds Gaussian noise to corrupt the ground-truth shape  $x_0$  into a random noise volume  $x_T$ , where  $T$  is the total number of time steps; and (ii) a backward process that employs a shape completion network, with learned parameters  $\theta$ , to iteratively remove the noise from  $x_T$  conditioned on the partial scan  $c$  and produce the final shape, denoted as  $p_\theta(x_{0:T}, c)$ . As both the forward and backward processes are governed by a discrete-time Markov chain with time steps  $\{0, ..., T\}$ , their Gaussian transition probabilities can be formulated as
$$
q \left(x _ {0: T}\right) = q \left(x _ {0}\right) \prod_ {t = 1} ^ {T} q \left(x _ {t} \mid x _ {t - 1}\right), \quad q \left(x _ {t} \mid x _ {t - 1}\right) := \mathcal {N} \left(\sqrt {1 - \beta_ {t}} x _ {t - 1}, \beta_ {t} \mathbf {I}\right) \tag {1}
$$
$$
\text {a n d} p _ {\theta} \left(x _ {0: T}, c\right) = p \left(x _ {T}\right) \prod_ {t = 1} ^ {T} p _ {\theta} \left(x _ {t - 1} \mid x _ {t}, c\right), \quad p _ {\theta} \left(x _ {t - 1} \mid x _ {t}\right) := \mathcal {N} \left(\mu_ {\theta} \left(x _ {t}, t, c\right), \sigma_ {t} ^ {2} \mathbf {I}\right). \tag {2}
$$
In the forward process (Eq. (1)), the scalars  $\beta_{t} \in [0,1]$  control a variance schedule that defines the amount of noise added in each step  $t$ . In the backward process (Eq. (2)),  $p(x_{T})$  is a Gaussian prior in time step  $t$ ,  $\mu_{\theta}$  represents the mean predicted from our network and  $\sigma_t^2$  is the variance. As suggested in DDPM [60], predicting  $\mu_{\theta}(x_{t},t,c)$  can be simplified to alternatively predicting  $\epsilon_{\theta}(x_{t},t,c)$ , which should approximate the noise used to corrupt  $x_{t-1}$  in the forward process, and  $\sigma_{t}$  can be replaced by the pre-defined  $\beta_{t}$ . With these modifications, we can optimize the network parameter  $\theta$  with a mean squared error loss to maximize the generation probability  $p_{\theta}(x_0)$ . The training objective is
$$
\underset {\theta} {\arg \min } E _ {t, x _ {0}, \epsilon , c} [ \| \epsilon - \epsilon_ {\theta} (x _ {t}, t, c) \| ^ {2} ], \quad \epsilon \in \mathcal {N} (0, \mathbf {I}) \tag {3}
$$
where  $\epsilon$  is the applied noise to corrupt  $x_0$  into  $x_{t}$  and  $\mathcal{N}(0,\mathbf{I})$  denotes a unit Gaussian distribution. We define all the diffusion processes in a volume space of a resolution  $S^3$ , i.e.,  $x_{0:T},c,\epsilon \in \mathbb{R}^{S\times S\times S}$ , where each voxel stores a scalar TSDF/TUDF value;  $S = 32$  in our experiments. Compared to previous latent diffusion models [11, 13, 66, 10] that require shape embedding first, we directly manipulate the shape with a better preservation of geometric features. Doing so naturally enables our hierarchical feature aggregation strategy (see Sec. 3.2). Next, we will introduce how to predict  $\epsilon_{\theta}(x_t,t,c)$ .
3.2 Shape Completion Pipeline
Overview. To enhance the completion accuracy, we encourage the incomplete scans to control completion behaviors. Inspired from the recent ControlNet [70], which shows great control ability given 2D conditions, we adopt a similar principle to train an additional control branch. To predict the noise  $\epsilon_{\theta}(x_t,t,c)$  in Eq. (3), we encode the corrupted ground-truth shape  $x_{t}$  by a main branch and the partial shape  $c$  by a control branch, where two branches have the same network structure without parameter sharing. Owing to the compact shape representation in 3D volume space, both complete and incomplete shapes are encoded into multi-resolution feature volumes with preserved spatial structures. Then, we hierarchically aggregate their features at each network level, as the sizes of two feature volumes are always aligned. In this work, we simply add up the two feature volumes to allow for a spatially-consistent feature aggregation, i.e., only features at the same 3D location are combined. Compared with frequently-used cross-attention technique [11, 13, 57], doing so can significantly reduce computation costs. By multi-scale feature interaction, the network can effectively correlate the difference in partialness between two shapes, both locally and globally, to learn the completion regularities. The final integrated features are then used to predict the noise volume  $\epsilon_{\theta}$ .
Network architecture. Fig. 2 provides an overview of our network architecture, which has a main branch and a control branch to respectively handle complete and incomplete shapes. The main branch (upper) is a 3D U-Net modified from a 2D version [59]. It takes as input an corrupted complete shape  $x_{t}$ , including (i) a pre-processing block  $\varepsilon_{x}(\cdot)$  of two convolution layers to project  $x_{t}$  into a high-dimension space, (ii)  $N$  subsequent encoder blocks, denoted as  $\{F_x^i (\cdot)\}_{i = 1}^N$ , to progressively encode and downsample the corrupted shape  $x_{t}$  into a collection of multi-scale features, (iii) a middle block  $M_{x}(\cdot)$  with a self-attention layer to incorporate non-local information into the encoded feature volume, and (iv)  $N$  decoder blocks  $\{D_x^i (\cdot)\}_{i = 1}^N$  that sequentially upsample features through an inversion convolution to produce a feature volume of the same size as the input  $x_{t}$ . The term
Figure 2: Given a corrupted complete shape  $x_{t}$  (diffused from  $x_{0}$ ) and an incomplete scan  $c$ , we first process them into  $\varepsilon_{x}(x_{t})$  and  $\varepsilon_{c}(c)$  to align the distributions. We employ a main branch to forward  $\varepsilon_{x}(x_{t})$ , and a control branch to propagate their fused features  $f$  into deep layers. Multi-level features of  $f$  are aggregated into the main branch for hierarchical control in predicting the diffusion noise. To support multiple partial scans as condition, e.g., two scans  $\{c_{1}, c_{2}\}$ , we switch on occupancy-aware fusion (Sec. 3.3). This strategy utilizes the occupancy masks to enable a weighted feature fusion for  $c_{1}$  and  $c_{2}$  by considering their geometry reliability before feeding them into the main branch.
“e(d)ncoder/middle block” denotes a group of neural layers commonly adopted as a unit in neural networks, e.g., a “resnet+downsample” block.
We use a control branch (bottom) to extract the features of incomplete scan(s). Likewise, we employ a pre-processing block,  $N$  encoder blocks, and a middle block to extract multi-scale features from the conditional input, denoted as  $\varepsilon_{c}(\cdot),\{F_{c}^{i}(\cdot)\}_{i = 1}^{N}$ , and  $M_{c}$ , respectively. They mirror the architecture of the corresponding blocks in the upper branch, yet operating with non-shared parameters, considering the different input data regularities (e.g., complete and sparse). Unlike the upper branch, there are no decoders for computation efficiency. Instead, we append a projection layer after each encoder/middle block  $F_{c}^{i} / M_{c}$  to forward multi-scale features and feed them into the decoder blocks of the main branch. As diffusion models are time-conditioned, we also convert the time step  $t$  into an embedding via two MLPs and add it to the volume features in each network block (the blue box in Fig. 2).
Hierarchical feature aggregation. To integrate the features of the complete and incomplete shapes for conditional control, we fuse them across multiple network levels. In this process, we first consider a single incomplete shape (denoted as  $c$ ) as the condition. Given the significant disparity between  $c$  and  $x_{t}$  (incomplete v.s. complete, TSDF v.s. TUDF), we employ the pre-processing layers  $(\varepsilon_{x}(\cdot)$  and  $\varepsilon_{c}(\cdot))$  to first empirically project the two different fields into more compatible feature space. Then we fuse them into feature volume  $f$ , i.e.,  $f = \varepsilon_{x}(x_{t}) + \varepsilon_{c}(c)$ . After that,  $f$  is passed through the control branch to better propagate the useful information to the deeper layers. In turn, we use multi-scale condition features from the control branch to guide each level of decoder blocks  $\{D_{x}^{i}(\cdot)\}_{i=1}^{N}$  in the main branch. The features that enter the  $i$ -th decoder block are denoted as
$$
d ^ {i} = \left[ D _ {x} ^ {i - 1} \left(x _ {t}\right), F _ {x} ^ {i} \left(x _ {t}\right) + \phi^ {i} \left(F _ {c} ^ {i} (f)\right) \right], \quad f = \varepsilon_ {x} \left(x _ {t}\right) + \varepsilon_ {c} (c) \tag {4}
$$
where  $[\cdot, \cdot]$  is the concatenation operation and  $\phi^i$  is one  $1 \times 1$  convolution layer for feature projection. For simplicity, we use  $F_x^i(x_0)$  to denote the features of the  $x_0$  after the  $i$ -th encoder block  $F_x^i(\cdot)$ , as  $x_0$  is not directly processed by  $F_x^i(\cdot)$ ; the same notation is used for  $F_c^i(f)$  and  $D_x^{i-1}(x_t)$ . The feature after the final decoder block is the network output  $\epsilon_\theta(x_t, t, c)$ .
By such a design, we can hierarchically incorporate conditional features, leveraging both their local and broader contexts to optimize the network outputs. Interestingly, we observe that adjusting the network level for feature aggregation can alter a trade-off between completion accuracy and diversity. Specifically, if we only aggregate features at the low-resolution network layers, we might miss finer geometry details present in the higher-resolution layers. This could reduce completion accuracy, as the model has less contextual information to work with. In contrast, if aggregating features at all levels, the effective control might make completion results closely resemble the ground truths, yet lowering the diversity. Further ablation study is provided in Sec. 4.5.
3.3 Occupancy-aware Fusion
Our framework can also take multiple incomplete scans as inputs. This option not only provides richer information to constrain the completed shape geometry but also enhances the approach's practicality, particularly in scenarios that a single-pass scan may not fully capture the entire object. Our critical design is to effectively fuse multiple partial shape features. As averaging the original TSDF volumes is sensitive to registration noise, we first register multiple partial shapes (e.g., using Fast-Robust-ICP [71]) and propose an occupancy-aware approach to fuse them in the feature space.
Given a set of incomplete shapes of the same object, denoted as  $\{c_1,\dots,c_M\}$ , we individually feed them into the control branch to produce feature volumes  $\{F_c^i (f_1),\dots,F_c^i (f_M)\}$  after the  $i$ -th encoder block. Before feeding them into the decoder  $D_{x}^{i}$ , we refer to their occupancy masks to perform a weighted feature average. Concretely, we first compute the original occupancy mask for each partial shape based on the TSDF values, i.e.,  $w_{j} = (abs(c_{j}) < \tau)\in \mathbf{B}^{S\times S\times S}$  for the  $j$ -th partial shape.  $\tau$  is a pre-defined threshold that assigns the volumes near the object surface as occupied and the rest as unoccupied;  $\mathbf{B}$  is a binary tensor. Then, we perform a pooling operation to resize  $w_{j}$  into  $w_{j}^{i}$  to align the resolution with feature volume  $F_{c}^{i}(f_{j})$  and normalize  $w_{j}^{i}$  by  $w_{j}^{i} = w_{j}^{i} / (w_{0}^{i} + . + w_{M}^{i})$ . For the voxels with zero values across all occupancy masks, we uniformly assign them a 1e-2 value to avoid the division-by-zero issue. We rewrite Eq. (4), which is for single partial shape condition, as
$$
d ^ {i} = \left[ D _ {x} ^ {i - 1} \left(x _ {t}\right), F _ {x} ^ {i} \left(x _ {t}\right) + \psi \left(\sum_ {j} w _ {j} ^ {i} \phi^ {i} \left(F _ {c} ^ {i} \left(f _ {j}\right)\right)\right) \right], \quad f _ {j} = \varepsilon_ {x} \left(x _ {t}\right) + \varepsilon_ {c} \left(c _ {j}\right) \tag {5}
$$