From b03ebdfd2d9cd6869d526a58159d2db1f0e52dcb Mon Sep 17 00:00:00 2001
From: Talmo Pereira <talmo@salk.edu>
Date: Fri, 1 Dec 2023 17:43:18 -0800
Subject: [PATCH 1/5] Initial commit

---
 environment.yml                    |   1 +
 sleap/nn/data/augmentation.py      | 232 ++++++++++++++++++++++++-----
 sleap/nn/data/pipelines.py         |  16 +-
 tests/nn/data/test_augmentation.py |   4 +-
 4 files changed, 208 insertions(+), 45 deletions(-)

diff --git a/environment.yml b/environment.yml
index 67ed39d01..3cf0086ae 100644
--- a/environment.yml
+++ b/environment.yml
@@ -13,6 +13,7 @@ dependencies:
     - conda-forge::attrs >=21.2.0  #,<=21.4.0
     - conda-forge::cattrs ==1.1.1
     - conda-forge::imgaug ==0.4.0
+    - conda-forge::albumentations
     - conda-forge::jsmin
     - conda-forge::jsonpickle ==1.2
     - conda-forge::networkx
diff --git a/sleap/nn/data/augmentation.py b/sleap/nn/data/augmentation.py
index 21dfb29e6..50073ef4d 100644
--- a/sleap/nn/data/augmentation.py
+++ b/sleap/nn/data/augmentation.py
@@ -1,19 +1,11 @@
 """Transformers for applying data augmentation."""
 
-# Monkey patch for: https://github.com/aleju/imgaug/issues/537
-# TODO: Fix when PyPI/conda packages are available for version fencing.
-import numpy
-
-if hasattr(numpy.random, "_bit_generator"):
-    numpy.random.bit_generator = numpy.random._bit_generator
-
 import sleap
 import numpy as np
 import tensorflow as tf
 import attr
 from typing import List, Text, Optional
-import imgaug as ia
-import imgaug.augmenters as iaa
+import albumentations as A
 from sleap.nn.config import AugmentationConfig
 from sleap.nn.data.instance_cropping import crop_bboxes
 
