d

RobinNet/allModelsCreatePredNpys.py

@@ -7,6 +7,7 @@ import matplotlib
 from tkinter import filedialog
 from skimage import transform
 import SimpleITK as sitk
+import nibabel as nib
 import argparse
 #import os
 #import pydot
@@ -235,10 +236,24 @@ def main():
                                 y_true3d = np.append(np.reshape(y_true_np, newshape=(1, 512, 512)), y_true3d, axis=0)
                         except:
                             pass
-                    npy_savepath = "finalResults/3dnpys/" + model_path
+                    npy_savepath = "finalResults/3dnpys/" + model_path + "P" + str(TP_num).zfill(2)
                     if not os.path.exists(npy_savepath):
                         os.makedirs(npy_savepath)
-                    np.save(npy_savepath + "P" + str(TP_num).zfill(2) + "_pred3D", y_pred3d)
+                    np.save(npy_savepath + "_pred3D", y_pred3d)
+
+                    output_volume = np.squeeze(y_pred3d)
+                    output_nifty = nib.Nifti1Image(output_volume, np.eye(4))
+
+                    fname = npy_savepath+'_Scan.nii.gz'
+                    nib.save(output_nifty, fname)
+
+                    output_label = np.argmax(output_volume, axis=-1)
+                    output_label = output_label.astype(np.uint8)
+                    output_label_nifty = nib.Nifti1Image(output_label, np.eye(4))
+
+                    fname = npy_savepath+'_Label.nii.gz'
+                    nib.save(output_label_nifty, fname)
+
 
           #####