add left-righ flip (#227)

paninski-lab · Dec 8, 2024 · 15d9fe2 · 15d9fe2
1 parent 0b26c0e
commit 15d9fe2
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 5 deletions.
diff --git a/docs/source/user_guide/config_file.rst b/docs/source/user_guide/config_file.rst
@@ -74,6 +74,7 @@ See the :ref:`FAQs <faq_oom>` for more information on memory management.
 
   * default: resizing only
   * dlc: imgaug pipeline implmented in DLC 2.0 package
+  * dlc-lr: dlc augmentations plus horizontal flips
   * dlc-top-down: dlc augmentations plus additional vertical and horizontal flips
 
 * ``training.train_batch_size`` (*int, default: 16*): batch size for labeled data during training

diff --git a/lightning_pose/data/augmentations.py b/lightning_pose/data/augmentations.py
@@ -36,15 +36,16 @@ def imgaug_transform(cfg: DictConfig) -> iaa.Sequential:
         # resizing happens below
         print("using default image augmentation pipeline (resizing only)")
 
-    elif kind == "dlc" or kind == "dlc-top-down":
+    elif kind in ["dlc", "dlc-lr", "dlc-top-down"]:
 
         print(f"using {kind} image augmentation pipeline")
 
-        # flip around horizontal/vertical axes first
-        if kind == "dlc-top-down":
-            # vertical axis
+        # flip horizontally/vertically
+        if kind in ["dlc-lr", "dlc-top-down"]:
+            # horizontal flip
             data_transform.append(iaa.Fliplr(p=0.5))
-            # horizontal axis
+        if kind == "dlc-top-down":
+            # vertical flip
             data_transform.append(iaa.Flipud(p=0.5))
 
         apply_prob_0 = 0.5

diff --git a/scripts/configs/config_default.yaml b/scripts/configs/config_default.yaml
@@ -29,6 +29,7 @@ training:
   # select from one of several predefined image/video augmentation pipelines
   # default: resizing only
   # dlc: imgaug pipeline implemented in DLC 2.0 package
+  # dlc-lr: dlc augmenation plus horizontal flips
   # dlc-top-down: dlc augmentations plus vertical and horizontal flips
   imgaug: dlc
   # batch size of labeled data during training