def __getitem__(self, index):
sid = index // self.scene_size
pid = index % self.scene_size
src, tgt, T = self._load_instance(sid, pid)
#################### Augmentation ###################
# hard rotation augmentation
if self.mode == 'train' and not self.opt.no_augmentation:
if self.opt.model.kpconv:
tgt, R_aug = pctk.rotate_point_cloud(tgt, max_degree=60)
else:
tgt, R_aug = pctk.rotate_point_cloud(tgt)
_, src = pctk.uniform_resample_np(src, self.opt.model.input_num)
_, tgt = pctk.uniform_resample_np(tgt, self.opt.model.input_num)
##########################################################
if self.opt.model.flag == 'rotation':
src = pctk.normalize_pc_np(src)
tgt = pctk.normalize_pc_np(tgt)
R, R_label = label_relative_rotation_np(self.anchors, T)
data = {'src': torch.from_numpy(src.astype(np.float32)),\
'tgt': torch.from_numpy(tgt.astype(np.float32)),\
'T': torch.from_numpy(T.astype(np.float32)),\
'R': torch.from_numpy(R.astype(np.float32)),\
'R_label': torch.Tensor([R_label]).long()}
else:
data = {'src': torch.from_numpy(src.astype(np.float32)),\
'tgt': torch.from_numpy(tgt.astype(np.float32)),\
'T': torch.from_numpy(T.astype(np.float32)),\
'fn': "None"}
return data
def __getitem__(self, index):
data = sio.loadmat(self.all_data[index])
if self.mode == 'train':
_, pc = pctk.uniform_resample_np(data['pc'], self.opt.model.input_num)
else:
pc = data['pc']
pc = p3dtk.normalize_np(pc.T)
pc = pc.T
R = np.eye(3)
R_label = 29
if not self.opt.no_augmentation:
if 'R' in data.keys() and self.mode != 'train':
pc, R = pctk.rotate_point_cloud(pc, data['R'])
else:
pc, R = pctk.rotate_point_cloud(pc)
_, R_label, R0 = rotation_distance_np(R, self.anchors)
if self.flag == 'rotation':
R = R0
return {'pc':torch.from_numpy(pc.astype(np.float32)),
'label':torch.from_numpy(data['label'].flatten()).long(),
'fn': data['name'][0],
'R': R,
'R_label': torch.Tensor([R_label]).long(),
}
def __getitem__(self, index):
meta = self.kptfiles[index]
choice = np.random.choice(np.arange(meta.indices.shape[0]), self.npt)
kpts = meta.indices[choice].astype(np.int32)
pcdA = o3d.io.read_point_cloud(meta.pathA)
pcdB = o3d.io.read_point_cloud(meta.pathB)
rawA, pcdA_down, normalA = radius_ball_search_o3d(
pcdA, kpts[:, 0], self.search_radius, self.voxel_size,
self.use_normals)
rawB, pcdB_down, normalB = radius_ball_search_o3d(
pcdB, kpts[:, 1], self.search_radius, self.voxel_size,
self.use_normals)
################################33
# pcdA = pcdA_down
# pcdB = pcdB_down
#####################################
pcdA = np.asarray(pcdA.points).astype(np.float32)
pcdB = np.asarray(pcdB.points).astype(np.float32)
# [nptxNx3]
inputA = []
inputB = []
# preprocessing and augmentation
T = RigidMatrix(meta.poseA).R.T @ RigidMatrix(meta.poseB).R
# T = T.R
# T = np.identity(4)
self.R_aug_src = None
self.R_aug_tgt = None
if not self.opt.no_augmentation:
_, self.R_aug_src = pctk.rotate_point_cloud(None, max_degree=30)
_, self.R_aug_tgt = pctk.rotate_point_cloud(None, max_degree=30)
_, pcdA = pctk.rotate_point_cloud(pcdA, self.R_aug_src)
_, pcdB = pctk.rotate_point_cloud(pcdB, self.R_aug_tgt)
# T = ?
# if self.use_normals:
# for pca,pcb,na,nb in zip(rawA, rawB, normalA, normalB):
# inputA.append(self._preprocess(pca,self.R_aug_src,na))
# inputB.append(self._preprocess(pcb,self.R_aug_tgt,nb))
# else:
for pca, pcb in zip(rawA, rawB):
inputA.append(self._preprocess(pca, self.R_aug_src))
inputB.append(self._preprocess(pcb, self.R_aug_tgt))
inputA = np.array(inputA)
inputB = np.array(inputB)
data = {'src': torch.from_numpy(inputA.astype(np.float32)),\
'tgt': torch.from_numpy(inputB.astype(np.float32)),\
'frag_src': pcdA,\
'frag_tgt': pcdB,\
'T': torch.from_numpy(T.astype(np.float32)),\
'fn': meta.id}
return data
def _preprocess(self, pc, R_aug=None, n=None):
idx, pc = pctk.uniform_resample_np(pc, self.input_num)
if n is not None:
n = n[idx]
if R_aug is not None:
# rotational augmentation
pc, _ = pctk.rotate_point_cloud(pc, R_aug)
if n is not None:
n, _ = pctk.rotate_point_cloud(n, R_aug)
if n is not None:
pc = np.concatenate([pc,n],axis=1)
return pc
def __getitem__(self, index):
# data = sio.loadmat(self.all_data[index])
data = np.load(self.all_data[index])
# _, pc = pctk.uniform_resample_np(data['pc'], self.opt.model.input_num)
# [10000, 3], the dimension should be correct according to the comments
_, pc = pctk.uniform_resample_np(data['points'], self.opt.model.input_num)
pc = p3dtk.normalize_np(pc.T)
pc = pc.T
R = np.eye(3)
R_label = 29
if not self.opt.no_augmentation:
# ###################### HACK #########################
# ridx = np.random.randint(0, high=self.anchors.shape[0])
# R = self.anchors[ridx]
# pcR, _ = pctk.rotate_point_cloud(pc, R)
#######################################################
if 'R' in data.keys() and self.mode != 'train':
pc, R = pctk.rotate_point_cloud(pc, data['R'])
else:
pc, R = pctk.rotate_point_cloud(pc)
# if self.mode == 'test':
# data['R'] = R
# output_path = os.path.join(self.dataset_path, data['cat'][0], 'testR')
# os.makedirs(output_path,exist_ok=True)
# sio.savemat(os.path.join(output_path, data['name'][0] + '.mat'), data)
_, R_label, R0 = rotation_distance_np(R, self.anchors)
if self.flag == 'rotation':
R = R0
return {'pc':torch.from_numpy(pc.astype(np.float32)),
'label':torch.from_numpy(data['label'].flatten()).long(),
# 'fn': data['name'][0],
'R': R,
'R_label': torch.Tensor([R_label]).long(),
}
def __getitem__(self, index):
# data = sio.loadmat(self.all_data[index])
data = np.load(self.all_data[index])
# _, pc = pctk.uniform_resample_np(data['pc'], self.opt.model.input_num)
# [10000, 3], the dimension should be correct according to the comments
_, pc = pctk.uniform_resample_np(data['points'], self.opt.model.input_num)
# normalization
pc = p3dtk.normalize_np(pc.T)
pc = pc.T
# R = np.eye(3)
# R_label = 29
# source shape
# if 'R' in data.keys() and self.mode != 'train':
# pc_src, R_src = pctk.rotate_point_cloud(pc, data['R'])
# else:
# pc_src, R_src = pctk.rotate_point_cloud(pc)
# Rotated point cloud, rotation matrix
pc_src, R_src = pctk.rotate_point_cloud(pc)
# target shape
# pc_tgt, R_tgt = pctk.rotate_point_cloud(pc)
pc_tgt = pc
# if self.mode == 'test':
# data['R'] = R
# output_path = os.path.join(self.dataset_path, data['cat'][0], 'testR')
# os.makedirs(output_path,exist_ok=True)
# sio.savemat(os.path.join(output_path, data['name'][0] + '.mat'), data)
# _, R_label, R0 = rotation_distance_np(R, self.anchors)
# T = R_src @ R_tgt.T
T = R_src # @ R_tgt.T
# RR_regress = np.einsum('abc,bj,ijk -> aick', self.anchors, T, self.anchors)
# R_label = np.argmax(np.einsum('abii->ab', RR_regress),axis=1)
# idxs = np.vstack([np.arange(R_label.shape[0]), R_label]).T
# R = RR_regress[idxs[:,0], idxs[:,1]]
R, R_label = label_relative_rotation_np(self.anchors, T)
pc_tensor = np.stack([pc_src, pc_tgt])
return {'pc':torch.from_numpy(pc_tensor.astype(np.float32)),
# 'fn': data['name'][0],
'T' : torch.from_numpy(T.astype(np.float32)),
'R': torch.from_numpy(R.astype(np.float32)),
'R_label': torch.Tensor([R_label]).long(),
}