def test_bunny(args, net):
pcd_src = o3d.io.read_point_cloud(
"CAD_models/office_1_chair_extracted1.pcd")
points_src = np.asarray(pcd_src.points)
points_src = helper.fit_in_m1_to_1(points_src)
pcd_target = o3d.io.read_point_cloud(
"CAD_models/office_1_chair_extracted1.pcd")
points_target = np.asarray(pcd_target.points)
points_target = helper.fit_in_m1_to_1(points_target)
points_rot, R_ab_gt = rotate_cloud(points_target)
points_cuda = Network_input_format(points_src)
points_rot_cuda = Network_input_format(points_rot)
src = points_cuda #
target = points_rot_cuda #
rotation_ab_pred, translation_ab_pred, rotation_ba_pred, translation_ba_pred, corr_mat_ab_pred = net(
src, target)
print(rotation_ab_pred, "rotation_ab_pred")
print(R_ab_gt, "R_ab_gt")
points_pred = rotation_ab_pred.detach().cpu().numpy().dot(points_src.T).T
helper.display_three_clouds(points_src,points_rot,points_pred,title="real world data" ,\
legend_list=[" source","target","prediction"] )
def create_partial_data(templates_ip):
import helper
templates = np.copy(templates_ip)
partial_templates = np.zeros(
(templates.shape[0], 1024, templates.shape[2]))
for i in range(templates.shape[0]):
templates[i] = templates[i] - np.array([1, 1, 1])
temp = remove_points(templates[i])
print(temp.shape)
helper.display_three_clouds(temp, templates[i], templates_ip[i], "")
def generate_icp_results(self, gt_pose): from icp import icp_test from scipy.spatial import KDTree M_given = self.templates[self.template_idx,:,:] S_given = helper.apply_transformation(M_given.reshape((1,-1,3)), gt_pose)[0] # S_given = S_given + np.random.normal(0,1,S_given.shape) # Noisy Data M_given = M_given[0:self.NUM_POINT,:] # template data S_given = S_given[0:self.NUM_POINT,:] # source data tree_M = KDTree(M_given) tree_M_sampled = KDTree(M_given[0:100,:]) final_pose, model_data, sensor_data, predicted_data, title, _, _ = icp_test(S_given[0:100,:], M_given, tree_M, M_given[0:100,:], tree_M_sampled, S_given, gt_pose.reshape((1,6)), self.MAX_LOOPS, self.ftol) self.find_errors(gt_pose, final_pose) helper.display_three_clouds(model_data, sensor_data, predicted_data, title)
def generate_results(self, ftol, gt_pose, swap_case): template_data = self.templates[self.template_idx,:,:].reshape((1,MAX_NUM_POINT,3)) # Extract the template and reshape it. template_data = template_data[:,0:self.NUM_POINT,:] source_data = helper.apply_transformation(template_data,gt_pose) # Generate Source Data. # source_data = source_data + np.random.normal(0,0.001,source_data.shape) # Noisy Data if swap_case: final_pose, TRANSFORMATIONS, loss_i, predicted_data, transformed_source_data, elapsed_time, itr = self.test_one_case(source_data, template_data) # Find final transformation by network. else: final_pose, TRANSFORMATIONS, loss_i, predicted_data, transformed_source_data, elapsed_time, itr = self.test_one_case(template_data, source_data) # Find final transformation by network. if not swap_case: title = "Actual T (Red->Green): " for i in range(len(gt_pose[0])): if i>2: title += str(round(gt_pose[0][i]*(180/np.pi),2)) else: title += str(gt_pose[0][i]) title += ', ' title += "\nPredicted T (Red->Blue): " for i in range(len(final_pose[0])): if i>2: title += str(round(final_pose[0,i]*(180/np.pi),3)) else: title += str(round(final_pose[0,i],3)) title += ', ' else: title = "Predicted Transformation: " for i in range(len(final_pose[0])): if i>2: title += str(round(final_pose[0,i]*(180/np.pi),3)) else: title += str(round(final_pose[0,i],3)) title += ', ' title += '\nElapsed Time: '+str(np.round(elapsed_time*1000,3))+' ms'+' & Iterations: '+str(itr) title += ' & Iterative Network' self.find_errors(gt_pose, final_pose) if swap_case: helper.display_three_clouds(source_data[0], template_data[0], transformed_source_data, title) else: helper.display_three_clouds(template_data[0], source_data[0], transformed_source_data, title)
def main():
parser = argparse.ArgumentParser(description='Point Cloud Registration')
# Settings for network.
parser.add_argument('--exp_name',
type=str,
default='exp',
metavar='N',
help='Name of the experiment')
parser.add_argument('--model',
type=str,
default='pcrnet',
metavar='N',
choices=['dcp'],
help='Model to use, [dcp]')
parser.add_argument('--emb_nn',
type=str,
default='pointnet',
metavar='N',
choices=['pointnet', 'dgcnn'],
help='Embedding nn to use, [pointnet, dgcnn]')
parser.add_argument(
'--pointer',
type=str,
default='identity',
metavar='N',
choices=['identity', 'transformer'],
help='Attention-based pointer generator to use, [identity, transformer]'
)
parser.add_argument('--head',
type=str,
default='mlp',
metavar='N',
choices=[
'mlp',
'svd',
],
help='Head to use, [mlp, svd]')
parser.add_argument('--iterations',
type=int,
default=8,
help='[No of iterations for PCRNet]')
# Settings for training
parser.add_argument('--emb_dims',
type=int,
default=1024,
metavar='N',
help='Dimension of embeddings')
parser.add_argument('--batch_size',
type=int,
default=32,
metavar='batch_size',
help='Size of batch)')
parser.add_argument('--test_batch_size',
type=int,
default=10,
metavar='batch_size',
help='Size of batch)')
parser.add_argument('--epochs',
type=int,
default=250,
metavar='N',
help='number of episode to train ')
parser.add_argument('--device',
action='store_true',
default=False,
help='enables CUDA training')
parser.add_argument('--seed',
type=int,
default=1234,
metavar='S',
help='random seed (default: 1)')
parser.add_argument('--eval',
action='store_true',
default=False,
help='evaluate the model')
# Settings for attention
parser.add_argument('--n_blocks',
type=int,
default=1,
metavar='N',
help='Num of blocks of encoder&decoder')
parser.add_argument('--n_heads',
type=int,
default=4,
metavar='N',
help='Num of heads in multiheadedattention')
parser.add_argument('--ff_dims',
type=int,
default=1024,
metavar='N',
help='Num of dimensions of fc in transformer')
parser.add_argument('--dropout',
type=float,
default=0.0,
metavar='N',
help='Dropout ratio in transformer')
parser.add_argument('--cycle',
type=bool,
default=False,
metavar='N',
help='Whether to use cycle consistency')
# Settings for dataset
parser.add_argument('--gaussian_noise',
type=bool,
default=False,
metavar='N',
help='Wheter to add gaussian noise')
parser.add_argument('--unseen',
type=bool,
default=False,
metavar='N',
help='Wheter to test on unseen category')
parser.add_argument('--num_points',
type=int,
default=1024,
metavar='N',
help='Num of points to use')
parser.add_argument('--dataset',
type=str,
default='modelnet40',
choices=['modelnet40'],
metavar='N',
help='dataset to use')
parser.add_argument('--factor',
type=float,
default=4,
metavar='N',
help='Divided factor for rotations')
parser.add_argument('--model_path',
type=str,
default='./pretrained/model.best.t7',
metavar='N',
help='Pretrained model path')
args = parser.parse_args()
use_cuda = torch.cuda.is_available()
args.device = torch.device("cuda" if use_cuda else "cpu")
torch.backends.cudnn.deterministic = True
torch.manual_seed(args.seed)
torch.cuda.manual_seed_all(args.seed)
np.random.seed(args.seed)
if args.model == 'pcrnet':
net = PCRNet(args).to(args.device)
model_path = args.model_path
net.load_state_dict(torch.load(model_path), strict=False)
else:
raise Exception('Not implemented')
source, template, transformed_src, _, _ = test_one_pair(args, net)
print(source.shape, template.shape, transformed_src.shape)
import helper
helper.display_three_clouds(template[0], source[0], transformed_src[0],
'Results')
print('FINISH')
def test_one_epoch(sess, ops_L, templates, shuffled_poses, saver, model_path):
# Arguments:
# sess: Tensorflow session to handle tensors.
# ops_L: Dictionary for tensors of Network_L
# ops19: Dictionary for tensors of Network19
# templates: Training Point Cloud data.
# poses: Training pose data.
# saver: To restore the weights.
# model_path: Path of log directory.
saver.restore(sess, model_path) # Restore the weights of trained network.
is_training = False
display_ptClouds = False
display_poses = False
display_poses_in_itr = False
display_ptClouds_in_itr = False
swap_case = False
templates = helper.process_templates('templates')
template_data = np.zeros((BATCH_SIZE, MAX_NUM_POINT,
3)) # Extract Templates for batch training.
for i in range(BATCH_SIZE):
template_data[i, :, :] = np.copy(templates[1, :, :])
batch_euler_poses = shuffled_poses[0].reshape(
(1, 6)) # Extract poses for batch training.
# Self defined test case.
batch_euler_poses[0] = [
0.5, 0.0, 0.2, 50 * (np.pi / 180), 0 * (np.pi / 180),
10 * (np.pi / 180)
]
source_data = helper.apply_transformation(
template_data, batch_euler_poses
) # Apply the poses on the templates to get source data.
# Only chose limited number of points from the source and template data.
template_data = template_data[:, 0:NUM_POINT, :]
source_data = source_data[:, 0:NUM_POINT, :]
if swap_case:
source_data, template_data = template_data, source_data # Swap the template and source.
transformation_template2source = helper.transformation(
batch_euler_poses)
transformation_source2template = np.linalg.inv(
transformation_template2source[0])
[euler_z, euler_y,
euler_x] = t3d.mat2euler(transformation_source2template[0:3, 0:3],
'szyx')
trans_x = transformation_source2template[0, 3]
trans_y = transformation_source2template[1, 3]
trans_z = transformation_source2template[2, 3]
pose_source2template = [
trans_x, trans_y, trans_z, euler_x * (18 / np.pi),
euler_y * (180 / np.pi), euler_z * (180 / np.pi)
]
batch_euler_poses[0] = pose_source2template
TEMPLATE_DATA = np.copy(
template_data) # Store the initial template to visualize results.
SOURCE_DATA = np.copy(
source_data) # Store the initial source to visualize results.
# To visualize the source and point clouds:
if display_ptClouds:
helper.display_clouds_data(source_data[0])
helper.display_clouds_data(template_data[0])
# Subtract the Centroids from the Point Clouds.
if centroid_subtraction_switch:
source_data, template_data, centroid_translation_pose = helper.centroid_subtraction(
source_data, template_data)
TRANSFORMATIONS = np.identity(
4) # Initialize identity transformation matrix.
TRANSFORMATIONS = np.matlib.repmat(TRANSFORMATIONS, BATCH_SIZE, 1).reshape(
BATCH_SIZE, 4,
4) # Intialize identity matrices of size equal to batch_size
# Feed the placeholders of Network_L with source data and template data obtained from N-Iterations.
feed_dict = {
ops_L['source_pointclouds_pl']: source_data,
ops_L['template_pointclouds_pl']: template_data,
ops_L['is_training_pl']: is_training
}
# Ask the network to predict transformation, calculate loss using distance between actual points.
import time
start = time.time()
step, predicted_transformation = sess.run(
[ops_L['step'], ops_L['predicted_transformation']],
feed_dict=feed_dict)
end = time.time()
print(end - start)
# Apply the final transformation on the template data and multiply it with the transformation matrix obtained from N-Iterations.
TRANSFORMATIONS, template_data = helper.transformation_quat2mat(
predicted_transformation, TRANSFORMATIONS, template_data)
if centroid_subtraction_switch: # If centroid is subtracted then apply the centorid translation back to point clouds.
TRANSFORMATIONS, template_data = helper.transformation_quat2mat(
centroid_translation_pose, TRANSFORMATIONS, template_data)
final_pose = helper.find_final_pose(TRANSFORMATIONS)
if not swap_case:
title = "Actual T (Red->Green): "
for i in range(len(batch_euler_poses[0])):
if i > 2:
title += str(round(batch_euler_poses[0][i] * (180 / np.pi), 2))
else:
title += str(batch_euler_poses[0][i])
title += ', '
title += "\nPredicted T (Red->Blue): "
for i in range(len(final_pose[0])):
if i > 2:
title += str(round(final_pose[0, i] * (180 / np.pi), 3))
else:
title += str(round(final_pose[0, i], 3))
title += ', '
else:
title = "Predicted T (Red->Blue): "
for i in range(len(final_pose[0])):
if i > 2:
title += str(round(final_pose[0, i] * (180 / np.pi), 3))
else:
title += str(round(final_pose[0, i], 3))
title += ', '
# Display the ground truth pose and predicted pose for first Point Cloud in batch
if display_poses:
print('Ground Truth Position: {}'.format(
batch_euler_poses[0, 0:3].tolist()))
print('Predicted Position: {}'.format(final_pose[0, 0:3].tolist()))
print('Ground Truth Orientation: {}'.format(
(batch_euler_poses[0, 3:6] * (180 / np.pi)).tolist()))
print('Predicted Orientation: {}'.format(
(final_pose[0, 3:6] * (180 / np.pi)).tolist()))
helper.display_three_clouds(TEMPLATE_DATA[0], SOURCE_DATA[0],
template_data[0], title)
print("Loss: {}".format(loss_val))
def run(self, source, template):
if not torch.cuda.is_available():
self.device = 'cpu'
self.device = torch.device(self.device)
model = self.create_model()
if self.model_path:
assert os.path.isfile(self.model_path)
model.load_state_dict(torch.load(self.model_path, map_location='cpu'))
model.to(self.device)
# testing
return (self.eval_1(model, source, template)).cpu().numpy()
if __name__ == '__main__':
act = Action(os.path.join('results_train_data','ex1_pointlk_0915_model_best.pth'))
templates, sources, poses = helper.read_partial_data('train_data', 'partial_data.h5')
template_data = templates[0].reshape(1,-1,3)
source_data = sources[0].reshape(1,-1,3)
#source_data = helper.apply_transformation(template_data, poses[0].reshape(1,-1))
transformation = act.run(source_data, template_data)
pred_data = np.matmul(transformation[0,0:3,0:3],source_data[0].T).T
pred_data = pred_data - np.mean(pred_data, axis=0, keepdims=True)
helper.display_three_clouds(template_data[0], source_data[0], pred_data, "PointNetLK Partial Source Result")
def test_one_epoch(sess, ops, templates, poses, saver, model_path):
# Arguments:
# sess: Tensorflow session to handle tensors.
# ops: Dictionary for tensors of Network
# templates: Training Point Cloud data.
# poses: Training pose data.
# saver: To restore the weights.
# model_path: Path of log directory.
saver.restore(sess, model_path) # Restore the weights of trained network.
is_training = False
display_ptClouds = False
display_poses = False
display_poses_in_itr = False
display_ptClouds_in_itr = False
swap_case = False
MAX_LOOPS = 4
template_data = np.zeros((BATCH_SIZE, MAX_NUM_POINT,
3)) # Extract Templates for batch training.
template_data[0] = np.copy(templates[FLAGS.template_idx, :, :])
batch_euler_poses = poses[0].reshape(
(1, 6)) # Extract poses for batch training.
# Define test case.
batch_euler_poses[0] = [
0.4, 0.5, 0.1, 10 * (np.pi / 180), 20 * (np.pi / 180),
20 * (np.pi / 180)
]
source_data = helper.apply_transformation(
template_data, batch_euler_poses
) # Apply the poses on the templates to get source data.
# Chose Random Points from point clouds for training.
if np.random.random_sample() < 0:
source_data = helper.select_random_points(
source_data, NUM_POINT
) # probability that source data has different points than template
else:
source_data = source_data[:, 0:NUM_POINT, :]
# Add noise to source point cloud.
if np.random.random_sample() < 1.0:
source_data = helper.add_noise(source_data)
# Only choose limited number of points from the source and template data.
source_data = source_data[:, 0:NUM_POINT, :]
template_data = template_data[:, 0:NUM_POINT, :]
TEMPLATE_DATA = np.copy(
template_data) # Store the initial template to visualize results.
SOURCE_DATA = np.copy(
source_data) # Store the initial source to visualize results.
# Subtract the Centroids from the Point Clouds.
if centroid_subtraction_switch:
source_data = source_data - np.mean(source_data, axis=1, keepdims=True)
template_data = template_data - np.mean(
template_data, axis=1, keepdims=True)
# To visualize the source and point clouds:
if display_ptClouds:
helper.display_clouds_data(source_data[0])
helper.display_clouds_data(template_data[0])
TRANSFORMATIONS = np.identity(
4) # Initialize identity transformation matrix.
TRANSFORMATIONS = npm.repmat(TRANSFORMATIONS, BATCH_SIZE, 1).reshape(
BATCH_SIZE, 4,
4) # Intialize identity matrices of size equal to batch_size
# Store the transformed point clouds after each iteration.
ITR = np.zeros((MAX_LOOPS, template_data.shape[0], template_data.shape[1],
template_data.shape[2]))
# Iterations for pose refinement.
for loop_idx in range(MAX_LOOPS - 1):
# 4a
# Feed the placeholders of Network with template data and source data.
feed_dict = {
ops['source_pointclouds_pl']: source_data,
ops['template_pointclouds_pl']: template_data,
ops['is_training_pl']: is_training
}
predicted_transformation = sess.run(
[ops['predicted_transformation']],
feed_dict=feed_dict) # Ask the network to predict the pose.
#print (predicted_transformation[0])
# 4b,4c
# Apply the transformation on the template data and multiply it to transformation matrix obtained in previous iteration.
TRANSFORMATIONS, source_data = helper.transformation_quat2mat(
predicted_transformation, TRANSFORMATIONS, source_data)
# Display Results after each iteration.
if display_poses_in_itr:
print(predicted_transformation[0, 0:3])
print(predicted_transformation[0, 3:7] * (180 / np.pi))
if display_ptClouds_in_itr:
helper.display_clouds_data(source_data[0])
ITR[loop_idx, :, :, :] = source_data
# Feed the placeholders of Network with source data and template data obtained from N-Iterations.
feed_dict = {
ops['source_pointclouds_pl']: source_data,
ops['template_pointclouds_pl']: template_data,
ops['is_training_pl']: is_training
}
# Ask the network to predict transformation, calculate loss using distance between actual points.
step, predicted_transformation = sess.run(
[ops['step'], ops['predicted_transformation']], feed_dict=feed_dict)
# Apply the final transformation on the template data and multiply it with the transformation matrix obtained from N-Iterations.
TRANSFORMATIONS, source_data = helper.transformation_quat2mat(
predicted_transformation, TRANSFORMATIONS, source_data)
final_pose = helper.find_final_pose_inv(TRANSFORMATIONS)
final_pose[0, 0:3] = final_pose[0, 0:3] + np.mean(SOURCE_DATA, axis=1)[0]
title = "Actual T (Red->Green): "
for i in range(len(batch_euler_poses[0])):
if i > 2:
title += str(round(batch_euler_poses[0][i] * (180 / np.pi), 2))
else:
title += str(batch_euler_poses[0][i])
title += ', '
title += "\nPredicted T (Red->Blue): "
for i in range(len(final_pose[0])):
if i > 2:
title += str(round(final_pose[0, i] * (180 / np.pi), 3))
else:
title += str(round(final_pose[0, i], 3))
title += ', '
# Display the ground truth pose and predicted pose for first Point Cloud in batch
if display_poses:
print('Ground Truth Position: {}'.format(
batch_euler_poses[0, 0:3].tolist()))
print('Predicted Position: {}'.format(final_pose[0, 0:3].tolist()))
print('Ground Truth Orientation: {}'.format(
(batch_euler_poses[0, 3:6] * (180 / np.pi)).tolist()))
print('Predicted Orientation: {}'.format(
(final_pose[0, 3:6] * (180 / np.pi)).tolist()))
helper.display_three_clouds(TEMPLATE_DATA[0], SOURCE_DATA[0],
source_data[0], title)
template_idx = 201
NUM_POINT = 1024
MAX_LOOPS = 500
M_given = templates[template_idx, :, :]
x_trans = 0.5
y_trans = 0.5
z_trans = 0.5
x_rot = 45 * (np.pi / 180)
y_rot = 45 * (np.pi / 180)
z_rot = 45 * (np.pi / 180)
gt_pose = np.array([[x_trans, y_trans, z_trans, x_rot, y_rot,
z_rot]]) # Pose: Source to Template
S_given = helper.apply_transformation(M_given.reshape((1, -1, 3)),
gt_pose)[0]
S_given = S_given + np.random.normal(0, 1, S_given.shape)
M_given = M_given[0:NUM_POINT, :] # template data
S_given = S_given[0:NUM_POINT, :] # source data
tree_M = KDTree(M_given)
tree_M_sampled = KDTree(M_given[0:100, :])
final_pose, model_data, sensor_data, predicted_data, title = icp_test(
S_given[0:100, :], M_given, tree_M, M_given[0:100, :], tree_M_sampled,
S_given, gt_pose.reshape((1, 6)), MAX_LOOPS, 1e-05)
helper.display_three_clouds(model_data, sensor_data, predicted_data, title)
