Dev

src/augmentations.py

@@ -1,5 +1,39 @@
+import math
 import random
 
+import numpy as np
+
+import math
+import torch
+
+def circle_intersection(x0, y0, r0, x1, y1, r1):
+    # circle 1: (x0, y0), radius r0
+    # circle 2: (x1, y1), radius r1
+
+    d=math.sqrt((x1-x0)**2 + (y1-y0)**2)
+    
+    # non intersecting
+    if d > r0 + r1 :
+        return None
+    # One circle within other
+    if d < abs(r0-r1):
+        return None
+    # coincident circles
+    if d == 0 and r0 == r1:
+        return None
+    else:
+        a=(r0**2-r1**2+d**2)/(2*d)
+        h=math.sqrt(r0**2-a**2)
+        x2=x0+a*(x1-x0)/d   
+        y2=y0+a*(y1-y0)/d   
+        x3=x2+h*(y1-y0)/d     
+        y3=y2-h*(x1-x0)/d 
+
+        x4=x2-h*(y1-y0)/d
+        y4=y2+h*(x1-x0)/d
+        
+        return (np.array([x3, y3]), np.array([x4, y4]))
+
 
 class MirrorKeypoints:
     def __call__(self, sample):    
@@ -8,4 +42,86 @@ class MirrorKeypoints:
         # flip the keypoints tensor
         sample = 1 -  sample
 
+        return sample
+
+class Z_augmentation:
+
+    def __init__(self, hand_side="left"):
+        self.hand_side = hand_side
+
+    def new_wrist(self, sample, hand_side="left", new_wrist=None):
+        if hand_side == "left":
+            wrist = sample[30:32]
+            shoulder = sample[22:24]
+            elbow = sample[26:28]
+        else:
+            wrist = sample[32:34]
+            shoulder = sample[24:26]
+            elbow = sample[28:30]
+
+        # calculate the length of the shoulder to elbow using math package
+        shoulder_elbow_length = math.sqrt((shoulder[0] - elbow[0])**2 + (shoulder[1] - elbow[1])**2)
+        # calculate the length of the wrist to elbow using math package
+        wrist_elbow_length = math.sqrt((wrist[0] - elbow[0])**2 + (wrist[1] - elbow[1])**2)
+
+        if shoulder_elbow_length == 0 or wrist_elbow_length == 0:
+            return sample, None
+
+        first_time = True
+        new_loc = False
+        while not new_loc:
+
+            if new_wrist is None or not first_time:
+                # get random new wrist point that is not too far from the elbow
+                new_wrist = [random.uniform(elbow[0] - 0.3, elbow[0] + 0.3), random.uniform(elbow[1] - 0.3, elbow[1] + 0.3)]
+
+            # get intersection points of the circles
+            c = circle_intersection(shoulder[0], shoulder[1], shoulder_elbow_length, new_wrist[0], new_wrist[1], wrist_elbow_length)
+            if c is not None:
+                (i1, i2) = c
+                new_loc = True
+            first_time = False
+
+        # get the point that is below the hand
+        if i1[1] > i2[1]:
+            new_elbow = i1
+        else:
+            new_elbow = i2
+        # new_elbow to shape (2,1)
+        new_elbow = np.array(new_elbow)
+        new_wrist = np.array(new_wrist)
+
+        # replace the keypoints in the sample
+        if hand_side == "left":
+            sample[26:28] = new_elbow
+            sample[30:32] = new_wrist
+        else:
+            sample[28:30] = new_elbow
+            sample[32:34] = new_wrist
+        return sample, new_wrist
+    
+    def __call__(self, samples):
+        # transform each sample in the batch
+        t_new = []
+
+        t = samples.numpy()
+        new_wrist = None
+        for t_i in t:
+            # if new_wrist is None:
+            #     new_t, w = self.new_wrist(t_i.reshape(-1), self.hand_side)
+            #     new_wrist = w
+            # else:
+            new_t, _ = self.new_wrist(t_i.reshape(-1), self.hand_side)
+            # reshape back to 2 dimensions
+            t_new.append(new_t.reshape(-1, 2))
+        return torch.tensor(np.array(t_new))
+    
+# augmentation to add little randow noise to the keypoints
+class NoiseAugmentation:
+    def __init__(self, noise=0.05):
+        self.noise = noise
+
+    def __call__(self, sample):
+        # add noise to the keypoints
+        sample = sample + torch.randn(sample.shape) * self.noise
         return sample