Saving predictions in NIFTI format from generative image-to-image translation model

[dubisx]

Hello,

I've trained a model by following the [https://github.com/Project-MONAI/GenerativeModels/blob/main/tutorials/generative/image_to_image_translation/tutorial_segmentation_with_ddpm.ipynb](Generative Image-to-Image translation tutorial). The tutorial gives the outputs as an ensemble.

for i in range(len(ensemble)):

    prediction = torch.where(ensemble[i] > 0.5, 1, 0).float()  # a binary mask is obtained via thresholding
    score = dice_coeff(
        prediction[0, 0].cpu(), inputlabel.cpu()
    )  # we compute the dice scores for all samples separately
    print("Dice score of sample" + str(i), score)

E = torch.where(torch.cat(ensemble) > 0.5, 1, 0).float()
var = torch.var(E, dim=0)  # pixel-wise variance map over the ensemble
mean = torch.mean(E, dim=0)  # pixel-wise mean map over the ensemble
mean_prediction = torch.where(mean > 0.5, 1, 0).float()

score = dice_coeff(mean_prediction[0, ...].cpu(), inputlabel.cpu())  # Here we predict the Dice score for the mean map
print("Dice score on the mean map", score)

plt.style.use("default")
plt.imshow(mean[0, ...].cpu(), vmin=0, vmax=1, cmap="gray")  # We plot the mean map
plt.tight_layout()
plt.axis("off")
plt.show()
plt.style.use("default")
plt.imshow(var[0, ...].cpu(), vmin=0, vmax=1, cmap="jet")  # We plot the variance map
plt.tight_layout()
plt.axis("off")
plt.show()`

I would like to threshold the combined output to a certain degree based on confidence displayed on the mean map, and then save the thresholded combined output of the model as a NifTi file. How would I go about achieving the desired result?

Many thanks for all the future responses.

[KumoLiu]

Hi @dubisx, SaveImage can help save images into specified output extensions.
https://github.com/Project-MONAI/MONAI/blob/0a3cdd9f1403cbad057b96230840ee17cf430dbf/monai/transforms/io/array.py#L300
Thanks!