Towards annotation-efficient segmentation via image-to-image translation

TL;DR

This paper introduces a semi-supervised image-to-image translation framework that leverages weak labels to improve tumor segmentation in medical images, reducing the need for extensive boundary annotations.

Contribution

It proposes a novel approach combining unpaired image translation with segmentation, enabling effective tumor detection using minimal annotations and synthetic data generation.

Findings

Achieved higher Dice scores on brain and liver datasets

Demonstrated improved performance over baseline semi-supervised methods

Validated on synthetic and real medical imaging datasets

Abstract

Often in medical imaging, it is prohibitively challenging to produce enough boundary annotations to train deep neural networks for accurate tumor segmentation. We propose the use of weak labels about whether an image presents tumor or whether it is absent to extend training over images that lack these annotations. Specifically, we propose a semi-supervised framework that employs unpaired image-to-image translation between two domains, presence vs. absence of cancer, as the unsupervised objective. We conjecture that translation helps segmentation -- both require the target to be separated from the background. We encode images into two codes: one that is common to both domains and one that is unique to the presence domain. Decoding from the common code yields healthy images; decoding with the addition of the unique code produces a residual change to this image that adds cancer. Translation proceeds from presence to absence and vice versa. In the first case, the tumor is re-added to the image and we successfully exploit the residual decoder to also perform segmentation. In the second case, unique codes are sampled, producing a distribution of possible tumors. To validate the method, we created challenging synthetic tasks and tumor segmentation datasets from public BRATS (brain, MRI) and LitS (liver, CT) datasets. We show a clear improvement (0.83 Dice on brain, 0.74 on liver) over baseline semi-supervised training with autoencoding (0.73, 0.66) and a mean teacher approach (0.75, 0.69), demonstrating the ability to generalize from smaller distributions of annotated samples.