Diffusion Models, in combination with so-called ControlNets, can generate labeled synthetic training data. To establish control over the segmentation label, the Control-Net is trained on the labeled data, which is also used for training the segmentation network. The resulting synthetic data, therefore, is drawn from the same distribution as the segmentation network’s training data and does not provide any novel information beyond it. This work aims to guide the generation process towards producing novel, meaningful training samples. We present UnIACorN, an approach that models the uncertainties of a given segmentation model. UnIACorN [1] is based on a Multi-ControlNet architecture that fuses the noise predictions of a semantic Control-Net and an uncertainty ControlNet. The semantic ControlNet is conditioned on the segmentation mask. Additionally, we introduce a novel uncertainty-guided Control-Net that generates images based on epistemic uncertainty. To train the uncertainty ControlNet, we first compute the segmentation network’s epistemic uncertainty on images, both from the labeled and unlabeled distributions. The epistemic uncertainty is high if the given image information was not part of the training distribution. We, therefore, obtain the relationship between uncertainty and image information. The uncertainty ControlNet is trained to predict the image given the corresponding epistemic uncertainty. During the diffusion process, the noise predictions of the Semantic ControlNet and the Uncertainty ControlNet are fused. Therefore, the stepwise generation process is guided by both conditions. We sample the uncertainty value from the Gaussian distribution of uncertainties obtained on the unlabeled set. The segmentation label is retained from the original labeled distribution. The segmentation network trained on a combination of synthetic and real data exhibits improved generalization to novel OCT image distributions. UnIACorN enables the segmentation network to learn from its own uncertainty.

错误:搜索内容不能为空,请输入英文关键词
错误:关键词超出字数限制,请精简
高级检索

Abstract: Uncertainty-aware ControlNet Bridging Domain Gaps with Synthetic Image Generation

  • Joshua Niemeijer,
  • Jan Ehrhardt,
  • Heinz Handels,
  • Hristina Uzunova

摘要

Diffusion Models, in combination with so-called ControlNets, can generate labeled synthetic training data. To establish control over the segmentation label, the Control-Net is trained on the labeled data, which is also used for training the segmentation network. The resulting synthetic data, therefore, is drawn from the same distribution as the segmentation network’s training data and does not provide any novel information beyond it. This work aims to guide the generation process towards producing novel, meaningful training samples. We present UnIACorN, an approach that models the uncertainties of a given segmentation model. UnIACorN [1] is based on a Multi-ControlNet architecture that fuses the noise predictions of a semantic Control-Net and an uncertainty ControlNet. The semantic ControlNet is conditioned on the segmentation mask. Additionally, we introduce a novel uncertainty-guided Control-Net that generates images based on epistemic uncertainty. To train the uncertainty ControlNet, we first compute the segmentation network’s epistemic uncertainty on images, both from the labeled and unlabeled distributions. The epistemic uncertainty is high if the given image information was not part of the training distribution. We, therefore, obtain the relationship between uncertainty and image information. The uncertainty ControlNet is trained to predict the image given the corresponding epistemic uncertainty. During the diffusion process, the noise predictions of the Semantic ControlNet and the Uncertainty ControlNet are fused. Therefore, the stepwise generation process is guided by both conditions. We sample the uncertainty value from the Gaussian distribution of uncertainties obtained on the unlabeled set. The segmentation label is retained from the original labeled distribution. The segmentation network trained on a combination of synthetic and real data exhibits improved generalization to novel OCT image distributions. UnIACorN enables the segmentation network to learn from its own uncertainty.