bug fix in left-right/top-down augmentation pipeline (#287)

Author: themattinthehatt

lightning_pose/data/augmentations.py

@@ -63,11 +63,15 @@ def imgaug_transform(params_dict: dict | DictConfig) -> iaa.Sequential:
         transform_args = args.get("args", ())
         transform_kwargs = args.get("kwargs", {})
 
-        # make sure any lists are converted to tuples; DictConfig cannot load tuples from yaml
-        # files, but no iaa args are lists
+        # DictConfig cannot load tuples from yaml files
+        # make sure any lists are converted to tuples
+        # unless the list contains a single item, then pass through the item (hack for Rot90)
         for kw, arg in transform_kwargs.items():
             if isinstance(arg, list) or isinstance(arg, ListConfig):
-                transform_kwargs[kw] = tuple(arg)
+                if len(arg) == 1:
+                    transform_kwargs[kw] = arg[0]
+                else:
+                    transform_kwargs[kw] = tuple(arg)
 
         # add transform to pipeline
         if apply_prob == 0.0:
@@ -91,13 +95,13 @@ def expand_imgaug_str_to_dict(params: str) -> dict[str, Any]:
         pass  # no augmentations
     elif params in ["dlc", "dlc-lr", "dlc-top-down"]:
 
-        # flip horizontally
-        if params in ["dlc-lr", "dlc-top-down"]:
-            params_dict["Fliplr"] = {"p": 1.0, "kwargs": {"p": 0.5}}
+        # rotate 0 or 180 degrees
+        if params in ["dlc-lr"]:
+            params_dict["Rot90"] = {"p": 1.0, "kwargs": {"k": [[0, 2]]}}
 
-        # flip vertically
+        # rotate 0, 90, 180, or 270 degrees
         if params in ["dlc-top-down"]:
-            params_dict["Flipud"] = {"p": 1.0, "kwargs": {"p": 0.5}}
+            params_dict["Rot90"] = {"p": 1.0, "kwargs": {"k": [[0, 1, 2, 3]]}}
 
         # rotate
         rotation = 25  # rotation uniformly sampled from (-rotation, +rotation)

tests/utils/test_scripts.py

@@ -122,6 +122,7 @@ def test_get_imgaug_transform_dlc(cfg):
     assert pipe.__str__().find("Resize") == -1
     assert pipe.__str__().find("Fliplr") == -1
     assert pipe.__str__().find("Flipud") == -1
+    assert pipe.__str__().find("Rot90") == -1
     assert pipe.__str__().find("Affine") != -1
     assert pipe.__str__().find("MotionBlur") != -1
     assert pipe.__str__().find("CoarseDropout") != -1
@@ -142,8 +143,9 @@ def test_get_imgaug_transform_dlc_lr(cfg):
     cfg_tmp.training.imgaug = "dlc-lr"
     pipe = get_imgaug_transform(cfg_tmp)
     assert pipe.__str__().find("Resize") == -1
-    assert pipe.__str__().find("Fliplr") != -1
+    assert pipe.__str__().find("Fliplr") == -1
     assert pipe.__str__().find("Flipud") == -1
+    assert pipe.__str__().find("Rot90(name=UnnamedRot90, parameters=[Choice(a=[0, 2]") != -1
     assert pipe.__str__().find("Affine") != -1
     assert pipe.__str__().find("MotionBlur") != -1
     assert pipe.__str__().find("CoarseDropout") != -1
@@ -164,8 +166,9 @@ def test_get_imgaug_transform_dlc_top_down(cfg):
     cfg_tmp.training.imgaug = "dlc-top-down"
     pipe = get_imgaug_transform(cfg_tmp)
     assert pipe.__str__().find("Resize") == -1
-    assert pipe.__str__().find("Fliplr") != -1
-    assert pipe.__str__().find("Flipud") != -1
+    assert pipe.__str__().find("Fliplr") == -1
+    assert pipe.__str__().find("Flipud") == -1
+    assert pipe.__str__().find("Rot90(name=UnnamedRot90, parameters=[Choice(a=[0, 1, 2, 3]") != -1
     assert pipe.__str__().find("Affine") != -1
     assert pipe.__str__().find("MotionBlur") != -1
     assert pipe.__str__().find("CoarseDropout") != -1