# Estimating Bayesian Uncertainty Using 3D Pre-trained TotalSegmentator Model

Cédric Hémon, Valentin Boussot, Jean-Louis Dillenseger, Jean-Claude Nunes

► To cite this version:

Cédric Hémon, Valentin Boussot, Jean-Louis Dillenseger, Jean-Claude Nunes. Estimating Bayesian Uncertainty Using 3D Pre-trained TotalSegmentator Model. CURVAS2025 Challenge, Modélisation bayésienne pour la segmentation du PDAC, MICCAI 2025, Sep 2025, Daejeon, South Korea. hal-05334639

HAL Id: hal-05334639

https://hal.science/hal-05334639v1

Submitted on 28 Oct 2025

HAL is a multi-disciplinary open access
archive for the deposit and dissemination of sci-
entific research documents, whether they are pub-
lished or not. The documents may come from
teaching and research institutions in France or
abroad, or from public or private research centers.

L'archive ouverte pluridisciplinaire HAL, est
destinée au dépôt et à la diffusion de documents
scientifiques de niveau recherche, publiés ou non,
émanant des établissements d'enseignement et de
recherche français ou étrangers, des laboratoires
publics ou privés.

<figure><img src="image_1.png" /></figure>

Distributed under a Creative Commons Attribution 4.0 International License
Estimating Bayesian Uncertainty Using 3D
Pre-trained TotalSegmentator Model *

Cédric Hémon¹[0009-0003-6669-5108], Valentin Boussot¹[0009-0003-2465-5458],
Jean-Louis Dillenseger¹, and Jean-Claude Nunes¹

Univ.Rennes, CLCC Eugène Marquis, INSERM, LTSI - UMR 1099, F-35000 Rennes,
France

https://ltsi.univ-rennes.fr/

**Abstract.** This paper presents a method for estimating Bayesian un-
certainty of medical image segmentation using pre-trained TotalSegmen-
tator models. By integrating the pre-trained nnU-Net M291 model and
employing the Adaptable Bayesian Neural Network (ABNN) strategy,
this approach addresses key challenges in uncertainty estimation for seg-
mentation. In contrast to conventional approaches, this method demands
only minimal fine-tuning of the training dataset, utilizes a very limited
number of trainable parameters without altering the pre-trained model
weights, and requires only a few training epochs. The results demon-
strate that this method provides a robust estimation of segmentation
uncertainty, effectively mitigating overfitting and enhancing generaliza-
tion.

**Keywords:** Bayesian Neural Network · Uncertainty · Segmentation

# 1 Method

Bayesian models are recognized for their ability to achieve both precise segmen-
tations and robust uncertainty quantification [3]. However, their application in
this context is particularly challenging due to the limited training data available
for the CURVAS challenge (only 20 CT images), each annotated by three inde-
pendent physicians to delineate 3 anatomical structures: the pancreas, kidneys,
and liver. This small dataset often leads to poor generalisation of the model and
low robustness of a 3D segmentation network created from scratch, particularly
in accurately quantifying uncertainty on new datasets. In the current landscape
of computer vision, with increasingly larger datasets, number of model param-
eters, and inference constraints, pre-trained models are becoming more popular
[4]. These models require post-hoc uncertainty estimation methods, which allow
for both leveraging their precision due to pre-training and to quantify the un-
certainty. We opted for this approach by selecting a pre-trained segmentation
network (TotalSegmentator [5]) and adding to it the ability to quantify these un-
certainties using the strategy proposed by Adaptable Bayesian Neural Network
(ABNN) [2].

* Supported by CominLabs CEMMTAUR 2022, by a PhD scholarship Grant from Elekta AB and by ANR as part of the VATSop project
## 1.1 From deterministic TotalSegmentator to Bayesian model

TotalSegmentator [5] is a powerful tool designed for the robust segmentation of numerous anatomical structures on CT scans. It is based on five nnU-Net models pre-trained on a large CT dataset. In particular, the M291 UNet model enables precise delineation of the three anatomical structures of interest (pancreas, kidneys, and liver) and 22 other labels. Our main challenge is to transform this deterministic network into a Bayesian one in order to quantify uncertainty from this segmentation model [5]. The choice of ABNN is motivated by the desire to avoid the complex and unstable optimization schemes typical of Bayesian methods, and to avoid overfitting on the dataset leading to a loss of generalization. In contrast to conventional approaches, this method demands only minimal fine-tuning of the training dataset, uses a limited number of trainable parameters without altering the initial model weights, and requires just a few training epochs. This allows the replacement of the UNet network's layer normalization with Bayesian Normalization Layers (BNL) [1]:

