add model

model_segvol_single.py

@@ -1,6 +1,7 @@
 from transformers import PreTrainedModel, PretrainedConfig
 import numpy as np
 import monai.transforms as transforms
+import nibabel as nib
 
 class SegVolConfig(PretrainedConfig):
     model_type = "segvol"
@@ -209,6 +210,12 @@ class SegVolProcessor():
             target = target.cpu()
         return dice
     
+    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
+
 class MinMaxNormalization(transforms.Transform):
     def __call__(self, data):
         d = dict(data)