[Paper Review] Learning to Rearrange Voxels in Binary Segmentation Masks for Smooth Manifold Triangulation

Jianning Li, Antonio Pepe+3 more|arXiv (Cornell University)|Aug 11, 2021

Anatomy and Medical Technology4 citations

TL;DR

This paper proposes a coarse-to-fine deep learning framework for high-resolution skull reconstruction from low-resolution inputs, using learned voxel rearrangement and hierarchical image synthesis to restore smooth, manifold-compatible surfaces. By leveraging binarized feature vectors and a hash table-based nearest neighbor search, the method achieves high-quality 3D-printable triangulations without direct high-memory processing of full-resolution volumes.

ABSTRACT

Medical images, especially volumetric images, are of high resolution and often exceed the capacity of standard desktop GPUs. As a result, most deep learning-based medical image analysis tasks require the input images to be downsampled, often substantially, before these can be fed to a neural network. However, downsampling can lead to a loss of image quality, which is undesirable especially in reconstruction tasks, where the fine geometric details need to be preserved. In this paper, we propose that high-resolution images can be reconstructed in a coarse-to-fine fashion, where a deep learning algorithm is only responsible for generating a coarse representation of the image, which consumes moderate GPU memory. For producing the high-resolution outcome, we propose two novel methods: learned voxel rearrangement of the coarse output and hierarchical image synthesis. Compared to the coarse output, the high-resolution counterpart allows for smooth surface triangulation, which can be 3D-printed in the highest possible quality. Experiments of this paper are carried out on the dataset of AutoImplant 2021 (https://autoimplant2021.grand-challenge.org/), a MICCAI challenge on cranial implant design. The dataset contains high-resolution skulls that can be viewed as 2D manifolds embedded in a 3D space. Codes associated with this study can be accessed at https://github.com/Jianningli/voxel_rearrangement.

Motivation & Objective

• Address the challenge of high memory consumption in 3D medical image processing, especially for high-resolution volumetric data like skulls.
• Preserve fine geometric details lost during standard downsampling for deep learning in medical image reconstruction.
• Enable high-quality 3D printing of reconstructed skulls by producing smooth, manifold-compatible triangulations.
• Develop memory-efficient methods to reconstruct high-resolution images indirectly via coarse representations and post-processing.
• Explore generalizable techniques applicable beyond skull reconstruction to other sparse, manifold-like medical image tasks.

Proposed method

• Train a deep neural network on downsampled (e.g., 128³) binary segmentation masks to generate a coarse 3D skull representation.
• Apply learned voxel rearrangement using a hash table to map coarse voxels to high-resolution positions via approximate nearest neighbor search on binarized feature vectors.
• Construct a pyramid of template images at multiple resolutions to guide hierarchical image synthesis for high-resolution output.
• Use sparse convolutional neural networks (Submanifold Sparse CNN) to exploit spatial sparsity and binariness of skull data, reducing FLOPs and memory usage.
• Implement a binary feature encoding and Hamming distance-based search to accelerate nearest neighbor lookup in the hash table, enabling fast, memory-efficient voxel mapping.
• Perform image synthesis using the nearest neighbor search results to reconstruct high-resolution voxels, with both hash table and kd-tree-based methods evaluated for accuracy and efficiency.

Experimental results

Research questions

• RQ1 Can a coarse-to-fine deep learning pipeline reconstruct high-resolution, smooth, and manifold-compatible skull surfaces without direct processing of full-resolution data?
• RQ2 How effective is learned voxel rearrangement via hash table-based nearest neighbor search in restoring surface smoothness and geometric fidelity?
• RQ3 What is the trade-off between reconstruction quality and computational efficiency when using approximate (hash table) versus exact (kd-tree) nearest neighbor search?
• RQ4 To what extent can binarized feature vectors and sparse representations reduce memory and FLOP requirements in 3D medical image reconstruction?
• RQ5 Can the proposed methods generalize to other medical imaging tasks involving sparse, manifold-like structures?

Key findings

• The learned voxel rearrangement method achieved high-quality, smooth, and manifold-compatible triangulations suitable for 3D printing, outperforming the coarse input in geometric fidelity.
• The hash table-based nearest neighbor search reduced CPU memory consumption to 12% compared to 59% for the kd-tree-based method, while maintaining acceptable reconstruction quality.
• The kd-tree-based method produced slightly higher-quality reconstructions due to exact nearest neighbor search, but at the cost of significantly higher memory usage.
• Both the voxel rearrangement and image synthesis methods completed the high-resolution reconstruction in under three minutes per level on CPUs, indicating feasibility for practical use.
• The method demonstrated that binarized feature vectors and sparse representations enable efficient, memory-saving processing of high-resolution medical images.
• The framework is generalizable and applicable to other medical image segmentation and reconstruction tasks involving spatially sparse, manifold-like structures.