add model

model_segvol_single.py

@@ -212,9 +212,16 @@ class SegVolProcessor():
     
     def save_preds(self, ct_path, save_path, logits_mask, start_coord, end_coord):
         ct = nib.load(ct_path)
-        print(ct.header, ct.affine, ct.shape)
-        print(start_coord, end_coord)
-        return
+        logits_mask = logits_mask.transpose(-1, -3)
+        start_coord[-1], start_coord[-3] = start_coord[-3], start_coord[-1]
+        end_coord[-1], end_coord[-3] = end_coord[-3], end_coord[-1]
+        preds_save = torch.zeros(ct.shape)
+        preds_save[start_coord[0]:end_coord[0], 
+                        start_coord[1]:end_coord[1], 
+                        start_coord[2]:end_coord[2]] = torch.sigmoid(logits_mask)
+        preds_save = torch.where(preds_save > 0.5, 1., 0.).numpy()
+        preds_nii = nib.Nifti1Image(preds_save, affine=ct.affine, header=ct.header)
+        nib.save(preds_nii, save_path)
 
 class MinMaxNormalization(transforms.Transform):
     def __call__(self, data):