update demo

app.py

@@ -16,7 +16,7 @@ from transformers import AutoModelForImageSegmentation
 from torchvision import transforms
 from PIL import Image
 import open3d as o3d
-from backend_utils import improved_multiway_registration
+from backend_utils import improved_multiway_registration, pts2normal, point2mesh, combine_and_clean_point_clouds
 
 
 # Default values
@@ -29,15 +29,45 @@ OPENGL = np.array([[1, 0, 0, 0],
                    [0, 0, -1, 0],
                    [0, 0, 0, 1]])
 
-def extract_frames(video_path: str) -> str:
+def export_geometry(geometry, as_pointcloud=False):
+    if as_pointcloud:
+        if not isinstance(geometry, o3d.geometry.PointCloud):
+            raise ValueError("Expected an Open3D PointCloud object when as_pointcloud is True")
+        output_path = tempfile.mktemp(suffix='.ply')
+    else:
+        if not isinstance(geometry, o3d.geometry.TriangleMesh):
+            raise ValueError("Expected an Open3D TriangleMesh object when as_pointcloud is False")
+        output_path = tempfile.mktemp(suffix='.obj')
+
+    # Apply rotation
+    rot = np.eye(4)
+    rot[:3, :3] = Rotation.from_euler('y', np.deg2rad(180)).as_matrix()
+    transform = np.linalg.inv(OPENGL @ rot)
+    geometry.transform(transform)
+
+    # Export the geometry
+    if as_pointcloud:
+        o3d.io.write_point_cloud(output_path, geometry, write_ascii=False, compressed=True)
+    else:
+        o3d.io.write_triangle_mesh(output_path, geometry, write_ascii=False, compressed=True)
+    
+    return output_path
+
+
+def extract_frames(video_path: str, duration: float = 20.0, fps: float = 3.0) -> str:
     temp_dir = tempfile.mkdtemp()
     output_path = os.path.join(temp_dir, "%03d.jpg")
+    
+    filter_complex = f"select='if(lt(t,{duration}),1,0)',fps={fps}"
+
     command = [
         "ffmpeg",
         "-i", video_path,
-        "-vf", "fps=1",
+        "-vf", filter_complex,
+        "-vsync", "0",
         output_path
     ]
+    
     subprocess.run(command, check=True)
     return temp_dir
 
@@ -144,9 +174,9 @@ def generate_mask(image: np.ndarray):
     # Convert mask to numpy array
     mask_np = np.array(mask) / 255.0
     return mask_np
-
 @torch.no_grad()
-def reconstruct(video_path, conf_thresh, kf_every, as_pointcloud=False, remove_background=False):
+def reconstruct(video_path, conf_thresh, kf_every, 
+                as_pointcloud=False, remove_background=False, refine=False):
     # Extract frames from video
     demo_path = extract_frames(video_path)
     
