Upload model

modeling.py

@@ -404,7 +404,9 @@ class TotalClassifierModel(PreTrainedModel):
             return slices, [dicom_files[idx] for idx in indices]
         return slices
 
-    def preprocess(self, x: np.ndarray, mode="2d") -> np.ndarray:
+    def preprocess(
+        self, x: np.ndarray, mode: str = "2d", torchify: bool = True
+    ) -> np.ndarray:
         mode = mode.lower()
         if mode == "2d":
             x = cv2.resize(x, self.image_size)
@@ -414,6 +416,11 @@ class TotalClassifierModel(PreTrainedModel):
             x = np.stack([cv2.resize(s, self.image_size) for s in x], axis=0)
             if x.ndim == 3:
                 x = x[:, :, :, np.newaxis]
+        if torchify:
+            if x.ndim == 3:
+                x = torch.from_numpy(x).float().permute(2, 0, 1)
+            elif x.ndim == 4:
+                x = torch.from_numpy(x).float().permute(3, 0, 1, 2)
         return x
 
     def crop_single_plane(