Label conditioned segmentation

Tianyu Ma, Benjamin C. Lee, Mert R. Sabuncu

Keywords: Semantic Segmentation

Abstract:

Semantic segmentation is an important task in computer vision that is often tackled with convolutional neural networks (CNNs).
A CNN learns to produce pixel-level predictions through training on pairs of images and their corresponding ground-truth segmentation labels.
For segmentation tasks with multiple classes, the standard approach is to use a network that computes a multi-channel probabilistic segmentation map, with each channel representing one class.
In applications where the image grid size (e.g., when it is a 3D volume) and/or the number of labels is relatively large, the standard (baseline) approach can become prohibitively expensive for our computational resources.
In this paper, we propose a simple yet effective method to address this challenge.
In our approach, the segmentation network produces a single-channel output, while being conditioned on a single class label, which determines the output class of the network.
Our method, called label conditioned segmentation (LCS), can be used to segment images with a very large number of classes, which might be infeasible for the baseline approach.
We also demonstrate in the experiments that label conditioning can improve the accuracy of a given backbone architecture, likely, thanks to its parameter efficiency.
Finally, as we show in our results, an LCS model can produce previously unseen fine-grained labels during inference time, when only coarse labels were available during training. We provide all of our code here: https://github.com/tym002/Label-conditioned-segmentation

Registration: I acknowledge that publication of this at MIDL and in the proceedings requires at least one of the authors to register and present the work during the conference.

Authorship: I confirm that I am the author of this work and that it has not been submitted to another publication before.

Paper Type: methodological development

Primary Subject Area: Segmentation

Secondary Subject Area: Application: Radiology

Confidentiality And Author Instructions: I read the call for papers and author instructions. I acknowledge that exceeding the page limit and/or altering the latex template can result in desk rejection.

Code And Data:

code: https://github.com/tym002/Label-conditioned-segmentation
data: http://brain-development.org/brain-atlases/adult-brain-atlases/

Community Implementations: [ 1 code implementation](https://www.catalyzex.com/paper/arxiv:2203.10091/code)