@@ -168,123 +198,143 @@ def reconstruct(video_path, conf_thresh, kf_every, as_pointcloud=False, remove_b
     print(f'Finished reconstruction for {demo_name}, FPS: {fps:.2f}')
     
     # Process results
-    pts_all, images_all, conf_all, mask_all = [], [], [], []
+    pcds = []
     for j, view in enumerate(batch):
         image = view['img'].permute(0, 2, 3, 1).cpu().numpy()[0]
+        image = (image + 1) / 2
         pts = preds[j]['pts3d' if j==0 else 'pts3d_in_other_view'].detach().cpu().numpy()[0]
+        pts_normal = pts2normal(preds[j]['pts3d' if j==0 else 'pts3d_in_other_view'][0]).cpu().numpy()
         conf = preds[j]['conf'][0].cpu().data.numpy()
-        
+        conf_sig = (conf - 1) / conf
         if remove_background:
             mask = generate_mask(image)
         else:
             mask = np.ones_like(conf)
+            
+        combined_mask = (conf_sig > conf_thresh) & (mask > 0.5)
         
-        images_all.append((image[None, ...] + 1.0)/2.0)
-        pts_all.append(pts[None, ...])
-        conf_all.append(conf[None, ...])
-        mask_all.append(mask[None, ...])
-    
-    images_all = np.concatenate(images_all, axis=0)
-    pts_all = np.concatenate(pts_all, axis=0) * 10
-    conf_all = np.concatenate(conf_all, axis=0)
-    mask_all = np.concatenate(mask_all, axis=0)
+        pcd = o3d.geometry.PointCloud()
+        pcd.points = o3d.utility.Vector3dVector(pts[combined_mask])
+        pcd.colors = o3d.utility.Vector3dVector(image[combined_mask])
+        pcd.normals = o3d.utility.Vector3dVector(pts_normal[combined_mask])
+        pcds.append(pcd)
     
-    # Create point cloud or mesh
-    conf_sig_all = (conf_all-1) / conf_all
-    combined_mask = (conf_sig_all > conf_thresh) & (mask_all > 0.5)
+    pcd_combined = combine_and_clean_point_clouds(pcds, voxel_size=0.001)
     
+    if as_pointcloud:
+        o3d_geometry = pcd_combined
+    else:
+        o3d_geometry = point2mesh(pcd_combined)
+        
     # Create coarse result
-    coarse_scene = create_scene(pts_all, images_all, combined_mask, as_pointcloud)
-    coarse_output_path = save_scene(coarse_scene, as_pointcloud)
-    
-    yield coarse_output_path, None, f"Reconstruction completed. FPS: {fps:.2f}"
+    coarse_output_path = export_geometry(o3d_geometry, as_pointcloud)
     
-    # Create point clouds for multiway registration
-    pcds = []
-    for j in range(len(pts_all)):
-        pcd = o3d.geometry.PointCloud()
-        mask = combined_mask[j]
-        pcd.points = o3d.utility.Vector3dVector(pts_all[j][mask])
-        pcd.colors = o3d.utility.Vector3dVector(images_all[j][mask])
-        pcds.append(pcd)
-
-    # Perform global optimization 
-    print("Performing global registration...")
-    transformed_pcds, pose_graph = improved_multiway_registration(pcds, voxel_size=0.01)
+    yield coarse_output_path, None
     
-    # Apply transformations from pose_graph to original pts_all
-    transformed_pts_all = np.zeros_like(pts_all)
-    for j in range(len(pts_all)):
-        # Get the transformation matrix from the pose graph
-        transformation = pose_graph.nodes[j].pose
-        
-        # Reshape pts_all[j] to (H*W, 3)
-        H, W, _ = pts_all[j].shape
-        pts_reshaped = pts_all[j].reshape(-1, 3)
+    if refine:
+        # Perform global optimization 
+        print("Performing global registration...")
+        transformed_pcds, _, _ = improved_multiway_registration(pcds, voxel_size=0.001)
         
-        # Apply transformation to all points
-        homogeneous_pts = np.hstack((pts_reshaped, np.ones((pts_reshaped.shape[0], 1))))
-        transformed_pts = (transformation @ homogeneous_pts.T).T[:, :3]
+        if as_pointcloud:
+            o3d_geometry = transformed_pcds
+        else:
+            o3d_geometry = point2mesh(transformed_pcds)
+            
+        # Create coarse result
+        refined_output_path = export_geometry(o3d_geometry, as_pointcloud)
         
-        # Reshape back to (H, W, 3) and store
-        transformed_pts_all[j] = transformed_pts.reshape(H, W, 3)
-
-    print(f"Original shape: {pts_all.shape}, Transformed shape: {transformed_pts_all.shape}")
-
-    # Create refined result
-    refined_scene = create_scene(transformed_pts_all, images_all, combined_mask, as_pointcloud)
-    refined_output_path = save_scene(refined_scene, as_pointcloud)
+        yield coarse_output_path, refined_output_path
 
     # Clean up temporary directory
     os.system(f"rm -rf {demo_path}")
-    
-    yield coarse_output_path, refined_output_path, f"Refinement completed. FPS: {fps:.2f}"
 
-def create_scene(pts_all, images_all, combined_mask, as_pointcloud):
-    scene = trimesh.Scene()
+# Update the Gradio interface with improved layout
+with gr.Blocks(
+        title="StableSpann3r: Making Spann3r stable with Odometry Backend",
+        css="""
+            #download {
+                height: 118px;
+            }
+            .slider .inner {
+                width: 5px;
+                background: #FFF;
+            }
+            .viewport {
+                aspect-ratio: 4/3;
+            }
+            .tabs button.selected {
+                font-size: 20px !important;
+                color: crimson !important;
+            }
+            h1 {
+                text-align: center;
+                display: block;
+            }
+            h2 {
+                text-align: center;
+                display: block;
+            }
+            h3 {
+                text-align: center;
+                display: block;
+            }
+            .md_feedback li {
+                margin-bottom: 0px !important;
+            }
+        """,
+        head="""
+            <script async src="https://www.googletagmanager.com/gtag/js?id=G-1FWSVCGZTG"></script>
+            <script>
+                window.dataLayer = window.dataLayer || [];
+                function gtag() {dataLayer.push(arguments);}
+                gtag('js', new Date());
+                gtag('config', 'G-1FWSVCGZTG');
+            </script>
+        """,
+    ) as iface:
+    gr.Markdown(
+        """
+        # StableSpann3r: Making Spann3r stable with Odometry Backend
+        <p align="center">
+            <a title="Website" href="https://stable-x.github.io/StableSpann3r/" target="_blank" rel="noopener noreferrer" style="display: inline-block;">
+                <img src="https://www.obukhov.ai/img/badges/badge-website.svg">
+            </a>
+            <a title="arXiv" href="https://arxiv.org/abs/XXXX.XXXXX" target="_blank" rel="noopener noreferrer" style="display: inline-block;">
+                <img src="https://www.obukhov.ai/img/badges/badge-pdf.svg">
+            </a>
+            <a title="Github" href="https://github.com/Stable-X/StableSpann3r" target="_blank" rel="noopener noreferrer" style="display: inline-block;">
+                <img src="https://img.shields.io/github/stars/Stable-X/StableSpann3r?label=GitHub%20%E2%98%85&logo=github&color=C8C" alt="badge-github-stars">
+            </a>
+            <a title="Social" href="https://x.com/ychngji6" target="_blank" rel="noopener noreferrer" style="display: inline-block;">
+                <img src="https://www.obukhov.ai/img/badges/badge-social.svg" alt="social">
+            </a>
+        </p>
+        """
+    )
+    with gr.Row():
+        with gr.Column(scale=1):
+            video_input = gr.Video(label="Input Video")
+            with gr.Row():
+                conf_thresh = gr.Slider(0, 1, value=1e-3, label="Confidence Threshold")
+                kf_every = gr.Slider(1, 30, step=1, value=1, label="Keyframe Interval")
+            with gr.Row():
+                remove_background = gr.Checkbox(label="Remove Background", value=False)
+                refine = gr.Checkbox(label="Enable Backend", value=False)
+                as_pointcloud = gr.Checkbox(label="As Pointcloud", value=False)
+            reconstruct_btn = gr.Button("Reconstruct")
+        
+        with gr.Column(scale=2):
+            with gr.Tab("Coarse Model"):
+                coarse_model = gr.Model3D(label="Coarse 3D Model", display_mode="solid", clear_color=[0.0, 0.0, 0.0, 0.0])
+            with gr.Tab("Refined Model"):
+                refined_model = gr.Model3D(label="Refined 3D Model", display_mode="solid", clear_color=[0.0, 0.0, 0.0, 0.0])
     
-    if as_pointcloud:
-        pcd = trimesh.PointCloud(
-            vertices=pts_all[combined_mask].reshape(-1, 3),
-            colors=images_all[combined_mask].reshape(-1, 3)
-        )
-        scene.add_geometry(pcd)
-    else:
-        meshes = []
-        for i in range(len(images_all)):
-            meshes.append(pts3d_to_trimesh(images_all[i], pts_all[i], combined_mask[i]))
-        mesh = trimesh.Trimesh(**cat_meshes(meshes))
-        scene.add_geometry(mesh)
-
-    rot = np.eye(4)
-    rot[:3, :3] = Rotation.from_euler('y', np.deg2rad(180)).as_matrix()
-    scene.apply_transform(np.linalg.inv(OPENGL @ rot))
-    return scene
-def save_scene(scene, as_pointcloud):
-    if as_pointcloud:
-        output_path = tempfile.mktemp(suffix='.ply')
-    else:
-        output_path = tempfile.mktemp(suffix='.obj')
-    scene.export(output_path)
-    return output_path
-
-# Update the Gradio interface
-iface = gr.Interface(
-    fn=reconstruct,
-    inputs=[
-        gr.Video(label="Input Video"),
-        gr.Slider(0, 1, value=1e-6, label="Confidence Threshold"),
-        gr.Slider(1, 30, step=1, value=5, label="Keyframe Interval"),
-        gr.Checkbox(label="As Pointcloud", value=False),
-        gr.Checkbox(label="Remove Background", value=False)
-    ],
-    outputs=[
-        gr.Model3D(label="Coarse 3D Model", display_mode="solid"),
-        gr.Model3D(label="Refined 3D Model", display_mode="solid"),
-        gr.Textbox(label="Status")
-    ],
-    title="3D Reconstruction with Spatial Memory, Background Removal, and Global Optimization",
-)
+    reconstruct_btn.click(
+        fn=reconstruct,
+        inputs=[video_input, conf_thresh, kf_every, as_pointcloud, remove_background, refine],
+        outputs=[coarse_model, refined_model]
+    )
 
 if __name__ == "__main__":
-    iface.launch(server_name="0.0.0.0",)
+    iface.launch(server_name="0.0.0.0")

backend_utils.py

@@ -1,90 +1,152 @@
 import numpy as np
 import open3d as o3d
+import torch
+from tqdm import tqdm
+import torch.nn.functional as F
 
-def improved_multiway_registration(pcds, voxel_size=0.05, max_correspondence_distance_coarse=None, max_correspondence_distance_fine=None, overlap=3, quadratic_overlap=True, use_colored_icp=True):
-    if max_correspondence_distance_coarse is None:
-        max_correspondence_distance_coarse = voxel_size * 15
-    if max_correspondence_distance_fine is None:
-        max_correspondence_distance_fine = voxel_size * 1.5
-
-    def preprocess_point_cloud(pcd, voxel_size):
-        pcd_down = pcd.voxel_down_sample(voxel_size)
-        pcd_down.estimate_normals(
-            o3d.geometry.KDTreeSearchParamHybrid(radius=voxel_size * 2, max_nn=30))
-        # Apply statistical outlier removal
-        cl, ind = pcd_down.remove_statistical_outlier(nb_neighbors=20, std_ratio=2.0)
-        pcd_down = pcd_down.select_by_index(ind)
-        return pcd_down
-
-    def pairwise_registration(source, target, use_colored_icp, voxel_size, max_correspondence_distance_coarse, max_correspondence_distance_fine):
-        current_transformation = np.identity(4)  # Start with identity matrix
-
-        if use_colored_icp:
-            print("Apply colored point cloud registration")
-            voxel_radius = [5*voxel_size, 3*voxel_size, voxel_size]
-            max_iter = [60, 35, 20]
-            
-            for scale in range(3):
-                iter = max_iter[scale]
-                radius = voxel_radius[scale]
+def pts2normal(pts):
+    h, w, _ = pts.shape
+    
+    # Compute differences in x and y directions
+    dx = torch.cat([pts[2:, 1:-1] - pts[:-2, 1:-1]], dim=0)
+    dy = torch.cat([pts[1:-1, 2:] - pts[1:-1, :-2]], dim=1)
+    
+    # Compute normal vectors using cross product
+    normal_map = F.normalize(torch.cross(dx, dy, dim=-1), dim=-1)
+    
+    # Create padded normal map
+    padded_normal_map = torch.zeros_like(pts)
+    padded_normal_map[1:-1, 1:-1, :] = normal_map
+
+    # Pad the borders
+    padded_normal_map[0, 1:-1, :] = normal_map[0, :, :]  # Top edge
+    padded_normal_map[-1, 1:-1, :] = normal_map[-1, :, :]  # Bottom edge
+    padded_normal_map[1:-1, 0, :] = normal_map[:, 0, :]  # Left edge
+    padded_normal_map[1:-1, -1, :] = normal_map[:, -1, :]  # Right edge
+    
+    # Pad the corners
+    padded_normal_map[0, 0, :] = normal_map[0, 0, :]  # Top-left corner
+    padded_normal_map[0, -1, :] = normal_map[0, -1, :]  # Top-right corner
+    padded_normal_map[-1, 0, :] = normal_map[-1, 0, :]  # Bottom-left corner
+    padded_normal_map[-1, -1, :] = normal_map[-1, -1, :]  # Bottom-right corner
+    
+    return padded_normal_map
+
+def point2mesh(pcd, depth=8, density_threshold=0.1, clean_mesh=True):
+    print("\nPerforming Poisson surface reconstruction...")
+    mesh, densities = o3d.geometry.TriangleMesh.create_from_point_cloud_poisson(
+        pcd, depth=depth, width=0, scale=1.1, linear_fit=False)
+
+    print(f"Reconstructed mesh has {len(mesh.vertices)} vertices and {len(mesh.triangles)} triangles")
+
+    # Normalize densities
+    densities = np.asarray(densities)
+    densities = (densities - densities.min()) / (densities.max() - densities.min())
+
+    # Remove low density vertices
+    print("\nPruning low-density vertices...")
+    vertices_to_remove = densities < np.quantile(densities, density_threshold)
+    mesh.remove_vertices_by_mask(vertices_to_remove)
+
+    print(f"Pruned mesh has {len(mesh.vertices)} vertices and {len(mesh.triangles)} triangles")
+
+    if clean_mesh:
+        print("\nCleaning the mesh...")
+        mesh.remove_degenerate_triangles()
+        mesh.remove_duplicated_triangles()
+        mesh.remove_duplicated_vertices()
+        mesh.remove_non_manifold_edges()
+
+        print(f"Final cleaned mesh has {len(mesh.vertices)} vertices and {len(mesh.triangles)} triangles")
 
-                source_down = source.voxel_down_sample(radius)
-                target_down = target.voxel_down_sample(radius)
+    mesh.compute_triangle_normals()
+    return mesh
 
-                source_down.estimate_normals(
-                    o3d.geometry.KDTreeSearchParamHybrid(radius=radius * 2, max_nn=30))
-                target_down.estimate_normals(
-                    o3d.geometry.KDTreeSearchParamHybrid(radius=radius * 2, max_nn=30))
+def combine_and_clean_point_clouds(pcds, voxel_size):
+    """
+    Combine, downsample, and clean a list of point clouds.
 
-                try:
-                    result_icp = o3d.pipelines.registration.registration_colored_icp(
-                        source_down, target_down, radius, current_transformation,
+    Parameters:
+    pcds (list): List of open3d.geometry.PointCloud objects to be processed.
+    voxel_size (float): The size of the voxel for downsampling.
+
+    Returns:
+    o3d.geometry.PointCloud: The cleaned and combined point cloud.
+    """
+    print("\nCombining point clouds...")
+    pcd_combined = o3d.geometry.PointCloud()
+    for p3d in pcds:
+        pcd_combined += p3d
+    
+    print("\nDownsampling the combined point cloud...")
+    pcd_combined = pcd_combined.voxel_down_sample(voxel_size)
+    print(f"Downsampled from {len(pcd_combined.points)} to {len(pcd_combined.points)} points")
+
+    print("\nCleaning the combined point cloud...")
+    cl, ind = pcd_combined.remove_statistical_outlier(nb_neighbors=20, std_ratio=2.0)
+    pcd_cleaned = pcd_combined.select_by_index(ind)
+
+    print(f"Cleaned point cloud contains {len(pcd_cleaned.points)} points.")
+    print(f"Removed {len(pcd_combined.points) - len(pcd_cleaned.points)} outlier points.")
+
+    return pcd_cleaned
+
+def improved_multiway_registration(pcds, descriptors=None, voxel_size=0.05, 
+                                   max_correspondence_distance_coarse=None, max_correspondence_distance_fine=None, 
+                                   overlap=5, quadratic_overlap=False, use_colored_icp=False):
+    if max_correspondence_distance_coarse is None:
+        max_correspondence_distance_coarse = voxel_size * 1.5
+    if max_correspondence_distance_fine is None:
+        max_correspondence_distance_fine = voxel_size * 0.15
+
+    def pairwise_registration(source, target, use_colored_icp, max_correspondence_distance_coarse, max_correspondence_distance_fine):
+        current_transformation = np.identity(4)
+        try:
+            if use_colored_icp:
+                icp_fine = o3d.pipelines.registration.registration_colored_icp(
+                        source, target, max_correspondence_distance_fine, current_transformation,
                         o3d.pipelines.registration.TransformationEstimationForColoredICP(),
                         o3d.pipelines.registration.ICPConvergenceCriteria(relative_fitness=1e-6,
-                                                                        relative_rmse=1e-6,
-                                                                        max_iteration=iter))
-                    current_transformation = result_icp.transformation
-                except RuntimeError as e:
-                    print(f"Colored ICP failed at scale {scale}: {str(e)}")
-                    print("Keeping the previous transformation")
-                    # We keep the previous transformation, no need to reassign
-
-            transformation_icp = current_transformation
-        else:
-            print("Apply point-to-plane ICP")
-            try:
-                icp_coarse = o3d.pipelines.registration.registration_icp(
-                    source, target, max_correspondence_distance_coarse, current_transformation,
-                    o3d.pipelines.registration.TransformationEstimationPointToPlane())
-                current_transformation = icp_coarse.transformation
-                
+                                                                      relative_rmse=1e-6,
+                                                                      max_iteration=100))
+            else:
                 icp_fine = o3d.pipelines.registration.registration_icp(
                     source, target, max_correspondence_distance_fine,
                     current_transformation,
                     o3d.pipelines.registration.TransformationEstimationPointToPlane())
