Calibrated Adversarial Refinement for Stochastic Semantic Segmentation

TL;DR

This paper introduces a two-stage calibrated adversarial refinement method for stochastic semantic segmentation, improving the modeling of prediction uncertainties and inter-pixel correlations, achieving state-of-the-art results on ambiguous datasets.

Contribution

A novel two-stage approach combining probabilistic segmentation and adversarial refinement to produce calibrated, coherent semantic maps with accurate uncertainty estimates.

Findings

Achieved state-of-the-art results on LIDC and Cityscapes datasets.

Effectively models inter-pixel correlations for coherent predictions.

Demonstrated adaptability to other tasks with calibrated predictive distributions.

Abstract

In semantic segmentation tasks, input images can often have more than one plausible interpretation, thus allowing for multiple valid labels. To capture such ambiguities, recent work has explored the use of probabilistic networks that can learn a distribution over predictions. However, these do not necessarily represent the empirical distribution accurately. In this work, we present a strategy for learning a calibrated predictive distribution over semantic maps, where the probability associated with each prediction reflects its ground truth correctness likelihood. To this end, we propose a novel two-stage, cascaded approach for calibrated adversarial refinement: (i) a standard segmentation network is trained with categorical cross entropy to predict a pixelwise probability distribution over semantic classes and (ii) an adversarially trained stochastic network is used to model the inter-pixel correlations to refine the output of the first network into coherent samples. Importantly, to calibrate the refinement network and prevent mode collapse, the expectation of the samples in the second stage is matched to the probabilities predicted in the first. We demonstrate the versatility and robustness of the approach by achieving state-of-the-art results on the multigrader LIDC dataset and on a modified Cityscapes dataset with injected ambiguities. In addition, we show that the core design can be adapted to other tasks requiring learning a calibrated predictive distribution by experimenting on a toy regression dataset. We provide an open source implementation of our method at https://github.com/EliasKassapis/CARSSS.