Post-processing model predictions with AUDIT
In the `LUMIERE preprocessing tutorial, we prepared and adapted the dataset LUMIERE to the file structure required to work with AUDIT. In this tutorial, we will see how, in addition to standardizing the organization of your project's files, it will be beneficial for your datasets to be standardized and follow the same naming conventions. This way, it'll be much easier to compare them with each other, evaluate models, and more.
Let's assume we've selected a pre-trained model from an open-source repository to generate predictions on your dataset. It is likely that the model wasn’t trained to predict the same labels that our segmentations use. In such cases, we’ll need to apply some post-processing to properly evaluate these predictions using AUDIT.
Fortunately, AUDIT provides users with tools to perform this post-processing and adapt the predictions as needed.
By the end of this tutorial, you will: - Organize predictions into a structured directory. - Rename files and labels to align with AUDIT's requirements. - Standardize labels for both ground truth and predictions, ensuring compatibility.
Let’s get started!
1. Prerequisites
This tutorial uses some utility functions to manipulate files and sequences.
from audit.utils.sequences.sequences import (
load_nii_by_subject_id,
iterative_labels_replacement,
count_labels
)
from audit.utils.commons.file_manager import (
list_dirs,
list_files,
organize_files_into_dirs,
add_suffix_to_files
)
2. Data understanding
Now, let's suppose we have a model that was pretrained on a brain MRI dataset. We will use this model to generate a series of predictions on our own dataset. In our case, we want to run inference on LUMIERE.
After running the inference, we store the predictions in the following directory:
root_path_sequences = "./datasets/LUMIERE/LUMIERE_images/"
root_path_predictions = "./datasets/LUMIERE/LUMIERE_seg/nnUnet"
Let’s check the contents of this directory:
[]
['Patient-001-week-044.nii.gz', 'Patient-001-week-056.nii.gz', 'Patient-002-week-000.nii.gz', 'Patient-002-week-003.nii.gz', 'Patient-002-week-021.nii.gz', 'Patient-002-week-037.nii.gz']
3. Organize folder
AUDIT is designed to work with multiple models and datasets, but it requires a specific organization of the data for processing. Currently, we have a single directory ".datasets/LUMIERE/LUMIERE_seg/". Have a look at the documentation to learn more about AUDIT project structure.
This directory contains the predictions generated by our model for the LUMIERE dataset. However, AUDIT requires that each segmentation be contained within a folder named after the subject ID. Since users may not structure their code to produce output in this specific format—like the pre-trained model we are using—AUDIT provides functions to facilitate the organization of files.
In this case, we can use the organize_files_into_folders function to create the necessary data structure without needing to perform complex manipulations.
[SAFE MODE] Would move: ./datasets/LUMIERE/LUMIERE_seg/nnUnet/Patient-058-week-002.nii.gz -> ./datasets/LUMIERE/LUMIERE_seg/nnUnet/Patient-058-week-002/Patient-058-week-002.nii.gz
[SAFE MODE] Would move: ./datasets/LUMIERE/LUMIERE_seg/nnUnet/Patient-018-week-063.nii.gz -> ./datasets/LUMIERE/LUMIERE_seg/nnUnet/Patient-018-week-063/Patient-018-week-063.nii.gz
...
Let's turn safe_mode into False to apply the changes.
The structure of the directory containing the predictions has changed. Instead of having all the files scattered within the directory, a separate folder has been created for each subject, containing their respective predictions.
This organized structure makes it much easier to manage and process the predictions for each subject in the dataset. Let’s take a moment to verify this new organization:
|--Patient-000-week-000
|---- Patient-000-week-000.nii.gz
|--Patient-000-week-001
|---- Patient-000-week-001.nii.gz
...
4. Rename files
Another important aspect of AUDIT is that it uses file extensions to distinguish between sequences (e.g., _t1, _t2, _t1ce, _flair), segmentations (_seg), and predictions (_pred). However, if we check the names of the files generated after inference, they did not contain any specific nomenclature.
To simplify the task for users and eliminate the need to modify their pipelines to ensure compatibility with AUDIT, the library provides methods to quickly and easily adapt the file names.
Let’s explore how to rename the files so they align with AUDIT’s requirements. Taking advantage of the function add_suffix_to_files, we can add the extension required by AUDIT.
add_suffix_to_files(
root_dir=root_path_predictions,
suffix='_pred',
ext='.nii.gz',
safe_mode=False
)
Now, all the files within the root_path contain the _pred extension.
5. Label replacement (Ground Truth)
We need to verify which labels were generated after the inference. The model used in this tutorial was pre-trained on the BraTS2020 dataset to predict the labels ENH, NEC, and EDE. However, these may not align with the labels used in the LUMIERE dataset. In fact, they are different and will need to be adjusted accordingly.
subject = "Patient-006-week-000"
seg = load_nii_by_subject_id(
root_dir=root_path_sequences,
subject_id=subject,
seq="_seg",
as_array=True
)
count_labels(seg)
The annotation protocol followed for BraTS 2020 dataset was the following:
- BKG: 0
- EDE: 3
- ENH: 1
- NEC: 2
In contrast, LUMIERE uses:
- BKG: 0
- EDE: 3
- ENH: 2
- NEC: 1
To resolve this mismatch, we will use the iterative_labels_replacement function. This function takes the old mapping and the new mapping as parameters and replaces the labels accordingly.
original_labels = [0, 1, 2, 3] # current mapping used by LUMIERE (BKG: 0 NEC: 1 ENH: 2 EDE: 3)
new_labels = [0, 2, 1, 3] # new mapping we want LUMIERE to use (BKG: 0 ENH: 1 NEC: 2 EDE: 3)
iterative_labels_replacement(
root_dir=root_path_sequences,
original_labels=original_labels,
new_labels=new_labels,
ext="_seg" # only files whose extension is '_seg' will be relabeled
)
Once we apply this function, we can confirm that the labels are now correct.
2025-01-20 13:52:47.302 | INFO | src.audit.utils.sequences.sequences:iterative_labels_replacement:216 - Iterative label replacement completed: 513 files processed, 2052 files skipped.
6. Label replacement (Predictions)
The model we used to generate the predictions was trained with the intention that tumor regions be labeled as follows:
- BKG: 0
- EDE: 2
- ENH: 4
- NEC: 1
As can be observed in the prediction generated for the same subject we have been working with.
pred = load_nii_by_subject_id(
root_dir=root_path_predictions,
subject_id=subject,
seq="_pred",
as_array=True
)
count_labels(pred)
Therefore, we will once again need to rename the labels to match the ground truth.
original_labels = [0, 1, 2, 4] # current mapping used in our predictions (BKG: 0 NEC: 1 EDE: 2 ENH: 4)
new_labels = [0, 2, 3, 1] # new mapping we want LUMIERE to use (BKG: 0 ENH: 1 NEC: 2 EDE: 3)
iterative_labels_replacement(
root_dir=root_path_predictions,
original_labels=original_labels,
new_labels=new_labels,
ext="_pred" # only files whose extension is '_seg' will be relabeled
)
2025-01-20 14:34:52.473 | INFO | src.audit.utils.sequences.sequences:iterative_labels_replacement:216 - Iterative label replacement completed: 507 files processed, 0 files skipped.
Finally, we have prepared the dataset to run the metric extraction module and start using AUDIT effectively.
seg = load_nii_by_subject_id(
root_dir=root_path_sequences,
subject_id=subject,
seq="_seg",
as_array=True
)
pred = load_nii_by_subject_id(
root_dir=root_path_predictions,
subject_id=subject,
seq="_pred",
as_array=True
)
# now they use the same labeling system
print(count_labels(seg))
print(count_labels(pred))