SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Wesley Khademi
khademiw@oregonstate.edu
Li Fuxin
lif@oregonstate.edu
Shape completion aims to recover the full 3D geometry of an object from a partial observation. This problem is inherently multi-modal since there can be many ways to plausibly complete the missing regions of a shape. Such diversity would be indicative of the underlying uncertainty of the shape and could be preferable for downstream tasks such as planning. In this paper, we propose a novel conditional generative adversarial network that can produce many diverse plausible completions of a partially observed point cloud. To enable our network to produce multiple completions for the same partial input, we introduce stochasticity into our network via style modulation. By extracting style codes from complete shapes during training, and learning a distribution over them, our style codes can explicitly carry shape category information leading to better completions. We further introduce diversity penalties and discriminators at multiple scales to prevent conditional mode collapse and to train without the need for multiple ground truth completions for each partial input. Evaluations across several synthetic and real datasets demonstrate that our method achieves significant improvements in respecting the partial observations while obtaining greater diversity in completions.
Figure 1: Given a partially observed point cloud (gray), our method is capable of producing many plausible completions (blue) of the missing regions.
With the rapid advancements in 3D sensing technologies, point clouds have emerged as a popular representation for capturing the geometry of real-world objects and scenes. Point clouds come from sensors such as LiDAR and depth cameras, and can find applications in various domains such as robotics, computer-aided design, augmented reality, and autonomous driving. However, the 3D geometry produced by such sensors is typically sparse, noisy, and incomplete, which hinders their effective utilization in many downstream tasks.
This motivates the task of 3D shape completion from a partially observed point cloud, which has seen significant research in the past few years [1, 2, 3, 4, 5, 6, 7]. Many early point cloud completion works mostly focus on generating a single completion that matches the ground truth in the training set, which does not take into account the potential uncertainty underlying the complete point cloud given the partial view. Ideally, an approach should correctly characterize such uncertainty – generating mostly similar completions when most of the object has been observed, and less similar completions when less of the object has been observed. A good characterization of uncertainty would be informative for downstream tasks such as planning or active perception to aim to reduce such uncertainty.
The task of completing a 3D point cloud with the shape uncertainty in mind is called multi-modal shape completion [8, 9], which aims to generate diverse point cloud completions (not to be confused with multi-modality of the input, e.g. text+image). A basic idea is to utilize the diversity coming from generative models such as generative adversarial networks (GAN), where diversity, or the avoidance of mode collapse to always generate the same output, has been studied extensively. However, early works [8, 9] often obtain diversity at a cost of poor fidelity to the partial observations due to their simplistic completion formulation that decodes from a single global latent vector. Alternatively, recent diffusion-based methods [10, 11, 12] and auto-regressive methods [13, 14] have shown greater generation capability, but suffer from slow inference time.
In this paper, we propose an approach to balance the diversity of the generated completions and fidelity to the input partial points. Our first novelty comes from the introduction of a style encoder to encode the global shape information of complete objects. During training, the ground truth shapes are encoded with this style encoder so that the completions match the ground truth. However, multiple ground truth completions are not available for each partial input; therefore, only using the ground truth style is not enough to obtain diversity in completions. To overcome this, we randomly sample style codes to provide diversity in the other generated completions. Importantly, we discovered that reducing the capacity of the style encoder and adding noise to the encoded ground truth shape leads to improved diversity of the generated shapes. We believe this avoids the ground truth from encoding too much content, which may lead the model to overfit to only reconstructing ground truth shapes.
Besides the style encoder, we also take inspiration from recent work SeedFormer [7] to adopt a coarse-to-fine completion architecture. SeedFormer has shown high-quality shape completion capabilities with fast inference time, but only provides deterministic completions. In our work, we make changes to the layers of SeedFormer, making it more suitable for the multi-modal completion task. Additionally, we utilize discriminators at multiple scales, which enable training without multiple ground truth completions and significantly improves completion quality. We further introduce a multi-scale diversity penalty that operates in the feature space of our discriminators. This added regularization helps ensure different sampled style codes produce diverse completions.
With these improvements, we build a multi-modal point cloud completion algorithm that outperforms state-of-the-art in both the fidelity to the input partial point clouds as well as the diversity in the generated shapes. Our method is capable of fast inference speeds since it does not rely on any iterative procedure, making it suitable for real-time applications such as in robotics.
- We design a novel conditional GAN for the task of diverse shape completion that achieves greater diversity along with higher fidelity partial reconstruction and completion quality.
- We introduce a style-based seed generator that produces diverse coarse shape completions via style modulation, where style codes are learned from a distribution of complete shapes.
- We propose a multi-scale discriminator and diversity penalty for training our diverse shape completion framework without access to multiple ground truth completions per partial input.
2.1 3D shape generation
The goal of 3D shape generation is to learn a generative model that can capture the distribution of 3D shapes. In recent years, generative modeling frameworks have been investigated for shape generation using different 3D representations, including voxels, point clouds, meshes, and neural fields.
One of the most popular frameworks for 3D shape generation has been generative adversarial networks (GANs). Most works have explored point cloud-based GANs [15, 16, 17, 18, 19], while recent works have shown that GANs can be trained on neural representations [20, 21, 22, 23]. To avoid the training instability and mode collapse in GANs, variational autoencoders have been used for learning 3D shapes [24, 25, 26], while other works have made use of Normalizing Flows [27, 28, 29]. Recently, diverse shape generation has been demonstrated by learning denoising diffusion models on point clouds [30], latent representations [31, 32], or neural fields [33].
2.2 Point cloud completion
Point cloud completion aims to recover the missing geometry of a shape while preserving the partially observed point cloud. PCN [1] was among the earliest deep learning-based methods that worked directly on point clouds using PointNet [34] for point cloud processing. Since then, other direct point cloud completion methods [35, 36, 4, 37, 6, 38, 7] have improved completion results by using local-to-global feature extractors and by producing completions through hierarchical decoding.
Recently, SeedFormer [7] has achieved state-of-the-art performance in point cloud completion. Their Patch Seeds representation has shown to be more effective than global shape representations due to carrying learned geometric features and explicit shape structure. Furthermore, their transformer-based upsampling layers enable reasoning about spatial relationships and aggregating local information in a coarse-to-fine manner, leading to improved recovery of fine geometric structures.
GAN-based completion networks have also been studied to enable point cloud completion learning in unpaired [39, 40, 41] or unsupervised [42] settings. To enhance completion quality, some works have leveraged adversarial training alongside explicit reconstruction losses [4, 43].
2.3 Multimodal shape completion
Most point cloud completion models are deterministic despite the ill-posed nature of shape completion. To address this, Wu et al. [8] proposed a GAN framework that learns a stochastic generator, conditioned on a partial shape code and noise sample, to generate complete shape codes in the latent space of a pre-trained autoencoder. Arora et al. [9] attempt to mitigate the mode collapse present in [8] by using implicit maximum likelihood estimation. These methods can only represent coarse geometry and struggle to respect the partial input due to decoding from a global shape latent vector.
ShapeFormer [13] and AutoSDF [14] explore auto-regressive approaches for probabilistic shape completion. Both methods propose compact discrete 3D representations for shapes and learn an auto-regressive transformer to model the distribution of object completions on such representation. However, these methods have a costly sequential inference process and rely on voxelization and quantization steps, potentially resulting in a loss of geometric detail.
Zhou et al. [10] propose a conditional denoising diffusion model that directly operates on point clouds to produce diverse shape completions. Alternatively, DiffusionSDF [11] and SDFusion [12] first learn a compact latent representation of neural SDFs and then learn a diffusion model over this latent space. These methods suffer from slow inference times due to the iterative denoising procedure, while [11, 12] have an additional costly dense querying of the neural SDF for extracting a mesh.
2.4 Diversity in GANs
Addressing diversity in GANs has also been extensively studied in the image domain. In particular, style-based generators have shown impressive capability in high-quality diverse image generation where several methods have been proposed for injecting style into generated images via adaptive instance normalization [44, 45, 46, 47] or weight modulation [46, 48]. In the conditional setting, diversity has been achieved by enforcing invertibility between output and latent codes [49] or by regularizing the generator to prevent mode collapse [50, 51, 52].
In this section, we present our conditional GAN framework for diverse point cloud completion. The overall architecture of our method is shown in Figure 2.
Our generator is tasked with producing high-quality shape completions when conditioned on a partial point cloud. To accomplish this, we first introduce a new partial shape encoder which extracts features from the partial input. We then follow SeedFormer [7] and utilize a seed generator to first propose a sparse set of points that represent a coarse completion of a shape given the extracted partial features. The coarse completion is then passed through a series of upsampling layers that utilize transformers with local attention to further refine and upsample the coarse completion into a dense completion.
To obtain diversity in our completions, we propose a style-based seed generator that introduces stochasticity into our completion network at the coarsest level. Our style-based seed generator
Figure 2: Overview of our diverse shape completion framework. A partial encoder is used to extract information from a partial point cloud. During training, a style encoder extracts style codes from complete point clouds, and at inference time style codes are randomly sampled from a normal distribution. Sampled style codes are injected into the partial information to produce diverse Patch Seeds in our style-based seed generator. The generated Patch Seeds are then upsampled into a dense completion through upsampling layers. Furthermore, discriminators and diversity penalties are used at every upsampling layer to train our model.
modulates the partial shape information with style codes before producing a sparse set of candidate points, enabling diverse coarse shape completions that propagate to dense completions through upsampling layers. The style codes used in modulating partial shape information are learned from an object category's complete shapes via a style encoder. Finally, we introduce discriminators and diversity penalties at multiple scales to train our model to produce diverse high-quality completions without having access to multiple ground truth completions that correspond to a partial observation.
3.1 Partial shape encoder
The goal of our partial encoder is to extract shape information in a local-to-global fashion, extracting local information that will be needed in the decoding stage to reconstruct fine geometric structures, while capturing global information needed to make sure a globally coherent shape is generated. An overview of the architecture for our proposed partial encoder is shown in Figure 3a.
Our partial encoder takes in a partially observed point cloud  $X_P$  and first applies a MLP to obtain a set of point-wise features  $F_0$ . To extract shape information in a local-to-global fashion,  $L$  consecutive downsampling blocks are applied to obtain a set of downsampled points  $X_L$  with local features  $F_L$ . In each down sampling block, a grid down sampling operation is performed followed by a series of PointConv [53] layers for feature interpolation and aggregation. Following the downsampling blocks, a global representation of the partial shape is additionally extracted by an MLP followed by max pooling, producing partial latent vector  $f_P$ .
3.2 Style encoder
To produce multi-modal completions, we need to introduce randomness into our completion model. One approach is to draw noise from a Gaussian distribution and combine it with partial features during the decoding phase. Another option is to follow StyleGAN [46, 54] and transform noise samples through a non-linear mapping to a latent space  $\mathcal{W}$  before injecting them into the partial features. However, these methods rely on implicitly learning a connection between latent samples and shape information. Instead, we propose to learn style codes from an object category's distribution of complete shapes and sample these codes to introduce stochasticity in our completion model. Our style codes explicitly carry information about ground truth shapes, leading to higher quality and more diverse completions.
To do so, we leverage the set of complete shapes we have access to during training. We introduce an encoder  $E$  that maps a complete shape  $X \in \mathbb{R}^{N \times 3}$  to a global latent vector via a 4-layer MLP followed by max pooling. We opt for a simple architecture as we would like the encoder to capture high level information about the distribution of shapes rather than fine-grained geometric structure.
Instead of assuming we have complete shapes to extract style codes from at inference time, we learn a distribution over style codes that we can sample from. Specifically, we define our style encoder as a learned Gaussian distribution  $E_S(z|X) = \mathcal{N}(z|\mu(E(X)), \sigma(E(X)))$  by adding two fully connected layers,  $\mu$  and  $\sigma$ , to the encoder  $E$  to predict the mean and standard deviation of a Gaussian to sample a style code from. Since our aim is to learn style codes that convey information about complete shapes that is useful for generating diverse completions, we train our style encoder with guidance from our
(b)
Figure 3: (a) Architecture of our partial shape encoder. (b) Overview of our style modulator network. For each style-modulated convolution (gray box),  $w_{i}$  and  $b_{i}$  are learned weights and biases of a convolution,  $w_{i}^{\prime \prime}$  are the weights after the modulation and demodulation process, and  $A$  is a learned Affine transformation.
completion network's losses. To enable sampling during inference, we minimize the KL-divergence between  $E_S(z|X)$  and a normal distribution during training. We additionally find that adding noise to our sampled style codes during training leads to higher fidelity and more diverse completions.
3.3 Style-based seed generator
We make use of the Patch Seeds representation proposed in SeedFormer [7], which enables faithfully completing unobserved regions while preserving partially observed structures. Patch Seeds are defined as a set of seed coordinates  $S \in \mathbb{R}^{N_S \times 3}$  and seed features  $\mathcal{F} \in \mathbb{R}^{N_S \times C_S}$  produced by a seed generator. In particular, a set of upsampled features  $F_{up} \in \mathbb{R}^{N_S \times C_S}$  are generated from the partial local features  $(X_L, F_L)$  via an Upsample Transformer [7]. Seed coordinates  $S$  and features  $\mathcal{F}$  are then produced from upsampled features  $F_{up}$  concatenated with partial latent code  $f_P$  via an MLP.
However, the described seed generator is deterministic, prohibiting the ability to generate diverse Patch Seeds that can then produce diverse completions through upsampling layers. We propose to incorporate stochasticity into the seed generator by injecting stochasticity into the partial latent vector  $f_{P}$ . We introduce a style modulator network  $M(f_{P},z)$ , shown in Figure 3b, that injects a style code  $z$  into a partial latent vector  $f_{P}$  to produce a styled partial shape latent vector  $f_{C}$ . Following [54], we use weight modulation to inject style into the activation outputs of a network layer, where the demodulated weights  $w^{\prime \prime}$  used in each convolution layer are computed as:
$$
s = A (z), \operatorname {m o d}: w _ {i j k} ^ {\prime} = s _ {i} \cdot w _ {i j k}, \operatorname {d e m o d}: w _ {i j k} ^ {\prime \prime} = w _ {i j k} ^ {\prime} / \sqrt {\sum_ {i , k} w _ {i j k} ^ {\prime 2} + \epsilon} \tag {1}
$$
where  $A$  is an Affine transformation, and  $w$  is the original convolution weights with  $i,j,k$  corresponding to the input channel, output channel, and spatial footprint of the convolution, respectively.
3.4 Coarse-to-fine decoder
Our decoder operates in a coarse-to-fine fashion, which has been shown to be effective in producing shapes with fine geometric structure for the task of point cloud completion [36, 38, 7, 4]. We treat our Patch Seed coordinates  $S$  as our coarsest completion  $\mathcal{G}_0$  and progressively upsample by a factor  $r$  to produce denser completions  $\mathcal{G}_i$  ( $i = 1, \dots, 3$ ). At each upsampling stage  $i$ , a set of seed features are first interpolated from the Patch Seeds. Seed features along with the previous layer's points and features are then used by an Upsample Transformer [7] where local self-attention is performed to produce a new set of points and features upsampled by a factor  $r$ . We replace the inverse distance weighted averaging used to interpolate seed features from Patch Seeds in SeedFormer with a PointConv interpolation. Since the importance of Patch Seed information may vary across the different layers, we believe a PointConv interpolation is more appropriate than a fixed weighted interpolation as it can learn the appropriate weighting of Patch Seed neighborhoods for each upsampling layer.
Unlike the fully-connected decoders in [8, 9], the coarse-to-fine decoder used in our method can reason about the structures of local regions, allowing us to generate cleaner shape surfaces. A coarse-to-fine design provides us with the additional benefit of having discriminators and diversity penalties at multiple resolutions, which we have found to lead to better completion fidelity and diversity.
3.5 Multi-scale discriminator
During training, we employ adversarial training to assist in learning realistic completions for any partial input and style code combination. We introduce a set of discriminators  $D_{i}$  for  $i = \{0,\dots,3\}$ , to discriminate against real and fake point clouds at each output level of our generator. Each discriminator follows a PointNet-Mix architecture proposed by Wang et al. [55]. In particular, an MLP first extracts a set of point features from a shape, which are max-pooled and average-pooled to produce  $f_{max}$  and  $f_{avg}$ , respectively. The features are then concatenated to produce mix-pooled feature  $f_{mix} = [f_{max},f_{avg}]$  before being passed through a fully-connected network to produce a final score about whether the point cloud is real or fake.
We also explored more complex discriminators that made use of PointConv or attention mechanisms, but we were unable to successfully train with any such discriminator. This is in line with the findings in [55], suggesting that more powerful discriminators may not guide the learning of point cloud shape generation properly. Thus, we instead use a weaker discriminator architecture but have multiple of them that operate at different scales to discriminate shape information at various feature levels.
For training, we use the WGAN loss [56] with R1 gradient penalty [57] and average over the losses at each output level  $i$ . We let  $\hat{X}_i = \mathcal{G}_i(X_P,z)$  be the completion output at level  $i$  for partial input  $X_P$  and sampled style code  $z\sim E_S(z|X)$ , and let  $X_{i}$  be a real point cloud of same resolution. Then our discriminator loss  $\mathcal{L}_D$  and generator loss  $\mathcal{L}_G$  are defined as:
$$
\mathcal {L} _ {D} = \frac {1}{4} \sum_ {i = 0} ^ {3} \left(\mathbb {E} _ {\hat {X} \sim P (\hat {X})} \left[ D _ {i} \left(\hat {X} _ {i}\right) \right] - \mathbb {E} _ {X \sim P (X)} \left[ D _ {i} \left(X _ {i}\right) \right] + \frac {\gamma}{2} \mathbb {E} _ {X \sim P (X)} \left[ \| \nabla D _ {i} \left(X _ {i}\right) \| ^ {2} \right]\right) \tag {2}
$$
$$
\mathcal {L} _ {G} = - \frac {1}{4} \sum_ {i = 0} ^ {3} \left(\mathbb {E} _ {\hat {X} \sim P (\hat {X})} \left[ D _ {i} \left(\hat {X} _ {i}\right) \right]\right) \tag {3}
$$
where  $\gamma$  is a hyperparameter ( $\gamma = 1$  in our experiments),  $P(\hat{X})$  is the distribution of generated shapes, and  $P(X)$  is the distribution of real shapes.
3.6 Diversity regularization
Despite introducing stochasticity into the partial latent vector, it is still possible for the network to learn to ignore the style code  $z$ , leading to mode collapse to a single completion. To address this, we propose a diversity penalty that operates in the feature space of our discriminator. Our key insight is that for a discriminator to be able to properly discriminate between real and fake point clouds, its extracted features should have learned relevant structural information. Then our assumption is that if two completions are structurally different, the discriminator's global mix-pooled features should be dissimilar as well, which we try to enforce through our diversity penalty.
Specifically, at every training iteration we sample two style codes  $z_{1} \sim E_{S}(z|X_{1})$  and  $z_{2} \sim \bar{E}_{S}(z|X_{2})$  from random complete shapes  $X_{1}$  and  $X_{2}$ . For a single partial input  $X_{P}$ , we produce two different completions  $\mathcal{G}_i(X_P,z_1)$  and  $\mathcal{G}_i(X_P,z_2)$ . We treat our discriminator  $D_{i}$  as a feature extractor and extract the mixed-pooled feature for both completions at every output level  $i$ . We denote the mixed-pooled feature corresponding to a completion conditioned on style code  $z$  at output level  $i$  by  $f_{mix_i}^z$ , then minimize:
$$
\mathcal {L} _ {d i v} = \sum_ {i = 0} ^ {3} \frac {1}{\left\| f _ {m i x _ {i}} ^ {z _ {1}} - f _ {m i x _ {i}} ^ {z _ {2}} \right\| _ {1}} \tag {4}
$$