--- a
+++ b/src/normalize_uncertainty.py
@@ -0,0 +1,62 @@
+import os
+import numpy as np
+from src.dataset.utils.nifi_volume import load_nifi_volume_return_nib, save_segmask_as_nifi_volume, load_nifi_volume
+from src.uncertainty.uncertainty import brats_normalize
+from tqdm import tqdm
+
+
+if __name__ == "__main__":
+
+    setx = "valid"
+    task = "uncertainty_task_tta/"
+    model_id = "model_1598640035"
+
+    gen_path = "/mnt/gpid07/users/laura.mora/"
+    model_path = f"{gen_path}/results/checkpoints/{model_id}/"
+    ground_truth_path = f"{gen_path}/datasets/2020/train/no_patch"
+    input_dir = os.path.join(model_path, f"{task}/{setx}")
+
+    output_dir = os.path.join(model_path, f"{task}/{setx}/normalized")
+    if not os.path.exists(output_dir):
+        os.makedirs(output_dir)
+
+    file_list = sorted([file for file in os.listdir(input_dir) if "unc" in file and "nii.gz"])
+    file_list_all = sorted([file for file in os.listdir(input_dir) if "nii.gz" in file])
+
+    max_uncertainty = 0
+    min_uncertainty = 10000
+
+    for uncertainty_map in tqdm(file_list, total=len(file_list), desc="Getting min and max"):
+
+        # Load Uncertainty maps
+        patient_name = uncertainty_map.split(".")[0].split("_unc")[0]
+        path_gt = os.path.join(ground_truth_path, patient_name, f"{patient_name}_flair.nii.gz")
+        flair = load_nifi_volume(path_gt, normalize=False)
+        brain_mask = np.zeros(flair.shape, np.float)
+        brain_mask[flair > 0] = 1
+
+        path = os.path.join(input_dir, uncertainty_map)
+        unc_map, _ = load_nifi_volume_return_nib(path, normalize=False)
+
+        tmp_max = np.max(unc_map[brain_mask == 1])
+        tmp_min = np.min(unc_map[brain_mask == 1])
+
+        if tmp_max > max_uncertainty:
+            max_uncertainty = tmp_max
+
+        if tmp_min < min_uncertainty:
+            min_uncertainty = tmp_min
+
+    for uncertainty_map_path in tqdm(file_list_all, total=len(file_list_all), desc="Normalizing.."):
+
+        path = os.path.join(input_dir, uncertainty_map_path)
+        output_path = os.path.join(output_dir, uncertainty_map_path)
+
+        unc_map, nib_data = load_nifi_volume_return_nib(path, normalize=False)
+
+        if "unc" in uncertainty_map_path:
+            uncertainty_map_normalized = brats_normalize(unc_map, max_unc=max_uncertainty, min_unc=min_uncertainty)
+            print(f"Saving to: {output_path}")
+            save_segmask_as_nifi_volume(uncertainty_map_normalized,  nib_data.affine, output_path)
+        else:
+            save_segmask_as_nifi_volume(unc_map, nib_data.affine, output_path)
\ No newline at end of file