139 lines
4.3 KiB
Python
139 lines
4.3 KiB
Python
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):
|
|
if random.random() > 0.5:
|
|
return sample
|
|
# 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
|
|
|
|
# augmentation to rotate all keypoints around 0,0
|
|
class RotateAugmentation:
|
|
def __call__(self, sample):
|
|
# generate a random angle between -13 and 13 degrees
|
|
angle_max = 13.0
|
|
angle = math.radians(random.uniform(a=-angle_max, b=angle_max))
|
|
# rotate the keypoints around 0.0
|
|
new_sample = sample
|
|
new_sample[:, :, 0] = sample[:, :, 0]*math.cos(angle) - sample[:, :, 1]*math.sin(angle)
|
|
new_sample[:, :, 1] = sample[:, :, 0]*math.sin(angle) + sample[:, :, 1]*math.cos(angle)
|
|
return new_sample |