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