@@ -111,15 +103,15 @@ def flip_instances_ud(
 
 
 @attr.s(auto_attribs=True)
-class ImgaugAugmenter:
-    """Data transformer based on the `imgaug` library.
+class AlbumentationsAugmenter:
+    """Data transformer based on the `albumentations` library.
 
     This class can generate a `tf.data.Dataset` from an existing one that generates
     image and instance data. Element of the output dataset will have a set of
     augmentation transformations applied.
 
     Attributes:
-        augmenter: An instance of `imgaug.augmenters.Sequential` that will be applied to
+        augmenter: An instance of `albumentations.Compose` that will be applied to
             each element of the input dataset.
         image_key: Name of the example key where the image is stored. Defaults to
             "image".
@@ -127,7 +119,7 @@ class ImgaugAugmenter:
             Defaults to "instances".
     """
 
-    augmenter: iaa.Sequential
+    augmenter: A.ReplayCompose
     image_key: str = "image"
     instances_key: str = "instances"
 
@@ -137,7 +129,7 @@ def from_config(
         config: AugmentationConfig,
         image_key: Text = "image",
         instances_key: Text = "instances",
-    ) -> "ImgaugAugmenter":
+    ) -> "AlbumentationsAugmenter":
         """Create an augmenter from a set of configuration parameters.
 
         Args:
@@ -148,52 +140,63 @@ def from_config(
                 Defaults to "instances".
 
         Returns:
-            An instance of `ImgaugAugmenter` with the specified augmentation
+            An instance of `AlbumentationsAugmenter` with the specified augmentation
             configuration.
         """
         aug_stack = []
         if config.rotate:
             aug_stack.append(
-                iaa.Affine(
-                    rotate=(config.rotation_min_angle, config.rotation_max_angle)
+                A.Rotate(
+                    limit=(config.rotation_min_angle, config.rotation_max_angle), p=1.0
                 )
             )
         if config.translate:
             aug_stack.append(
-                iaa.Affine(
+                A.Affine(
                     translate_px={
                         "x": (config.translate_min, config.translate_max),
                         "y": (config.translate_min, config.translate_max),
-                    }
+                    },
+                    p=1.0,
                 )
             )
         if config.scale:
-            aug_stack.append(iaa.Affine(scale=(config.scale_min, config.scale_max)))
-        if config.uniform_noise:
             aug_stack.append(
-                iaa.AddElementwise(
-                    value=(config.uniform_noise_min_val, config.uniform_noise_max_val)
-                )
+                A.Affine(scale=(config.scale_min, config.scale_max), p=1.0)
             )
+        if config.uniform_noise:
+
+            def uniform_noise(image, **kwargs):
+                return image + np.random.uniform(
+                    config.uniform_noise_min_val, config.uniform_noise_max_val
+                )
+
+            aug_stack.append(A.Lambda(image=uniform_noise))
         if config.gaussian_noise:
             aug_stack.append(
-                iaa.AdditiveGaussianNoise(
-                    loc=config.gaussian_noise_mean, scale=config.gaussian_noise_stddev
+                A.GaussNoise(
+                    mean=config.gaussian_noise_mean,
+                    var_limit=config.gaussian_noise_stddev,
                 )
             )
         if config.contrast:
             aug_stack.append(
-                iaa.GammaContrast(
-                    gamma=(config.contrast_min_gamma, config.contrast_max_gamma)
+                A.RandomGamma(
+                    gamma_limit=(config.contrast_min_gamma, config.contrast_max_gamma),
+                    p=1.0,
                 )
             )
         if config.brightness:
             aug_stack.append(
-                iaa.Add(value=(config.brightness_min_val, config.brightness_max_val))
+                A.RandomBrightness(
+                    limit=(config.brightness_min_val, config.brightness_max_val), p=1.0
+                )
             )
 
         return cls(
-            augmenter=iaa.Sequential(aug_stack),
+            augmenter=A.ReplayCompose(
+                aug_stack, keypoint_params=A.KeypointParams(format="xy")
+            ),
             image_key=image_key,
             instances_key=instances_key,
         )
@@ -228,20 +231,25 @@ def py_augment(image, instances):
             """Local processing function that will not be autographed."""
             # Ensure that the transformations applied to all data within this
             # example are kept consistent.
-            aug_det = self.augmenter.to_deterministic()
+            # aug_det = self.augmenter.to_deterministic()
 
             # Augment the image.
-            aug_img = aug_det.augment_image(image.numpy())
+            aug_res = self.augmenter(image=image.numpy(), keypoints=[])
+            aug_img = aug_res["image"]
 
             # This will get converted to a rank 3 tensor (n_instances, n_nodes, 2).
             aug_instances = np.full_like(instances, np.nan)
 
             # Augment each set of points for each instance.
             for i, instance in enumerate(instances):
-                kps = ia.KeypointsOnImage.from_xy_array(
-                    instance.numpy(), tuple(image.shape)
+                # kps = ia.KeypointsOnImage.from_xy_array(
+                #     instance.numpy(), tuple(image.shape)
+                # )
+                # aug_instances[i] = aug_det.augment_keypoints(kps).to_xy_array()
+                aug_res_i = A.ReplayCompose.replay(
+                    aug_res["replay"], image=[], keypoints=instance.numpy()
                 )
-                aug_instances[i] = aug_det.augment_keypoints(kps).to_xy_array()
+                aug_instances[i] = aug_res_i["keypoints"]
 
             return aug_img, aug_instances
 
@@ -264,6 +272,160 @@ def augment(frame_data):
         return output_ds
 
 
+# @attr.s(auto_attribs=True)
+# class ImgaugAugmenter:
+#     """Data transformer based on the `imgaug` library.
+
+#     This class can generate a `tf.data.Dataset` from an existing one that generates
+#     image and instance data. Element of the output dataset will have a set of
+#     augmentation transformations applied.
+
+#     Attributes:
+#         augmenter: An instance of `imgaug.augmenters.Sequential` that will be applied to
+#             each element of the input dataset.
+#         image_key: Name of the example key where the image is stored. Defaults to
+#             "image".
+#         instances_key: Name of the example key where the instance points are stored.
+#             Defaults to "instances".
+#     """
+
+#     augmenter: iaa.Sequential
+#     image_key: str = "image"
+#     instances_key: str = "instances"
+
+#     @classmethod
+#     def from_config(
+#         cls,
+#         config: AugmentationConfig,
+#         image_key: Text = "image",
+#         instances_key: Text = "instances",
+#     ) -> "ImgaugAugmenter":
+#         """Create an augmenter from a set of configuration parameters.
+
+#         Args:
+#             config: An `AugmentationConfig` instance with the desired parameters.
+#             image_key: Name of the example key where the image is stored. Defaults to
+#                 "image".
+#             instances_key: Name of the example key where the instance points are stored.
+#                 Defaults to "instances".
+
+#         Returns:
+#             An instance of `ImgaugAugmenter` with the specified augmentation
+#             configuration.
+#         """
+#         aug_stack = []
+#         if config.rotate:
+#             aug_stack.append(
+#                 iaa.Affine(
+#                     rotate=(config.rotation_min_angle, config.rotation_max_angle)
+#                 )
+#             )
+#         if config.translate:
+#             aug_stack.append(
+#                 iaa.Affine(
+#                     translate_px={
+#                         "x": (config.translate_min, config.translate_max),
+#                         "y": (config.translate_min, config.translate_max),
+#                     }
+#                 )
+#             )
+#         if config.scale:
+#             aug_stack.append(iaa.Affine(scale=(config.scale_min, config.scale_max)))
+#         if config.uniform_noise:
+#             aug_stack.append(
+#                 iaa.AddElementwise(
+#                     value=(config.uniform_noise_min_val, config.uniform_noise_max_val)
+#                 )
+#             )
+#         if config.gaussian_noise:
+#             aug_stack.append(
+#                 iaa.AdditiveGaussianNoise(
+#                     loc=config.gaussian_noise_mean, scale=config.gaussian_noise_stddev
+#                 )
+#             )
+#         if config.contrast:
+#             aug_stack.append(
+#                 iaa.GammaContrast(
+#                     gamma=(config.contrast_min_gamma, config.contrast_max_gamma)
+#                 )
+#             )
+#         if config.brightness:
+#             aug_stack.append(
+#                 iaa.Add(value=(config.brightness_min_val, config.brightness_max_val))
+#             )
+
+#         return cls(
+#             augmenter=iaa.Sequential(aug_stack),
+#             image_key=image_key,
+#             instances_key=instances_key,
+#         )
+
+#     @property
+#     def input_keys(self) -> List[Text]:
+#         """Return the keys that incoming elements are expected to have."""
+#         return [self.image_key, self.instances_key]
+
+#     @property
+#     def output_keys(self) -> List[Text]:
+#         """Return the keys that outgoing elements will have."""
+#         return self.input_keys
+
+#     def transform_dataset(self, input_ds: tf.data.Dataset) -> tf.data.Dataset:
+#         """Create a `tf.data.Dataset` with elements containing augmented data.
+
+#         Args:
+#             input_ds: A dataset with elements that contain the keys "image" and
+#                 "instances". This is typically raw data from a data provider.
+
+#         Returns:
+#             A `tf.data.Dataset` with the same keys as the input, but with images and
+#             instance points updated with the applied augmentations.
+
+#         Notes:
+#             The "scale" key in examples are not modified when scaling augmentation is
+#             applied.
+#         """
+#         # Define augmentation function to map over each sample.
+#         def py_augment(image, instances):
+#             """Local processing function that will not be autographed."""
+#             # Ensure that the transformations applied to all data within this
+#             # example are kept consistent.
+#             aug_det = self.augmenter.to_deterministic()
+
+#             # Augment the image.
+#             aug_img = aug_det.augment_image(image.numpy())
+
+#             # This will get converted to a rank 3 tensor (n_instances, n_nodes, 2).
+#             aug_instances = np.full_like(instances, np.nan)
+
+#             # Augment each set of points for each instance.
+#             for i, instance in enumerate(instances):
+#                 kps = ia.KeypointsOnImage.from_xy_array(
+#                     instance.numpy(), tuple(image.shape)
+#                 )
+#                 aug_instances[i] = aug_det.augment_keypoints(kps).to_xy_array()
+
+#             return aug_img, aug_instances
+
+#         def augment(frame_data):
+#             """Wrap local processing function for dataset mapping."""
+#             image, instances = tf.py_function(
+#                 py_augment,
+#                 [frame_data["image"], frame_data["instances"]],
+#                 [frame_data["image"].dtype, frame_data["instances"].dtype],
+#             )
+#             image.set_shape(frame_data["image"].get_shape())
+#             instances.set_shape(frame_data["instances"].get_shape())
+#             frame_data.update({"image": image, "instances": instances})
+#             return frame_data
+
+#         # Apply the augmentation to each element.
+#         # Note: We map sequentially since imgaug gets slower with tf.data parallelism.
+#         output_ds = input_ds.map(augment)
+
+#         return output_ds
+
+
 @attr.s(auto_attribs=True)
 class RandomCropper:
     """Data transformer for applying random crops to input images.
diff --git a/sleap/nn/data/pipelines.py b/sleap/nn/data/pipelines.py
index b0892f8a1..2e334456a 100644
--- a/sleap/nn/data/pipelines.py
+++ b/sleap/nn/data/pipelines.py
@@ -18,7 +18,7 @@
 from sleap.nn.data.providers import LabelsReader, VideoReader
 from sleap.nn.data.augmentation import (
     AugmentationConfig,
-    ImgaugAugmenter,
+    AlbumentationsAugmenter,
     RandomCropper,
     RandomFlipper,
 )
@@ -68,7 +68,7 @@
 
 PROVIDERS = (LabelsReader, VideoReader)
 TRANSFORMERS = (
-    ImgaugAugmenter,
+    AlbumentationsAugmenter,
     RandomCropper,
     Normalizer,
     Resizer,
@@ -406,7 +406,7 @@ def make_training_pipeline(self, data_provider: Provider) -> Pipeline:
                 self.data_config.labels.skeletons[0],
                 horizontal=self.optimization_config.augmentation_config.flip_horizontal,
             )
-        pipeline += ImgaugAugmenter.from_config(
+        pipeline += AlbumentationsAugmenter.from_config(
             self.optimization_config.augmentation_config
         )
         if self.optimization_config.augmentation_config.random_crop:
@@ -550,7 +550,7 @@ def make_training_pipeline(self, data_provider: Provider) -> Pipeline:
                 self.data_config.labels.skeletons[0],
                 horizontal=self.optimization_config.augmentation_config.flip_horizontal,
             )
-        pipeline += ImgaugAugmenter.from_config(
+        pipeline += AlbumentationsAugmenter.from_config(
             self.optimization_config.augmentation_config
         )
         if self.optimization_config.augmentation_config.random_crop:
@@ -713,7 +713,7 @@ def make_training_pipeline(self, data_provider: Provider) -> Pipeline:
                 self.data_config.labels.skeletons[0],
                 horizontal=self.optimization_config.augmentation_config.flip_horizontal,
             )
-        pipeline += ImgaugAugmenter.from_config(
+        pipeline += AlbumentationsAugmenter.from_config(
             self.optimization_config.augmentation_config
         )
         pipeline += Normalizer.from_config(self.data_config.preprocessing)
@@ -863,7 +863,7 @@ def make_training_pipeline(self, data_provider: Provider) -> Pipeline:
                 self.data_config.labels.skeletons[0],
                 horizontal=aug_config.flip_horizontal,
             )
-        pipeline += ImgaugAugmenter.from_config(aug_config)
+        pipeline += AlbumentationsAugmenter.from_config(aug_config)
         if aug_config.random_crop:
             pipeline += RandomCropper(
                 crop_height=aug_config.random_crop_height,
@@ -1028,7 +1028,7 @@ def make_training_pipeline(self, data_provider: Provider) -> Pipeline:
                 horizontal=aug_config.flip_horizontal,
             )
 
-        pipeline += ImgaugAugmenter.from_config(aug_config)
+        pipeline += AlbumentationsAugmenter.from_config(aug_config)
         if aug_config.random_crop:
             pipeline += RandomCropper(
                 crop_height=aug_config.random_crop_height,
@@ -1186,7 +1186,7 @@ def make_training_pipeline(self, data_provider: Provider) -> Pipeline:
                 config=self.data_config.preprocessing,
                 provider=data_provider,
             )
-        pipeline += ImgaugAugmenter.from_config(
+        pipeline += AlbumentationsAugmenter.from_config(
             self.optimization_config.augmentation_config
         )
         pipeline += Normalizer.from_config(self.data_config.preprocessing)
diff --git a/tests/nn/data/test_augmentation.py b/tests/nn/data/test_augmentation.py
index d2b468522..fe84aea81 100644
--- a/tests/nn/data/test_augmentation.py
+++ b/tests/nn/data/test_augmentation.py
@@ -14,7 +14,7 @@ def test_augmentation(min_labels):
     ds = labels_reader.make_dataset()
     example_preaug = next(iter(ds))
 
-    augmenter = augmentation.ImgaugAugmenter.from_config(
+    augmenter = augmentation.AlbumentationsAugmenter.from_config(
         augmentation.AugmentationConfig(
             rotate=True, rotation_min_angle=-90, rotation_max_angle=-90
         )
@@ -52,7 +52,7 @@ def test_augmentation_with_no_instances(min_labels):
     )
 
     p = min_labels.to_pipeline(user_labeled_only=False)
-    p += augmentation.ImgaugAugmenter.from_config(
+    p += augmentation.AlbumentationsAugmenter.from_config(
         augmentation.AugmentationConfig(rotate=True)
     )
     exs = p.run()

From dd5a59bcd9f05eee7ed44f2040c3f745cbe2f90d Mon Sep 17 00:00:00 2001
From: Talmo Pereira <talmo@salk.edu>
Date: Fri, 1 Dec 2023 18:26:14 -0800
Subject: [PATCH 2/5] Fix augmentation

---
 environment.yml                    |   1 -
 sleap/nn/data/augmentation.py      | 190 ++---------------------------
 tests/nn/data/test_augmentation.py |   2 +-
 3 files changed, 13 insertions(+), 180 deletions(-)

diff --git a/environment.yml b/environment.yml
index 3cf0086ae..4004f4694 100644
--- a/environment.yml
+++ b/environment.yml
@@ -12,7 +12,6 @@ dependencies:
     # Packages SLEAP uses directly
     - conda-forge::attrs >=21.2.0  #,<=21.4.0
     - conda-forge::cattrs ==1.1.1
-    - conda-forge::imgaug ==0.4.0
     - conda-forge::albumentations
     - conda-forge::jsmin
     - conda-forge::jsonpickle ==1.2
diff --git a/sleap/nn/data/augmentation.py b/sleap/nn/data/augmentation.py
index 50073ef4d..9c8c10842 100644
--- a/sleap/nn/data/augmentation.py
+++ b/sleap/nn/data/augmentation.py
@@ -119,7 +119,7 @@ class AlbumentationsAugmenter:
             Defaults to "instances".
     """
 
-    augmenter: A.ReplayCompose
+    augmenter: A.Compose
     image_key: str = "image"
     instances_key: str = "instances"
 
@@ -194,7 +194,7 @@ def uniform_noise(image, **kwargs):
             )
 
         return cls(
-            augmenter=A.ReplayCompose(
+            augmenter=A.Compose(
                 aug_stack, keypoint_params=A.KeypointParams(format="xy")
             ),
             image_key=image_key,
@@ -229,27 +229,16 @@ def transform_dataset(self, input_ds: tf.data.Dataset) -> tf.data.Dataset:
         # Define augmentation function to map over each sample.
         def py_augment(image, instances):
             """Local processing function that will not be autographed."""
-            # Ensure that the transformations applied to all data within this
-            # example are kept consistent.
-            # aug_det = self.augmenter.to_deterministic()
-
-            # Augment the image.
-            aug_res = self.augmenter(image=image.numpy(), keypoints=[])
-            aug_img = aug_res["image"]
-
-            # This will get converted to a rank 3 tensor (n_instances, n_nodes, 2).
-            aug_instances = np.full_like(instances, np.nan)
-
-            # Augment each set of points for each instance.
-            for i, instance in enumerate(instances):
-                # kps = ia.KeypointsOnImage.from_xy_array(
-                #     instance.numpy(), tuple(image.shape)
-                # )
-                # aug_instances[i] = aug_det.augment_keypoints(kps).to_xy_array()
-                aug_res_i = A.ReplayCompose.replay(
-                    aug_res["replay"], image=[], keypoints=instance.numpy()
-                )
-                aug_instances[i] = aug_res_i["keypoints"]
+            # Convert to numpy arrays.
+            img = image.numpy()
+            kps = instances.numpy()
+            original_shape = kps.shape
+            kps = kps.reshape(-1, 2)
+
+            # Augment.
+            augmented = self.augmenter(image=img, keypoints=kps)
+            aug_img = augmented["image"]
+            aug_instances = np.array(augmented["keypoints"]).reshape(original_shape)
 
             return aug_img, aug_instances
 
@@ -266,166 +255,11 @@ def augment(frame_data):
             return frame_data
 
         # Apply the augmentation to each element.
-        # Note: We map sequentially since imgaug gets slower with tf.data parallelism.
         output_ds = input_ds.map(augment)
 
         return output_ds
 
 
-# @attr.s(auto_attribs=True)
-# class ImgaugAugmenter:
-#     """Data transformer based on the `imgaug` library.
-
-#     This class can generate a `tf.data.Dataset` from an existing one that generates
-#     image and instance data. Element of the output dataset will have a set of
-#     augmentation transformations applied.
-
-#     Attributes:
-#         augmenter: An instance of `imgaug.augmenters.Sequential` that will be applied to
-#             each element of the input dataset.
-#         image_key: Name of the example key where the image is stored. Defaults to
-#             "image".
-#         instances_key: Name of the example key where the instance points are stored.
-#             Defaults to "instances".
-#     """
-
-#     augmenter: iaa.Sequential
-#     image_key: str = "image"
-#     instances_key: str = "instances"
-
-#     @classmethod
-#     def from_config(
-#         cls,
-#         config: AugmentationConfig,
-#         image_key: Text = "image",
-#         instances_key: Text = "instances",
-#     ) -> "ImgaugAugmenter":
-#         """Create an augmenter from a set of configuration parameters.
-
-#         Args:
-#             config: An `AugmentationConfig` instance with the desired parameters.
-#             image_key: Name of the example key where the image is stored. Defaults to
-#                 "image".
-#             instances_key: Name of the example key where the instance points are stored.
-#                 Defaults to "instances".
-
-#         Returns:
-#             An instance of `ImgaugAugmenter` with the specified augmentation
-#             configuration.
-#         """
-#         aug_stack = []
-#         if config.rotate:
-#             aug_stack.append(
-#                 iaa.Affine(
-#                     rotate=(config.rotation_min_angle, config.rotation_max_angle)
-#                 )
-#             )
-#         if config.translate:
-#             aug_stack.append(
-#                 iaa.Affine(
-#                     translate_px={
-#                         "x": (config.translate_min, config.translate_max),
-#                         "y": (config.translate_min, config.translate_max),
-#                     }
-#                 )
-#             )
-#         if config.scale:
-#             aug_stack.append(iaa.Affine(scale=(config.scale_min, config.scale_max)))
-#         if config.uniform_noise:
-#             aug_stack.append(
-#                 iaa.AddElementwise(
-#                     value=(config.uniform_noise_min_val, config.uniform_noise_max_val)
-#                 )
-#             )
-#         if config.gaussian_noise:
-#             aug_stack.append(
-#                 iaa.AdditiveGaussianNoise(
-#                     loc=config.gaussian_noise_mean, scale=config.gaussian_noise_stddev
-#                 )
-#             )
-#         if config.contrast:
-#             aug_stack.append(
-#                 iaa.GammaContrast(
-#                     gamma=(config.contrast_min_gamma, config.contrast_max_gamma)
-#                 )
-#             )
-#         if config.brightness:
-#             aug_stack.append(
-#                 iaa.Add(value=(config.brightness_min_val, config.brightness_max_val))
-#             )
-
-#         return cls(
-#             augmenter=iaa.Sequential(aug_stack),
-#             image_key=image_key,
-#             instances_key=instances_key,
-#         )
-
-#     @property
-#     def input_keys(self) -> List[Text]:
-#         """Return the keys that incoming elements are expected to have."""
-#         return [self.image_key, self.instances_key]
-
-#     @property
-#     def output_keys(self) -> List[Text]:
-#         """Return the keys that outgoing elements will have."""
-#         return self.input_keys
-
-#     def transform_dataset(self, input_ds: tf.data.Dataset) -> tf.data.Dataset:
-#         """Create a `tf.data.Dataset` with elements containing augmented data.
-
-#         Args:
-#             input_ds: A dataset with elements that contain the keys "image" and
-#                 "instances". This is typically raw data from a data provider.
-
-#         Returns:
-#             A `tf.data.Dataset` with the same keys as the input, but with images and
-#             instance points updated with the applied augmentations.
-
-#         Notes:
-#             The "scale" key in examples are not modified when scaling augmentation is
-#             applied.
-#         """
-#         # Define augmentation function to map over each sample.
-#         def py_augment(image, instances):
-#             """Local processing function that will not be autographed."""
-#             # Ensure that the transformations applied to all data within this
-#             # example are kept consistent.
-#             aug_det = self.augmenter.to_deterministic()
-
-#             # Augment the image.
-#             aug_img = aug_det.augment_image(image.numpy())
-
-#             # This will get converted to a rank 3 tensor (n_instances, n_nodes, 2).
-#             aug_instances = np.full_like(instances, np.nan)
-
-#             # Augment each set of points for each instance.
-#             for i, instance in enumerate(instances):
-#                 kps = ia.KeypointsOnImage.from_xy_array(
-#                     instance.numpy(), tuple(image.shape)
-#                 )
-#                 aug_instances[i] = aug_det.augment_keypoints(kps).to_xy_array()
-
-#             return aug_img, aug_instances
-
-#         def augment(frame_data):
-#             """Wrap local processing function for dataset mapping."""
-#             image, instances = tf.py_function(
-#                 py_augment,
-#                 [frame_data["image"], frame_data["instances"]],
-#                 [frame_data["image"].dtype, frame_data["instances"].dtype],
-#             )
-#             image.set_shape(frame_data["image"].get_shape())
-#             instances.set_shape(frame_data["instances"].get_shape())
-#             frame_data.update({"image": image, "instances": instances})
-#             return frame_data
-
-#         # Apply the augmentation to each element.
-#         # Note: We map sequentially since imgaug gets slower with tf.data parallelism.
-#         output_ds = input_ds.map(augment)
-
-#         return output_ds
-
-
 @attr.s(auto_attribs=True)
 class RandomCropper:
     """Data transformer for applying random crops to input images.
diff --git a/tests/nn/data/test_augmentation.py b/tests/nn/data/test_augmentation.py
index fe84aea81..dd2306f5a 100644
--- a/tests/nn/data/test_augmentation.py
+++ b/tests/nn/data/test_augmentation.py
@@ -16,7 +16,7 @@ def test_augmentation(min_labels):
 
     augmenter = augmentation.AlbumentationsAugmenter.from_config(
         augmentation.AugmentationConfig(
-            rotate=True, rotation_min_angle=-90, rotation_max_angle=-90
+            rotate=True, rotation_min_angle=90, rotation_max_angle=90
         )
     )
     ds = augmenter.transform_dataset(ds)

From 301f28946f50e01c6221af9202a86dc05ee9ccf0 Mon Sep 17 00:00:00 2001
From: Talmo Pereira <talmo@salk.edu>
Date: Mon, 4 Dec 2023 10:56:53 -0800
Subject: [PATCH 3/5] Update more deps requirements

---
 .conda/meta.yaml        | 4 ++--
 environment_mac.yml     | 2 +-
 environment_no_cuda.yml | 2 +-
 3 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/.conda/meta.yaml b/.conda/meta.yaml
index caffe9fcb..bb957fd87 100644
--- a/.conda/meta.yaml
+++ b/.conda/meta.yaml
@@ -33,7 +33,7 @@ requirements:
     - conda-forge::attrs ==21.4.0
     - conda-forge::cattrs ==1.1.1
     - conda-forge::h5py ==3.1  # [not win]
-    - conda-forge::imgaug ==0.4.0
+    - conda-forge::albumentations
     - conda-forge::jsmin
     - conda-forge::jsonpickle ==1.2
     - conda-forge::networkx
@@ -61,7 +61,7 @@ requirements:
     - conda-forge::cudnn=8.2.1
     - nvidia::cuda-nvcc=11.3
     - conda-forge::h5py ==3.1  # [not win]
-    - conda-forge::imgaug ==0.4.0
+    - conda-forge::albumentations
     - conda-forge::jsmin
     - conda-forge::jsonpickle ==1.2
     - conda-forge::networkx
diff --git a/environment_mac.yml b/environment_mac.yml
index 85ef7d3b9..d32cc03e6 100644
--- a/environment_mac.yml
+++ b/environment_mac.yml
@@ -11,7 +11,7 @@ dependencies:
     - conda-forge::attrs >=21.2.0
     - conda-forge::cattrs ==1.1.1
     - conda-forge::h5py
-    - conda-forge::imgaug ==0.4.0
+    - conda-forge::albumentations
     - conda-forge::jsmin
     - conda-forge::jsonpickle ==1.2
     - conda-forge::keras <2.10.0,>=2.9.0rc0  # Required by tensorflow-macos
diff --git a/environment_no_cuda.yml b/environment_no_cuda.yml
index 7e384b5f9..e51ce61db 100644
--- a/environment_no_cuda.yml
+++ b/environment_no_cuda.yml
@@ -13,7 +13,7 @@ dependencies:
     # Packages SLEAP uses directly
     - conda-forge::attrs >=21.2.0  #,<=21.4.0
     - conda-forge::cattrs ==1.1.1
-    - conda-forge::imgaug ==0.4.0
+    - conda-forge::albumentations
     - conda-forge::jsmin
     - conda-forge::jsonpickle ==1.2
     - conda-forge::networkx

From 70f2d5a040261c15adfaeb09f7004d9b6fede5cc Mon Sep 17 00:00:00 2001
From: Talmo Pereira <talmo@salk.edu>
Date: Thu, 14 Dec 2023 15:42:09 -0800
Subject: [PATCH 4/5] Use pip for installing albumentations and avoid
 reinstalling OpenCV

---
 .conda/bld.bat          | 2 +-
 .conda/build.sh         | 2 +-
 .conda/meta.yaml        | 2 --
 .conda_mac/build.sh     | 2 +-
 .conda_mac/meta.yaml    | 2 --
 environment.yml         | 3 +--
 environment_mac.yml     | 1 +
 environment_no_cuda.yml | 1 +
 requirements.txt        | 2 ++
 9 files changed, 8 insertions(+), 9 deletions(-)

diff --git a/.conda/bld.bat b/.conda/bld.bat
index 22b63e50a..65cbec481 100644
--- a/.conda/bld.bat
+++ b/.conda/bld.bat
@@ -7,7 +7,7 @@ set PIP_IGNORE_INSTALLED=False
 
 @REM Install the pip dependencies. Note: Using urls to wheels might be better: 
 @REM https://docs.conda.io/projects/conda-build/en/stable/user-guide/wheel-files.html)
-pip install --no-cache-dir -r .\requirements.txt
+pip install --no-cache-dir -r .\requirements.txt --no-binary qudida,albumentations
 
 @REM Install sleap itself. This does not install the requirements, but will list which 
 @REM requirements are missing (see "install_requires") when user attempts to install.
diff --git a/.conda/build.sh b/.conda/build.sh
index 86ab5af73..5a4e0bf87 100644
--- a/.conda/build.sh
+++ b/.conda/build.sh
@@ -7,7 +7,7 @@ export PIP_IGNORE_INSTALLED=False
 
 # Install the pip dependencies. Note: Using urls to wheels might be better: 
 # https://docs.conda.io/projects/conda-build/en/stable/user-guide/wheel-files.html)
-pip install --no-cache-dir -r ./requirements.txt
+pip install --no-cache-dir -r ./requirements.txt --no-binary qudida,albumentations
 
 
 # Install sleap itself. This does not install the requirements, but will list which 
diff --git a/.conda/meta.yaml b/.conda/meta.yaml
index bb957fd87..5707a3a2c 100644
--- a/.conda/meta.yaml
+++ b/.conda/meta.yaml
@@ -33,7 +33,6 @@ requirements:
     - conda-forge::attrs ==21.4.0
     - conda-forge::cattrs ==1.1.1
     - conda-forge::h5py ==3.1  # [not win]
-    - conda-forge::albumentations
     - conda-forge::jsmin
     - conda-forge::jsonpickle ==1.2
     - conda-forge::networkx
@@ -61,7 +60,6 @@ requirements:
     - conda-forge::cudnn=8.2.1
     - nvidia::cuda-nvcc=11.3
     - conda-forge::h5py ==3.1  # [not win]
-    - conda-forge::albumentations
     - conda-forge::jsmin
     - conda-forge::jsonpickle ==1.2
     - conda-forge::networkx
diff --git a/.conda_mac/build.sh b/.conda_mac/build.sh
index 2036035f6..85c488903 100644
--- a/.conda_mac/build.sh
+++ b/.conda_mac/build.sh
@@ -7,6 +7,6 @@ export PIP_NO_INDEX=False
 export PIP_NO_DEPENDENCIES=False
 export PIP_IGNORE_INSTALLED=False
 
-pip install --no-cache-dir -r requirements.txt
+pip install --no-cache-dir -r requirements.txt --no-binary qudida,albumentations
 
 python setup.py install --single-version-externally-managed --record=record.txt
\ No newline at end of file
diff --git a/.conda_mac/meta.yaml b/.conda_mac/meta.yaml
index 7496f2057..141d5784e 100644
--- a/.conda_mac/meta.yaml
+++ b/.conda_mac/meta.yaml
@@ -34,7 +34,6 @@ requirements:
     - conda-forge::attrs >=21.2.0
     - conda-forge::cattrs ==1.1.1
     - conda-forge::h5py
-    - conda-forge::imgaug ==0.4.0
     - conda-forge::jsmin
     - conda-forge::jsonpickle ==1.2
     - conda-forge::keras <2.10.0,>=2.9.0rc0  # Required by tensorflow-macos
@@ -61,7 +60,6 @@ requirements:
     - conda-forge::attrs >=21.2.0
     - conda-forge::cattrs ==1.1.1
     - conda-forge::h5py
-    - conda-forge::imgaug ==0.4.0
     - conda-forge::jsmin
     - conda-forge::jsonpickle ==1.2
     - conda-forge::keras <2.10.0,>=2.9.0rc0  # Required by tensorflow-macos
diff --git a/environment.yml b/environment.yml
index 4004f4694..f4426cfee 100644
--- a/environment.yml
+++ b/environment.yml
@@ -12,7 +12,6 @@ dependencies:
     # Packages SLEAP uses directly
     - conda-forge::attrs >=21.2.0  #,<=21.4.0
     - conda-forge::cattrs ==1.1.1
-    - conda-forge::albumentations
     - conda-forge::jsmin
     - conda-forge::jsonpickle ==1.2
     - conda-forge::networkx
@@ -46,4 +45,4 @@ dependencies:
 
     - pip:  
         - "--editable=.[conda_dev]"
-        
\ No newline at end of file
+        - "--no-binary qudida,albumentations"
diff --git a/environment_mac.yml b/environment_mac.yml
index d32cc03e6..bb540bd51 100644
--- a/environment_mac.yml
+++ b/environment_mac.yml
@@ -38,3 +38,4 @@ dependencies:
     - conda-forge::tensorflow-hub
     - pip:
         - "--editable=.[conda_dev]"
+        - "--no-binary qudida,albumentations"
\ No newline at end of file
diff --git a/environment_no_cuda.yml b/environment_no_cuda.yml
index e51ce61db..752c80a81 100644
--- a/environment_no_cuda.yml
+++ b/environment_no_cuda.yml
@@ -41,3 +41,4 @@ dependencies:
 
     - pip:  
         - "--editable=.[conda_dev]"
+        - "--no-binary qudida,albumentations"
\ No newline at end of file
diff --git a/requirements.txt b/requirements.txt
index cb0ef45c5..50aec76b6 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -12,3 +12,5 @@ tensorflow-metal==0.5.0; sys_platform == 'darwin' and platform_machine == 'arm64
 # Conda installing results in https://github.com/h5py/h5py/issues/2037
 h5py<3.2; sys_platform == 'win32'  # Newer versions result in error above, linking issue in Linux
 pynwb>=2.3.3  # 2.0.0 required by ndx-pose, 2.3.3 fixes importlib-metadata incompatibility
+
+albumentations
\ No newline at end of file

From 364da53559ee659d43a4a465c283787a066a9cd5 Mon Sep 17 00:00:00 2001
From: Talmo Pereira <talmo@salk.edu>
Date: Fri, 15 Dec 2023 15:47:21 -0800
Subject: [PATCH 5/5] Update other conda envs

---
 environment_mac.yml     | 1 -
 environment_no_cuda.yml | 1 -
 2 files changed, 2 deletions(-)

diff --git a/environment_mac.yml b/environment_mac.yml
index bb540bd51..c292c40e0 100644
--- a/environment_mac.yml
+++ b/environment_mac.yml
@@ -11,7 +11,6 @@ dependencies:
     - conda-forge::attrs >=21.2.0
     - conda-forge::cattrs ==1.1.1
     - conda-forge::h5py
-    - conda-forge::albumentations
     - conda-forge::jsmin
     - conda-forge::jsonpickle ==1.2
     - conda-forge::keras <2.10.0,>=2.9.0rc0  # Required by tensorflow-macos
diff --git a/environment_no_cuda.yml b/environment_no_cuda.yml
index 752c80a81..a02c21222 100644
--- a/environment_no_cuda.yml
+++ b/environment_no_cuda.yml
@@ -13,7 +13,6 @@ dependencies:
     # Packages SLEAP uses directly
     - conda-forge::attrs >=21.2.0  #,<=21.4.0
     - conda-forge::cattrs ==1.1.1
-    - conda-forge::albumentations
     - conda-forge::jsmin
     - conda-forge::jsonpickle ==1.2
     - conda-forge::networkx