NAISR: A 3D Neural Additive Model for Interpretable Shape Representation

`NAISR`: A 3D Neural Additive Model for Interpretable Shape Representation

Yining Jiao, Carlton Zdanski, Julia Kimbell, Andrew Prince, Cameron Worden, Samuel Kirse, Christopher Rutter, Benjamin Shields, William Dunn, Jisan Mahmud, Marc Niethammer

University of North Carolina at Chapel Hill
ICLR 2024, Spotlight Top 5%

Paper Code Colab Video

Motivation of `NAISR`

For scientific and medical discovery, it is very important to capture the dependencies of shapes on covariates, e.g., growth charts are to describe how children's height and weight develop with age. Despite the growth charts, a 3D shape representation is needed to describe how different covariates affect shapes. NAISR is such a shape representation method to extrapolate the complete, disentangled template shape space from incomplete, entangled and individual observations.

Conceptual comparison between our approach NAISR and prior state-of-the-art for shape analysis tasks. NAISR is the first implicit shape representation method to investigate an atlas-based representation of 3D shapes in a deformable, disentangleable, transferable and evolvable way.

Abstract

Deep implicit functions (DIFs) have emerged as a powerful paradigm for many computer vision tasks such as 3D shape reconstruction, generation, registration, completion, editing, and understanding. However, given a set of 3D shapes with associated covariates there is at present no shape representation method which allows to precisely represent the shapes while capturing the individual dependencies on each covariate. Such a method would be of high utility to researchers to discover knowledge hidden in a population of shapes. For scientific shape discovery, we propose a 3D Neural Additive Model for Interpretable Shape Representation (NAISR) which describes individual shapes by deforming a shape atlas in accordance to the effect of disentangled covariates. Our approach captures shape population trends and allows for patient-specific predictions through shape transfer. NAISR is the first approach to combine the benefits of deep implicit shape representations with an atlas deforming according to specified covariates. We evaluate NAISR with respect to shape reconstruction, shape disentanglement, shape evolution, and shape transfer on three datasets: 1) Starman, a simulated 2D shape dataset; 2) the ADNI hippocampus 3D shape dataset; and 3) a pediatric airway 3D shape dataset. Our experiments demonstrate that NAISR achieves excellent shape reconstruction performance while retaining interpretability.

Video

Method

Neural Additive Implicit Shape Representation. During training we learn the template $\mathcal{T}$ and the multi-layer perceptrons (MLPs) $\{g_i\}$ predicting the covariate-wise displacement fields $\{\mathbf{d}_i\}$. The displacement fields are added to obtain the overall displacement field $\mathbf{d}$ defined in the target space; $\mathbf{d}$ provides the displacement between the deformed template shape $\mathcal{T}$ and the target shape. Specifically the template shape is queried not at its original coordinates $\mathbf{p}$, but at $\tilde{\mathbf{p}}=\mathbf{p}+\mathbf{d}$ effectively spatially deforming the template. At test time we evaluate the trained MLPs for shape reconstruction, evolution, disentanglement, and shape transfer.

Disentangled Shape Evolutions

Disentangled Shape Evolution for Starman. The blue shapes are the groundtruth shapes and the red shapes are the reconstructions.

Arm Movement

Leg Movement

Disentangled Shape Evolution for Pediatric Airways.

Age: from 0 - 240 months

Weight: from 0 - 160 kgs

Sex: from F to M

Disentangled Shape Evolution for ADNI Hippocampus.

Age: from 50 - 100 yrs old

AD: from health to AD

Edu.: from 0 - 30 years

Sex

Shape Extrapolation in Template Space

Template Shape Space for Starman.

Template Shape Space for ADNI Hippocampus.

Template Shape Space for Pediatric Airways.

BibTeX

@inproceedings{
jiao2024naisr,
title={\texttt{NAISR}: A 3D Neural Additive Model for Interpretable Shape Representation},
author={Yining Jiao and Carlton Jude ZDANSKI and Julia S Kimbell and Andrew Prince and Cameron P Worden and Samuel Kirse and Christopher Rutter and Benjamin Shields and William Alexander Dunn and Jisan Mahmud and Marc Niethammer},
booktitle={The Twelfth International Conference on Learning Representations},
year={2024},
url={https://openreview.net/forum?id=wg8NPfeMF9}
}