Upload model

modeling.py

@@ -113,7 +113,7 @@ class CTCropModel(PreTrainedModel):
         fname: str = "",
         sort_by_instance_number: bool = False,
         exclude_invalid_dicoms: bool = False,
-    ):
+    ) -> bool:
         attributes = [
             "pixel_array",
             "RescaleSlope",
@@ -281,7 +281,7 @@ class CTCropModel(PreTrainedModel):
         coords: torch.Tensor,
         buffer: float | tuple[float, float] = 0.05,
         empty_threshold: float = 1e-4,
-    ):
+    ) -> torch.Tensor:
         coords = coords.clone()
         empty = (coords < empty_threshold).all(dim=1)
         # assumes coords is a torch.Tensor of shape (N, 4) containing
@@ -352,7 +352,12 @@ class CTCropModel(PreTrainedModel):
         raw_hu: bool = False,
         remove_empty_slices: bool = False,
         add_buffer: float | tuple[float, float] | None = None,
-    ) -> np.ndarray | tuple[np.ndarray, list[int]]:
+        return_coords: bool = False,
+    ) -> (
+        np.ndarray
+        | tuple[np.ndarray, list[int]]
+        | tuple[np.ndarray, list[int], list[int]]
+    ):
         assert mode in ["2d", "3d"]
         if device is None:
             device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -393,5 +398,11 @@ class CTCropModel(PreTrainedModel):
             empty_indices = list(torch.where(empty)[0].cpu().numpy())
             print(f"removing {empty.sum()} empty slices ...")
             cropped = cropped[~empty.cpu().numpy()]
-            return cropped, empty_indices
+            if not isinstance(cropped, tuple):
+                cropped = (cropped,)
+            cropped = cropped + (empty_indices,)
+        if return_coords:
+            if not isinstance(cropped, tuple):
+                cropped = (cropped,)
+            cropped = cropped + ([x1, y1, x2, y2],)
         return cropped