U-SEG: Uncertainty in SEGmentation - A systematic multi-variable exploration

This repository contains the code for the U-SEG paper (CVPR Findings 2026) by Michael Smith and Frank Ferrie. It is based on the Mask2Former project from Facebook. Thus, there is a lot of inherited code not used in this project.

See installation instructions.

• Where possible, we used models provided in the original Model Zoo from Facebook.
• In some cases, when training a model we swapped out the backbone but kept the Mask2Former head for transfer learning. Those models are suffixed with backbone_stripped e.g model_final_54b88a_backbone_stripped.pkl. In those cases, the original model (e.g. model_final_54b88a.pkl) can be found from the Facebook Model Zoo and the strip-backbone-weights.py script in tools/ can be used to generate the weights file we used.
• All models trained by us are available on HuggingFace and mirrored at McGill University.

See the README.

The basic process for training and model if necessary (as we do for datasets such as VIPER) and running inference follows the standard Mask2Former format, documented on their page.

All of the predictions made by Mask2Former models are done through configs in the configs directory. They are in one of four directories:

• cityscapes/panoptic-segmentation
• cityscapes_VIPERformat/panoptic-segmentation
• viper/panoptic-segmentation
• VIPER_cityscapesformat/panoptic-segmentation

All others are inherited from Mask2Former.

To generate the results used in the paper, the general procedure is:

1. Download the datasets and set them up using the scripts in the datasets folder and listed in the datasets README.

2. Acquire the weights or train a model. For example, with Cityscapes, we used the model trained by the Mask2Former authors, but with VIPER we needed to train a model ourselves. This was done via the following command:

   python train_net.py --num-gpus="${SLURM_GPUS_ON_NODE}" --config-file "${MASK2FORMER_CONFIG_FILE}" OUTPUT_DIR ${OUTPUT_DIR}

   An example config file for training is configs/viper/panoptic-segmentation/maskformer2_R50_bs1_300ep.yaml. Another example is configs/viper/panoptic-segmentation/timm/maskformer2_vit-b_dinov2_bs1_300ep.yaml, which loads a pretrained DINOv2 ViT-B model from timm and then trains from there.

3. Run inference on the dataset e.g.

   python train_net.py --eval-only --num-gpus="${SLURM_GPUS_ON_NODE}" --config-file "${MASK2FORMER_CONFIG_FILE}" OUTPUT_DIR "${OUTPUT_DIR}$" SEED "${SLURM_ARRAY_TASK_ID}"

   Relevant config files can be found in the uncert_eval folder e.g. configs/viper/panoptic-segmentation/uncert_eval/r50/maskformer2_R50_MC3_maskdist.yaml for 3 MC Dropout samples. Effectively all experiment settings are defined in these config files, and they can inherit from each other.

   Running inference with different models, datasets, and seeds will generate a lot of data, especially as we save a lot of "raw" results to disk via the saveResults "evaluator" class. Some metrics, such as those for calibration, are calculated on the fly. Others, however, including those for out-of-distribution detection and failure detection, are calculated later (we call this offline) as they are somewhat slow to calculate and would waste time better spent sending new images to the GPU. Thus, we save the data we need for these calculations, such as per-pixel predictions and uncertainty information, to disk. Do note that this will very quickly consume terabytes of disk space, as there are multiple models, datasets, seeds and other variables.

4. These offline calculations are done using the offline_preprocess.py script:

   python offline_preprocess.py "${EXPERIMENTS_FOLDER}" --delete --yes

   which, with the option --delete specified, will process the potentially terabytes of data, save the processed results to disk, and delete the original files.

5. Once the offline processing is done, there is still a lot of manipulation required to create the figures in the paper. These are all done using Jupyter notebooks in the notebooks folder.

The notebooks are used for the final figure generation for the paper. Much of the analysis that resulted in the paper were done while writing them, so they are disorganized but are provided for completeness for anyone building off this work.

• process_results.ipynb - The main notebook for most of the plots in the paper. Run this first as it generates files used by some other ones.
• process_results_ood.ipynb - Similar to the main notebook, but for the OOD experiments. Run this after process_results.ipynb.
• process_aggcompare_results.ipynb - Handles the comparison of different aggregation methods supporting the use in the paper of the introduced Mask Distance method.
• diagram.ipynb - Used to get some prediction images for the overview diagram in the paper.
• timing.ipynb - Used to calculate some of the runtime comparison results in the Supplementary Material.

We also provide CSV versions of the DataFrames created in the notebooks; these contain all the key information but are significantly more compact that even the processed results from which they are derived. Simply find the cells where they are saved (e.g. df.to_csv('all_uncert_data_exceptood.csv')) and load them instead.

• aggcompare_results.csv - The DataFrame used in process_aggcompare_results.ipynb.
• all_uncert_data_exceptood.csv - The main DataFrame used in process_results.ipynb.
• uncert_data_ood_only.csv - The DataFrame used in process_results_ood.ipynb.

If you make use of any code, data or results, please cite as follows:

@InProceedings{Smith_2026_CVPR,
    author    = {Smith, Michael and Ferrie, Frank P.},
    title     = {U-SEG: Uncertainty in SEGmentation - A systematic multi-variable exploration},
    booktitle = {Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR) Findings},
    month     = {June},
    year      = {2026},
    pages     = {1850-1861}
}

You may also want to cite the original Mask2Former paper if appropriate.