-                transformation_icp = icp_fine.transformation
-            except RuntimeError as e:
-                print(f"Point-to-plane ICP failed: {str(e)}")
-                print("Keeping the best available transformation")
-                transformation_icp = current_transformation
+            
+            
+            fitness = icp_fine.fitness
+            FITNESS_THRESHOLD = 0.01
 
-        try:
-            information_icp = o3d.pipelines.registration.get_information_matrix_from_point_clouds(
-                source, target, max_correspondence_distance_fine,
-                transformation_icp)
+            if fitness >= FITNESS_THRESHOLD:
+                current_transformation = icp_fine.transformation
+            
+                information_icp = o3d.pipelines.registration.get_information_matrix_from_point_clouds(
+                    source, target, max_correspondence_distance_fine,
+                    current_transformation)
+                return current_transformation, information_icp, True
+            else:
+                print(f"Registration failed. Fitness {fitness} is below threshold {FITNESS_THRESHOLD}")
+                return None, None, False
+        
         except RuntimeError as e:
-            print(f"Failed to compute information matrix: {str(e)}")
-            print("Using identity information matrix")
-            information_icp = np.identity(6)
+            print(f"  ICP registration failed: {str(e)}")
+            return None, None, False
 
-        return transformation_icp, information_icp
-
-    def full_registration(pcds_down):
-        pose_graph = o3d.pipelines.registration.PoseGraph()
-        odometry = np.identity(4)
-        pose_graph.nodes.append(o3d.pipelines.registration.PoseGraphNode(odometry))
-        n_pcds = len(pcds_down)
+    def detect_loop_closure(descriptors, min_interval=3, similarity_threshold=0.9):
+        n_pcds = len(descriptors)
+        loop_edges = []
+        
+        for i in range(n_pcds):
+            for j in range(i + min_interval, n_pcds):
+                similarity = torch.dot(descriptors[i], descriptors[j])
+                if similarity > similarity_threshold:
+                    loop_edges.append((i, j))
         
