Creating soft segmentations using predictions' probability maps.

[plbenveniste]

Previously, we created soft segmentations by using multiple folds such as described in this figure:

However, when producing predictions, nnUNet also produces .npz files which are probability maps of the prediction. Here is the comparison of the binary prediction and the probability map of the prediction.

To maximize the "softness" of the predictions which are later used to train a model, I suggest combining the probability maps of the 5 folds.

Note:
The .npz files need to be reoriented. Here is the code I used to do it for this example.

# Load image 
image = nib.load(image_file)
# Save image to output folder
nib.save(image, out_folder + "/image.nii.gz")
# Load seg 
seg = nib.load(seg_file)
# Save seg to output folder
nib.save(seg, out_folder + "/seg.nii.gz")
# Load npz
npz = np.load(npz_file)
probability = npz["probabilities"][1]
# Threshold probability at 1e-5
probability[probability < 1e-5] = 0
print(probability.shape)
# rotate probability matrix from 100x110x120 to 120x100x110
probability = np.transpose(probability, (2, 1, 0))
print(probability.shape)
# Save probability to output folder
nib.save(nib.Nifti1Image(probability, seg.affine), out_folder + "/probability.nii.gz")