$$ \mathrm{BNL}(h_j) = \frac{h_j - \hat{\mu}_j}{\hat{\sigma}_j} \times \gamma_j(1 + \epsilon_j) + \beta_j. \quad (1) $$

BNL($h_j$) represents an affine function with respect to the layer normalization (LN) of the layer $j$. Our normalization layers consist of trainable parameters $\beta_j$ and $\gamma_j$ of the layer $j$. The parameters $\hat{\mu}_j$ and $\hat{\sigma}_j$ denote the empirical mean and variance, respectively, of the layer $h_j$. The term $\epsilon \sim N(0, 1)$ refers to a reduced centered Gaussian perturbation applied to these parameters. These Gaussian perturbations are applied to the normalization layer weights before training begins, effectively introducing randomness that allows the model to approximate the posterior distribution of the weights. This step is essential for enabling the Bayesian approach and for capturing the uncertainty in the network's predictions [2].

In this method, uncertainty is introduced by sampling the $\epsilon$ from a reduced centered Gaussian distribution and using the different $\gamma$ and $\beta$ values. These values are not predetermined; instead, they are empirically estimated through fine-tuning on the training dataset. In section 1.2, we describe the process for determining the parameters of the distributions, and in section 1.3, we explain how sampling from these distributions generates our uncertainty maps and the mean segmentation.

## 1.2 Fine-tuning procedure

To estimate the $\gamma$ and $\beta$ values in each Bayesian normalization layer, we retrained the model 20 times over 10 epochs, keeping all model weights fixed, except for the $\gamma$ and $\beta$ parameters. The empirical values of these parameters were then derived from the final 20 configurations. This process provides a comprehensive estimate of the uncertainty associated with the network's predictions [2].

To achieve a reliable estimation of the distributions, several factors are crucial, starting with the diversity of the datasets used for each re-training. To
ensure diversity and robustness in the uncertainty estimation, we used a struc-
tured strategy based on five-fold cross-validation. Specifically, five different data
splits were created, and for each split, the model was retrained using four distinct
annotation configurations: three runs using a single physician annotation (one
per annotator), and one run using all three annotations combined. Together,
these 20 training runs enabled a reliable empirical estimation of the posterior
distributions of the Bayesian normalization parameters, enhancing the model’s
capacity to express prediction uncertainty [2]. The network is retrained while
adhering to the inference constraints and optimization scheme (loss, optimizer,
preprocessing ...) of the original TotalSegmentator [5]. An exponentially decay-
ing scheduler was employed during fine-tuning to minimize the loss as effectively
as possible. Additionally, the weights of the classes of interest in the segmen-
tation loss vary uniformly between [0.5, 1] at each epoch in order to introduce
more diversity into the uncertainty estimation, while the weights for the other
22 labels are fixed at 0.2.

TotalSegmentator includes the segmentation of 22 additional classes beyond
those specified in the challenge, and these have been incorporated into the chal-
lenge annotations for training. This expansion introduces additional constraints
on the segmentation process, as the model must now consider and accurately ac-
count for these extended anatomical features throughout the rest of the patient’s
anatomy. The inclusion of these extra classes allows for a more comprehensive
segmentation and extends the model’s capabilities beyond the original scope of
this segmentation challenge.

1.3 Inference

The incorporation of prior knowledge about the weights completes the likeli-
hood and leads to the search for the maximum a posteriori (MAP). Specifically,
a regularization term derived from a Gaussian prior is added to the loss func-
tion, where the prior is modeled as a normal distribution with mean zero and
a predefined standard deviation. This encourages the weights to remain close
to zero, effectively acting as L2 regularization. The prediction y for a new sam-
ple x is ideally computed as the expected outcome from an infinite ensemble,
encompassing all possible weights sampled from the posterior distribution:

$$
P(y | \mathbf{x}, \mathcal{D}) = \int_{\omega \in \Omega} P(y | \mathbf{x}, \omega) P(\omega | \mathcal{D}) d\omega. \quad (2)
$$

where $P(y | \mathbf{x}, \omega)$ represents the network's prediction given the input $\mathbf{x}$ and
weight configuration $\omega$, and $P(\omega | \mathcal{D})$ is the posterior distribution over the
weights given the training dataset $\mathcal{D}$.

