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)