Post-processing model predictions with AUDIT
In the Preprocessing dataset tutorial, we prepared and adapted the LUMIERE dataset 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. Load functions
This tutorial uses some utility functions from the audit.utils.commons.file_manager and audit.utils.sequences.sequences
modules 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_folders,
add_suffix_to_files
)
2. Understanding the data
Now, let's suppose we have a model that was trained on a brain MRI dataset that we don't know in advance (in this case, it was BraTS2020). We use this model to generate a series of predictions on our own dataset. In our scenario, 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/"
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. Add extension name
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. As mentioned, the model used 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 BraTS and UCSF datasets provided by AUDIT use the following mapping (after preprocessing that we performed earlier).
- 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.py 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))