Upload model

modeling.py

@@ -349,6 +349,7 @@ class CTCropModel(PreTrainedModel):
         mode: str,
         device: str | None = None,
         raw_hu: bool = False,
+        remove_empty_slices: bool = False,
         add_buffer: float | tuple[float, float] | None = None,
     ) -> np.ndarray:
         assert mode in ["2d", "3d"]
@@ -375,7 +376,8 @@ class CTCropModel(PreTrainedModel):
         coords = self.forward(x, img_shape=img_shapes, add_buffer=add_buffer)
         # get the union of all slice-wise bounding boxes
         # exclude empty boxes
-        coords = coords[coords.sum(dim=1) != 0]
+        empty = coords.sum(dim=1) == 0
+        coords = coords[~empty]
         # if all empty, return original input
         if coords.shape[0] == 0:
             print("no foreground detected, returning original input ...")
@@ -386,4 +388,9 @@ class CTCropModel(PreTrainedModel):
         x1, y1 = x1.min().item(), y1.min().item()
         x2, y2 = x2.max().item(), y2.max().item()
         cropped = x0[:, y1:y2, x1:x2] if mode == "3d" else x0[y1:y2, x1:x2]
+        if remove_empty_slices and empty.sum() > 0:
+            empty_indices = list(torch.where(empty)[0].cpu().numpy())
+            print(f"removing {empty.sum()} empty slices ...")
+            cropped = cropped[~empty]
+            return cropped, empty_indices
         return cropped