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)