Label conditioned segmentation

TL;DR

This paper introduces label conditioned segmentation (LCS), a method that enables efficient multi-class segmentation with large label sets by conditioning a single-channel network on class labels, improving accuracy and flexibility.

Contribution

The paper proposes a novel label conditioned segmentation approach that reduces computational costs and enhances accuracy for large label sets in semantic segmentation tasks.

Findings

LCS can segment images with many classes more efficiently than baseline methods.

Label conditioning improves segmentation accuracy of backbone architectures.

LCS can generate fine-grained labels during inference from coarse training labels.

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