Fix for albumentations transform

v7labs · Sep 13, 2024 · e325ecb · e325ecb
1 parent 943684c
commit e325ecb
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 4 deletions.
diff --git a/darwin/future/data_objects/properties.py b/darwin/future/data_objects/properties.py
@@ -84,8 +84,6 @@ class FullProperty(DefaultDarwin):
     options: Optional[List[PropertyValue]] = None
     granularity: PropertyGranularity = PropertyGranularity("section")
 
-    # model_config = ConfigDict(use_enum_values=True)
-
     def to_create_endpoint(
         self,
     ) -> dict:

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