Fix explicit augmentations passing in the fit method

weigertlab · Nov 22, 2024 · 86af2a3 · 86af2a3
1 parent f6c7bba
commit 86af2a3
Showing 1 changed file with 2 additions and 4 deletions.
diff --git a/spotiflow/model/spotiflow.py b/spotiflow/model/spotiflow.py
@@ -415,11 +415,9 @@ def fit(
                 transforms.Crop if not self.config.is_3d else transforms3d.Crop3D
             )
             assert any(
-                isinstance(p, _crop_cls) for p in augment_train.transforms
+                isinstance(p, _crop_cls) for p in augment_train.augmentations
             ), "Custom augmenter must contain a cropping transform!"
-            tr_augmenter = self.build_image_augmenter(
-                crop_size, point_priority=point_priority
-            )
+            tr_augmenter = augment_train
         elif augment_train:
             tr_augmenter = self.build_image_augmenter(
                 crop_size, point_priority=point_priority