def pc_perspective_transform(cfg, point_cloud,
transform, predicted_translation=None,
focal_length=None):
"""
:param cfg:
:param point_cloud: [B, N, 3]
:param transform: [B, 4] if quaternion or [B, 4, 4] if camera matrix
:param predicted_translation: [B, 3] translation vector
:return:
"""
camera_distance = cfg.camera_distance
if focal_length is None:
focal_length = cfg.focal_length
else:
focal_length = tf.expand_dims(focal_length, axis=-1)
if cfg.pose_quaternion:
pc2 = quaternion_rotate(point_cloud, transform)
if predicted_translation is not None:
predicted_translation = tf.expand_dims(predicted_translation, axis=1)
pc2 += predicted_translation
xs = tf.slice(pc2, [0, 0, 2], [-1, -1, 1])
ys = tf.slice(pc2, [0, 0, 1], [-1, -1, 1])
zs = tf.slice(pc2, [0, 0, 0], [-1, -1, 1])
# translation part of extrinsic camera
zs += camera_distance
# intrinsic transform
xs *= focal_length
ys *= focal_length
else:
xyz1 = tf.pad(point_cloud, tf.constant([[0, 0], [0, 0], [0, 1]]), "CONSTANT", constant_values=1.0)
extrinsic = transform
intr = intrinsic_matrix(cfg, dims=4)
intrinsic = tf.convert_to_tensor(intr)
intrinsic = tf.expand_dims(intrinsic, axis=0)
intrinsic = tf.tile(intrinsic, [tf.shape(extrinsic)[0], 1, 1])
full_cam_matrix = tf.matmul(intrinsic, extrinsic)
pc2 = tf.matmul(xyz1, tf.transpose(full_cam_matrix, [0, 2, 1]))
# TODO unstack instead of split
xs = tf.slice(pc2, [0, 0, 2], [-1, -1, 1])
ys = tf.slice(pc2, [0, 0, 1], [-1, -1, 1])
zs = tf.slice(pc2, [0, 0, 0], [-1, -1, 1])
xs /= zs
ys /= zs
zs -= camera_distance
if predicted_translation is not None:
zt = tf.slice(predicted_translation, [0, 0, 0], [-1, -1, 1])
zs -= zt
xyz2 = tf.concat([zs, ys, xs], axis=2)
return xyz2
def compute_alignment_candidates(cfg, dataset, all_rotations_file):
exp_dir = cfg.checkpoint_dir
save_pred_name = cfg.save_predictions_dir
save_dir = os.path.join(exp_dir, save_pred_name)
gt_dir = os.path.join(cfg.gt_pc_dir, cfg.synth_set)
num_models = dataset.num_samples()
num_views = cfg.num_views
num_show = 5
rmse = np.ones((num_models, num_views), np.float32)
rotations = np.zeros((num_models, num_views, 4), np.float32)
for model_idx in range(num_models):
sample = dataset.data[model_idx]
print(f"{model_idx}/{num_models}")
print(sample.name)
gt_filename = f"{gt_dir}/{sample.name}.mat"
if not os.path.isfile(gt_filename):
continue
obj = scipy.io.loadmat(gt_filename)
gt = obj["points"]
data = scipy.io.loadmat(f"{save_dir}/{sample.name}_pc.mat")
all_pcs = np.squeeze(data["points"])
all_cameras = data["camera_pose"]
for view_idx in range(num_views):
pred_pc_ref = all_pcs[view_idx, :, :]
quat_pred = all_cameras[view_idx, :]
cam_pos = sample.cam_pos[view_idx, :]
quat_gt = quaternion_from_campos(cam_pos)
quat, rmse_view = alignment_to_ground_truth(
pred_pc_ref, quat_pred, gt, quat_gt)
rmse[model_idx, view_idx] = rmse_view
rotations[model_idx, view_idx, :] = np.squeeze(quat)
if cfg.vis_voxels and view_idx < num_show:
rotated_with_quat = quaternion_rotate(
np.expand_dims(pred_pc_ref, 0), quat)
draw_registration_result(np.squeeze(rotated_with_quat), gt)
print("rmse:", rmse[model_idx, :])
scipy.io.savemat(all_rotations_file,
mdict={
"rotations": rotations,
"rmse": rmse
})
def model_unrotate_points(cfg):
"""
un_q = quat_gt^(-1) * predicted_quat
pc_unrot = un_q * pc_np * un_q^(-1)
"""
from util.quaternion import quaternion_normalise, quaternion_conjugate, \
quaternion_rotate, quaternion_multiply
gt_quat = tf.placeholder(dtype=tf.float32, shape=[1, 4])
pred_quat_n = quaternion_normalise(pred_quat)
gt_quat_n = quaternion_normalise(gt_quat)
un_q = quaternion_multiply(quaternion_conjugate(gt_quat_n), pred_quat_n)
pc_unrot = quaternion_rotate(input_pc, un_q)
return input_pc, pred_quat, gt_quat, pc_unrot
has_number = True
else:
has_number = False
obj = scipy.io.loadmat(gt_filename)
Vgt = obj["points"]
for i in range(num_views):
chamfer_dists_current = np.zeros((2), dtype=np.float64)
pred = all_pcs[i, :, :]
if has_number:
pred = pred[0:all_pcs_nums[i], :]
if eval_unsup:
pred = np.expand_dims(pred, 0)
pred = quaternion_rotate(torch.from_numpy(pred).cuda(),
torch.from_numpy(reference_rotation).cuda()).cpu().numpy()
pred = np.squeeze(pred)
pred_to_gt, idx_np = compute_distance(cfg, pred, Vgt)
gt_to_pred, _ = compute_distance(cfg, Vgt, pred)
chamfer_dists_current[0] = np.mean(pred_to_gt)
chamfer_dists_current[1] = np.mean(gt_to_pred)
is_nan = np.isnan(pred_to_gt)
assert (not np.any(is_nan))
campos = sample['cam_pos'][i]
images = sample['image'][i].transpose(1,2,0)
g = get_group(campos, divs)
# chamfer_dict[g] = np.concatenate((chamfer_dict[g], np.expand_dims(chamfer_dists_current, 0)))
def compute_predictions():
cfg = app_config
setup_environment(cfg)
exp_dir = cfg.checkpoint_dir
cfg.batch_size = 1
cfg.step_size = 1
pc_num_points = cfg.pc_num_points
vox_size = cfg.vox_size
save_pred = cfg.save_predictions
save_voxels = cfg.save_voxels
fast_conversion = True
pose_student = cfg.pose_predictor_student and cfg.predict_pose
g = tf.Graph()
with g.as_default():
model = model_pc.ModelPointCloud(cfg)
out = build_model(model)
input_image = out["inputs"]
cam_matrix = out["camera_extr_src"]
cam_quaternion = out["cam_quaternion"]
point_cloud = out["points_1"]
rgb = out["rgb_1"] if cfg.pc_rgb else tf.no_op()
projs = out["projs"]
projs_rgb = out["projs_rgb"]
projs_depth = out["projs_depth"]
cam_transform = out["cam_transform"]
z_latent = out["z_latent"]
if pose_student:
proj_student, camera_pose_student = model_student(
input_image, model)
input_pc = tf.placeholder(tf.float32, [cfg.batch_size, None, 3])
if save_voxels:
if fast_conversion:
voxels, _ = pointcloud2voxels3d_fast(cfg, input_pc, None)
voxels = tf.expand_dims(voxels, axis=-1)
voxels = smoothen_voxels3d(cfg, voxels, model.gauss_kernel())
else:
voxels = pointcloud2voxels(cfg, input_pc, model.gauss_sigma())
q_inp = tf.placeholder(tf.float32, [1, 4])
q_matrix = as_rotation_matrix(q_inp)
input_pc, pred_quat, gt_quat, pc_unrot = model_unrotate_points(cfg)
pc_rot = quaternion_rotate(input_pc, pred_quat)
config = tf.ConfigProto(device_count={'GPU': 1})
config.gpu_options.per_process_gpu_memory_fraction = cfg.per_process_gpu_memory_fraction
sess = tf.Session(config=config)
sess.run(tf.global_variables_initializer())
sess.run(tf.local_variables_initializer())
variables_to_restore = slim.get_variables_to_restore(exclude=["meta"])
restorer = tf.train.Saver(variables_to_restore)
checkpoint_file = tf.train.latest_checkpoint(exp_dir)
print("restoring checkpoint", checkpoint_file)
restorer.restore(sess, checkpoint_file)
save_dir = os.path.join(exp_dir,
'{}_vis_proj'.format(cfg.save_predictions_dir))
mkdir_if_missing(save_dir)
save_pred_dir = os.path.join(exp_dir, cfg.save_predictions_dir)
mkdir_if_missing(save_pred_dir)
vis_size = cfg.vis_size
dataset = Dataset3D(cfg)
pose_num_candidates = cfg.pose_predict_num_candidates
num_views = cfg.num_views
plot_h = 4
plot_w = 6
num_views = int(min(num_views, plot_h * plot_w / 2))
if cfg.models_list:
model_names = parse_lines(cfg.models_list)
else:
model_names = [sample.name for sample in dataset.data]
num_models = len(model_names)
for k in range(num_models):
model_name = model_names[k]
sample = dataset.sample_by_name(model_name)
images = sample.image
masks = sample.mask
if cfg.saved_camera:
cameras = sample.camera
cam_pos = sample.cam_pos
if cfg.vis_depth_projs:
depths = sample.depth
if cfg.variable_num_views:
num_views = sample.num_views
print("{}/{} {}".format(k, num_models, model_name))
if pose_num_candidates == 1:
grid = np.empty((plot_h, plot_w), dtype=object)
else:
plot_w = pose_num_candidates + 1
if pose_student:
plot_w += 1
grid = np.empty((num_views, plot_w), dtype=object)
if save_pred:
all_pcs = np.zeros((num_views, pc_num_points, 3))
all_cameras = np.zeros((num_views, 4))
all_voxels = np.zeros((num_views, vox_size, vox_size, vox_size))
all_z_latent = np.zeros((num_views, cfg.fc_dim))
for view_idx in range(num_views):
input_image_np = images[[view_idx], :, :, :]
gt_mask_np = masks[[view_idx], :, :, :]
if cfg.saved_camera:
extr_mtr = cameras[view_idx, :, :]
cam_quaternion_np = quaternion_from_campos(
cam_pos[view_idx, :])
cam_quaternion_np = np.expand_dims(cam_quaternion_np, axis=0)
else:
extr_mtr = np.zeros((4, 4))
if cfg.pc_rgb:
proj_tensor = projs_rgb
elif cfg.vis_depth_projs:
proj_tensor = projs_depth
else:
proj_tensor = projs
(pc_np, rgb_np, proj_np, cam_transf_np, z_latent_np) = sess.run(
[point_cloud, rgb, proj_tensor, cam_transform, z_latent],
feed_dict={
input_image: input_image_np,
cam_matrix: extr_mtr,
cam_quaternion: cam_quaternion_np
})
if pose_student:
(proj_student_np, camera_student_np) = sess.run(
[proj_student, camera_pose_student],
feed_dict={input_image: input_image_np})
predicted_camera = camera_student_np
else:
predicted_camera = cam_transf_np
if cfg.vis_depth_projs:
proj_np = normalise_depthmap(proj_np)
if depths is not None:
depth_np = depths[view_idx, :, :, :]
depth_np = normalise_depthmap(depth_np)
else:
depth_np = 1.0 - np.squeeze(gt_mask_np)
if pose_student:
proj_student_np = normalise_depthmap(proj_student_np)
if cfg.predict_pose:
if cfg.save_rotated_points:
ref_rot = scipy.io.loadmat(
"{}/final_reference_rotation.mat".format(exp_dir))
ref_rot = ref_rot["rotation"]
pc_np_unrot = sess.run(pc_rot,
feed_dict={
input_pc: pc_np,
pred_quat: ref_rot
})
pc_np = pc_np_unrot
if cfg.pc_rgb:
gt_image = input_image_np
elif cfg.vis_depth_projs:
gt_image = depth_np
else:
gt_image = gt_mask_np
if pose_num_candidates == 1:
view_j = view_idx * 2 // plot_w
view_i = view_idx * 2 % plot_w
gt_image = np.squeeze(gt_image)
grid[view_j, view_i] = mask4vis(cfg, gt_image, vis_size)
curr_img = np.squeeze(proj_np)
grid[view_j, view_i + 1] = mask4vis(cfg, curr_img, vis_size)
if cfg.save_individual_images:
curr_dir = os.path.join(save_dir, sample.name)
if not os.path.exists(curr_dir):
os.makedirs(curr_dir)
imageio.imwrite(
os.path.join(curr_dir,
'{}_{}.png'.format(view_idx, 'rgb_gt')),
mask4vis(cfg, np.squeeze(input_image_np), vis_size))
imageio.imwrite(
os.path.join(curr_dir,
'{}_{}.png'.format(view_idx,
'mask_pred')),
mask4vis(cfg, np.squeeze(proj_np), vis_size))
else:
view_j = view_idx
gt_image = np.squeeze(gt_image)
grid[view_j, 0] = mask4vis(cfg, gt_image, vis_size)
for kk in range(pose_num_candidates):
curr_img = np.squeeze(proj_np[kk, :, :, :])
grid[view_j, kk + 1] = mask4vis(cfg, curr_img, vis_size)
if cfg.save_individual_images:
curr_dir = os.path.join(save_dir, sample.name)
if not os.path.exists(curr_dir):
os.makedirs(curr_dir)
imageio.imwrite(
os.path.join(
curr_dir,
'{}_{}_{}.png'.format(view_idx, kk,
'mask_pred')),
mask4vis(cfg, np.squeeze(curr_img), vis_size))
if cfg.save_individual_images:
imageio.imwrite(
os.path.join(curr_dir,
'{}_{}.png'.format(view_idx, 'mask_gt')),
mask4vis(cfg, np.squeeze(gt_mask_np), vis_size))
if pose_student:
grid[view_j,
-1] = mask4vis(cfg, np.squeeze(proj_student_np),
vis_size)
if save_pred:
all_pcs[view_idx, :, :] = np.squeeze(pc_np)
all_z_latent[view_idx] = z_latent_np
if cfg.predict_pose:
all_cameras[view_idx, :] = predicted_camera
if save_voxels:
# multiplying by two is necessary because
# pc->voxel conversion expects points in [-1, 1] range
pc_np_range = pc_np
if not fast_conversion:
pc_np_range *= 2.0
voxels_np = sess.run(voxels,
feed_dict={input_pc: pc_np_range})
all_voxels[view_idx, :, :, :] = np.squeeze(voxels_np)
vis_view = view_idx == 0 or cfg.vis_all_views
if cfg.vis_voxels and vis_view:
rgb_np = np.squeeze(rgb_np) if cfg.pc_rgb else None
vis_pc(np.squeeze(pc_np), rgb=rgb_np)
grid_merged = merge_grid(cfg, grid)
imageio.imwrite("{}/{}_proj.png".format(save_dir, sample.name),
grid_merged)
if save_pred:
if cfg.save_as_mat:
save_dict = {"points": all_pcs, "z_latent": all_z_latent}
if cfg.predict_pose:
save_dict["camera_pose"] = all_cameras
scipy.io.savemat("{}/{}_pc".format(save_pred_dir, sample.name),
mdict=save_dict)
else:
np.savez("{}/{}_pc".format(save_pred_dir, sample.name),
all_pcs)
if save_voxels:
np.savez("{}/{}_vox".format(save_pred_dir, sample.name),
all_voxels)
sess.close()
def pc_perspective_transform(cfg,
point_cloud,
transform,
predicted_translation=None,
focal_length=None):
"""
:param cfg:
:param point_cloud: [B, N, 3]
:param transform: [B, 4] if quaternion or [B, 4, 4] if camera matrix
:param predicted_translation: [B, 3] translation vector
:return:
"""
camera_distance = cfg.camera_distance
if focal_length is None:
focal_length = cfg.focal_length
else:
focal_length = focal_length.unsqueeze(-1)
if cfg.pose_quaternion:
pc2 = quaternion_rotate(point_cloud, transform)
if predicted_translation is not None:
predicted_translation = predicted_translation.unsqueeze(1)
pc2 += predicted_translation
xs = pc2[:, :, 2:3]
ys = pc2[:, :, 1:2]
zs = pc2[:, :, 0:1]
# translation part of extrinsic camera
zs += camera_distance
# intrinsic transform
xs *= focal_length
ys *= focal_length
else:
xyz1 = F.pad(point_cloud, (0, 1))
extrinsic = transform
if instr_global is None:
intr = intrinsic_matrix(cfg, dims=4)
intrinsic = torch.from_numpy(intr).cuda()
intrinsic = intrinsic.reshape(1, intrinsic.shape[0],
intrinsic.shape[1])
instr_global = intrinsic
intrinsic = instr_global
intrinsic = intrinsic.repeat(extrinsic.shape[0], 1, 1)
full_cam_matrix = torch.matmul(intrinsic, extrinsic)
pc2 = torch.matmul(xyz1, full_cam_matrix.permute(0, 2, 1))
# TODO unstack instead of split
xs = pc2[:, :, 2:3]
ys = pc2[:, :, 1:2]
zs = pc2[:, :, 0:1]
xs = torch.div(xs, zs)
ys = torch.div(ys, zs)
zs = zs - camera_distance
if predicted_translation is not None:
zt = predicted_translation[:, :, 0:1]
zs -= zt
xyz2 = torch.cat([zs, ys, xs], dim=2)
return xyz2
def run_eval():
cfg = app_config
exp_dir = cfg.checkpoint_dir
num_views = cfg.num_views
eval_unsup = cfg.eval_unsupervised_shape
dataset_folder = cfg.inp_dir
gt_dir = os.path.join(cfg.gt_pc_dir, cfg.synth_set)
#g = tf.Graph()
#with g.as_default():
# source_pc = tf.placeholder(dtype=tf.float64, shape=[None, 3])
# target_pc = tf.placeholder(dtype=tf.float64, shape=[None, 3])
# quat_tf = tf.placeholder(dtype=tf.float64, shape=[1, 4])
# _, min_dist, min_idx = point_cloud_distance(source_pc, target_pc)
# source_pc_2 = tf.placeholder(dtype=tf.float64, shape=[1, None, 3])
# rotated_pc = quaternion_rotate(source_pc_2, quat_tf)
# sess = tf.Session(config=config)
# sess.run(tf.global_variables_initializer())
# sess.run(tf.local_variables_initializer())
save_pred_name = "{}_{}".format(cfg.save_predictions_dir, cfg.eval_split)
save_dir = os.path.join(exp_dir, cfg.save_predictions_dir)
device = 'cuda' if torch.cuda.is_available() else 'cpu'
if eval_unsup:
reference_rotation = scipy.io.loadmat(
"{}/final_reference_rotation.mat".format(exp_dir))["rotation"]
dataset = ShapeRecords(dataset_folder, cfg, 'test')
if cfg.models_list:
model_names = parse_lines(cfg.models_list)
else:
model_names = dataset.file_names
num_models = len(model_names)
chamfer_dists = np.zeros((0, num_views, 2), dtype=np.float64)
for k in range(num_models):
sample = dataset.__getitem__(k)
print("{}/{}".format(k, num_models))
print(model_names[k])
gt_filename = "{}/{}.mat".format(gt_dir, model_names[k]).replace(
'_features.p', '')
mat_filename = "{}/{}_pc.pkl".format(save_dir, model_names[k])
if not os.path.isfile(gt_filename) or not os.path.isfile(mat_filename):
continue
with open(mat_filename, 'rb') as handle:
data = pickle.load(handle)
all_pcs = np.squeeze(data["points"])
if "num_points" in data:
all_pcs_nums = np.squeeze(data["num_points"])
has_number = True
else:
has_number = False
obj = scipy.io.loadmat(gt_filename)
Vgt = obj["points"]
chamfer_dists_current = np.zeros((num_views, 2), dtype=np.float64)
for i in range(num_views):
pred = all_pcs[i, :, :]
if has_number:
pred = pred[0:all_pcs_nums[i], :]
if eval_unsup:
pred = np.expand_dims(pred, 0)
pred = quaternion_rotate(
torch.from_numpy(pred).to(device),
torch.from_numpy(reference_rotation).to(
device)).cpu().numpy()
pred = np.squeeze(pred)
pred_to_gt, idx_np = compute_distance(cfg, pred, Vgt)
gt_to_pred, _ = compute_distance(cfg, Vgt, pred)
chamfer_dists_current[i, 0] = np.mean(pred_to_gt)
chamfer_dists_current[i, 1] = np.mean(gt_to_pred)
is_nan = np.isnan(pred_to_gt)
assert (not np.any(is_nan))
print(i, ":", chamfer_dists_current)
current_mean = np.mean(chamfer_dists_current, 0)
print("total:", current_mean)
chamfer_dists = np.concatenate(
(chamfer_dists, np.expand_dims(chamfer_dists_current, 0)))
final = np.mean(chamfer_dists, axis=(0, 1)) * 100
print(final)
scipy.io.savemat(
os.path.join(exp_dir, "chamfer_{}.mat".format(save_pred_name)), {
"chamfer": chamfer_dists,
"model_names": to_np_object(model_names)
})
file = open(os.path.join(exp_dir, "chamfer_{}.txt".format(save_pred_name)),
"w")
file.write("{} {}\n".format(final[0], final[1]))
file.close()
def run_eval():
divs = 2
cfg = app_config
exp_dir = cfg.checkpoint_dir
num_views = cfg.num_views
eval_unsup = cfg.eval_unsupervised_shape
dataset_folder = cfg.inp_dir
gt_dir = os.path.join(cfg.gt_pc_dir, cfg.synth_set)
# g = tf.Graph()
# with g.as_default():
# source_pc = tf.placeholder(dtype=tf.float64, shape=[None, 3])
# target_pc = tf.placeholder(dtype=tf.float64, shape=[None, 3])
# quat_tf = tf.placeholder(dtype=tf.float64, shape=[1, 4])
# _, min_dist, min_idx = point_cloud_distance(source_pc, target_pc)
# source_pc_2 = tf.placeholder(dtype=tf.float64, shape=[1, None, 3])
# rotated_pc = quaternion_rotate(source_pc_2, quat_tf)
# sess = tf.Session(config=config)
# sess.run(tf.global_variables_initializer())
# sess.run(tf.local_variables_initializer())
save_pred_name = "{}_{}".format(cfg.save_predictions_dir, cfg.eval_split)
save_dir = os.path.join(exp_dir, cfg.save_predictions_dir)
if eval_unsup:
reference_rotation = scipy.io.loadmat("{}/final_reference_rotation.mat".format(exp_dir))["rotation"]
dataset = ShapeRecords(dataset_folder, cfg, 'test')
if cfg.models_list:
model_names = parse_lines(cfg.models_list)
else:
model_names = dataset.file_names
num_models = len(model_names)
# making groups for samples and views according to 8 groups of yaw, pitch, roll
chamfer_dict = {}
for j in range(np.power(divs, 3)):
chamfer_dict[j] = np.zeros((0,2))
for k in range(num_models):
sample = dataset.__getitem__(k)
print("{}/{}".format(k, num_models))
print(model_names[k])
gt_filename = "{}/{}.mat".format(gt_dir, model_names[k]).replace('_features.p', '')
mat_filename = "{}/{}_pc.pkl".format(save_dir, model_names[k])
if not os.path.isfile(gt_filename) or not os.path.isfile(mat_filename):
continue
with open(mat_filename, 'rb') as handle:
data = pickle.load(handle)
all_pcs = np.squeeze(data["points"])
if "num_points" in data:
all_pcs_nums = np.squeeze(data["num_points"])
has_number = True
else:
has_number = False
obj = scipy.io.loadmat(gt_filename)
Vgt = obj["points"]
for i in range(num_views):
chamfer_dists_current = np.zeros((2), dtype=np.float64)
pred = all_pcs[i, :, :]
if has_number:
pred = pred[0:all_pcs_nums[i], :]
if eval_unsup:
pred = np.expand_dims(pred, 0)
pred = quaternion_rotate(torch.from_numpy(pred).cuda(),
torch.from_numpy(reference_rotation).cuda()).cpu().numpy()
pred = np.squeeze(pred)
pred_to_gt, idx_np = compute_distance(cfg, pred, Vgt)
gt_to_pred, _ = compute_distance(cfg, Vgt, pred)
chamfer_dists_current[0] = np.mean(pred_to_gt)
chamfer_dists_current[1] = np.mean(gt_to_pred)
is_nan = np.isnan(pred_to_gt)
assert (not np.any(is_nan))
campos = sample['cam_pos'][i]
g = get_group(campos)
chamfer_dict[g] = np.concatenate((chamfer_dict[g], np.expand_dims(chamfer_dists_current, 0)))
print(i, ":", chamfer_dists_current)
# current_mean = np.mean(chamfer_dists_current, 0)
# print("total:", current_mean)
for key in chamfer_dict:
print(key, np.mean(chamfer_dict[key],0)*100)
def run_eval():
config = tf.ConfigProto(device_count={'GPU': 1})
cfg = app_config
exp_dir = cfg.checkpoint_dir
num_views = cfg.num_views
eval_unsup = cfg.eval_unsupervised_shape
gt_dir = os.path.join(cfg.gt_pc_dir, cfg.synth_set)
g = tf.Graph()
with g.as_default():
source_pc = tf.placeholder(dtype=tf.float64, shape=[None, 3])
target_pc = tf.placeholder(dtype=tf.float64, shape=[None, 3])
quat_tf = tf.placeholder(dtype=tf.float64, shape=[1, 4])
_, min_dist, min_idx = point_cloud_distance(source_pc, target_pc)
source_pc_2 = tf.placeholder(dtype=tf.float64, shape=[1, None, 3])
rotated_pc = quaternion_rotate(source_pc_2, quat_tf)
sess = tf.Session(config=config)
sess.run(tf.global_variables_initializer())
sess.run(tf.local_variables_initializer())
save_pred_name = "{}_{}".format(cfg.save_predictions_dir, cfg.eval_split)
save_dir = os.path.join(exp_dir, cfg.save_predictions_dir)
if eval_unsup:
reference_rotation = scipy.io.loadmat(
"{}/final_reference_rotation.mat".format(exp_dir))["rotation"]
dataset = Dataset3D(cfg)
num_models = dataset.num_samples()
model_names = []
chamfer_dists = np.zeros((0, num_views, 2), dtype=np.float64)
for k in range(num_models):
sample = dataset.data[k]
print("{}/{}".format(k, num_models))
print(sample.name)
gt_filename = "{}/{}.mat".format(gt_dir, sample.name)
if not os.path.isfile(gt_filename):
continue
model_names.append(sample.name)
mat_filename = "{}/{}_pc.mat".format(save_dir, sample.name)
if os.path.isfile(mat_filename):
data = scipy.io.loadmat(mat_filename)
all_pcs = np.squeeze(data["points"])
if "num_points" in data:
all_pcs_nums = np.squeeze(data["num_points"])
has_number = True
else:
has_number = False
else:
data = np.load("{}/{}_pc.npz".format(save_dir, sample.name))
all_pcs = np.squeeze(data["arr_0"])
if 'arr_1' in list(data.keys()):
all_pcs_nums = np.squeeze(data["arr_1"])
has_number = True
else:
has_number = False
obj = scipy.io.loadmat(gt_filename)
Vgt = obj["points"]
chamfer_dists_current = np.zeros((num_views, 2), dtype=np.float64)
for i in range(num_views):
pred = all_pcs[i, :, :]
if has_number:
pred = pred[0:all_pcs_nums[i], :]
if eval_unsup:
pred = np.expand_dims(pred, 0)
pred = sess.run(rotated_pc,
feed_dict={
source_pc_2: pred,
quat_tf: reference_rotation
})
pred = np.squeeze(pred)
pred_to_gt, idx_np = compute_distance(cfg, sess, min_dist, min_idx,
source_pc, target_pc, pred,
Vgt)
gt_to_pred, _ = compute_distance(cfg, sess, min_dist, min_idx,
source_pc, target_pc, Vgt, pred)
chamfer_dists_current[i, 0] = np.mean(pred_to_gt)
chamfer_dists_current[i, 1] = np.mean(gt_to_pred)
is_nan = np.isnan(pred_to_gt)
assert (not np.any(is_nan))
current_mean = np.mean(chamfer_dists_current, 0)
print("total:", current_mean)
chamfer_dists = np.concatenate(
(chamfer_dists, np.expand_dims(chamfer_dists_current, 0)))
final = np.mean(chamfer_dists, axis=(0, 1)) * 100
print(final)
scipy.io.savemat(
os.path.join(exp_dir, "chamfer_{}.mat".format(save_pred_name)), {
"chamfer": chamfer_dists,
"model_names": to_np_object(model_names)
})
file = open(os.path.join(exp_dir, "chamfer_{}.txt".format(save_pred_name)),
"w")
file.write("{} {}\n".format(final[0], final[1]))
file.close()
def compute_predictions():
cfg = app_config
setup_environment(cfg)
exp_dir = cfg.checkpoint_dir
cfg.batch_size = 1
cfg.step_size = 1
pc_num_points = cfg.pc_num_points
vox_size = cfg.vox_size
save_pred = cfg.save_predictions
save_voxels = cfg.save_voxels
fast_conversion = True
pose_student = cfg.pose_predictor_student and cfg.predict_pose
device = 'cuda' if torch.cuda.is_available() else 'cpu'
model = model_pc.ModelPointCloud(cfg)
model = model.to(device)
log_dir = '../../dpc/run/model_run_data/'
learning_rate = 1e-4
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate, weight_decay = cfg.weight_decay)
global_step = 100000
if global_step>0:
checkpoint_path = os.path.join(log_dir,'model.ckpt_{}.pth'.format(global_step))
print("Loading from path:",checkpoint_path)
checkpoint = torch.load(checkpoint_path)
global_step_val = checkpoint['global_step']
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
loss = checkpoint['loss']
else:
global_step_val = global_step
print('Restored checkpoint at {} with loss {}'.format(global_step, loss))
save_dir = os.path.join(exp_dir, '{}_vis_proj'.format(cfg.save_predictions_dir))
mkdir_if_missing(save_dir)
save_pred_dir = os.path.join(exp_dir, cfg.save_predictions_dir)
mkdir_if_missing(save_pred_dir)
vis_size = cfg.vis_size
split_name = "val"
dataset_folder = cfg.inp_dir
dataset = ShapeRecords(dataset_folder, cfg, split_name)
dataset_loader = torch.utils.data.DataLoader(dataset,
batch_size=cfg.batch_size, shuffle=cfg.shuffle_dataset,
num_workers=4,drop_last=True)
pose_num_candidates = cfg.pose_predict_num_candidates
num_views = cfg.num_views
plot_h = 4
plot_w = 6
num_views = int(min(num_views, plot_h * plot_w / 2))
if cfg.models_list:
model_names = parse_lines(cfg.models_list)
else:
model_names = dataset.file_names
num_models = len(model_names)
for k in range(num_models):
model_name = model_names[k]
sample = dataset.__getitem__(k)
images = sample['image']
masks = sample['mask']
if cfg.saved_camera:
cameras = sample['extrinsic']
cam_pos = sample['cam_pos']
if cfg.vis_depth_projs:
depths = sample['depth']
if cfg.variable_num_views:
num_views = sample['num_views']
print("{}/{} {}".format(k, num_models, model_name))
if pose_num_candidates == 1:
grid = np.empty((plot_h, plot_w), dtype=object)
else:
plot_w = pose_num_candidates + 1
if pose_student:
plot_w += 1
grid = np.empty((num_views, plot_w), dtype=object)
if save_pred:
all_pcs = np.zeros((num_views, pc_num_points, 3))
all_cameras = np.zeros((num_views, 4))
#all_voxels = np.zeros((num_views, vox_size, vox_size, vox_size))
#all_z_latent = np.zeros((num_views, cfg.fc_dim))
for view_idx in range(num_views):
input_image_np = images[[view_idx], :, :, :]
gt_mask_np = masks[[view_idx], :, :, :]
if cfg.saved_camera:
extr_mtr = cameras[view_idx, :, :]
cam_quaternion_np = quaternion_from_campos(cam_pos[view_idx, :])
cam_quaternion_np = np.expand_dims(cam_quaternion_np, axis=0)
else:
extr_mtr = np.zeros((4, 4))
code = 'images' if cfg.predict_pose else 'images_1'
input = {code: input_image_np,
'matrices': extr_mtr,
'camera_quaternion': cam_quaternion_np}
out = build_model(model, input, global_step)
input_image = out["inputs"]
cam_matrix = out["camera_extr_src"]
cam_quaternion = out["cam_quaternion"]
point_cloud = out["points_1"]
#gb = out["rgb_1"] if cfg.pc_rgb else None
#rojs = out["projs"]
#rojs_rgb = out["projs_rgb"]
#rojs_depth = out["projs_depth"]
cam_transform = out["cam_transform"]
#_latent = out["z_latent"]
#if cfg.pc_rgb:
# proj_tensor = projs_rgb
#elif cfg.vis_depth_projs:
# proj_tensor = projs_depth
#else:
# proj_tensor = projs
if pose_student:
camera_student_np = out["pose_student"]
predicted_camera = camera_student_np
else:
predicted_camera = cam_transf_np
#if cfg.vis_depth_projs:
# proj_np = normalise_depthmap(out["projs"])
# if depths is not None:
# depth_np = depths[view_idx, :, :, :]
# depth_np = normalise_depthmap(depth_np)
# else:
# depth_np = 1.0 - np.squeeze(gt_mask_np)
# if pose_student:
# proj_student_np = normalise_depthmap(proj_student_np)
#if save_voxels:
# if fast_conversion:
# voxels, _ = pointcloud2voxels3d_fast(cfg, input_pc, None)
# voxels = tf.expand_dims(voxels, axis=-1)
# voxels = smoothen_voxels3d(cfg, voxels, model.gauss_kernel())
# else:
# voxels = pointcloud2voxels(cfg, input_pc, model.gauss_sigma())
if cfg.predict_pose:
if cfg.save_rotated_points:
ref_rot = scipy.io.loadmat("{}/final_reference_rotation.mat".format(exp_dir))
ref_rot = ref_rot["rotation"]
pc_unrot = quaternion_rotate(input_pc, ref_quat)
point_cloud = pc_np_unrot
if cfg.pc_rgb:
gt_image = input_image_np
elif cfg.vis_depth_projs:
gt_image = depth_np
else:
gt_image = gt_mask_np
# if pose_num_candidates == 1:
# view_j = view_idx * 2 // plot_w
# view_i = view_idx * 2 % plot_w
# gt_image = np.squeeze(gt_image)
# grid[view_j, view_i] = mask4vis(cfg, gt_image, vis_size)
# curr_img = np.squeeze(out[projs])
# grid[view_j, view_i + 1] = mask4vis(cfg, curr_img, vis_size)
# if cfg.save_individual_images:
# curr_dir = os.path.join(save_dir, model_names[k])
# if not os.path.exists(curr_dir):
# os.makedirs(curr_dir)
# imageio.imwrite(os.path.join(curr_dir, '{}_{}.png'.format(view_idx, 'rgb_gt')),
# mask4vis(cfg, np.squeeze(input_image_np), vis_size))
# imageio.imwrite(os.path.join(curr_dir, '{}_{}.png'.format(view_idx, 'mask_pred')),
# mask4vis(cfg, np.squeeze(proj_np), vis_size))
# else:
# view_j = view_idx
# gt_image = np.squeeze(gt_image)
# grid[view_j, 0] = mask4vis(cfg, gt_image, vis_size)
# for kk in range(pose_num_candidates):
# curr_img = np.squeeze(out["projs"][kk, :, :, :].detach().cpu())
# grid[view_j, kk + 1] = mask4vis(cfg, curr_img, vis_size)
# if cfg.save_individual_images:
# curr_dir = os.path.join(save_dir, model_names[k])
# if not os.path.exists(curr_dir):
# os.makedirs(curr_dir)
# imageio.imwrite(os.path.join(curr_dir, '{}_{}_{}.png'.format(view_idx, kk, 'mask_pred')),
# mask4vis(cfg, np.squeeze(curr_img), vis_size))
# if cfg.save_individual_images:
# imageio.imwrite(os.path.join(curr_dir, '{}_{}.png'.format(view_idx, 'mask_gt')),
# mask4vis(cfg, np.squeeze(gt_mask_np), vis_size))
# if pose_student:
# grid[view_j, -1] = mask4vis(cfg, np.squeeze(proj_student_np.detach().cpu()), vis_size)
if save_pred:
#pc_np = pc_np.detach().cpu().numpy()
all_pcs[view_idx, :, :] = np.squeeze(point_cloud.detach().cpu())
#all_z_latent[view_idx] = z_latent.detach().cpu()
if cfg.predict_pose:
all_cameras[view_idx, :] = predicted_camera.detach().cpu()
# if save_voxels:
# # multiplying by two is necessary because
# # pc->voxel conversion expects points in [-1, 1] range
# pc_np_range = pc_np
# if not fast_conversion:
# pc_np_range *= 2.0
# voxels_np = sess.run(voxels, feed_dict={input_pc: pc_np_range})
# all_voxels[view_idx, :, :, :] = np.squeeze(voxels_np)
# vis_view = view_idx == 0 or cfg.vis_all_views
# if cfg.vis_voxels and vis_view:
# rgb_np = np.squeeze(rgb_np) if cfg.pc_rgb else None
# vis_pc(np.squeeze(pc_np), rgb=rgb_np)
#grid_merged = merge_grid(cfg, grid)
#imageio.imwrite("{}/{}_proj.png".format(save_dir, sample.file_names), grid_merged)
if save_pred:
if 0:
save_dict = {"points": all_pcs}
if cfg.predict_pose:
save_dict["camera_pose"] = all_cameras
scipy.io.savemat("{}/{}_pc.mat".format(save_pred_dir, model_names[k]),
mdict=save_dict)
else:
save_dict = {"points": all_pcs}
if cfg.predict_pose:
save_dict["camera_pose"] = all_cameras
with open("{}/{}_pc.pkl".format(save_pred_dir, model_names[k]), 'wb') as handle:
pickle.dump(save_dict, handle, protocol=pickle.HIGHEST_PROTOCOL)
