def __getitem__(self, idx): anno_pair, perm_mat = self.ds.get_pair(self.cls) if perm_mat.size <= 2 * 2: return self.__getitem__(idx) P1_gt = [(kp['x'], kp['y']) for kp in anno_pair[0]['keypoints']] P2_gt = [(kp['x'], kp['y']) for kp in anno_pair[1]['keypoints']] n1_gt, n2_gt = len(P1_gt), len(P2_gt) P1_gt = np.array(P1_gt) P2_gt = np.array(P2_gt) G1_gt, H1_gt, e1_gt = build_graphs(P1_gt, n1_gt, stg=cfg.PAIR.GT_GRAPH_CONSTRUCT, k=cfg.PAIR.NUM_NEIGHBOURS) if cfg.PAIR.REF_GRAPH_CONSTRUCT == 'same': G2_gt = perm_mat.transpose().dot(G1_gt) H2_gt = perm_mat.transpose().dot(H1_gt) e2_gt = e1_gt else: G2_gt, H2_gt, e2_gt = build_graphs( P2_gt, n2_gt, stg=cfg.PAIR.REF_GRAPH_CONSTRUCT, k=cfg.PAIR.NUM_NEIGHBOURS) ret_dict = { 'Ps': [torch.Tensor(x) for x in [P1_gt, P2_gt]], 'ns': [torch.tensor(x) for x in [n1_gt, n2_gt]], 'es': [torch.tensor(x) for x in [e1_gt, e2_gt]], 'gt_perm_mat': perm_mat, 'Gs': [torch.Tensor(x) for x in [G1_gt, G2_gt]], 'Hs': [torch.Tensor(x) for x in [H1_gt, H2_gt]] } imgs = [anno['image'] for anno in anno_pair] if imgs[0] is not None: trans = transforms.Compose([ transforms.ToTensor(), transforms.Normalize(cfg.NORM_MEANS, cfg.NORM_STD) ]) imgs = [trans(img) for img in imgs] ret_dict['images'] = imgs elif 'feat' in anno_pair[0]['keypoints'][0]: feat1 = np.stack([kp['feat'] for kp in anno_pair[0]['keypoints']], axis=-1) feat2 = np.stack([kp['feat'] for kp in anno_pair[1]['keypoints']], axis=-1) ret_dict['features'] = [torch.Tensor(x) for x in [feat1, feat2]] return ret_dict
def __getitem__(self, item): file = os.path.join(self.root, self.partition, self.files[item]) data = np.load(file) sample = {'points':data['x'], 'R':data['R'], 't':np.expand_dims(data['t'], -1),'idx': np.array(item, dtype=np.int32)} # 'label': self.files[idx].split('/')[0] if self.transform: sample = self.transform(sample) T_ab = sample['transform_gt'] T_ba = np.concatenate((T_ab[:, :3].T, np.expand_dims(-(T_ab[:, :3].T).dot(T_ab[:, 3]), axis=1)), axis=-1) n1_gt, n2_gt = sample['perm_mat'].shape A1_gt, e1_gt = build_graphs(sample['points_src'], sample['src_inlier'], n1_gt, stg=cfg.PAIR.GT_GRAPH_CONSTRUCT) if cfg.PAIR.REF_GRAPH_CONSTRUCT == 'same': A2_gt = A1_gt.transpose().contiguous() e2_gt = e1_gt else: A2_gt, e2_gt = build_graphs(sample['points_ref'], sample['ref_inlier'], n2_gt, stg=cfg.PAIR.REF_GRAPH_CONSTRUCT) if cfg.DATASET.NOISE_TYPE != 'clean': src_o3 = open3d.geometry.PointCloud() ref_o3 = open3d.geometry.PointCloud() src_o3.points = open3d.utility.Vector3dVector(sample['points_src'][:, :3]) ref_o3.points = open3d.utility.Vector3dVector(sample['points_ref'][:, :3]) src_o3.estimate_normals(search_param=open3d.geometry.KDTreeSearchParamHybrid(radius=0.1, max_nn=30)) ref_o3.estimate_normals(search_param=open3d.geometry.KDTreeSearchParamHybrid(radius=0.1, max_nn=30)) sample['points_src'] = np.concatenate([sample['points_src'],src_o3.normals], axis=1).astype(np.float32) sample['points_ref'] = np.concatenate([sample['points_ref'],ref_o3.normals], axis=1).astype(np.float32) ret_dict = {'Ps': [torch.Tensor(x) for x in [sample['points_src'], sample['points_ref']]], 'ns': [torch.tensor(x) for x in [n1_gt, n2_gt]], 'es': [torch.tensor(x) for x in [e1_gt, e2_gt]], 'gt_perm_mat': torch.tensor(sample['perm_mat'].astype('float32')), 'As': [torch.Tensor(x) for x in [A1_gt, A2_gt]], 'Ts': [torch.Tensor(x) for x in [T_ab.astype('float32'), T_ba.astype('float32')]], 'Ins': [torch.Tensor(x) for x in [sample['src_inlier'], sample['ref_inlier']]], 'label':torch.tensor(1) } return ret_dict
def __getitem__(self, item): sample = {'points': self.data[item, :, :], 'label': self.label[item], 'idx': np.array([item], dtype=np.int32)} if self.transform: sample = self.transform(sample) # if item==139: # np.save(cfg.DATASET.NOISE_TYPE+'sample'+str(item),sample) # samplecrop = np.load('/home/science/code/python/PGM/cropsample139.npy', allow_pickle=True).item() T_ab = sample['transform_gt'] T_ba = np.concatenate((T_ab[:,:3].T, np.expand_dims(-(T_ab[:,:3].T).dot(T_ab[:,3]), axis=1)), axis=-1) n1_gt, n2_gt = sample['perm_mat'].shape A1_gt, e1_gt = build_graphs(sample['points_src'], sample['src_inlier'], n1_gt, stg=cfg.PAIR.GT_GRAPH_CONSTRUCT) if cfg.PAIR.REF_GRAPH_CONSTRUCT == 'same': A2_gt = A1_gt.transpose().contiguous() e2_gt= e1_gt else: A2_gt, e2_gt = build_graphs(sample['points_ref'], sample['ref_inlier'], n2_gt, stg=cfg.PAIR.REF_GRAPH_CONSTRUCT) if cfg.DATASET.NOISE_TYPE != 'clean': src_o3 = open3d.geometry.PointCloud() ref_o3 = open3d.geometry.PointCloud() src_o3.points = open3d.utility.Vector3dVector(sample['points_src'][:, :3]) ref_o3.points = open3d.utility.Vector3dVector(sample['points_ref'][:, :3]) src_o3.estimate_normals(search_param=open3d.geometry.KDTreeSearchParamHybrid(radius=0.1, max_nn=30)) ref_o3.estimate_normals(search_param=open3d.geometry.KDTreeSearchParamHybrid(radius=0.1, max_nn=30)) sample['points_src'][:, 3:6] = src_o3.normals sample['points_ref'][:, 3:6] = ref_o3.normals ret_dict = {'Ps': [torch.Tensor(x) for x in [sample['points_src'], sample['points_ref']]], 'ns': [torch.tensor(x) for x in [n1_gt, n2_gt]], 'es': [torch.tensor(x) for x in [e1_gt, e2_gt]], 'gt_perm_mat': torch.tensor(sample['perm_mat'].astype('float32')), 'As': [torch.Tensor(x) for x in [A1_gt, A2_gt]], 'Ts': [torch.Tensor(x) for x in [T_ab.astype('float32'), T_ba.astype('float32')]], 'Ins': [torch.Tensor(x) for x in [sample['src_inlier'], sample['ref_inlier']]], 'label': torch.tensor(sample['label']), 'raw': torch.Tensor(sample['points_raw']), } return ret_dict
def __getitem__(self, idx): sampling_strategy = cfg.train_sampling if self.ds.sets == "train" else cfg.eval_sampling if self.num_graphs_in_matching_instance is None: raise ValueError("Num_graphs has to be set to an integer value.") idx = idx if self.true_epochs else None anno_list, perm_mat_list = self.ds.get_k_samples(idx, k=self.num_graphs_in_matching_instance, cls=self.cls, mode=sampling_strategy) for perm_mat in perm_mat_list: if ( not perm_mat.size or (perm_mat.size < 2 * 2 and sampling_strategy == "intersection") and not self.true_epochs ): # 'and not self.true_epochs' because we assume all data is valid when sampling a true epoch next_idx = None if idx is None else idx + 1 return self.__getitem__(next_idx) points_gt = [np.array([(kp["x"], kp["y"]) for kp in anno_dict["keypoints"]]) for anno_dict in anno_list] n_points_gt = [len(p_gt) for p_gt in points_gt] graph_list = [] for p_gt, n_p_gt in zip(points_gt, n_points_gt): edge_indices, edge_features = build_graphs(p_gt, n_p_gt) # Add dummy node features so the __slices__ of them is saved when creating a batch pos = torch.tensor(p_gt).to(torch.float32) / 256.0 assert (pos > -1e-5).all(), p_gt graph = Data( edge_attr=torch.tensor(edge_features).to(torch.float32), edge_index=torch.tensor(edge_indices, dtype=torch.long), x=pos, pos=pos, ) graph.num_nodes = n_p_gt graph_list.append(graph) ret_dict = { "Ps": [torch.Tensor(x) for x in points_gt], "ns": [torch.tensor(x) for x in n_points_gt], "gt_perm_mat": perm_mat_list, "edges": graph_list, } imgs = [anno["image"] for anno in anno_list] if imgs[0] is not None: trans = transforms.Compose([transforms.ToTensor(), transforms.Normalize(cfg.NORM_MEANS, cfg.NORM_STD)]) imgs = [trans(img) for img in imgs] ret_dict["images"] = imgs elif "feat" in anno_list[0]["keypoints"][0]: feat_list = [np.stack([kp["feat"] for kp in anno_dict["keypoints"]], axis=-1) for anno_dict in anno_list] ret_dict["features"] = [torch.Tensor(x) for x in feat_list] return ret_dict
def construct_universe_graph(self, sets): if cfg.EDGE_PREDICTION.train_graph_only: graph_build_path = "data/cache/" + "voc_graph_build_" + "train" else: graph_build_path = "data/cache/" + "voc_graph_build_" + sets if cfg.EDGE_PREDICTION.regression: graph_build_path += "_regression.pk" else: graph_build_path += "_classification.pk" if os.path.exists(graph_build_path): with open(graph_build_path, "rb") as f: self.adj_list = pickle.load(f) print("load {} successfully.".format(graph_build_path)) else: for cls in self.classes: print(cls) cls_idx = self.classes.index(cls) n_max = len(KPT_NAMES[cls]) adj = np.zeros((n_max, n_max)) adj_full = np.zeros((n_max, n_max)) for xml_name in self.xml_list[cls_idx]: anno_dict = self.__get_anno_dict(xml_name, cls_idx) nodes = np.array([(kp["x"], kp["y"]) for kp in anno_dict["keypoints"]]) num_nodes = len(nodes) if num_nodes == 0: continue edge_indices, _ = build_graphs(nodes, num_nodes) for i in range(edge_indices.shape[1]): x, y = edge_indices[0][i], edge_indices[1][i] xname, yname = anno_dict["keypoints"][x][ "name"], anno_dict["keypoints"][y]["name"] adj[KPT_NAMES[cls].index(xname), KPT_NAMES[cls].index(yname)] += 1 node_indices = np.array([ KPT_NAMES[cls].index(kp["name"]) for kp in anno_dict["keypoints"] ]) nodex_indices = np.tile(np.expand_dims(node_indices, 0), (num_nodes, 1)).reshape(-1) nodey_indices = np.tile(np.expand_dims(node_indices, 1), (1, num_nodes)).reshape(-1) adj_full[nodex_indices, nodey_indices] += 1 adj_full = adj_full + 1e-8 assert np.sum(np.abs(adj_full == adj_full.transpose)) == 0 assert np.sum(np.abs(adj_full < adj)) == 0 if cfg.EDGE_PREDICTION.regression: adj = adj / adj_full else: adj = (adj / adj_full > 0.5).astype(float) self.adj_list.append(adj) with open(graph_build_path, "wb") as f: pickle.dump(self.adj_list, f)
def __getitem__(self, idx): anno_pair, perm_mat = self.ds.get_pair(self.cls) if perm_mat.size <= 2 * 2: return self.__getitem__(idx) P1_gt = [(kp['x'], kp['y']) for kp in anno_pair[0]['keypoints']] P2_gt = [(kp['x'], kp['y']) for kp in anno_pair[1]['keypoints']] ## outlier generation # for jj in range(perm_mat[0]): # num_out = 4 # num_row = perm_mat.shape[0] # perm_mat = np.concatenate((np.array(perm_mat), np.zeros([num_out, num_row])), axis=0) # perm_mat = np.concatenate((np.array(perm_mat), np.zeros([num_out + num_row, num_out])), axis=1) # P1_gt = np.concatenate((np.array(P1_gt), np.random.normal(0, 10, [2, 2])), axis=0) # P2_gt = np.concatenate((np.array(P2_gt), np.random.normal(0, 10, [2, 2])), axis=0) # if len(P1_gt) == 0 or len(P2_gt) == 0 or len(P1_gt) != len(P2_gt): # P1_gt = np.array(np.random.normal(0, 10, [3,2])) # P2_gt = np.array(np.random.normal(0, 10, [3,2])) P1_gt = np.array(P1_gt) P2_gt = np.array(P2_gt) n1_gt, n2_gt = len(P1_gt), len(P2_gt) G1_gt, H1_gt, e1_gt, edge_indices1, edge_feat1 = build_graphs( P1_gt, n1_gt, stg=cfg.PAIR.GT_GRAPH_CONSTRUCT) if cfg.PAIR.REF_GRAPH_CONSTRUCT == 'same': G2_gt = perm_mat.transpose().dot(G1_gt) H2_gt = perm_mat.transpose().dot(H1_gt) e2_gt = e1_gt else: G2_gt, H2_gt, e2_gt, edge_indices2, edge_feat2 = build_graphs( P2_gt, n2_gt, stg=cfg.PAIR.REF_GRAPH_CONSTRUCT) ret_dict = { 'Ps': [torch.Tensor(x) for x in [P1_gt, P2_gt]], 'ns': [torch.tensor(x) for x in [n1_gt, n2_gt]], 'es': [torch.tensor(x) for x in [e1_gt, e2_gt]], 'gt_perm_mat': perm_mat, 'Gs': [torch.Tensor(x) for x in [G1_gt, G2_gt]], 'Hs': [torch.Tensor(x) for x in [H1_gt, H2_gt]], "edge_src": [torch.Tensor(x) for x in [edge_indices1]], "edge_tgt": [torch.Tensor(x) for x in [edge_indices2]], "edge_feat1": [torch.Tensor(x) for x in [edge_feat1]], "edge_feat2": [torch.Tensor(x) for x in [edge_feat2]] } imgs = [anno['image'] for anno in anno_pair] if imgs[0] is not None: trans = transforms.Compose([ transforms.ToTensor(), transforms.Normalize(cfg.NORM_MEANS, cfg.NORM_STD) ]) imgs = [trans(img) for img in imgs] ret_dict['images'] = imgs elif 'feat' in anno_pair[0]['keypoints'][0]: feat1 = np.stack([kp['feat'] for kp in anno_pair[0]['keypoints']], axis=-1) feat2 = np.stack([kp['feat'] for kp in anno_pair[1]['keypoints']], axis=-1) ret_dict['features'] = [torch.Tensor(x) for x in [feat1, feat2]] return ret_dict