Decomposing 3D Scenes into Objects via Unsupervised Volume Segmentation

1 INTRODUCTION . The ability to recognize and reason about 3D geometry is key to a wide variety of important robotics and AI tasks , such as dynamics modelling , rapid physical inference , robot grasping , or autonomous driving ( Battaglia et al. , 2013 ; Chen et al. , 2016 ; Mahler et al. , 2019 ; Driess et al. , 2020 ; Li et al. , 2020 ) . While progress has been made on integrating traditional representations of geometry with machine learning systems ( Wang et al. , 2018 ; Qi et al. , 2017 ; Nash et al. , 2020 ) , they remain difficult to work with due to either undesirable scaling behavior ( voxels , point clouds ) or discrete structures and an inability to train on observational data ( polygonal meshes ) . Recent work has revealed coordinate-based neural networks as an alternative . These representations are continuous , independent of resolution or scene dimensions , and can directly map 3D coordinates to binary occupancy indicators ( Mescheder et al. , 2019 ) , signed distances to the shape ( Park et al. , 2019 ) , or volume radiances ( Mildenhall et al. , 2020 ) . When used with an encoder network , such functions can facilitate learning low-dimensional implicit representations of geometry ( Mescheder et al. , 2019 ; Kosiorek et al. , 2021 ) . They can also render high quality images ( Martin-Brualla et al. , 2021 ; Barron et al. , 2021 ) , and synthesize novel scenes ( Guo et al. , 2020 ; Niemeyer & Geiger , 2020 ) . In this paper , we investigate a different aspect . We ask whether a latent-variable model built around implicit representations may be used to infer semantically meaningful information from a given image , without human supervision . Such representations may then be used for downstream reasoning tasks . In particular , we focus on obtaining object-based representations from multi-object scenes in an unsupervised way . Representations factored into objects are beneficial for dynamics modelling , visual question answering , and many other tasks ( Battaglia et al. , 2016 ; Santoro et al. , 2017 ; Battaglia et al. , 2018 ) . Previous works in unsupervised object-based representation learning have mostly focused on segmenting 2D images ( Eslami et al. , 2016 ; Greff et al. , 2020 ; Locatello et al. , 2020 ) . Such methods have remained limited to visually simple , mostly synthetic data , whereas more realistic scenes with complex textures and geometries have remained out of reach ( Weis et al. , 2020 ) . To move towards more complex scenarios , we present ObSuRF , a model which learns to decompose scenes consisting of multiple Objects into a Superposition of Radiance Fields . We first encode the input image with a slot-based encoder similar to ( Locatello et al. , 2020 ) , but use the resulting set of latent codes to condition continuous 3D scene functions ( Mildenhall et al. , 2020 ; Kosiorek et al. , 2021 ) instead of directly generating 2D images . For training the model in 3D , we provide three RGB-D views of each scene , and optimize the model to match the observed depths and colors . To do so , we reframe NeRF ’ s volumetric rendering as a Poisson process and derive a novel training objective , which allows for more efficient training when depth information is available as supervision . After confirming that the resulting model is capable of segmenting 2D images as well or better as previous approaches , we test it on two new 3D benchmarks : A 3D version of CLEVR ( Johnson et al. , 2017 ) featuring multiple viewpoints , camera positions , and depth information , and MultiShapeNet , a novel multiobject dataset in which the objects are shapes from the ShapeNet dataset ( Chang et al. , 2015 ) . We will publish both of these benchmarks and the code for our model with this paper . 2 RELATED WORK . Learned Implicit Representations of 3D Shapes . Early work focused on predicting the geometry of shapes using occupancy fields ( Mescheder et al. , 2019 ) or signed distance functions ( SDFs ) ( Park et al. , 2019 ) . These models were trained using ground-truth geometry , for instance by sampling points from the meshes in the ShapeNet dataset ( Chang et al. , 2015 ) . Mildenhall et al . ( 2020 ) subsequently introduced Neural Radiance Fields ( NeRF ) , which predict density and color of points in a scene . This admits rendering via differentiable ray tracing through the scene . The required supervision is thereby reduced to a small set of posed images . However , ray tracing makes this approach very computationally expensive . Therefore , in this work we opt for a pragmatic middle ground : we assume access to views of the training scenes enriched with depth information ( RGB-D ) , and show how this allows for training NeRFs without expensive raymarching . We note that unlike ground-truth meshes , such data may be obtained at scale in the real world with reasonable accuracy using either low-cost camera systems ( Horaud et al. , 2016 ) or by estimating depth from stereo vision or motion , see e.g . ( Chang & Chen , 2018 ; Teed & Deng , 2020 ) . In its original form , NeRF is fitted to a single scene , which makes it too slow for real-time applications , and also ill-suited for representation learning , as the scene is encoded in neural network weights . Recent work has used encoded representations of input scenes to condition a decoder shared between scenes ( Yu et al. , 2021a ; Trevithick & Yang , 2020 ; Kosiorek et al. , 2021 ) . We take a similar approach , but put a greater emphasis on compact and factored representations . Using depth supervision to improve NeRF training has concurrently been proposed by Deng et al . ( 2021 ) , without however avoiding ray marching and the corresponding performance costs . Unsupervised Image Segmentation . The work in this area can be broadly grouped in two categories . Patch-based approaches represent scenes as a collection of object bounding boxes ( Eslami et al. , 2016 ; Kosiorek et al. , 2018 ; Stelzner et al. , 2019 ; Crawford & Pineau , 2019 ) . This encourages spatial consistency in objects , but limits the model ’ s ability to explain complicated objects . Scenemixture models represent scenes as pixelwise mixture models , with each mixture component corresponding to a single object ( Greff et al. , 2016 ; 2017 ; van Steenkiste et al. , 2018 ; Burgess et al. , 2019 ; Greff et al. , 2020 ; Locatello et al. , 2020 ; Engelcke et al. , 2020 ; 2021 ) . Since spatial consistency is not enforced , these models sometimes segment images by color rather than by object ( Weis et al. , 2020 ) , motivating combinations of the two concepts ( Lin et al. , 2020 ) . We mitigate these issues by employing NeRF . Explicitly formulated as an MLP of Fourier-encoded spatial coordinates ( instead of a CNN , similarly to e.g . Kabra et al . ( 2021 ) ) , it is naturally biased toward spatial coherence . We note that no model in this line of research has been shown to work well on general natural images . Instead they have targeted synthetic benchmarks or constrained real world environments . Some recent models use viewpoint-dependent decoders ( Nanbo et al. , 2020 ; Chen et al. , 2021 ; Kabra et al. , 2021 ) , which allows generating images from novel viewpoints . Since there is nothing to enforce multi-view consistency , these models are unlikely to generalize to novel viewpoints . Unsupervised 3D Segmentation . Most work in this area focuses on segmenting single shapes into parts , and is trained on ground truth geometry . BAE-NET ( Chen et al. , 2019 ) reconstructs voxel inputs as a union of occupancy functions . CvxNet ( Deng et al. , 2020 ) and BSP-Net ( Chen et al. , 2020 ) represent shapes as a union of convex polytopes , with the latter capable of segmenting them into meaningful parts . Similarly , UCSG-Net ( Kania et al. , 2020 ) combines SDF primitives via boolean logic operations , yielding rigid but interpretable shape descriptions . Only a small number of publications attempt to segment full scenes into objects . Ost et al . ( 2021 ) train a model to segment a complex real-world video into objects , but doing so requires a full scene graph , including manually annotated tracking information such as object positions . In GIRAFFE ( Niemeyer & Geiger , 2020 ) object-centric representations emerge when a set of NeRFs conditioned on latent codes is trained as a GAN . However , the focus lies on image synthesis and not inference . Consequently , it is not straightforward to obtain representations for a given scene . BlockGAN ( Nguyen-Phuoc et al. , 2020 ) achieves similar results , but uses latent-space perspective projections and CNNs instead of NeRFs . RELATE ( Ehrhardt et al. , 2020 ) builds on this to allow the manipulation of scenes . Elich et al . ( 2021 ) present a model which encodes a single input image into a set of deep SDFs representing the objects in the scene . In contrast to our work , theirs requires pretraining on ground-truth shapes and operates on simpler scenes without shadows or reflections . Closest to our work is a system concurrently developed by Yu et al . ( 2021b ) , which also learns to segment scenes into a set of NeRFs using a slot based encoder . However , they train using NeRF rendering and mitigate its computational cost by adjusting the rendering resolution . Set Prediction . Objects present in a scene are interchangeable and form a set . Any attempt to assign a fixed ordering to the objects ( such as left-to-right ) will introduce discontinuities as they switch positions ( e.g . as they pass behind each other ) ( Zhang et al. , 2020 ) . It is therefore essential to use permutation equivariant prediction methods in such settings . We build on recent works which construct permutation equivariant set prediction architectures ( Locatello et al. , 2020 ; Goyal et al. , 2021 ; Kosiorek et al. , 2020 ; Stelzner et al. , 2020 ) . 3 METHODS . ObSuRF decomposes 3D scenes into multiple NeRFs , each representing a separate object . We derive a principled way of composing multiple NeRFs into a single , albeit more expensive ( by a factor of n ) , scene function . To decrease this cost , we show how to make NeRF integration cheaper when ground-truth depth is available at training time . 3.1 NEURAL RADIANCE FIELDS ( NERFS ) . NeRFs ( Mildenhall et al. , 2020 ) represent the geometry and appearance of a scene as a neural network f : ( x , d ) → ( c , σ ) mapping world coordinates x and viewing direction d to a color value c and a density value σ . To guarantee coherent geometry , the architecture is chosen such that the density σ is independent of the viewing direction d. To simplify notation , we refer to the individual outputs of f ( · ) as σ ( x ) and c ( x , d ) . NeRFs allow rendering images using classic ray marching techniques for volume rendering ( Blinn , 1982 ; Kajiya & von Herzen , 1984 ) . Specifically , the color Ĉ ( r ) corresponding to the ray r is given by Ĉ ( r ) = ∫ ∞ 0 T ( t ) σ ( r ( t ) ) c ( r ( t ) , d ) dt , ( 1 ) with transmittance T ( t ) = exp ( − ∫ t 0 σ ( r ( t′ ) ) dt′ ) . Since this procedure is differentiable , NeRFs are trained by minimizing the L2 loss between the rendered colors Ĉ ( r ) and the colors of the training image C ( r ) , i.e . LNeRF ( r ) = ‖Ĉ ( r ) −C ( r ) ‖22 . This is expensive , however , as approximating Eq . ( 1 ) requires evaluating f many times per pixel ( 256 in Mildenhall et al . ( 2020 ) ) . We will show that , if depth is available , we can instead train NeRFs with just two evaluations per pixel , reducing the required computation by a factor of 128 per iteration1 .

Summary: This paper is about unsupervised segmentation both of 3D volumes and of images. The idea is to learn a set of Neural Rendering Fields (NeRFs) which explain a given image. To reduce computationally effort, the authors assume RGB-D data and show how that reduces the computational burden by two magnitudes. The proposed algorithm can learn segmentation masks and a set of NeRFs that generate each individual object and the background.