def get_shape_chamfer_loss(self, pts, quat1, quat2, center1, center2, part_cnt): num_point = pts.shape[1] part_pcs1 = qrot(quat1.unsqueeze(1).repeat(1, num_point, 1), pts) + center1.unsqueeze(1).repeat(1, num_point, 1) part_pcs2 = qrot(quat2.unsqueeze(1).repeat(1, num_point, 1), pts) + center2.unsqueeze(1).repeat(1, num_point, 1) t = 0 shape_pcs1 = [] shape_pcs2 = [] for cnt in part_cnt: cur_shape_pc1 = part_pcs1[t:t + cnt].view(1, -1, 3) cur_shape_pc2 = part_pcs2[t:t + cnt].view(1, -1, 3) with torch.no_grad(): idx1 = furthest_point_sample(cur_shape_pc1, 2048).long()[0] idx2 = furthest_point_sample(cur_shape_pc2, 2048).long()[0] shape_pcs1.append(cur_shape_pc1[:, idx1]) shape_pcs2.append(cur_shape_pc2[:, idx2]) t += cnt shape_pcs1 = torch.cat(shape_pcs1, dim=0) # numshapes x 2048 x 3 shape_pcs2 = torch.cat(shape_pcs2, dim=0) # numshapes x 2048 x 3 dist1, dist2 = chamfer_distance(shape_pcs1, shape_pcs2, transpose=False) loss_per_data = torch.mean(dist1, dim=1) + torch.mean(dist2, dim=1) return loss_per_data
def get_quat_loss(self, pts, quat1, quat2): num_point = pts.shape[1] pts1 = qrot(quat1.unsqueeze(1).repeat(1, num_point, 1), pts) pts2 = qrot(quat2.unsqueeze(1).repeat(1, num_point, 1), pts) dist1, dist2 = chamfer_distance(pts1, pts2, transpose=False) loss_per_data = torch.mean(dist1, dim=1) + torch.mean(dist2, dim=1) return loss_per_data
def get_l2_rotation_loss(self, pts, quat1, quat2): num_point = pts.shape[1] pts1 = qrot(quat1.unsqueeze(1).repeat(1, num_point, 1), pts) pts2 = qrot(quat2.unsqueeze(1).repeat(1, num_point, 1), pts) loss_per_data = (pts1 - pts2).pow(2).sum(dim=2).mean(dim=1) return loss_per_data
def forward_kinematics(self, rotations, root_positions): """ Perform forward kinematics using the given trajectory and local rotations. Arguments (where N = batch size, L = sequence length, J = number of joints): -- rotations: (N, L, J, 4) tensor of unit quaternions describing the local rotations of each joint. -- root_positions: (N, L, 3) tensor describing the root joint positions. """ assert len(rotations.shape) == 4 assert rotations.shape[-1] == 4 positions_world = [] rotations_world = [] expanded_offsets = self._offsets.expand(rotations.shape[0], rotations.shape[1], self._offsets.shape[0], self._offsets.shape[1]) # Parallelize along the batch and time dimensions for i in range(self._offsets.shape[0]): if self._parents[i] == -1: positions_world.append(root_positions) rotations_world.append(rotations[:, :, 0]) else: positions_world.append(qrot(rotations_world[self._parents[i]], expanded_offsets[:, :, i]) \ + positions_world[self._parents[i]]) if self._has_children[i]: rotations_world.append( qmul(rotations_world[self._parents[i]], rotations[:, :, i])) else: # This joint is a terminal node -> it would be useless to compute the transformation rotations_world.append(None) return torch.stack(positions_world, dim=3).permute(0, 1, 3, 2)
def render_boxes(out_fn, boxes): boxes = torch.from_numpy(boxes) cmd = 'cp %s %s' % (os.path.join(BASE_DIR, 'cube.mtl'), out_fn + '.mtl') call(cmd, shell=True) with open(out_fn + '.obj', 'w') as fout: fout.write('mtllib %s\n' % (out_fn.split('/')[-1] + '.mtl')) for j in range(boxes.shape[0]): v = (qrot(boxes[j, 6:].unsqueeze(dim=0).repeat(8, 1), cube_v_torch * boxes[j, 3:6].unsqueeze(dim=0)) + boxes[j, :3].unsqueeze(dim=0)).numpy() for i in range(8): fout.write('v %f %f %f\n' % (v[i, 0], v[i, 1], v[i, 2])) for i in range(6): fout.write('usemtl f%d\n' % i) fout.write('f %d %d %d\n' % (cube_f[2 * i, 0] + 8 * j, cube_f[2 * i, 1] + 8 * j, cube_f[2 * i, 2] + 8 * j)) fout.write( 'f %d %d %d\n' % (cube_f[2 * i + 1, 0] + 8 * j, cube_f[2 * i + 1, 1] + 8 * j, cube_f[2 * i + 1, 2] + 8 * j)) cmd = 'cd %s && blender -noaudio --background camera_centered.blend --python render_blender.py %s %s > /dev/null' \ % (os.path.join(os.path.dirname(os.path.abspath(__file__))), \ out_fn+'.obj', out_fn) call(cmd, shell=True)
def transform_pc_batch(pc, feat, anchor=False): batch_size = feat.size(0) num_point = pc.size(0) pc = pc.repeat(batch_size, 1, 1) center = feat[:, :3].unsqueeze(dim=1).repeat(1, num_point, 1) shape = feat[:, 3:6].unsqueeze(dim=1).repeat(1, num_point, 1) quat = feat[:, 6:].unsqueeze(dim=1).repeat(1, num_point, 1) if not anchor: pc = pc * shape pc = qrot(quat.view(-1, 4), pc.view(-1, 3)).view(batch_size, num_point, 3) if not anchor: pc = pc + center return pc
def sensorFromWorld(self, pts_world): X1 = pts_world quat = torch.cat((torch.ones(1), self.R1)) quat = quat / quat.norm() quat = quat.expand(X1.shape[0], 4) X1 = quaternion.qrot(quat, X1) #X1_p = torch.cat((X1,torch.zeros(X1.shape[0],1)),dim=1) #X2 = quaternion.qmul(quat,X1_p) #qp = quat * torch.tensor([1.0,-1.0,-1.0,-1.0]) #X3 = quaternion.qmul(X2,qp) #print pts_world, X1, quat X1.add_(self.T1) ab1 = X1.div(X1[:, 2].view(X1.shape[0], 1)) ab1 = ab1[:, :2] r1 = ab1.mul(ab1) r1 = r1.sum(dim=1, keepdim=True) r1.sqrt_() theta1 = r1.atan() theta1_sq = theta1.mul(theta1) theta1_d = torch.zeros(theta1.shape) theta1_d.add_(self.K1[3]) theta1_d.mul_(theta1_sq) theta1_d.add_(self.K1[2]) theta1_d.mul_(theta1_sq) theta1_d.add_(self.K1[1]) theta1_d.mul_(theta1_sq) theta1_d.add_(self.K1[0]) theta1_d.mul_(theta1_sq) theta1_d.add_(1) theta1_d.mul_(theta1) #theta1_scale = theta1_d.div(r1) theta1_scale = theta1_d.div(theta1) x1_p = ab1.mul(theta1_scale) uv_1 = x1_p.mul(self.F1) uv_1.add_(self.C1) return uv_1
def get_adj_loss(self, pts, quat, center, adjs): num_point = pts.shape[1] part_pcs = qrot(quat.unsqueeze(1).repeat(1, num_point, 1), pts) + center.unsqueeze(1).repeat(1, num_point, 1) loss = [] for cur_shape_adj in adjs: cur_shape_loss = [] for adj in cur_shape_adj: idx1, idx2 = adj dist1, dist2 = chamfer_distance(part_pcs[idx1].unsqueeze(0), part_pcs[idx2].unsqueeze(0), transpose=False) cur_loss = torch.min(dist1, dim=1)[0] + torch.min(dist2, dim=1)[0] cur_shape_loss.append(cur_loss) loss.append(torch.stack(cur_shape_loss).mean()) return loss
def other_forward(t, r, t0_, r0_, t1_, r1_): t0 = t0_.view(-1, 3) r0 = r0_.view(-1, 4) t1 = t1_.view(-1, 3) r1 = r1_.view(-1, 4) r0_inv = qinv(r0) t_diff = qrot(r0_inv, (t1 - t0)) r_diff = qmul(r0_inv.clone(), r1) self_t = t.repeat(t0.shape[0], 1) self_r = r.repeat(t0.shape[0], 1) t_loss = t_diff - self_t self_r_inv = qinv(self_r) r_loss = qmul(self_r_inv, r_diff) norm_r_loss = r_loss / qlen(r_loss) cut_norm_r_loss = norm_r_loss[:, 1:] return torch.cat([t_loss, cut_norm_r_loss], dim=1)
def linear_assignment(self, mask1, mask2, similar_cnt, pts, centers1, quats1, centers2, quats2): ''' mask1, mask 2: # part_cnt x 224 x 224 similar cnt # shape: max_mask x 2 # first index is the index of parts without similar parts 1, 2, 3, 4, 4, 5, 5, 6, 6, 6, 6, .... # second index is the number of similar part count: 1, 1, 1, 2, 2, 2, 2, 4, 4, 4, 4, .... ''' bids = [] ids1 = [] ids2 = [] inds1 = [] inds2 = [] img_size = mask1.shape[-1] t = 0 num_point = pts.shape[1] max_num_part = centers1.shape[1] # part_cnt = [item for sublist in part_cnt for item in sublist] with torch.no_grad(): while t < similar_cnt.shape[0]: cnt = similar_cnt[t].item() bids = [t] * cnt cur_mask1 = mask1[t:t + cnt].unsqueeze(1).repeat( 1, cnt, 1, 1).view(-1, img_size, img_size) cur_mask2 = mask2[t:t + cnt].unsqueeze(0).repeat( cnt, 1, 1, 1).view(-1, img_size, img_size) dist_mat_mask = self.get_mask_loss(cur_mask1, cur_mask2).view(cnt, cnt) dist_mat_mask = torch.clamp(dist_mat_mask, max=-0.1) cur_pts = pts[t:t + cnt] cur_quats1 = quats1[t:t + cnt].unsqueeze(1).repeat( 1, num_point, 1) cur_centers1 = centers1[t:t + cnt].unsqueeze(1).repeat( 1, num_point, 1) cur_pts1 = qrot(cur_quats1, cur_pts) + cur_centers1 cur_quats2 = quats2[t:t + cnt].unsqueeze(1).repeat( 1, num_point, 1) cur_centers2 = centers2[t:t + cnt].unsqueeze(1).repeat( 1, num_point, 1) cur_pts2 = qrot(cur_quats2, cur_pts) + cur_centers2 cur_pts1 = cur_pts1.unsqueeze(1).repeat(1, cnt, 1, 1).view( -1, num_point, 3) cur_pts2 = cur_pts2.unsqueeze(0).repeat(cnt, 1, 1, 1).view( -1, num_point, 3) dist1, dist2 = chamfer_distance(cur_pts1, cur_pts2, transpose=False) dist_mat_pts = (dist1.mean(1) + dist2.mean(1)).view(cnt, cnt) dist_mat_pts = torch.clamp(dist_mat_pts, max=1) * 0.1 dist_mat = torch.add(dist_mat_mask, dist_mat_pts) t += cnt rind, cind = linear_sum_assignment(dist_mat.cpu().numpy()) ids1 = list(rind) ids2 = list(cind) inds1 += [bids[i] + ids1[i] for i in range(len(ids1))] inds2 += [bids[i] + ids2[i] for i in range(len(ids2))] return inds1, inds2
def camera_to_world(X, R, t): Q = np.tile(R, (*X.shape[:-1], 1)) V = X return qrot(Q, V) + t
def world_to_camera(X, R, t): Rt = qinverse(R) Q = np.tile(Rt, (*X.shape[:-1], 1)) V = X - t return qrot(Q, V)
a_hist = np.zeros((max_iteration, 3)) p_hist[0,0] = p_c[0,0] p_hist[0,1] = p_c[1,0] p_hist[0,2] = p_c[2,0] a_hist[0,0] = pitch_c a_hist[0,1] = yaw_c a_hist[0,2] = roll_c #qr = qn.angle_to_q(roll_c, pitch_c, yaw_c) qr = qn.angle_to_q(roll, pitch, yaw) qa = np.array([0, 1, 0, 0]) qb = qn.qrot(qr, qa) Rq = qn.q_to_r(qr) print("original R") print(R) print("quaternion to R") print(Rq) for i in range(max_iteration): d_pitch = p_gain * (pitch_t - pitch_c) d_yaw = p_gain * (yaw_t - yaw_c) d_roll = p_gain * (roll_t - roll_c) #pitch_c += d_pitch #yaw_c += d_yaw
for j in range(cnt): cur_part_id = cur_parts[j] cur_part_pose = np.hstack([ matched_pred_center2_all[t + j].cpu().numpy(), matched_pred_quat2_all[t + j].cpu().numpy() ]) to_save[cur_part_id] = cur_part_pose fn = input_shape_id[t] + '-' + input_view_id[t] + '.npy' np.save(os.path.join(test_res_matched_dir, fn), to_save) print('saved', fn) t += cnt # get metric stats pred_pts = qrot( matched_pred_quat2_all.unsqueeze(1).repeat(1, num_point, 1), input_pts) + matched_pred_center2_all.unsqueeze(1).repeat( 1, num_point, 1) gt_pts = qrot( matched_gt_quat_all.unsqueeze(1).repeat(1, num_point, 1), input_pts) + matched_gt_center_all.unsqueeze(1).repeat( 1, num_point, 1) dist1, dist2 = chamfer_distance(gt_pts, pred_pts, transpose=False, sqrt=True) dist = torch.mean(dist1, dim=1).cpu().numpy() + torch.mean( dist2, dim=1).cpu().numpy() dist /= 2.0 batch_size = input_box_size.shape[0] correct = 0
def __getitem__(self, index): shape_id = self.data[index] cur_data_fn = os.path.join(self.data_dir, self.category, '%s.npy' % (shape_id)) cur_data = np.load(cur_data_fn, allow_pickle=True).item() view_id = self.convert_view(np.random.choice(np.arange(24))) if self.data_split == "val" or self.data_split == "vis" or self.data_split == "test": view_id = "00" cur_data_view_fn = os.path.join(self.data_dir, self.category, shape_id, '%s.npy' % (view_id)) cur_data_view = np.load(cur_data_view_fn, allow_pickle=True).item() part_id = np.random.choice(cur_data['all_parts']) equiv_edge_indices = [] # sort part order to put the single part in front and similar parts in the back single_parts = [ x for x in cur_data['all_parts'] if len(cur_data[x]['similar']) == 1 ] # similar_part_inds = [i+1 for i in range(len(single_parts))] other_parts = [ x for x in cur_data['all_parts'] if len(cur_data[x]['similar']) > 1 ] part_order = sorted(single_parts) t = len(part_order) for part in other_parts: if part not in part_order: cur_cnt = len(cur_data[part]['similar']) part_order += sorted(cur_data[part]['similar']) for t1 in range(t, t + cur_cnt): for t2 in range(t, t + cur_cnt): if t1 != t2: equiv_edge_indices.append([t1, t2]) t += cur_cnt if len(part_order) != len(cur_data['all_parts']): print('check data ', shape_id, view_id) data_feats = () for feat in self.data_features: if feat == 'img': img_fn = os.path.join(self.img_path, shape_id, 'view-{}'.format(view_id), 'shape-rgb.png') if not os.path.exists(img_fn): return None with Image.open(img_fn) as fimg: out = np.array(fimg, dtype=np.float32) / 255 white_img = np.ones((self.img_size, self.img_size, 3), dtype=np.float32) mask = np.tile(out[:, :, 3:4], [1, 1, 3]) out = out[:, :, :3] * mask + white_img * (1 - mask) out = torch.from_numpy(out).permute(2, 0, 1).unsqueeze(0) data_feats = data_feats + (out, ) elif feat == 'mask': out = np.zeros( (self.max_num_mask, self.img_size, self.img_size), dtype=np.float32) if len(cur_data['all_parts']) > self.max_num_mask: return None if len(cur_data['all_parts']) <= 1: return None for i in range(len(part_order)): part_id = part_order[i] mask_ids = np.array(cur_data_view[part_id]['mask']) if len(mask_ids) > 0: out[i, mask_ids[:, 0], mask_ids[:, 1]] = 1 # real mask starts from [:, 1:, 1:] of the tensor out = torch.from_numpy(out).unsqueeze(0) data_feats = data_feats + ( out, ) # output tenshor shape [1, max mask, 225, 225] elif feat == 'pts': out = torch.zeros( self.max_num_mask, self.num_points, 3).float() # first-dim is 0/1, meaning exist or not if len(cur_data['all_parts']) > self.max_num_mask: return None if len(cur_data['all_parts']) <= 1: return None for i in range(len(part_order)): part_id = part_order[i] if len(cur_data[part_id] ['similar']) > self.max_num_similar_parts: return None out[i, :, :] = torch.from_numpy( cur_data[part_id]['pts']).float() out = out.unsqueeze(0) data_feats = data_feats + (out, ) elif feat == 'box_size': out = torch.zeros( self.max_num_mask, 3).float() # first-dim is 0/1, meaning exist or not if len(cur_data['all_parts']) > self.max_num_mask: return None if len(cur_data['all_parts']) <= 1: return None for i in range(len(part_order)): part_id = part_order[i] if len(cur_data[part_id] ['similar']) > self.max_num_similar_parts: return None out[i, :] = torch.from_numpy( cur_data[part_id]['bbox_size']).float() out = out.unsqueeze(0) data_feats = data_feats + (out, ) elif feat == 'anchor': out = torch.zeros( self.max_num_mask, 6, 3).float() # first-dim is 0/1, meaning exist or not if len(cur_data['all_parts']) > self.max_num_mask: return None if len(cur_data['all_parts']) <= 1: return None for i in range(len(part_order)): part_id = part_order[i] if len(cur_data[part_id] ['similar']) > self.max_num_similar_parts: return None xyz = cur_data[part_id]['bbox_size'] neg_xyz = -1.0 * cur_data[part_id]['bbox_size'] out[i, 0, :] = torch.tensor([xyz[0], 0, 0]).float() out[i, 1, :] = torch.tensor([0, xyz[1], 0]).float() out[i, 2, :] = torch.tensor([0, 0, xyz[2]]).float() out[i, 3, :] = torch.tensor([neg_xyz[0], 0, 0]).float() out[i, 4, :] = torch.tensor([0, neg_xyz[1], 0]).float() out[i, 5, :] = torch.tensor([0, 0, neg_xyz[2]]).float() out = out.unsqueeze(0) data_feats = data_feats + (out, ) elif feat == 'sem_one_hot': out = torch.zeros(self.max_num_mask, len(self.part_sems)).float() if len(cur_data['all_parts']) > self.max_num_mask: return None if len(cur_data['all_parts']) <= 1: return None for i in range(len(part_order)): part_id = part_order[i] if len(cur_data[part_id] ['similar']) > self.max_num_similar_parts: return None out[i, self.part_sem2id[cur_data[part_id]['sem']]] = 1 out = out.unsqueeze(0) data_feats = data_feats + (out, ) elif feat == 'ins_one_hot': # instance one hot: shape : max parts allowed x max_similar number of parts # out [i, :] is the instance one hot of the parts # [1, 0 .... ], [1, 0, ..... ] [1, 0, ..... ] [0, 1, ..... ] out = torch.zeros(self.max_num_mask, self.max_num_similar_parts).float() if len(cur_data['all_parts']) > self.max_num_mask: return None if len(cur_data['all_parts']) <= 1: return None for i in range(len(part_order)): part_id = part_order[i] if len(cur_data[part_id] ['similar']) > self.max_num_similar_parts: return None counter = 0 for part in sorted(cur_data[part_id]['similar']): ind = part_order.index(part) out[ind, counter] = 1 counter += 1 out = out.unsqueeze(0) data_feats = data_feats + (out, ) elif feat == 'shape_id': if len(cur_data['all_parts']) > self.max_num_mask: return None data_feats = data_feats + (shape_id, ) elif feat == 'part_id': if len(cur_data['all_parts']) > self.max_num_mask: return None data_feats = data_feats + (part_order, ) elif feat == 'view_id': if len(cur_data['all_parts']) > self.max_num_mask: return None data_feats = data_feats + (view_id, ) elif feat == 'adj': adj_pairs = [] with open( os.path.join(self.data_dir, self.category, shape_id + '-adj.txt'), 'r') as f: for line in f: l = line.strip() adj_pairs.append((l.split()[0], l.split()[1])) out = [] for pair in adj_pairs: out.append( (part_order.index(pair[0]), part_order.index(pair[1]))) data_feats = data_feats + (out, ) elif feat == 'adj_pt': adj_idx_pairs = [] with open( os.path.join(self.data_dir, self.category, shape_id + '-adj-ptid.txt'), 'r') as f: for line in f: l = line.strip() adj_idx_pairs.append((l.split()[0], l.split()[1], l.split()[2], l.split()[3])) out = [] for i in range(len(adj_idx_pairs)): pair = adj_idx_pairs[i] pts1 = torch.from_numpy(cur_data[pair[0]]['pts']).float() pts2 = torch.from_numpy(cur_data[pair[1]]['pts']).float() t1, rot1 = cur_data_view[pair[0]][ 'cam_dof'][:3], cur_data_view[pair[0]]['cam_dof'][3:] t2, rot2 = cur_data_view[pair[1]][ 'cam_dof'][:3], cur_data_view[pair[1]]['cam_dof'][3:] t1 = torch.from_numpy(t1).float() t2 = torch.from_numpy(t2).float() rot1 = torch.from_numpy(rot1).float() rot2 = torch.from_numpy(rot2).float() num_point = 1000 pc1 = qrot(rot1.unsqueeze(0).repeat(num_point, 1), pts1) + t1.unsqueeze(0).repeat(num_point, 1) pc2 = qrot(rot2.unsqueeze(0).repeat(num_point, 1), pts2) + t2.unsqueeze(0).repeat(num_point, 1) xyz1 = pc1[int(pair[2])] xyz2 = pc2[int(pair[3])] xyz1 = xyz1.float() xyz2 = xyz2.float() cur_pair_xyz = torch.zeros(2, 3) cur_pair_xyz[0] = xyz1 cur_pair_xyz[1] = xyz2 out.append(cur_pair_xyz) out = torch.stack(out) data_feats = data_feats + (out, ) elif feat == 'total_parts_cnt': if len(cur_data['all_parts']) > self.max_num_mask: return None if len(cur_data['all_parts']) <= 1: return None out = len(cur_data['all_parts']) data_feats = data_feats + (out, ) elif feat == 'similar_parts_cnt': # shape: max_mask x 2 # first index is the index of parts without similar parts 1, 2, 3, 4, 4, 5, 5, 6, 6, 6, 6, .... # second index is the number of similar part count: 1, 1, 1, 2, 2, 2, 2, 4, 4, 4, 4, .... out = torch.zeros(self.max_num_mask, 1, dtype=torch.long) if len(cur_data['all_parts']) > self.max_num_mask: return None if len(cur_data['all_parts']) <= 1: return None total_part_cnt = len(cur_data['all_parts']) counter = 0 for i in range(len(part_order)): part_id = part_order[i] if len(cur_data[part_id] ['similar']) > self.max_num_similar_parts: return None out[i, 0] = len(cur_data[part_id]['similar']) out = out.unsqueeze(0) data_feats = data_feats + (out, ) elif feat == 'similar_parts_edge_indices': out = torch.Tensor(equiv_edge_indices).long() data_feats = data_feats + (out, ) elif feat == 'parts_cam_dof': out = torch.zeros(self.max_num_mask, 7).float() if len(cur_data[part_id] ['similar']) > self.max_num_similar_parts: return None if len(cur_data['all_parts']) <= 1: return None for i in range(len(part_order)): part_id = part_order[i] if len(cur_data[part_id] ['similar']) > self.max_num_similar_parts: return None out[i, :] = torch.from_numpy( cur_data_view[part_id]['cam_dof']) data_feats = data_feats + (out, ) else: raise ValueError('ERROR: unknown feat type %s!' % feat) return data_feats
n_diff = np.linalg.norm(diff) print("|diff|: {0}".format(n_diff)) # angle estimation using quaternion print("Quaternion") qa = np.array([0, p_ref[0][0], p_ref[1][0], p_ref[2][0]]) qb = np.array([0, p_c[0][0], p_c[1][0], p_c[2][0]]) qr = qn.solve_qr_lm(qa, qb) roll_q, pitch_q, yaw_q = qn.q_to_angle(qr) print("pitch: {0} --> {1}".format(pitch, pitch_q)) print("yaw: {0} --> {1}".format(yaw, yaw_q)) print("roll: {0} --> {1}".format(roll, roll_q)) inv_qr = qn.qinv(qr) qa_reproj = qn.qrot(inv_qr, qb) p_reprojq = qa_reproj[1:] print("inv_qr: {0}".format(inv_qr)) diffq = np.array(p_ref).flatten() - p_reprojq n_diffq = np.linalg.norm(diffq) print("|diffq|: {0}".format(n_diffq)) print("Euler") roll_e, pitch_e, yaw_e = ea.solve_angle(np.array(p_ref).flatten(), np.array(p_c).flatten()) print("pitch: {0} --> {1}".format(pitch, pitch_e)) print("yaw: {0} --> {1}".format(yaw, yaw_e)) print("roll: {0} --> {1}".format(roll, roll_e))
def forward(batch, data_features, network, conf, \ is_val=False, step=None, epoch=None, batch_ind=0, num_batch=1, start_time=0, \ log_console=False, log_tb=False, tb_writer=None, lr=None): # prepare input batch_index = 1 if len(batch) == 0: return None cur_batch_size = len(batch[data_features.index('total_parts_cnt')]) total_part_cnt = batch[data_features.index('total_parts_cnt')][0] input_total_part_cnt = batch[data_features.index('total_parts_cnt')][ 0] # 1 input_img = batch[data_features.index('img')][0] # 3 x H x W input_img = input_img.repeat(input_total_part_cnt, 1, 1, 1) # part_cnt 3 x H x W input_pts = batch[data_features.index('pts')][0].squeeze( 0)[:input_total_part_cnt] input_ins_one_hot = batch[data_features.index('ins_one_hot')][0].squeeze( 0)[:input_total_part_cnt] # part_cnt x max_similar_parts input_similar_part_cnt = batch[data_features.index('similar_parts_cnt')][ 0].squeeze(0)[:input_total_part_cnt] # part_cnt x 1 input_shape_id = [batch[data_features.index('shape_id')][0] ] * input_total_part_cnt input_view_id = [batch[data_features.index('view_id')][0] ] * input_total_part_cnt # prepare gt: gt_cam_dof = batch[data_features.index('parts_cam_dof')][0].squeeze( 0)[:input_total_part_cnt] gt_mask = [ batch[data_features.index('mask')][0].squeeze(0) [:input_total_part_cnt].to(conf.device) ] input_total_part_cnt = [batch[data_features.index('total_parts_cnt')][0]] input_similar_parts_edge_indices = [ batch[data_features.index('similar_parts_edge_indices')][0].to( conf.device) ] while total_part_cnt < 70 and batch_index < cur_batch_size: cur_input_cnt = batch[data_features.index( 'total_parts_cnt')][batch_index] total_part_cnt += cur_input_cnt if total_part_cnt > 90: total_part_cnt -= cur_input_cnt batch_index += 1 continue cur_batch_img = batch[data_features.index('img')][batch_index].repeat( cur_input_cnt, 1, 1, 1) input_img = torch.cat((input_img, cur_batch_img), dim=0) input_pts = torch.cat((input_pts, batch[data_features.index('pts')] [batch_index].squeeze(0)[:cur_input_cnt]), dim=0) input_ins_one_hot = torch.cat( (input_ins_one_hot, batch[data_features.index('ins_one_hot')] [batch_index].squeeze(0)[:cur_input_cnt]), dim=0) # B x max_parts x max_similar_parts input_total_part_cnt.append( batch[data_features.index('total_parts_cnt')][batch_index]) # 1 input_similar_part_cnt = torch.cat( (input_similar_part_cnt, batch[data_features.index( 'similar_parts_cnt')][batch_index].squeeze(0)[:cur_input_cnt]), dim=0) # B x max_parts x 2 input_shape_id += [ batch[data_features.index('shape_id')][batch_index] ] * cur_input_cnt input_view_id += [batch[data_features.index('view_id')][batch_index] ] * cur_input_cnt gt_cam_dof = torch.cat((gt_cam_dof, batch[data_features.index( 'parts_cam_dof')][batch_index].squeeze(0)[:cur_input_cnt]), dim=0) # prepare gt gt_mask.append(batch[data_features.index('mask')][batch_index].squeeze( 0)[:cur_input_cnt].to(conf.device)) input_similar_parts_edge_indices.append(batch[data_features.index( 'similar_parts_edge_indices')][batch_index].to(conf.device)) batch_index += 1 input_img = input_img.to(conf.device) input_pts = input_pts.to(conf.device) input_similar_part_cnt = input_similar_part_cnt.to(conf.device) input_ins_one_hot = input_ins_one_hot.to(conf.device) gt_cam_dof = gt_cam_dof.to(conf.device) # prepare gt gt_center = gt_cam_dof[:, :3] # B x 3 gt_quat = gt_cam_dof[:, 3:] # B x 4 batch_size = input_img.shape[0] num_point = input_pts.shape[1] # forward through the network pred_masks, pred_center, pred_quat, pred_center2, pred_quat2 = network( input_img - 0.5, input_pts, input_ins_one_hot, input_total_part_cnt, input_similar_parts_edge_indices) mask_loss_per_data = [] t = 0 matched_pred_mask_all = [] matched_gt_mask_all = [] matched_mask_loss_per_data_all = [] matched_pred_center_all = [] matched_gt_center_all = [] matched_pred_center2_all = [] matched_pred_quat_all = [] matched_gt_quat_all = [] matched_pred_quat2_all = [] matched_ins_onehot_all = [] for i in range(len(input_total_part_cnt)): total_cnt = input_total_part_cnt[i] matched_gt_ids, matched_pred_ids = network.linear_assignment(gt_mask[i], pred_masks[i], \ input_similar_part_cnt[t:t+total_cnt], input_pts[t:t+total_cnt], gt_center[t:t+total_cnt], \ gt_quat[t:t+total_cnt], pred_center[t:t+total_cnt], pred_quat[t:t+total_cnt]) # select the matched data matched_pred_mask = pred_masks[i][matched_pred_ids] matched_gt_mask = gt_mask[i][matched_gt_ids] matched_pred_center = pred_center[t:t + total_cnt][matched_pred_ids] matched_pred_center2 = pred_center2[t:t + total_cnt][matched_pred_ids] matched_gt_center = gt_center[t:t + total_cnt][matched_gt_ids] matched_pred_quat = pred_quat[t:t + total_cnt][matched_pred_ids] matched_pred_quat2 = pred_quat2[t:t + total_cnt][matched_pred_ids] matched_gt_quat = gt_quat[t:t + total_cnt][matched_gt_ids] matched_ins_onehot = input_ins_one_hot[t:t + total_cnt][matched_pred_ids] matched_ins_onehot_all.append(matched_ins_onehot) matched_gt_mask_all.append(matched_gt_mask) matched_pred_mask_all.append(matched_pred_mask) matched_pred_center_all.append(matched_pred_center) matched_pred_center2_all.append(matched_pred_center2) matched_gt_center_all.append(matched_gt_center) matched_pred_quat_all.append(matched_pred_quat) matched_pred_quat2_all.append(matched_pred_quat2) matched_gt_quat_all.append(matched_gt_quat) # for computing mask soft iou loss matched_mask_loss_per_data = network.get_mask_loss( matched_pred_mask, matched_gt_mask) matched_mask_loss_per_data_all.append(matched_mask_loss_per_data) mask_loss_per_data.append(matched_mask_loss_per_data.mean()) t += total_cnt matched_ins_onehot_all = torch.cat(matched_ins_onehot_all, dim=0) matched_pred_mask_all = torch.cat(matched_pred_mask_all, dim=0) matched_gt_mask_all = torch.cat(matched_gt_mask_all, dim=0) matched_mask_loss_per_data_all = torch.cat(matched_mask_loss_per_data_all, dim=0) matched_pred_quat_all = torch.cat(matched_pred_quat_all, dim=0) matched_pred_quat2_all = torch.cat(matched_pred_quat2_all, dim=0) matched_gt_quat_all = torch.cat(matched_gt_quat_all, dim=0) matched_pred_center_all = torch.cat(matched_pred_center_all, dim=0) matched_pred_center2_all = torch.cat(matched_pred_center2_all, dim=0) matched_gt_center_all = torch.cat(matched_gt_center_all, dim=0) center_loss_per_data = network.get_center_loss(matched_pred_center_all, matched_gt_center_all) + \ network.get_center_loss(matched_pred_center2_all, matched_gt_center_all) quat_loss_per_data = network.get_quat_loss(input_pts, matched_pred_quat_all, matched_gt_quat_all) + \ network.get_quat_loss(input_pts, matched_pred_quat2_all, matched_gt_quat_all) l2_rot_loss_per_data = network.get_l2_rotation_loss(input_pts, matched_pred_quat_all, matched_gt_quat_all) + \ network.get_l2_rotation_loss(input_pts, matched_pred_quat2_all, matched_gt_quat_all) whole_shape_cd_per_data = network.get_shape_chamfer_loss(input_pts, matched_pred_quat_all, matched_gt_quat_all, matched_pred_center_all, matched_gt_center_all, input_total_part_cnt) + \ network.get_shape_chamfer_loss(input_pts, matched_pred_quat2_all, matched_gt_quat_all, matched_pred_center2_all, matched_gt_center_all, input_total_part_cnt) # for each type of loss, compute avg loss per batch mask_loss_per_data = torch.stack(mask_loss_per_data) center_loss = center_loss_per_data.mean() quat_loss = quat_loss_per_data.mean() mask_loss = mask_loss_per_data.mean() l2_rot_loss = l2_rot_loss_per_data.mean() shape_chamfer_loss = whole_shape_cd_per_data.mean() # compute total loss total_loss = \ center_loss * conf.loss_weight_center + \ quat_loss * conf.loss_weight_quat + \ l2_rot_loss * conf.loss_weight_l2_rot + \ shape_chamfer_loss * conf.loss_weight_shape_chamfer # display information data_split = 'train' if is_val: data_split = 'val' with torch.no_grad(): # log to console if log_console: utils.printout(conf.flog, \ f'''{strftime("%H:%M:%S", time.gmtime(time.time()-start_time)):>9s} ''' f'''{epoch:>5.0f}/{conf.epochs:<5.0f} ''' f'''{data_split:^10s} ''' f'''{batch_ind:>5.0f}/{num_batch:<5.0f} ''' f'''{100. * (1+batch_ind+num_batch*epoch) / (num_batch*conf.epochs):>9.1f}% ''' f'''{lr:>5.2E} ''' f'''{mask_loss.item():>10.5f}''' f'''{center_loss.item():>10.5f}''' f'''{quat_loss.item():>10.5f}''' f'''{l2_rot_loss.item():>10.5f}''' f'''{shape_chamfer_loss.item():>10.5f}''' f'''{total_loss.item():>10.5f}''') conf.flog.flush() # log to tensorboard if log_tb and tb_writer is not None: tb_writer.add_scalar('mask_loss', mask_loss.item(), step) tb_writer.add_scalar('center_loss', center_loss.item(), step) tb_writer.add_scalar('quat_loss', quat_loss.item(), step) tb_writer.add_scalar('l2 rotation_loss', l2_rot_loss.item(), step) tb_writer.add_scalar('shape_chamfer_loss', shape_chamfer_loss.item(), step) tb_writer.add_scalar('total_loss', total_loss.item(), step) tb_writer.add_scalar('lr', lr, step) # gen visu if is_val and ( not conf.no_visu) and epoch % conf.num_epoch_every_visu == 0: visu_dir = os.path.join(conf.exp_dir, 'val_visu') out_dir = os.path.join(visu_dir, 'epoch-%04d' % epoch) input_img_dir = os.path.join(out_dir, 'input_img') input_pts_dir = os.path.join(out_dir, 'input_pts') gt_mask_dir = os.path.join(out_dir, 'gt_mask') pred_mask_dir = os.path.join(out_dir, 'pred_mask') gt_dof_dir = os.path.join(out_dir, 'gt_dof') pred_dof_dir = os.path.join(out_dir, 'pred_dof') pred_dof2_dir = os.path.join(out_dir, 'pred_dof2') info_dir = os.path.join(out_dir, 'info') if batch_ind == 0: # create folders os.mkdir(out_dir) os.mkdir(input_img_dir) os.mkdir(input_pts_dir) os.mkdir(gt_mask_dir) os.mkdir(pred_mask_dir) os.mkdir(gt_dof_dir) os.mkdir(pred_dof_dir) os.mkdir(pred_dof2_dir) os.mkdir(info_dir) if batch_ind < conf.num_batch_every_visu: utils.printout(conf.flog, 'Visualizing ...') # compute pred_pts and gt_pts pred_pts = qrot( matched_pred_quat_all.unsqueeze(1).repeat(1, num_point, 1), input_pts) + matched_pred_center_all.unsqueeze(1).repeat( 1, num_point, 1) pred2_pts = qrot( matched_pred_quat2_all.unsqueeze(1).repeat( 1, num_point, 1), input_pts) + matched_pred_center2_all.unsqueeze(1).repeat( 1, num_point, 1) gt_pts = qrot( matched_gt_quat_all.unsqueeze(1).repeat(1, num_point, 1), input_pts) + matched_gt_center_all.unsqueeze(1).repeat( 1, num_point, 1) t = 0 for i in range(batch_size): fn = 'data-%03d.png' % (batch_ind * batch_size + i) cur_input_img = ( input_img[i].permute(1, 2, 0).cpu().numpy() * 255).astype(np.uint8) Image.fromarray(cur_input_img).save( os.path.join(input_img_dir, fn)) cur_input_pts = input_pts[i].cpu().numpy() render_utils.render_pts(os.path.join( BASE_DIR, input_pts_dir, fn), cur_input_pts, blender_fn='object_centered.blend') cur_gt_mask = (matched_gt_mask_all[i].cpu().numpy() > 0.5).astype(np.uint8) * 255 Image.fromarray(cur_gt_mask).save( os.path.join(gt_mask_dir, fn)) cur_pred_mask = (matched_pred_mask_all[i].cpu().numpy() > 0.5).astype(np.uint8) * 255 Image.fromarray(cur_pred_mask).save( os.path.join(pred_mask_dir, fn)) cur_pred_pts = pred_pts[i].cpu().numpy() render_utils.render_pts(os.path.join( BASE_DIR, pred_dof_dir, fn), cur_pred_pts, blender_fn='camera_centered.blend') cur_pred_pts = pred2_pts[i].cpu().numpy() render_utils.render_pts(os.path.join( BASE_DIR, pred_dof2_dir, fn), cur_pred_pts, blender_fn='camera_centered.blend') cur_gt_pts = gt_pts[i].cpu().numpy() render_utils.render_pts(os.path.join( BASE_DIR, gt_dof_dir, fn), cur_gt_pts, blender_fn='camera_centered.blend') with open( os.path.join(info_dir, fn.replace('.png', '.txt')), 'w') as fout: fout.write('shape_id: %s, view_id: %s\n' % (\ input_shape_id[i],\ input_view_id[i])) fout.write('ins onehot %s\n' % matched_ins_onehot_all[i]) fout.write('mask_loss: %f\n' % matched_mask_loss_per_data_all[i].item()) fout.write('center_loss: %f\n' % center_loss_per_data[i].item()) fout.write('quat_loss: %f\n' % quat_loss_per_data[i].item()) fout.write('l2_rot_loss: %f\n' % l2_rot_loss_per_data[i].item()) fout.write('shape_chamfer_loss %f\n' % shape_chamfer_loss.item()) return total_loss
def forward(self, img, pc, pred_masks, part_feat, part_cnt, equiv_edge_indices): # get image feature img = img.clone() img[:, 0] = (img[:, 0] - 0.485) / 0.229 img[:, 1] = (img[:, 1] - 0.456) / 0.224 img[:, 2] = (img[:, 2] - 0.406) / 0.225 img_feat = torch.relu(self.resnet_final_bn(self.resnet(img))) # get mask feature all_masks = (torch.cat(pred_masks, dim=0).unsqueeze(dim=1) > 0.5).float() mask_feat = torch.relu( self.mask_resnet_final_bn(self.mask_resnet(all_masks))) # get global feature per shape part_feat_5d = torch.cat((img_feat, mask_feat, part_feat), dim=1) part_feat_5d = torch.relu(self.mlp3(part_feat_5d)) # perform graph-conv t = 0 i = 0 output_part_feat_5d = [] for cnt in part_cnt: child_feats = part_feat_5d[t:t + cnt] iter_feats = [child_feats] cur_equiv_edge_indices = equiv_edge_indices[i] for j in range(self.num_iteration): if cur_equiv_edge_indices.shape[0] > 0: node_edge_feats = torch.cat([ child_feats[cur_equiv_edge_indices[:, 0], :], child_feats[cur_equiv_edge_indices[:, 1], :] ], dim=1) node_edge_feats = torch.relu( self.node_edge_op[j](node_edge_feats)) new_child_feats = child_feats.new_zeros(cnt, 256) new_child_feats = torch_scatter.scatter_mean( node_edge_feats, cur_equiv_edge_indices[:, 0], dim=0, out=new_child_feats) child_feats = new_child_feats iter_feats.append(child_feats) all_iter_feat = torch.cat(iter_feats, dim=1) output_part_feat_5d.append(all_iter_feat) t += cnt i += 1 output_part_feat_5d = torch.cat(output_part_feat_5d, dim=0) # decode pose center, quat = self.pose_decoder(output_part_feat_5d) # refine final poses num_point = pc.shape[1] with torch.no_grad(): pc2 = qrot(quat.unsqueeze(1).repeat(1, num_point, 1), pc) + center.unsqueeze(1).repeat(1, num_point, 1) feat = torch.cat([output_part_feat_5d, center, quat], dim=1) center2, quat2 = self.refine_net(pc2, feat, part_cnt, equiv_edge_indices) return center, quat, center2, quat2
def forward(self, vert): vert = quat.qrot( (self.rotation / (self.rotation.norm())).repeat(vert.shape[0], 1), vert - 0.5) vert = vert + self.translation.unsqueeze(0) + 0.5 return vert
gt_w[20:40, :] = [0.00, 0.05, 0.0] gt_w[50:60, :] = [0.01, -0.05, 0.01] gt_w[70:90, :] = [-0.01, -0.00, 0.0] for i in range(2, seq_len): gt_angle[i] = gt_angle[i - 1] + gt_w[i] # generate accelerometer G.T. gn = [0, 0, 0, 1] gt_a = np.zeros((seq_len, 3)) # gx, gy, gz # ya,t = -Rt gn + e_a,t for i in range(seq_len): qr = qt.angle_to_q(gt_angle[i, 0], gt_angle[i, 1], gt_angle[i, 2]) qa = qt.qrot(qr, gn) gt_a[i, 0] = -qa[1] gt_a[i, 1] = -qa[2] gt_a[i, 2] = -qa[3] # generate magnetometer G.T. gt_m = np.zeros((seq_len, 3)) mn = np.array([0, 1, 0, 0]) for i in range(seq_len): qr = qt.angle_to_q(gt_angle[i, 0], gt_angle[i, 1], gt_angle[i, 2]) qm = qt.qrot(qr, mn) b = np.zeros((3, 1)) b[0, 0] = qm[1] b[1, 0] = qm[2]