Uncertainty-Aware ControlNet: Bridging Domain Gaps with Synthetic Image Generation

AI summary

Overview

• Research area: Computer vision / generative modeling — controlled diffusion-based image synthesis (ControlNet), uncertainty estimation, and domain adaptation for semantic segmentation, with a primary focus on medical imaging (retinal OCT).
• Technical level: Advanced. It assumes familiarity with latent diffusion models, ControlNet conditioning, epistemic/entropy-based uncertainty, and semantic segmentation metrics (mIoU, FID).
• Scope: The paper proposes UnIA CorN (an uncertainty-aware ControlNet) that trains a second ControlNet on the uncertainty of a frozen segmentation model over an unlabeled domain, then fuses it with a semantic ControlNet at inference to synthesize labeled images from unseen domains.

What This Paper Is About

Generative models such as ControlNet can create labeled synthetic images to augment segmentation training data, but they tend to reproduce the distribution they were trained on, so they add little value for test data that lies far from that distribution. The authors want to generate labeled images that are deliberately "relevant" to a downstream segmentation network, i.e., images that carry the high uncertainty the network exhibits on recorded-but-unlabeled data from a different domain. They do this by learning control over uncertainty separately from control over semantic labels, and combining both during sampling.

Key Contributions

1. An uncertainty-based conditioning mechanism for ControlNets that enables training on existing unlabeled datasets with large domain gaps, without requiring any semantic control information for those unlabeled images.
2. A dual-control generation scheme that combines semantic label control (trained on labeled data) with uncertainty control (trained on labeled plus unlabeled data), so that diverse labeled datasets can be synthesized — unlike a conventional ControlNet, which can only reproduce the labeled distribution.
3. Demonstration that this is not just style transfer: because control is exerted through sampled uncertainty levels rather than a learned image style, arbitrary domain shifts can be represented, as illustrated in a traffic-scene experiment as well as the retinal OCT scenario.
4. Sampling-based dataset construction that allows multiple synthetic images per input label map by sampling uncertainty values from a Gaussian, a property most style-transfer models lack (though it was not exploited in the reported comparisons, to keep them fair).

Main Findings

• Retinal OCT domain gap: In a dataset of 50 healthy subjects imaged with both a Spectralis OCT device and a prototype self-examination low-cost full-field (HOME) OCT device, training a segmentation U-Net on real Spectralis data yields a mean Jaccard index (mIoU) of 0.52, while the upper bound of training on real HOME-OCT images is 0.77.
• Uncertainty-guided data helps: Training on synthetic images from UnIA CorN alone gives mIoU 0.53; combining those synthetic images with the real Spectralis data gives 0.65, the best training scenario reported. Excluding the background labels (mean over labels 1–8), the corresponding values are 0.46 and 0.59.
• Comparison with style transfer: CycleGAN-generated data scores 0.40 mIoU alone and 0.57 with real Spectralis; CUT scores 0.51 alone and 0.59 with real Spectralis. CycleGAN is the worst performer because it lacks semantic consistency, producing realistic HOME-OCT looks that do not align with the shapes and labels of the source Spectralis images.
• Uncertainty matches the target domain: The generated UnIA CorN distribution has a mean epistemic uncertainty of 12.85, compared with 7.27 for real Spectralis data and 10.89 for real HOME-OCT data — indicating the generated labeled data resembles the unlabeled target domain in difficulty without learning its style.
• FID behaves unexpectedly: CycleGAN achieves the lowest FID against real HOME-OCT (68.9), but the authors attribute this to CycleGAN mapping large spatial shifts rather than producing label-consistent images. UnIA CorN reaches FID 132.5 against real HOME-OCT, versus 152.6 for real Spectralis; CUT gives 173.5, and the real Spectralis-to-real HOME-OCT FID is 152.6 versus 151.9 between generated data and real Spectralis.
• Traffic scenes generalize: Using Cityscapes (DCS) as the labeled distribution and ACDC (rain, snow, fog, night) as unlabeled, with uncertainty sampled from a Gaussian N(9.6, 2.4) and a generated dataset of 2975 images, training a SegFormer-based model improved mIoU by around 2% over training on Cityscapes only, and by around 1% over an equivalent ALDM-based prompted generation strategy.
• Qualitative domain effects: With texts prompts, UnIA CorN makes generated scenes visibly "harder" for the segmentation network (fog becomes foggier, night becomes darker), and its outputs are closer in style to the original ACDC data.

Methodology in Plain English

The pipeline has three stages.

Pretraining. First, a latent diffusion model (the "backbone" DDPM) is trained or reused so that it can generate images from both the labeled and the unlabeled distributions. Then two ControlNets are trained on top of it. The Semantic-ControlNet sees only labeled pairs (image, segmentation mask) — in the medical case, Spectralis OCT images with retinal layer segmentations — so it reproduces that labeled distribution. Separately, a segmentation network is trained on the labeled data only.

