OOD-SEG: Exploiting out-of-distribution detection techniques for learning image segmentation from sparse multi-class positive-only annotations

Abstract

Despite significant advancements, segmentation based on deep neural networks in medical and surgical imaging faces several challenges, two of which we aim to address in this work. First, acquiring complete pixel-level segmentation labels for medical images is time-consuming and requires domain expertise. Second, typical segmentation pipelines cannot detect out-of-distribution (OOD) pixels, leaving them prone to spurious outputs during deployment. In this work, we propose a novel segmentation approach which broadly falls within the positive-unlabelled (PU) learning paradigm and exploits tools from OOD detection techniques. Our framework learns only from sparsely annotated pixels from multiple positive-only classes and does not use any annotation for the background class. These multi-class positive annotations naturally fall within the in-distribution (ID) set. Unlabelled pixels may contain positive classes but also negative ones, including what is typically referred to as \emph{background} in standard segmentation formulations. To the best of our knowledge, this work is the first to formulate multi-class segmentation with sparse positive-only annotations as a pixel-wise PU learning problem and to address it using OOD detection techniques. Here, we forgo the need for background annotation and consider these together with any other unseen classes as part of the OOD set. Our framework can integrate, at a pixel-level, any OOD detection approaches designed for classification tasks. To address the lack of existing OOD datasets and established evaluation metric for medical image segmentation, we propose a cross-validation strategy that treats held-out labelled classes as OOD. Extensive experiments on both multi-class hyperspectral and RGB surgical imaging datasets demonstrate the robustness and generalisation capability of our proposed framework.

Figures (13)

• Figure 1: Decision boundaries under different learning settings for a multi-class problem with annotations from only positive classes. Coloured circle markers ( / / ) denote positively labelled samples for each class. Black question marks (?) indicate unused unlabelled training data, while coloured question marks (?/?/?) represent used unlabelled training data. Unlabelled data could be stemming from a positive class or from the background class. Our proposed framework adapts a multi-class positive-only learning setting (as shown on the right hand side), forming distinct decision boundaries that enclose positively labelled data for each class. Unlabelled data points outside these boundaries are detected by OOD techniques and flagged as background.
• Figure 2: Overview of the proposed OOD-SEG framework. During the training stage, only annotated pixels for the multiple positive classes are used to update the model weights. We define a confidence score $\boldsymbol{S}$ to correlate probability distribution for ID classes. $\boldsymbol{S}$ can be replaced by multiple OOD detection methods (See bottom left block). At the inference stage, we compute the maximum probability of $\boldsymbol{S}$ from $c$ classes followed by thresholding from a pre-selected threshold $\tau_m$ to obtain the predicted mask.
• Figure 3: Graphical representation of the proposed OOD-focused two-level cross-validation strategy. For simplicity, only the first fold is shown in detail. In this example, the number of subject partitions (SP) and class partitions (CP) are set to $4$, resulting in a total of $16$ partitions. Subject-Class Partitions (SCP) marked in red, blue and grey respectively highlight training, testing or untouched data for a particular cross-validation fold.
• Figure 4: Graphical illustration of confusion matrix incorporating multi-class ID and OOD data. Left: with actual OOD data as negative class. Right: In the case without actual OOD data and with class $2$ considered as positive while others are negative in a one-vs-rest approach.
• Figure 5: Example confusion matrix at threshold $\tau_0=0$ (left) and $\tau_m$ (right) for the ODIN method on the Heiporspectral dataset with a specific held-out class partition. The first row and column represent the negative / outlier class.