Upload app_fine.py

app_fine.py

@@ -263,7 +263,9 @@ def refine_pose_mast3r(rend_image_pil, target_image_pil, original_size, fxy, tar
     target_extrinsic_final = torch.tensor(predict_c2w_refine).inverse().cuda()[None].float()
     return target_extrinsic_final
 
-def pointcloud_registration(rend_image_pil, target_image_pil, original_size, fxy, target_extrinsic, rend_depth, target_pointmap):
+def pointcloud_registration(rend_image_pil, target_image_pil, original_size, 
+                            fxy, target_extrinsic, rend_depth, target_pointmap, 
+                            down_pcd, pcd):
     images_mast3r = load_images_new([rend_image_pil, target_image_pil], size=512, square_ok=True)
     with torch.no_grad():
         output = inference([tuple(images_mast3r)], mast3r_model, device, batch_size=1, verbose=False)
@@ -359,7 +361,10 @@ def pointcloud_registration(rend_image_pil, target_image_pil, original_size, fxy
     )
     transformation_matrix = result.transformation.copy()
     transformation_matrix[:3,:3] = transformation_matrix[:3,:3] * (scale_1 / scale_2)
-    return transformation_matrix, result.fitness
+    evaluation = o3d.pipelines.registration.evaluate_registration(
+        down_pcd, pcd, 0.02, transformation_matrix
+    )
+    return transformation_matrix, evaluation.fitness
 
 @spaces.GPU(duration=120)
 def generate_and_extract_glb(
@@ -465,7 +470,9 @@ def generate_and_extract_glb(
                 cl, ind = pcd.remove_statistical_outlier(nb_neighbors=30, std_ratio=3.0)
                 inlier_cloud = pcd.select_by_index(ind)
                 outlier_cloud = pcd.select_by_index(ind, invert=True)
-                voxel_size = 1/64
+                distance = np.array(inlier_cloud.points) - np.array(inlier_cloud.points).mean(axis=0)[None]
+                scale = np.percentile(np.linalg.norm(distance, axis=1), 97)
+                voxel_size = 1/64*scale*2
                 down_pcd = inlier_cloud.voxel_down_sample(voxel_size)
             torch.cuda.empty_cache()
 
@@ -476,6 +483,16 @@ def generate_and_extract_glb(
             target_intrinsics = []
             target_transforms = []
             target_fitnesses = []   
+            pcd = o3d.geometry.PointCloud()
+            mesh = outputs['mesh'][0]
+            idxs = torch.randperm(mesh.vertices.shape[0])[:min(50000, pointcloud.shape[0])]
+            pcd.points = o3d.utility.Vector3dVector(mesh.vertices[idxs].cpu().numpy())
+            distance = np.array(pcd.points) - np.array(pcd.points).mean(axis=0)[None]
+            scale = np.linalg.norm(distance, axis=1).max()
+            voxel_size = 1/64*scale*2
+            pcd = pcd.voxel_down_sample(voxel_size)
+            # pcd.points = o3d.utility.Vector3dVector((coords[:,1:].cpu().numpy() + 0.5) / 64 - 0.5)
+
             for k in range(len(image_files)):
                 images = torch.stack([TF.ToTensor()(render_image) for render_image in video['color']] + [TF.ToTensor()(image_files[k].convert("RGB"))], dim=0)
                 # if len(images) == 0:
@@ -536,7 +553,8 @@ def generate_and_extract_glb(
                 
                 idx = iou_list.index(max(iou_list))
                 target_extrinsic[j:j+1] = target_extrinsic_list[idx]
-                target_transform, fitness = pointcloud_registration(rend_image_pil, target_image_pil, original_size, fxy[j:j+1], target_extrinsic[j:j+1], rend_depth_eroded, point_map_perframe[k].cpu().numpy())
+                target_transform, fitness = pointcloud_registration(rend_image_pil, target_image_pil, original_size, fxy[j:j+1], target_extrinsic[j:j+1], \
+                                                                    rend_depth_eroded, point_map_perframe[k].cpu().numpy(), down_pcd, pcd)
                 target_transforms.append(target_transform)
                 target_fitnesses.append(fitness)
                 
@@ -549,15 +567,15 @@ def generate_and_extract_glb(
             idx = target_fitnesses.index(max(target_fitnesses_filtered))
             target_transform = target_transforms[idx]
             down_pcd_align = copy.deepcopy(down_pcd).transform(target_transform)
-            pcd = o3d.geometry.PointCloud()
-            pcd.points = o3d.utility.Vector3dVector(coords[:,1:].cpu().numpy() / 64 - 0.5)
+            # pcd = o3d.geometry.PointCloud()
+            # pcd.points = o3d.utility.Vector3dVector(coords[:,1:].cpu().numpy() / 64 - 0.5)
             reg_p2p = o3d.pipelines.registration.registration_icp(
-                down_pcd_align, pcd, 0.01, np.eye(4),
+                down_pcd_align, pcd, 0.02, np.eye(4),
                 o3d.pipelines.registration.TransformationEstimationPointToPoint(with_scaling=True),
                 o3d.pipelines.registration.ICPConvergenceCriteria(max_iteration = 10000))
             down_pcd_align_2 = copy.deepcopy(down_pcd_align).transform(reg_p2p.transformation)
             input_points = torch.tensor(np.asarray(down_pcd_align_2.points)).to(extrinsic.device).float()
-            input_points = ((input_points + 0.5).clip(0, 1) * 63).to(torch.int32)
+            input_points = ((input_points + 0.5).clip(0, 1) * 64 - 0.5).to(torch.int32)
             
             outputs = pipeline.run_refine(
                 image=image_files,