Resolve WES-155 "Mirror augmentation"

src/augmentations.py

@@ -36,11 +36,28 @@ def circle_intersection(x0, y0, r0, x1, y1, r1):
 
 
 class MirrorKeypoints:
-    def __call__(self, sample):    
+    def __call__(self, sample):
+        if sample.shape[0] == 0:
+            return sample    
         if random.random() > 0.5:
             return sample
-        # flip the keypoints tensor
-        sample = 1 -  sample
+        
+        # flip the x coordinates
+        sample[:, :, 0] *= -1
+
+        # switch hands (left becomes right and vice versa)
+        left, right, n = 12, 33, 21
+        if isinstance(sample, np.ndarray): # For testing purposes only
+            sample[:, left:left+n, :], sample[:, right:right+n, :] = sample[: , right:right+n, :], sample[:, left:left+n, :].copy()
+        else:
+            sample[:, left:left+n, :], sample[:, right:right+n, :] = sample[: , right:right+n, :], sample[:, left:left+n, :].clone()
+
+        # switch pose keypoints
+        sample[:, [1, 2], :] = sample[:, [2, 1], :] #eye
+        sample[:, [3, 4], :] = sample[:, [4, 3], :] #ear
+        sample[:, [6, 7], :] = sample[:, [7, 6], :] #shoulder
+        sample[:, [8, 9], :] = sample[:, [9, 8], :] #elbow
+        sample[:, [10, 11], :] = sample[:, [11, 10], :] #wrist
 
         return sample