ToNNO: Tomographic Reconstruction of a Neural Network's Output for Weakly Supervised Segmentation of 3D Medical Images

TL;DR

ToNNO introduces a tomographic reconstruction approach for weakly supervised 3D medical image segmentation, leveraging 2D encoders and inverse Radon transform to produce dense 3D predictions from image-level labels.

Contribution

The paper presents a novel tomographic reconstruction method, ToNNO, enabling dense 3D segmentation using 2D encoders and inverse Radon transform, outperforming existing CAM-based methods.

Findings

Outperforms 2D CAM methods on four datasets

Combining tomographic reconstruction with CAM improves results

Applicable to various 2D encoders for 3D segmentation

Abstract

Annotating lots of 3D medical images for training segmentation models is time-consuming. The goal of weakly supervised semantic segmentation is to train segmentation models without using any ground truth segmentation masks. Our work addresses the case where only image-level categorical labels, indicating the presence or absence of a particular region of interest (such as tumours or lesions), are available. Most existing methods rely on class activation mapping (CAM). We propose a novel approach, ToNNO, which is based on the Tomographic reconstruction of a Neural Network's Output. Our technique extracts stacks of slices with different angles from the input 3D volume, feeds these slices to a 2D encoder, and applies the inverse Radon transform in order to reconstruct a 3D heatmap of the encoder's predictions. This generic method allows to perform dense prediction tasks on 3D volumes using any 2D image encoder. We apply it to weakly supervised medical image segmentation by training the 2D encoder to output high values for slices containing the regions of interest. We test it on four large scale medical image datasets and outperform 2D CAM methods. We then extend ToNNO by combining tomographic reconstruction with CAM methods, proposing Averaged CAM and Tomographic CAM, which obtain even better results.