However, due to the computational challenges of calculating this exact integral,
it is approximated using Monte Carlo sampling methods. This involves calculat-
ing the average of the predictions after using the argmax function from a finite
set of *N* weight configurations sampled from the posterior distribution. These
configurations are derived from the estimated parameters *γ* and *μ* in the nor-
malization layers and the Gaussian perturbation $\epsilon$.

During the inference process, the standard TotalSegmentator model [5] is employed to predict the average segmentation of the CT images. Simultaneously, our Bayesian TotalSegmentator model is used to estimate the uncertainty by generating 10 different predictions. The CT intensities are adjusted using the normalization function specific to TotalSegmentator [5]. The CT images undergo a reshaping process that includes resampling to a consistent voxel size, applying a canonical function to standardize image orientation, and padding to ensure uniformity across all slices. Subsequently, the CT slices are segmented into distinct 3D patches of size 128×128×128, with an overlap of 10 voxels for generating mean predictions and 20 voxels for uncertainty predictions.

For the final 3D segmentation, the results from the individual 3D patches are combined using a cosine function applied after the argmax operation. This method helps in smooth integration of the segmentation patches by considering the spatial coherence of the predictions. The cosine function effectively combines the patch-wise predictions to produce a coherent final segmentation map. This method promotes smooth spatial transitions between neighboring patches and avoids abrupt edges that can arise from hard decisions.

The uncertainty map is derived by averaging the segmentation score maps from the 10 predictions for each region of interest. This results in an uncertainty value between 0 and 1, with increments of 0.1, indicating the confidence level in the segmentation predictions.

Post-processing is performed to retain only the largest connected volume in both the mean and uncertainty maps for each label. This step ensures that the final segmentation result is not fragmented and is represented by the most significant anatomical structures. After completing these steps, all preprocessing operations, such as resampling and padding, are reversed to return the CT images and segmentation maps to their original format for consistent evaluation and interpretation.

## Implementation and Code Availability

All experiments were implemented using our framework *KonfAI* [1], an in-house deep learning framework built on top of PyTorch. KonfAI provides a modular and fully configurable training and inference pipeline, driven by YAML-based experiment management. The code and trained models are openly available at https://github.com/vboussot/Curvas.

## References

1. Boussot, V., Dillenseger, J.: Konfai: A modular and fully configurable framework for deep learning in medical imaging. arXiv preprint arXiv:2508.09823 (2025), https://arxiv.org/abs/2508.09823
2. Franchi, G., Laurent, O., Leguéry, M., Bursuc, A., Pilzer, A., Yao, A.: Make Me a BNN: A Simple Strategy for Estimating Bayesian Uncertainty from Pre-trained Models (Dec 2023). https://doi.org/10.48550/arXiv.2312.15297, http://arxiv.org/abs/2312.15297, arXiv:2312.15297 [cs, stat]

3. Hiasa Y., Otake, Y., Takao, M., Ogawa, T., Sugano, N., Sato, Y.: Automated muscle segmentation from clinical ct using bayesian u-net for personalized musculoskeletal modeling. IEEE Transactions on Medical Imaging 39(4), 1030–1040 (Apr 2020). https://doi.org/10.1109/tmi.2019.2940555, http://dx.doi.org/10.1109/TMI.2019.2940555

4. Kolides, A., Nawaz, A., Rathor, A., Beeman, D., Hashmi, M., Fatima, S., Berdik, D., Al-Ayyoub, M., Jararweh, Y.: Artificial intelligence foundation and pre-trained models: Fundamentals, applications, opportunities, and social impacts. Simulation Modelling Practice and Theory 126, 102754 (Jul 2023). https://doi.org/10.1016/j.simpat.2023.102754, https://www.sciencedirect.com/science/article/pii/S1569190X2300031X

5. Wasserthal, J., Breit, H.C., Meyer, M.T., Pradella, M., Hinck, D., Sauter, A.W., Heye, T., Boll, D.T., Cyriac, J., Yang, S., Bach, M., Segeroth, M.: TotalSegmentator: Robust Segmentation of 104 Anatomic Structures in CT Images. Radiology: Artificial Intelligence 5(5), e230024 (Jul 2023). https://doi.org/10.1148/ryai.230024, https://www.ncbi.nlm.nih.gov/pmc/articles/PMC10546353/