Benefits of Linear Conditioning with Metadata for Image Segmentation

TL;DR

This paper demonstrates that linear conditioning with metadata, via FiLM, improves medical image segmentation performance, robustness, and adaptability, especially with limited or missing annotations.

Contribution

It adapts FiLM for segmentation tasks, integrating metadata into neural networks to enhance accuracy and robustness in medical imaging.

Findings

Average Dice score increase of 5.1% with metadata incorporation.

FiLMed U-Net outperforms single-task U-Net with fewer labels.

Robustness to missing labels and improved multi-task adaptation.

Abstract

Medical images are often accompanied by metadata describing the image (vendor, acquisition parameters) and the patient (disease type or severity, demographics, genomics). This metadata is usually disregarded by image segmentation methods. In this work, we adapt a linear conditioning method called FiLM (Feature-wise Linear Modulation) for image segmentation tasks. This FiLM adaptation enables integrating metadata into segmentation models for better performance. We observed an average Dice score increase of 5.1% on spinal cord tumor segmentation when incorporating the tumor type with FiLM. The metadata modulates the segmentation process through low-cost affine transformations applied on feature maps which can be included in any neural network's architecture. Additionally, we assess the relevance of segmentation FiLM layers for tackling common challenges in medical imaging: multi-class training with missing segmentations, model adaptation to multiple tasks, and training with a limited or unbalanced number of annotated data. Our results demonstrated the following benefits of FiLM for segmentation: FiLMed U-Net was robust to missing labels and reached higher Dice scores with few labels (up to 16.7%) compared to single-task U-Net. The code is open-source and available at www.ivadomed.org.