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: The paper presents an array of techniques for solving the problem of decomposing a single image of a scene into a set of NeRFs, each of which represents an object in the scene. The decomposition requires accurate estimation of scene depth and the segmentation of the scene into objects. To achieve this, the authors propose an enhanced version of NeRF that operates on RGB-D images where its rendering only requires two evaluations per each camera ray. The authors show that the proposed method works effectively on CLEVR-3D and a custom built dataset named MultiShapeNet.