Ensure Keypoint Visibility Is Updated After Spatial Transforms (#346)

Author: JSabadin

luxonis_ml/data/__main__.py

@@ -1,4 +1,3 @@
-import random
 import shutil
 from collections.abc import Iterator
 from pathlib import Path
@@ -334,9 +333,6 @@ def inspect(
 ):
     """Inspects images and annotations in a dataset."""
     check_exists(name, bucket_storage)
-    if deterministic:
-        np.random.seed(42)
-        random.seed(42)
 
     view = view or ["train"]
     dataset = LuxonisDataset(name, bucket_storage=bucket_storage)
@@ -349,7 +345,12 @@ def inspect(
     if aug_config is not None:
         h, w, _ = loader[0][0].shape
         loader.augmentations = loader._init_augmentations(
-            "albumentations", aug_config, h, w, not ignore_aspect_ratio
+            augmentation_engine="albumentations",
+            augmentation_config=aug_config,
+            height=h,
+            width=w,
+            keep_aspect_ratio=not ignore_aspect_ratio,
+            seed=42 if deterministic else None,
         )
 
     if len(dataset) == 0:

luxonis_ml/data/augmentations/albumentations_engine.py

@@ -6,6 +6,7 @@
 
 import albumentations as A
 import numpy as np
+from albumentations.core.composition import TransformsSeqType
 from loguru import logger
 from typing_extensions import override
 
@@ -396,6 +397,20 @@ def __init__(
                     f"Only subclasses of `A.BasicTransform` are allowed. "
                 )
 
+        wrapped_spatial_ops: TransformsSeqType = []
+        if "keypoints" in targets.values():
+            for op in spatial_transforms:
+                wrapped_spatial_ops.append(op)
+                wrapped_spatial_ops.append(
+                    A.Lambda(
+                        image=lambda img, **kw: img,
+                        keypoints=self._mark_invisible_keypoints,
+                        p=1.0,
+                    )
+                )
+        else:
+            wrapped_spatial_ops = spatial_transforms
+
         if resize_transform is None:
             if keep_aspect_ratio:
                 resize_transform = LetterboxResize(height=height, width=width)
@@ -424,7 +439,7 @@ def get_params(is_custom: bool = False) -> dict[str, Any]:
                 batch_transforms, **get_params(is_custom=True)
             )
             self.spatial_transform = wrap_transform(
-                A.Compose(spatial_transforms, **get_params())
+                A.Compose(wrapped_spatial_ops, **get_params())
             )
             self.pixel_transform = wrap_transform(
                 A.Compose(pixel_transforms), is_pixel=True
@@ -625,6 +640,26 @@ def postprocess(
 
         return out_image, out_labels
 
+    @staticmethod
+    def _mark_invisible_keypoints(
+        keypoints: np.ndarray, **kwargs
+    ) -> np.ndarray:
+        """
+        keypoints: np.ndarray of shape (N,6) columns = [x, y, z, a, s, v]
+        Zeroes out the visibility (last) column if (x,y) is out of image bounds.
+        """
+        shape = kwargs.get("shape")
+        if shape is None:
+            raise ValueError(
+                "Shape must be provided in kwargs to mark invisible keypoints."
+            )
+        h, w = shape[:2]
+        kps = keypoints.copy()
+        xs, ys = kps[:, 0], kps[:, 1]
+        oob = (xs < 0) | (ys < 0) | (xs >= w) | (ys >= h)
+        kps[oob, -1] = 0.0
+        return kps
+
     @staticmethod
     def create_transformation(
         config: AlbumentationConfigItem,

tests/test_data/test_augmentations/test_special.py

@@ -77,5 +77,10 @@ def test_skip_augmentations():
         t.__class__.__name__ for t in augmentations.batch_transform.transforms
     ]
 
-    assert spatial_transform_names == ["Perspective", "Rotate"]
+    assert spatial_transform_names == [
+        "Perspective",
+        "Lambda",
+        "Rotate",
+        "Lambda",
+    ]
     assert batched_transform_names == ["Mosaic4"]

tests/test_data/test_export.py

@@ -41,6 +41,9 @@ def test_dir_parser(
         .to_dict(as_series=False)
     )
     anns = {k: sorted(v) for k, v in anns.items()}
+    splits = dataset.get_splits()
+    assert splits is not None
+    splits = {split: sorted(files) for split, files in splits.items()}
 
     zip_result = dataset.export(tempdir / "exported")
     zip_path = zip_result[0] if isinstance(zip_result, list) else zip_result
@@ -67,6 +70,12 @@ def test_dir_parser(
         .to_dict(as_series=False)
     )
     imported_anns = {k: sorted(v) for k, v in imported_anns.items()}
+    imported_splits = exported_dataset.get_splits()
+    assert imported_splits is not None
+    imported_splits = {
+        split: sorted(files) for split, files in imported_splits.items()
+    }
+    assert imported_splits == splits
     del imported_metadata["tasks"]
     del imported_metadata["skeletons"]
     assert imported_metadata == metadata

tests/test_data/test_loader.py

@@ -545,7 +545,7 @@ def round_nested_list(
     new_aug_annotations = [convert_annotation(ann) for _, ann in loader_aug]
 
     original_aug_annotations = load_annotations(
-        "test_augmentation_reproducibility_labels.json"
+        "test_augmentation_reproducibility_annotations.json"
     )
 
     for orig_ann, new_ann in zip(