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

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.

• The paper introduces an uncertainty-based method that preserves high-uncertainty regions while regenerating context with a diffusion inpainting model.
• It employs an iterative, active learning-like process that significantly improves mIoU scores, especially for rare classes like 'bus' and 'moving_car'.
• The approach ensures strict label validity through a pixel-exact paste-back mechanism, outperforming traditional spatial heuristic augmentation techniques.

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

Motivation and Problem Setting

Semantic segmentation for complex real-world imagery is fundamentally constrained by sparse and imbalanced annotations, with rare classes and small objects receiving little supervision. Prior approaches leveraging synthetic augmentation via diffusion models attempt to inject diversity but suffer from label-content misalignment or reliance on coarse, spatial heuristics (foreground vs. background regeneration). This work critically examines why such spatial strategies are insufficient and proposes an uncertainty-guided augmentation method that explicitly maximizes the informativeness of synthetic samples while strictly preserving label validity.

Methodological Contributions

The method is characterized by a pipeline that identifies per-sample, high-uncertainty regions of an image—as determined by a baseline segmenter’s per-pixel predictive entropy—and strictly preserves these regions. The complementary context is regenerated via a black-box diffusion inpainting model (SDXL-Inpaint-1.0), ensuring that only informative, challenging pixels are retained and labeled, while context diversity is synthetically injected around them. Fine-tuning excludes synthetic pixels from the loss, eliminating label noise by construction.

Figure 1: Uncertainty-based region selection transcends the foreground-background contradiction, demonstrating augmentation should be based on predictive entropy rather than spatial heuristics.

Class-level predictive entropy is computed for each image, and the most uncertain classes are selected until a target preserve-area threshold $\tau$ is satisfied. This yields a binary mask marking hard regions for preservation. Inpainting, performed only on the mask complement, supplies novel visual context.

To enforce strict label preservation post-inpaint, a pixel-exact paste-back restores the original content within the preserve mask, mitigating reconstruction drift and guaranteeing correspondence between preserved pixels and label targets.

Figure 2: Pixel-exact paste-back mechanism ensures hard regions retain exact original pixel values, strictly maintaining label validity.

The method is naturally iterative: entropy and preserve regions are recomputed using the improved segmenter after each augmentation round, generating new synthetic samples tailored to the segmenter's current failure modes, akin to active learning cycles.

Empirical Results

Experiments on Cityscapes, UAVID, and BDD100K systematically demonstrate superior mIoU over state-of-the-art instance-level and background-based augmentation baselines. Gains concentrate on rare or difficult classes, with substantial improvement on underrepresented categories (e.g., +11.94 IoU for "bus" on Cityscapes, +8.75 IoU for "moving_car" on UAVID). Iterative refinement compounds class-specific gains, evidenced by the "train" class on Cityscapes, where accuracy improves from +1.1 to +9.2 IoU after multiple iterations.

Figure 3: Example from Cityscapes validation illustrating the preserved mask, entropy map, and the resulting uncertainty-guided augmented sample.

Qualitative analysis confirms that uncertainties align with genuinely challenging regions (e.g., ambiguous boundaries, small objects), and regenerated contexts show consistent diversity in lighting, season, and scene geometry without label drift.

Figure 4: Additional uncurated qualitative samples on Cityscapes demonstrate robustness of uncertainty-guided augmentation across various real scenarios.

Figure 5: Additional uncurated qualitative samples on UAVID validate method reliability in aerial urban scenes.

Ablation studies isolate the informational axis (uncertainty of regions) as the major driver of improvement, substantially outperforming random or spatial heuristics. The pixel-exact paste-back and synthetic-region ignore masking both contribute to abating label noise—without either, numerical results degrade.

The approach generalizes across segmentation architectures (DINOv2-ViT, SegFormer) and inpainting models (SDXL, FLUX).

Theoretical and Practical Implications

This work establishes the model’s predictive uncertainty, rather than spatial semantics, as the optimum axis for allocating synthetic-data augmentation. By focusing training on the most informative, hard regions in novel contexts, while rigorously avoiding label-corrupting synthetic targets, the method realizes significant mIoU gains with tightly controlled compute overhead (dominated by diffusion inpainting). The scheme is model-agnostic, amenable to any off-the-shelf segmentation backbone or high-fidelity diffusion inpainting model.

Crucially, this moves beyond the field's dominant focus on spatial heuristics, opening paths for principled, entropy-driven allocation of augmentation resources—analogous to automated, model-driven active learning cycles but without human annotation. The methodology is especially pertinent for safety-critical applications (autonomous driving, aerial mapping), where rare class accuracy and label correctness are paramount.

Future Directions

Immediate extensions include the incorporation of domain-adaptive or conditional inpainting models to further improve synthetic context in specialized domains (e.g., medical imagery), and exploration of more granular region selection strategies (e.g., uncertainty clustering) for finer-grained control. The active-learning analogy suggests potential for budget-aware or cost-sensitive augmentation cycles.

As diffusion models continue to scale and improve in inpainting fidelity, methodologies that tightly couple model-driven selection and label-preserving augmentation—as pioneered here—will become central to leveraging synthetic context in limited-annotation regimes.

Conclusion

Uncertainty-guided context augmentation, via iterative, entropy-based region selection and diffusion inpainting with label-preserving constraints, outperforms spatially-driven augmentation baselines in semantic segmentation. Gains are largest for rare classes, and the entire pipeline is model-, backbone-, and inpainting-agnostic with minimal added computational cost relative to off-the-shelf diffusion pipelines. The results advocate for an informational, rather than spatial, allocation of synthetic augmentation budgets in data-sparse settings, providing a robust methodology for enhancing segmentation in challenging domains (2606.31603).