+        return loop_edges
+
+    def generate_pairs(n_pcds, overlap, quadratic_overlap):
         pairs = []
         for i in range(n_pcds - 1):
             for j in range(i + 1, min(i + overlap + 1, n_pcds)):
@@ -93,52 +155,122 @@ def improved_multiway_registration(pcds, voxel_size=0.05, max_correspondence_dis
                     q = 2**(j-i)
                     if q > overlap and i + q < n_pcds:
                         pairs.append((i, i + q))
+        return pairs
+
+    def full_registration(pcds_down, pairs, loop_edges):
+        pose_graph = o3d.pipelines.registration.PoseGraph()
+        n_pcds = len(pcds_down)
         
-        for source_id, target_id in pairs:
-            transformation_icp, information_icp = pairwise_registration(
+        for i in range(n_pcds):
+            pose_graph.nodes.append(o3d.pipelines.registration.PoseGraphNode(np.identity(4)))
+
+        print("\nPerforming pairwise registration:")
+        for source_id, target_id in tqdm(pairs):
+            transformation_icp, information_icp, success = pairwise_registration(
                 pcds_down[source_id], pcds_down[target_id], use_colored_icp,
-                voxel_size, max_correspondence_distance_coarse, max_correspondence_distance_fine)
-            print(f"Build PoseGraph: {source_id} -> {target_id}")
-            
-            if target_id == source_id + 1:
-                odometry = np.dot(transformation_icp, odometry)
-                pose_graph.nodes.append(
-                    o3d.pipelines.registration.PoseGraphNode(
-                        np.linalg.inv(odometry)))
+                max_correspondence_distance_coarse, max_correspondence_distance_fine)
             
-            pose_graph.edges.append(
-                o3d.pipelines.registration.PoseGraphEdge(source_id,
-                                                         target_id,
-                                                         transformation_icp,
-                                                         information_icp,
-                                                         uncertain=False))
+            if success:            
+                uncertain = abs(target_id - source_id) == 1
+                pose_graph.edges.append(
+                    o3d.pipelines.registration.PoseGraphEdge(source_id,
+                                                            target_id,
+                                                            transformation_icp,
+                                                            information_icp,
+                                                            uncertain=uncertain))
+            else:
+                print(f"  Skipping edge between {source_id} and {target_id} due to ICP failure")
+
+        # Add loop closure edges
+        print("\nAdding loop closure edges:")
+        for source_id, target_id in tqdm(loop_edges):
+            transformation_icp, information_icp, success = pairwise_registration(
+                pcds_down[source_id], pcds_down[target_id], use_colored_icp,
+                max_correspondence_distance_coarse, max_correspondence_distance_fine)
+
+            if success:
+                pose_graph.edges.append(
+                    o3d.pipelines.registration.PoseGraphEdge(source_id,
+                                                            target_id,
+                                                            transformation_icp,
+                                                            information_icp,
+                                                            uncertain=True))
+            else:
+                print(f"  Skipping loop closure edge between {source_id} and {target_id} due to ICP failure")
+
         return pose_graph
 
-    # Preprocess point clouds
-    print("Preprocessing point clouds...")
-    pcds_down = [preprocess_point_cloud(pcd, voxel_size) for pcd in pcds]
+    print("\n--- Improved Multiway Registration Process ---")
+    print(f"Number of point clouds: {len(pcds)}")
+    print(f"Voxel size: {voxel_size}")
+    print(f"Max correspondence distance (coarse): {max_correspondence_distance_coarse}")
+    print(f"Max correspondence distance (fine): {max_correspondence_distance_fine}")
+    print(f"Overlap: {overlap}")
+    print(f"Quadratic overlap: {quadratic_overlap}")
 
-    print("Full registration ...")
-    pose_graph = full_registration(pcds_down)
-    
-    print("Optimizing PoseGraph ...")
+    print("\nPreprocessing point clouds...")
+    pcds_down = pcds
+    print(f"Preprocessing complete. {len(pcds_down)} point clouds processed.")
+
+    print("\nGenerating initial graph pairs...")
+    pairs = generate_pairs(len(pcds), overlap, quadratic_overlap)
+    print(f"Generated {len(pairs)} pairs for initial graph.")
+
+    if descriptors is None:
+        print("\nNo descriptors provided. Skipping loop closure detection.")
+        loop_edges = []
+    else:
+        print(descriptors[0].shape)
+        print("\nDetecting loop closures...")
+        loop_edges = detect_loop_closure(descriptors)
+        print(f"Detected {len(loop_edges)} loop closures.")
+
+    print("\nPerforming full registration...")
+    pose_graph = full_registration(pcds_down, pairs, loop_edges)
+
+    print("\nOptimizing PoseGraph...")
     option = o3d.pipelines.registration.GlobalOptimizationOption(
         max_correspondence_distance=max_correspondence_distance_fine,
         edge_prune_threshold=0.25,
         reference_node=0)
+    o3d.pipelines.registration.global_optimization(
+        pose_graph,
+        o3d.pipelines.registration.GlobalOptimizationLevenbergMarquardt(),
+        o3d.pipelines.registration.GlobalOptimizationConvergenceCriteria(),
+        option)
+
+    # Count edges for each node
+    edge_count = {i: 0 for i in range(len(pcds))}
+    for edge in pose_graph.edges:
+        edge_count[edge.source_node_id] += 1
+        edge_count[edge.target_node_id] += 1
+
+    # Filter nodes with more than 3 edges
+    valid_nodes = [count > 3 for count in edge_count.values()]
     
-    with o3d.utility.VerbosityContextManager(o3d.utility.VerbosityLevel.Debug) as cm:
-        o3d.pipelines.registration.global_optimization(
-            pose_graph,
-            o3d.pipelines.registration.GlobalOptimizationLevenbergMarquardt(),
-            o3d.pipelines.registration.GlobalOptimizationConvergenceCriteria(),
-            option)
-
-    print("Transform points and combine")
+    print("\nTransforming and combining point clouds...")
     pcd_combined = o3d.geometry.PointCloud()
-    for point_id in range(len(pcds)):
-        print(pose_graph.nodes[point_id].pose)
-        pcds[point_id].transform(pose_graph.nodes[point_id].pose)
-        pcd_combined += pcds[point_id]
+
+    for point_id, is_valid in enumerate(valid_nodes):
+        if is_valid:
+            pcds[point_id].transform(pose_graph.nodes[point_id].pose)
+            pcd_combined += pcds[point_id]
+        else:
+            print(f"Skipping point cloud {point_id} as it has {edge_count[point_id]} edges (<=3)")
+            
+    print("\nDownsampling the combined point cloud...")
+    # pcd_combined.orient_normals_consistent_tangent_plane(k=30)
+    pcd_combined = pcd_combined.voxel_down_sample(voxel_size * 0.1)
+    print(f"Downsampled from {len(pcd_combined.points)} to {len(pcd_combined.points)} points")
+
+    print("\nCleaning the combined point cloud...")
+    cl, ind = pcd_combined.remove_statistical_outlier(nb_neighbors=20, std_ratio=2.0)
+    pcd_cleaned = pcd_combined.select_by_index(ind)
+
+    print(f"Cleaned point cloud contains {len(pcd_cleaned.points)} points.")
+    print(f"Removed {len(pcd_combined.points) - len(pcd_cleaned.points)} outlier points.")
+
+    print("\nMultiway registration complete.")
+    print(f"Included {len(valid_nodes)} out of {len(pcds)} point clouds (with >3 edges).")
     
-    return pcd_combined, pose_graph
+    return pcd_cleaned, pose_graph, valid_nodes