Fix for albumentations transform

darwin/exceptions.py

@@ -392,13 +392,17 @@ def __init__(
         super().__init__("Complex polygons not yet supported for dataloop import")
 
 
-class ExportException(DarwinException): ...
+class ExportException(DarwinException):
+    ...
 
 
-class ExportException_CouldNotAssembleOutputPath(ExportException): ...
+class ExportException_CouldNotAssembleOutputPath(ExportException):
+    ...
 
 
-class ExportException_CouldNotBuildOutput(ExportException): ...
+class ExportException_CouldNotBuildOutput(ExportException):
+    ...
 
 
-class ExportException_CouldNotWriteFile(ExportException): ...
+class ExportException_CouldNotWriteFile(ExportException):
+    ...

darwin/future/core/types/common.py

@@ -9,7 +9,8 @@
 
 
 class Implements_str(Protocol):
-    def __str__(self) -> str: ...
+    def __str__(self) -> str:
+        ...
 
 
 Stringable = Union[str, Implements_str]

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:

darwin/future/tests/data_objects/workflow/invalidvaluefortest.py

@@ -1 +1,2 @@
-class InvalidValueForTest: ...
+class InvalidValueForTest:
+    ...

darwin/importer/importer.py

@@ -283,9 +283,9 @@ def _get_team_properties_annotation_lookup(client, team_slug):
     team_properties = client.get_team_properties(team_slug)
 
     # (property-name, annotation_class_id): FullProperty object
-    team_properties_annotation_lookup: Dict[Tuple[str, Optional[int]], FullProperty] = (
-        {}
-    )
+    team_properties_annotation_lookup: Dict[
+        Tuple[str, Optional[int]], FullProperty
+    ] = {}
     for prop in team_properties:
         team_properties_annotation_lookup[(prop.name, prop.annotation_class_id)] = prop
 
@@ -1407,9 +1407,9 @@ def _import_annotations(
         # Insert the default slot name if not available in the import source
         annotation = _handle_slot_names(annotation, dataset.version, default_slot_name)
 
-        annotation_class_ids_map[(annotation_class.name, annotation_type)] = (
-            annotation_class_id
-        )
+        annotation_class_ids_map[
+            (annotation_class.name, annotation_type)
+        ] = annotation_class_id
         serial_obj = {
             "annotation_class_id": annotation_class_id,
             "data": data,

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