Fix topdown aug

talmolab · Aug 19, 2024 · d8f49a8 · d8f49a8
1 parent 8704e22
commit d8f49a8
Showing 1 changed file with 15 additions and 15 deletions.
diff --git a/sleap_nn/data/pipelines.py b/sleap_nn/data/pipelines.py
@@ -66,13 +66,21 @@ def make_training_pipeline(
             provider=provider,
         )
 
-        if use_augmentations and "intensity" in self.data_config.augmentation_config:
-            datapipe = KorniaAugmenter(
-                datapipe,
-                **dict(self.data_config.augmentation_config.intensity),
-                image_key="image",
-                instance_key="instances",
-            )
+        if use_augmentations:
+            if "intensity" in self.data_config.augmentation_config:
+                datapipe = KorniaAugmenter(
+                    datapipe,
+                    **dict(self.data_config.augmentation_config.intensity),
+                    image_key="image",
+                    instance_key="instances",
+                )
+            if "geometric" in self.data_config.augmentation_config:
+                datapipe = KorniaAugmenter(
+                    datapipe,
+                    **dict(self.data_config.augmentation_config.geometric),
+                    image_key="image",
+                    instance_key="instances",
+                )
 
         datapipe = InstanceCentroidFinder(
             datapipe, anchor_ind=self.confmap_head.anchor_part
@@ -83,14 +91,6 @@ def make_training_pipeline(
             self.data_config.preprocessing.crop_hw,
         )
 
-        if use_augmentations and "geometric" in self.data_config.augmentation_config:
-            datapipe = KorniaAugmenter(
-                datapipe,
-                **dict(self.data_config.augmentation_config.geometric),
-                image_key="instance_image",
-                instance_key="instance",
-            )
-
         datapipe = Resizer(
             datapipe,
             scale=self.data_config.preprocessing.scale,