Fix for albumentations transform

v7labs · Sep 13, 2024 · 93d1f4b · 93d1f4b
1 parent 943684c
commit 93d1f4b
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/darwin/torch/transforms.py b/darwin/torch/transforms.py
@@ -368,8 +368,11 @@ def _pre_process(self, image: np.ndarray, annotation: dict) -> dict:
         if (
             masks is not None and masks.numel() > 0
         ):  # using numel() to check if tensor is non-empty
-            print("WE GOT MASKS")
-            albumentation_dict["masks"] = masks.numpy()
+            if isinstance(masks, torch.Tensor):
+                masks = masks.numpy()
+            if masks.ndim == 3:  # Ensure masks is a list of numpy arrays
+                masks = [masks[i] for i in range(masks.shape[0])]
+            albumentation_dict["masks"] = masks
 
         return albumentation_dict