Building the uncertainty signal. The frozen segmentation network is applied to images from both domains. For each pixel, the authors compute the entropy of the predicted class distribution (equation 3), then average it over all pixels and normalize it to the range [0, 100] (equation 4), yielding one uncertainty value per image. That value is broadcast into a constant-valued "control image" fed to the Uncertainty-ControlNet, which is trained on the combined labeled and unlabeled datasets. Because a single scalar is used rather than the full pixelwise entropy map, the network does not learn spatial structure — the Semantic-ControlNet handles image structure, while uncertainty conveys how confident the segmentation model is about the content.

Inference and retraining. At each of 1000 diffusion steps, both ControlNets predict noise, and these predictions are fused as a weighted sum with weight α = 0.4 for the uncertainty branch, then combined with the backbone's noise prediction through a 1×1 convolution. To generate a dataset, uncertainty values are drawn from a Gaussian fitted to the unlabeled domain's uncertainty (N(10.89, 2.71) for OCT) while label maps are taken from labeled images, so the output is a labeled image with target-domain-like difficulty. Finally, the segmentation network is retrained on these synthetic images (optionally together with the real labeled data) and evaluated on real target-domain images.

Why This Matters

• Research impact: It reframes data augmentation for segmentation as uncertainty-targeted generation rather than style transfer, showing that unlabeled recorded data can be exploited for ControlNet training without manual annotation, and that generative models can be steered outside their training distribution while keeping labels valid.
• Medical imaging with home devices: Lower-quality HOME-OCT scans from self-examination can be segmented using models trained only with high-quality clinical Spectralis scans and their ground truth.
• Rare disease and clinical shift: New acquisition techniques, rare conditions, or scanner changes can be handled without a costly expert labeling round for each new device.
• Autonomous driving: Vehicle fleets collect huge unlabeled volumes across weather and lighting conditions; uncertainty-guided synthesis could convert these into labeled training material, which the authors explicitly name as an automotive use case.
• Industry relevance: The method needs no explicit knowledge of the target domain's style, which lowers the engineering effort for companies deploying perception models across heterogeneous sensors, and the code is released at https://github.com/JNiemeijer/UnAICorN.git.

Future Directions

• Augmented generation per label map: Sampling multiple uncertainties per label map was deliberately not used in the baseline comparison to avoid an unfair advantage; the authors plan to explore this augmented generation in future work, noting its value for automotive fleets with large unlabeled data pools.
• Iterative refinement loops: After retraining the segmentation model, its uncertainty changes, so the generative model can be updated in turn — a circular dependency the authors flag as a door to further advancements.
• Broader and more diverse unlabeled distributions: The authors suggest intentionally presenting varied domain shifts in the unlabeled dataset to study the uncertainty-guidance property in more depth.
• Foundation-model substitution: The OCT backbone DDPM was trained from scratch because of data specificity, but the authors state this step may be omitted in the future as foundation models advance, and they plan to release the trained diffusion model.

Target Audience

Researchers and practitioners working on generative data augmentation, controlled diffusion models, domain adaptation, and uncertainty-aware segmentation — particularly those in medical image analysis dealing with device-induced domain shifts — as well as machine learning engineers in automotive or other industrial perception settings who have abundant unlabeled data and scarce labels. Readers should already be comfortable with diffusion models and segmentation evaluation to follow the method and results fully.

Authors’ abstract

Generative Models are a valuable tool for the controlled creation of high-quality image data. Controlled diffusion models like the ControlNet have allowed the creation of labeled distributions. Such synthetic datasets can augment the original training distribution when discriminative models, like semantic segmentation, are trained. However, this augmentation effect is limited since ControlNets tend to reproduce the original training distribution. This work introduces a method to utilize data from unlabeled domains to train ControlNets by introducing the concept of uncertainty into the control mechanism. The uncertainty indicates that a given image was not part of the training distribution of a downstream task, e.g., segmentation. Thus, two types of control are engaged in the final network: an uncertainty control from an unlabeled dataset and a semantic control from the labeled dataset. The resulting ControlNet allows us to create annotated data with high uncertainty from the target domain, i.e., synthetic data from the unlabeled distribution with labels. In our scenario, we consider retinal OCTs, where typically high-quality Spectralis images are available with given ground truth segmentations, enabling the training of segmentation networks. The recent development in Home-OCT devices, however, yields retinal OCTs with lower quality and a large domain shift, such that out-of-the-pocket segmentation networks cannot be applied for this type of data. Synthesizing annotated images from the Home-OCT domain using the proposed approach closes this gap and leads to significantly improved segmentation results without adding any further supervision. The advantage of uncertainty-guidance becomes obvious when compared to style transfer: it enables arbitrary domain shifts without any strict learning of an image style. This is also demonstrated in a traffic scene experiment.