feat: Update demo

app.py

@@ -18,6 +18,9 @@ from PIL import Image
 import open3d as o3d
 from backend_utils import improved_multiway_registration, pts2normal, point2mesh, combine_and_clean_point_clouds
 from gs_utils import point2gs
+from gradio.helpers import Examples as GradioExamples
+from gradio.utils import get_cache_folder
+from pathlib import Path
 
 # Default values
 DEFAULT_CKPT_PATH = './checkpoints/spann3r.pth'
@@ -29,6 +32,14 @@ OPENGL = np.array([[1, 0, 0, 0],
                    [0, 0, -1, 0],
                    [0, 0, 0, 1]])
 
+class Examples(GradioExamples):
+    def __init__(self, *args, directory_name=None, **kwargs):
+        super().__init__(*args, **kwargs, _initiated_directly=False)
+        if directory_name is not None:
+            self.cached_folder = get_cache_folder() / directory_name
+            self.cached_file = Path(self.cached_folder) / "log.csv"
+            self.create()
+
 def export_geometry(geometry):
     output_path = tempfile.mktemp(suffix='.obj')
 
@@ -163,6 +174,41 @@ def generate_mask(image: np.ndarray):
     # Convert mask to numpy array
     mask_np = np.array(mask) / 255.0
     return mask_np
+
+def center_pcd(pcd: o3d.geometry.PointCloud, normalize=False) -> o3d.geometry.PointCloud:
+    # Convert to numpy array
+    points = np.asarray(pcd.points)
+    
+    # Compute centroid
+    centroid = np.mean(points, axis=0)
+    
+    # Center the point cloud
+    centered_points = points - centroid
+    
+    if normalize:
+         # Compute the maximum distance from the center
+        max_distance = np.max(np.linalg.norm(centered_points, axis=1))
+        
+        # Normalize the point cloud
+        normalized_points = centered_points / max_distance
+        
+        # Create a new point cloud with the normalized points
+        normalized_pcd = o3d.geometry.PointCloud()
+        normalized_pcd.points = o3d.utility.Vector3dVector(normalized_points)
+        
+        # If the original point cloud has colors, normalize them too
+        if pcd.has_colors():
+            normalized_pcd.colors = pcd.colors
+        
+        # If the original point cloud has normals, copy them
+        if pcd.has_normals():
+            normalized_pcd.normals = pcd.normals
+        
+        return normalized_pcd
+    else:
+        pcd.points = o3d.utility.Vector3dVector(centered_points)
+        return pcd
+
 @torch.no_grad()
 def reconstruct(video_path, conf_thresh, kf_every, 
                 remove_background=False):
@@ -209,6 +255,7 @@ def reconstruct(video_path, conf_thresh, kf_every,
         pcds.append(pcd)
     
     pcd_combined = combine_and_clean_point_clouds(pcds, voxel_size=0.001)
+    pcd_combined = center_pcd(pcd_combined, normalize=True)
     o3d_geometry = point2mesh(pcd_combined)
         
     # Create coarse result
@@ -216,10 +263,9 @@ def reconstruct(video_path, conf_thresh, kf_every,
     
     yield coarse_output_path, None
     
-    # Perform global optimization 
-    print("Performing global registration...")
     transformed_pcds, _, _ = improved_multiway_registration(pcds, voxel_size=0.01)
-
+    transformed_pcds = center_pcd(transformed_pcds)
+    
     # Create coarse result
     refined_output_path = tempfile.mktemp(suffix='.ply')
     point2gs(refined_output_path, transformed_pcds)
@@ -228,9 +274,11 @@ def reconstruct(video_path, conf_thresh, kf_every,
     # Clean up temporary directory
     os.system(f"rm -rf {demo_path}")
 
+example_videos = [os.path.join('./examples', f) for f in os.listdir('./examples') if f.endswith(('.mp4', '.webm'))]
+
 # Update the Gradio interface with improved layout
 with gr.Blocks(
-        title="StableSpann3r: Making Spann3r stable with Odometry Backend",
+        title="StableSpann3r: 3D Reconstruction from Video",
         css="""
             #download {
                 height: 118px;
@@ -276,12 +324,6 @@ with gr.Blocks(
         """
         # 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>
@@ -293,26 +335,72 @@ with gr.Blocks(
     )
     with gr.Row():
         with gr.Column(scale=1):
-            video_input = gr.Video(label="Input Video")
+            video_input = gr.Video(label="Input Video", sources=["upload"])
             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)
-            reconstruct_btn = gr.Button("Reconstruct")
+            reconstruct_btn = gr.Button("Start Reconstruction")
         
         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 Gaussian Splatting", display_mode="solid", 
-                                           clear_color=[0.0, 0.0, 0.0, 0.0])
+            with gr.Tab("3D Models"):
+                with gr.Group():
+                    initial_model = gr.Model3D(label="Initial 3D Model", display_mode="solid", 
+                                               clear_color=[0.0, 0.0, 0.0, 0.0])
+                    gr.Markdown(
+                        """
+                        <div class="model-description">
+                        This is the initial 3D model generated from the video. Finish within 10 seconds.
+                        </div>
+                        """
+                    )
+                
+                with gr.Group():
+                    optimized_model = gr.Model3D(label="Optimized 3D Model", display_mode="solid", 
+                                                 clear_color=[0.0, 0.0, 0.0, 0.0])
+                    gr.Markdown(
+                        """
+                        <div class="model-description">
+                        This is the optimized 3D model with improved accuracy and detail using Gaussian Splatting. Finish within 60 seconds.
+                        </div>
+                        """
+                    )
+            
+            with gr.Tab("Help"):
+                gr.Markdown(
+                    """
+                    ## How to use this tool:
+                    1. Upload a video of the object you want to reconstruct.
+                    2. Adjust the Confidence Threshold and Keyframe Interval if needed.
+                    3. Choose whether to remove the background.
+                    4. Click "Start Reconstruction" to begin the process.
+                    5. The Initial 3D Model will appear first, giving you a quick preview.
+                    6. Once processing is complete, the Optimized 3D Model will show the final result.
+                    
+                    ### Tips:
+                    - For best results, ensure your video captures the object from multiple angles.
+                    - If the model appears noisy, try increasing the Confidence Threshold.
+                    - Experiment with different Keyframe Intervals to balance speed and accuracy.
+                    """
+                )
+    
+    Examples(
+        fn=reconstruct,
+        examples=sorted([
+            os.path.join("examples", name) 
+            for name in os.listdir(os.path.join("examples")) if name.endswith('.webm')
+        ]),
+        inputs=[video_input],
+        outputs=[initial_model, optimized_model],
+        directory_name="examples_video",
+        cache_examples=False,
+    )
     
     reconstruct_btn.click(
         fn=reconstruct,
         inputs=[video_input, conf_thresh, kf_every, remove_background],
-        outputs=[coarse_model, refined_model]
+        outputs=[initial_model, optimized_model]
     )
 
 if __name__ == "__main__":