(source_path, model_path, ckpt_path, device, batch_size, image_size, schedule, lr, niter,
min_conf_thr, llffhold, n_views, co_vis_dsp, depth_thre, conf_aware_ranking=False, focal_avg=True, infer_video=False)
| 22 | |
| 23 | |
| 24 | def main(source_path, model_path, ckpt_path, device, batch_size, image_size, schedule, lr, niter, |
| 25 | min_conf_thr, llffhold, n_views, co_vis_dsp, depth_thre, conf_aware_ranking=False, focal_avg=True, infer_video=False): |
| 26 | |
| 27 | # ---------------- (1) Load model and images ---------------- |
| 28 | save_path, sparse_0_path, sparse_1_path = init_filestructure(Path(source_path), n_views) |
| 29 | model = AsymmetricMASt3R.from_pretrained(ckpt_path).to(device) |
| 30 | image_dir = Path(source_path) / 'images' |
| 31 | image_files, image_suffix = get_sorted_image_files(image_dir) |
| 32 | if infer_video: |
| 33 | train_img_files = image_files |
| 34 | else: |
| 35 | train_img_files, test_img_files = split_train_test(image_files, llffhold, n_views, verbose=True) |
| 36 | |
| 37 | # when init test pose, use all images |
| 38 | image_files = train_img_files + test_img_files |
| 39 | images, org_imgs_shape = load_images(image_files, size=image_size) |
| 40 | |
| 41 | start_time = time() |
| 42 | print(f'>> Making pairs...') |
| 43 | pairs = make_pairs(images, scene_graph='complete', prefilter=None, symmetrize=True) |
| 44 | print(f'>> Inference...') |
| 45 | output = inference(pairs, model, device, batch_size=1, verbose=True) |
| 46 | |
| 47 | # Load Focal |
| 48 | train_pts_all_path = sparse_0_path / 'points3D_all.npy' |
| 49 | train_pts_all = np.load(train_pts_all_path) |
| 50 | train_pts3d_m1 = train_pts_all |
| 51 | if args.focal_avg: |
| 52 | focals_file = sparse_0_path / 'non_scaled_focals.npy' |
| 53 | preset_focals = np.load(focals_file) |
| 54 | preset_focal = np.mean(preset_focals) |
| 55 | print(f">> preset_focal: {preset_focal}") |
| 56 | |
| 57 | print(f'>> Global alignment...') |
| 58 | scene = global_aligner(output, device=args.device, mode=GlobalAlignerMode.PointCloudOptimizer) |
| 59 | loss = scene.compute_global_alignment(init="mst", niter=500, schedule=schedule, lr=lr, focal_avg=args.focal_avg, known_focal=preset_focal) |
| 60 | imgs = np.array(scene.imgs) |
| 61 | focals = np.repeat(preset_focal, len(test_img_files)) |
| 62 | |
| 63 | all_poses = to_numpy(scene.get_im_poses()) |
| 64 | all_pts3d = to_numpy(scene.get_pts3d()) |
| 65 | train_pts3d_n1 = all_pts3d[:n_views] |
| 66 | test_poses_n1 = all_poses[n_views:] |
| 67 | train_pts3d_n1 = np.array(to_numpy(train_pts3d_n1)).reshape(-1,3) |
| 68 | test_poses_n1 = np.array(to_numpy(test_poses_n1)) # test_pose_n1: c2w |
| 69 | |
| 70 | #---------------- (4) Applying pointcloud registration & Calculate transform_matrix & Save initial_test_pose---------------- |
| 71 | # compute transform that goes from cam to world |
| 72 | train_pts3d_n1 = torch.from_numpy(train_pts3d_n1) |
| 73 | train_pts3d_m1 = torch.from_numpy(train_pts3d_m1) |
| 74 | scale, R, T = rigid_points_registration(train_pts3d_n1, train_pts3d_m1, conf=None) |
| 75 | |
| 76 | transform_matrix = torch.eye(4) |
| 77 | transform_matrix[:3, :3] = R |
| 78 | transform_matrix[:3, 3] = T |
| 79 | transform_matrix[:3, 3] *= scale |
| 80 | transform_matrix = transform_matrix.numpy() |
| 81 | test_poses_m1 = transform_matrix @ test_poses_n1 |
no test coverage detected