def img_warp(img, mask, theta_max=15, offset_max=0, scale_max=1.1, shear_max=0.1):
"""Training data augmentation with random affine transformation"""
# Rotation
vec = np.random.normal(0, 1, 3)
vec /= np.sqrt(np.sum(vec ** 2))
theta = np.random.uniform(- theta_max, theta_max, 1) * pi / 180
R = euler.axangle2mat(vec, theta)
# Scale/zoom
sign = -1 if np.random.random() < 0.5 else 1
Z = np.ones(3) * np.random.uniform(1, scale_max, 1) ** sign
# Translation
c_in = np.array(img.shape[1:]) // 2
offset = np.random.uniform(- offset_max, offset_max, 3)
T = - (c_in).dot((R * Z).T) + c_in + offset
# Shear
S = np.random.uniform(- shear_max, shear_max, 3)
# Compose affine
mat = affines.compose(T, R, Z, S)
# Apply warp
img_warped = np.zeros_like(img)
mask_warped = np.zeros_like(mask)
for i in range(len(img)):
img_warped[i,] = affine_transform(img[i,], mat, order=1) # Trilinear
mask_warped[0,] = affine_transform(mask[0,], mat, order=0) # Nearest neighbor
return img_warped, mask_warped
def get_augmented_pcs(pcs):
batch = pcs.shape[0]
size = pcs.shape[1]
scale = torch.rand(batch) * 0.2 + 0.9 #scale 0.7 to 1.1
scale = scale.cuda() #batch
#print (scale)
T = (torch.rand(batch, 3) - 0.5) * 0.2 #-10cm to 10cm
T = T.cuda() #batch*3
R = []
for i in range(batch):
axis = np.random.rand(3)
axis = axis / np.sqrt(pow(axis, 2).sum())
theta = (np.random.rand() - 0.5) * np.pi * 2
mat = euler.axangle2mat(axis, theta) #3,3
R += [mat]
R = torch.FloatTensor(np.array(R)).cuda() #batch*3*3
new_pcs = torch.einsum('b,bsc->bsc', [scale, pcs])
new_pcs = torch.einsum('bdc,bsc->bsd', [R, new_pcs])
new_pcs = new_pcs + T.view(batch, 1, 3).repeat(1, size, 1)
return new_pcs
def generate_fibonacci_spiral(structure, phi, theta, psi, max_theta, _):
# Vogel's method -> disk. Modify -> sphere surface
zone_to_rotation = np.identity(3)
lattice_to_zone = np.identity(3)
lattice_to_zone = euler2mat(phi, -theta, psi, 'rzxz')
n = 100
golden_angle = np.pi * (3 - np.sqrt(5))
theta = golden_angle * np.arange(n)
z = np.linspace(1, np.cos(max_theta), n)
radius = np.sqrt(1 - z**2)
points = np.zeros((n, 3))
points[:, 0] = radius * np.cos(theta)
points[:, 1] = radius * np.sin(theta)
points[:, 2] = z
rotations = np.empty((n, 3, 3))
for i, point in enumerate(points):
# Simplifications to cos angle formula since one of the directions is (0, 0, 1)
point_angle = math.acos(point[2] / np.linalg.norm(point))
point_axis = np.cross(np.array([0, 0, 1]), point)
if np.count_nonzero(point_axis) == 0:
point_axis = np.array([0, 0, 1])
zone_to_rotation = axangle2mat(point_axis, point_angle)
rotations[i] = lattice_to_zone @ zone_to_rotation
return rotations.reshape(-1, 3, 3)
def augment_train_frame_data(train_frame_data, T, axisR):
hip_index = joint_index['hip']
hip_pos = train_frame_data[hip_index * 3:hip_index * 3 + 3]
for i in range(int(len(train_frame_data) / 3)):
if (i != hip_index):
train_frame_data[i * 3:i * 3 +
3] = train_frame_data[i * 3:i * 3 + 3] + hip_pos
mat_r_augment = euler.axangle2mat(axisR[0:3], axisR[3])
n = int(len(train_frame_data) / 3)
for i in range(n):
raw_data = train_frame_data[i * 3:i * 3 + 3]
new_data = np.dot(mat_r_augment, raw_data) + T
train_frame_data[i * 3:i * 3 + 3] = new_data
hip_pos = train_frame_data[hip_index * 3:hip_index * 3 + 3]
for i in range(int(len(train_frame_data) / 3)):
if (i != hip_index):
train_frame_data[i * 3:i * 3 +
3] = train_frame_data[i * 3:i * 3 + 3] - hip_pos
return train_frame_data
def get_augmented_pc(pc):
size = pc.shape[0]
scale = np.random.rand() * 0.2 + 0.9 #scale 0.9 to 1.1
new_pc = pc * scale
#axis = np.array([0,1,0])
axis = np.random.rand(3)
axis = axis / np.sqrt(pow(axis, 2).sum())
#theta=(np.random.randint(0,30)-15)/180.0*np.pi
theta = (np.random.rand() - 0.5) * np.pi * 2
R = euler.axangle2mat(axis, theta).reshape((1, 3, 3)).repeat(size, 0)
#T = np.random.randn(1,3,1)/50
T = (np.random.rand(1, 3, 1) - 0.5) * 0.2 #-10cm to 10cm
#T[:,1] = T[:,1]*0
T = T.repeat(size, 0)
R = torch.FloatTensor(R).cuda()
T = torch.FloatTensor(T).cuda()
new_pc = torch.matmul(R, new_pc.view(size, 3, 1)) + T
#new_pc = np.matmul(R, new_pc.reshape((size,3,1))) +T
return new_pc.view(size, 3)
def augment_train_data_torch(train_data, T, axisR):
batch = train_data.shape[0]
frame_num = train_data.shape[1]
joint_num = int(train_data.shape[2] / 3)
mat_r = torch.autograd.Variable(
torch.FloatTensor(euler.axangle2mat(axisR[0:3], axisR[3])).cuda())
t = torch.autograd.Variable(torch.FloatTensor(T).cuda())
seq = train_data.view(batch, frame_num, joint_num, 3)
hip_index = joint_index['hip']
hip_seq = seq[:, :, hip_index].clone()
abs_seq = seq + hip_seq.view(batch, frame_num, 1, 3).expand(
-1, -1, joint_num, 3)
abs_seq[:, :, hip_index] = hip_seq
abs_seq = abs_seq.view(batch * frame_num * joint_num, 3, 1)
mat_r_batch = mat_r.unsqueeze(0).expand(batch * frame_num * joint_num, -1,
-1)
rotated_abs_seq = torch.bmm(mat_r_batch,
abs_seq).view(batch, frame_num, joint_num, 3)
final_abs_seq = rotated_abs_seq + t
new_hip_seq = final_abs_seq[:, :, hip_index].clone()
new_seq = final_abs_seq - new_hip_seq.view(batch, frame_num, 1, 3).expand(
-1, -1, joint_num, 3)
new_seq[:, :, hip_index] = new_hip_seq
return new_seq.view(batch, frame_num, -1)
def get_augmented_dance_seq_from_quat_bvh(dance_quat_seq):
#dance_quat_seq = np.load(quat_npy_fn)
dance_quat_seq = dance_quat_seq[Skip_initial_frame_num: dance_quat_seq.shape[0]] #seq_len*(3+joint_num*4)
seq_len=dance_quat_seq.shape[0]
dance_quat_seq_cuda = torch.autograd.Variable(torch.FloatTensor(dance_quat_seq).cuda()) #seq_len*(3+joint_num*4)
#augment the rotation
hip_quat_seq_cuda = dance_quat_seq_cuda[:, 3:7]
hip_matrix_seq_cuda = tools.compute_rotation_matrix_from_quaternion(hip_quat_seq_cuda)#seq_len*3*3
##generate a random rotation matrix
axisR= [0,1,0,0]#[0,1,0,np.random.randint(0,360)/180.0*np.pi]
mat_r=torch.autograd.Variable(torch.FloatTensor(euler.axangle2mat(axisR[0:3], axisR[3])).cuda())#3*3
mat_r = mat_r.view(1, 3,3).repeat(hip_matrix_seq_cuda.shape[0],1,1) #seq_len*3*3
new_mat_r = torch.matmul(mat_r,hip_matrix_seq_cuda,) #seq_len*3*3
##get the quaternion seq
new_hip_quat_seq_cuda = tools.compute_quaternions_from_rotation_matrices(new_mat_r) #seq_len*4
new_hip_pose_seq_cuda = torch.matmul(mat_r, dance_quat_seq_cuda[:,0:3].unsqueeze(2)).view(seq_len,3)
new_hip_pose_seq_cuda[:,0]=new_hip_pose_seq_cuda[:,0]+np.random.randn(1)[0]
new_hip_pose_seq_cuda[:,2]=new_hip_pose_seq_cuda[:,2]+np.random.randn(1)[0]
new_dance_quat_seq_cuda = dance_quat_seq_cuda#seq_len*(3+joint_num*4)
new_dance_quat_seq_cuda[:,0:3]=new_hip_pose_seq_cuda #fix the hip global pose to be at 0
new_dance_quat_seq_cuda[:,3:7]=new_hip_quat_seq_cuda
return new_dance_quat_seq_cuda
def augment_train_frame_data(train_frame_data, T, axisR):
#train_Frame_Data=(102,171)
print('before train_frame_data:', train_frame_data.shape)
mat_r_augment = euler.axangle2mat(axisR[0:3], axisR[3])
print(train_frame_data.shape[0])
n = int(train_frame_data[0] / 3)
for i in range(n):
raw_data = train_frame_data[i * 3:i * 3 + 3]
new_data = np.dot(mat_r_augment, raw_data) + T
train_frame_data[i * 3:i * 3 + 3] = new_data
return train_frame_data
def get_augmented_pc_ret(pc):
size = pc.shape[0]
new_pc = pc
axis = np.random.rand(3)
axis = axis / np.sqrt(pow(axis, 2).sum())
theta = (np.random.rand() - 0.5) * np.pi * 2
Rorgmat = euler.axangle2mat(axis, theta)
R = Rorgmat.reshape((1, 3, 3)).repeat(size, 0)
Torg = (np.random.rand(1, 3, 1) - 0.5) * 0.2 #-10cm to 10cm
T = Torg.repeat(size, 0)
new_pc = np.matmul(R, new_pc.reshape((size, 3, 1))) + T
return new_pc.reshape((size, 3)), Rorgmat, Torg[..., 0]
def rotmat_from_vecs(v1, v2=np.asarray([0, 0, 1])):
"""
Returns a rotation matrix R_1_2
:param v1: vector in frame 1
:param v2: vector in frame 2
:return:
"""
v1 = v1 / np.linalg.norm(v1)
v2 = v2 / np.linalg.norm(v2)
v = np.cross(v2, v1)
vx = np.asarray([[0, -v[2], +v[1], 0], [+v[2], 0, -v[0], 0],
[-v[1], +v[0], 0, 0], [0, 0, 0, 0]])
dotp = np.dot(v1, v2)
if np.abs(dotp + 1) < 1e-3:
R = np.eye(4)
x = np.cross(v2, [1, 0, 0])
R[:3, :3] = txe.axangle2mat(x, np.pi)
else:
R = np.eye(4) + vx + np.dot(vx, vx) / (1 + dotp)
return R
def get_augmented_gt_seq_for_training(dances_lst,
target_framerate,
min_seq_len=16,
max_seq_len=128):
dance_id = np.random.randint(0, len(dances_lst))
dance_quat_seq, frame_rate = dances_lst[dance_id]
dance_quat_seq = dance_quat_seq[Skip_initial_frame_num:dance_quat_seq.
shape[0]]
random_length = np.random.randint(
min_seq_len, min(max_seq_len, dance_quat_seq.shape[0]))
random_start_id = np.random.randint(
0, dance_quat_seq.shape[0] - random_length)
dance_quat_seq = dance_quat_seq[random_start_id:random_start_id +
random_length]
dance_quat_seq = dance_quat_seq[0:int(dance_quat_seq.shape[0] -
dance_quat_seq.shape[0] % 16)]
stride = int(frame_rate / target_framerate)
dance_quat_seq = dance_quat_seq.reshape(
int(dance_quat_seq.shape[0] / stride), stride,
dance_quat_seq.shape[1])[:, np.random.randint(0, stride)]
seq_len = dance_quat_seq.shape[0]
dance_quat_seq_cuda = torch.autograd.Variable(
torch.FloatTensor(dance_quat_seq).cuda()) #seq_len*(3+joint_num*4)
#augment the rotation
hip_quat_seq_cuda = dance_quat_seq_cuda[:, 3:7]
hip_matrix_seq_cuda = tools.compute_rotation_matrix_from_quaternion(
hip_quat_seq_cuda) #seq_len*3*3
##generate a random rotation matrix
axisR = [0, 1, 0, np.random.randint(0, 360) / 180 * np.pi]
mat_r = torch.autograd.Variable(
torch.FloatTensor(euler.axangle2mat(axisR[0:3],
axisR[3])).cuda()) #3*3
mat_r = mat_r.view(1, 3, 3).repeat(hip_matrix_seq_cuda.shape[0], 1,
1) #seq_len*3*3
new_mat_r = torch.matmul(
mat_r,
hip_matrix_seq_cuda,
) #seq_len*3*3
##get the quaternion seq
new_hip_quat_seq_cuda = tools.compute_quaternions_from_rotation_matrices(
new_mat_r) #seq_len*4
new_hip_pose_seq_cuda = torch.matmul(
mat_r, dance_quat_seq_cuda[:, 0:3].unsqueeze(2)).view(seq_len, 3)
new_hip_pose_seq_cuda[:,
0] = new_hip_pose_seq_cuda[:,
0] + np.random.randn(1)[0]
new_hip_pose_seq_cuda[:,
2] = new_hip_pose_seq_cuda[:,
2] + np.random.randn(1)[0]
new_dance_quat_seq_cuda = dance_quat_seq_cuda #seq_len*(3+joint_num*4)
new_dance_quat_seq_cuda[:, 0:
3] = new_hip_pose_seq_cuda #fix the hip global pose to be at 0
new_dance_quat_seq_cuda[:, 3:7] = new_hip_quat_seq_cuda
return new_dance_quat_seq_cuda
def generate_complete_rotation_list(structure, corner_a, corner_b, corner_c,
inplane_rotations, resolution):
"""Generate a rotation list covering the inverse pole figure specified by three
corners in cartesian coordinates.
Arguments:
structure: diffpy.structure.Structure, used for calculating angles
corner_a, corner_b, corner_c: The three corners of the inverse pole
figure given by three coordinates in the coordinate system
specified by the structure lattice (crystal directions).
resolution: Angular resolution in radians of the generated rotation
list.
inplane_rotations: List of angles in radians for in-plane rotation
of the diffraction pattern. This corresponds to the third Euler
angle rotation. The rotation list will be generated for each of
these angles, and combined. This should be done automatically, but
by including all possible rotations in the rotation list, it
becomes too large.
Returns:
Rotations covering the inverse pole figure given as a `np.array` of Euler
angles in degress. This `np.array` can be passed directly to pyxem.
"""
# Convert the crystal directions to cartesian vectors, and then normalize
# to get the directions.
if len(corner_a) == 4:
corner_a = uvtw_to_uvw(*corner_a)
if len(corner_b) == 4:
corner_b = uvtw_to_uvw(*corner_b)
if len(corner_c) == 4:
corner_c = uvtw_to_uvw(*corner_c)
corner_a = direction_to_cartesian(structure, corner_a)
corner_b = direction_to_cartesian(structure, corner_b)
corner_c = direction_to_cartesian(structure, corner_c)
corner_a /= np.linalg.norm(corner_a)
corner_b /= np.linalg.norm(corner_b)
corner_c /= np.linalg.norm(corner_c)
angle_a_to_b = angle_between_cartesian(corner_a, corner_b)
angle_a_to_c = angle_between_cartesian(corner_a, corner_c)
angle_b_to_c = angle_between_cartesian(corner_b, corner_c)
axis_a_to_b = np.cross(corner_a, corner_b)
axis_a_to_c = np.cross(corner_a, corner_c)
# Input validation. The corners have to define a non-degenerate triangle
if np.count_nonzero(axis_a_to_b) == 0:
raise ValueError('Directions a and b are parallel')
if np.count_nonzero(axis_a_to_c) == 0:
raise ValueError('Directions a and c are parallel')
# Find the maxiumum number of points we can generate, given by the
# resolution, then allocate storage for them. For the theta direction,
# ensure that we keep the resolution also along the direction to the corner
# b or c farthest away from a.
theta_count = math.ceil(max(angle_a_to_b, angle_a_to_c) / resolution)
phi_count = math.ceil(angle_b_to_c / resolution)
inplane_rotation_count = len(inplane_rotations)
rotations = np.zeros(
(theta_count, phi_count, inplane_rotation_count, 3, 3))
# Generate a list of theta_count evenly spaced angles theta_b in the range
# [0, angle_a_to_b] and an equally long list of evenly spaced angles
# theta_c in the range[0, angle_a_to_c].
for i, (theta_b, theta_c) in enumerate(
zip(np.linspace(0, angle_a_to_b, theta_count),
np.linspace(0, angle_a_to_c, theta_count))):
# Define the corner local_b at a rotation theta_b from corner_a toward
# corner_b on the circle surface. Similarly, define the corner local_c
# at a rotation theta_c from corner_a toward corner_c.
rotation_a_to_b = axangle2mat(axis_a_to_b, theta_b)
rotation_a_to_c = axangle2mat(axis_a_to_c, theta_c)
local_b = np.dot(rotation_a_to_b, corner_a)
local_c = np.dot(rotation_a_to_c, corner_a)
# Then define an axis and a maximum rotation to create a great cicle
# arc between local_b and local_c. Ensure that this is not a degenerate
# case where local_b and local_c are coincident.
angle_local_b_to_c = angle_between_cartesian(local_b, local_c)
axis_local_b_to_c = np.cross(local_b, local_c)
if np.count_nonzero(axis_local_b_to_c) == 0:
# Theta rotation ended at the same position. First position, might
# be other cases?
axis_local_b_to_c = corner_a
axis_local_b_to_c /= np.linalg.norm(axis_local_b_to_c)
# Generate points along the great circle arc with a distance defined by
# resolution.
phi_count_local = math.ceil(angle_local_b_to_c / resolution)
for j, phi in enumerate(
np.linspace(0, angle_local_b_to_c, phi_count_local)):
rotation_phi = axangle2mat(axis_local_b_to_c, phi)
for k, psi in enumerate(inplane_rotations):
rotation_psi = axangle2mat((0, 0, 1), psi)
# Combine the rotations. Order is important. The matrix is
# applied from the left, and we rotate by theta first toward
# local_b, then across the triangle toward local_c, and finally
# an inplane rotation psi
rotations[i, j,
k] = rotation_psi @ (rotation_phi @ rotation_a_to_b)
return rotations
def _axangle2mat_(ax, a):
a = numpy.deg2rad(a)
return euler.axangle2mat(ax, a)
def generate_complete_rotation_list(structure, corner_a, corner_b, corner_c,
resolution):
# Start defining some angles and normals from the given corners
angle_a_to_b = angle_between_directions(structure, corner_a, corner_b)
angle_a_to_c = angle_between_directions(structure, corner_a, corner_c)
angle_b_to_c = angle_between_directions(structure, corner_b, corner_c)
axis_a_to_b = np.cross(corner_a, corner_b)
axis_a_to_c = np.cross(corner_a, corner_c)
# Input validation. The corners have to define a non-degenerate triangle
if np.count_nonzero(axis_a_to_b) == 0:
raise ValueError("Directions a and b are parallel")
if np.count_nonzero(axis_a_to_c) == 0:
raise ValueError("Directions a and c are parallel")
# Find the maxiumum number of points we can generate, given by the
# resolution, then allocate storage for them. For the theta direction,
# ensure that we keep the resolution also along the direction to the corner
# b or c farthest away from a.
theta_count = math.ceil(max(angle_a_to_b, angle_a_to_c) / resolution)
psi_count = math.ceil(angle_b_to_c / resolution)
rotations = np.zeros((theta_count, psi_count, 3, 3))
local_cs = []
# For each theta_count angle theta, evenly spaced
for i, (theta_b, theta_c) in enumerate(
zip(np.linspace(0, angle_a_to_b, theta_count),
np.linspace(0, angle_a_to_c, theta_count))):
# Define the corner local_b at a rotation theta from corner_a toward
# corner_b on the circle surface
# Similarly, define the corner local_c at a rotation theta from
# corner_a toward corner_c
# The rotation angle is negated, since we rotate the observer instead
# of the structure itself.
rotation_a_to_b = axangle2mat(axis_a_to_b, -theta_b)
rotation_a_to_c = axangle2mat(axis_a_to_c, -theta_c)
local_b = np.dot(corner_a, rotation_a_to_b)
local_c = np.dot(corner_a, rotation_a_to_c)
local_cs.append(local_c)
# Then define an axis and a maximum rotation to create a great cicle
# arc between local_b and local_c. Ensure that this is not a degenerate
# case where local_b and local_c are coincident.
angle_local_b_to_c = angle_between_directions(structure_zb, local_b,
local_c)
axis_local_b_to_c = np.cross(local_b, local_c)
if np.count_nonzero(axis_local_b_to_c) == 0:
# Theta rotation ended at the same position. First position, might
# be other cases?
axis_local_b_to_c = corner_a
axis_local_b_to_c /= np.linalg.norm(axis_local_b_to_c)
# Generate points along the great circle arc with a distance defined by
# resolution.
psi_count_local = math.ceil(angle_local_b_to_c / resolution)
for j, psi in enumerate(
np.linspace(0, angle_local_b_to_c, psi_count_local)):
# The angle psi is negated as for theta, then the two rotations are
# combined to give the final rotation matrix.
rotation_psi = axangle2mat(axis_local_b_to_c, -psi)
rotations[i, j] = np.matmul(rotation_a_to_b, rotation_psi)
# We also remove duplicates before returning. This eliminates the unused rotations.
return np.unique(rotations.reshape(-1, 3, 3), axis=0)
