Preserve the Hard, Regenerate the Rest: Uncertainty-Guided Synthetic Training Data Augmentation with Diffusion Models

2026-06-30

AI Summary

This work addresses the performance limitations of semantic segmentation models on sparse and visually diverse rare regions—such as small objects in aerial imagery or autonomous driving scenes—where conventional synthetic augmentation often causes pixel-label misalignment and inefficient computation. The authors propose an uncertainty-guided contextual augmentation method that identifies uncertain regions via prediction entropy and selectively applies diffusion-based inpainting only to their complementary context. During fine-tuning, loss is computed exclusively on original pixels to preserve label consistency while maximizing information gain. Notably, this approach requires no external models or heuristic rules and represents the first integration of predictive uncertainty with diffusion-based restoration for targeted enhancement of challenging regions. Experiments on Cityscapes, UAVID, and BDD100K demonstrate significant mIoU improvements, with the largest gains observed in difficult classes such as buses, trains, and aerial-view cars.

Abstract

Semantic segmentation models struggle with data sparsity and rare or visually diverse regions, e.g., dense regions or small objects in aerial or autonomous mobility data. While synthetic augmentation is an appealing solution, directly generating new labeled data risks misalignment of labels and generated pixels. Existing solutions to this problem often rely on external models, or employ coarse heuristics such as indiscriminately augmenting all foreground objects or entire backgrounds, which wastes capacity on uninformative pixels. To address this, we propose an uncertainty-guided synthetic context augmentation strategy that strictly preserves label validity and efficiently maximizes pixel informativeness per synthetic sample - no external guardrails required. Using a baseline segmenter's predictive entropy, we identify uncertain semantic regions and inpaint only the complementary visual context. When fine-tuning the segmenter on this synthetic data, we compute the loss only over the original pixels, excluding inpainted regions. This focuses learning on the unmodified, uncertain regions while presenting them in novel contexts. We demonstrate substantial mIoU gains on Cityscapes, UAVID, and BDD100K with the largest gains on rare and difficult classes such as buses, trains, or (from the aerial perspective) cars. Our results demonstrate that uncertainty-guided context augmentation is a highly effective lever to improve segmentation performance on complex datasets, with code provided at https://github.com/XITASO/Preserve-the-Hard-Regenerate-the-Rest.

Problem

• semantic segmentation
• data sparsity
• synthetic data augmentation
• label misalignment
• rare classes

Innovation

• uncertainty-guided augmentation
• diffusion models
• semantic segmentation
• synthetic data
• context inpainting