Change style a bit, add new whl. Add examples.

app.py

@@ -2,31 +2,25 @@ import sys
 sys.path.append('wild-gaussian-splatting/mast3r/')
 sys.path.append('demo/')
 
-import os
-os.environ['CUDA_LAUNCH_BLOCKING'] = '1'
-
-
 import gradio as gr
 import torch
 from mast3r.demo import get_args_parser
 from mast3r_demo import mast3r_demo_tab
 from gs_demo import gs_demo_tab
 
-torch.backends.cuda.matmul.allow_tf32 = True
-
 if __name__ == '__main__':
     with gr.Blocks() as demo:
         gr.HTML('''
         <div style="text-align: center; padding: 20px; background-color: #f9f9f9; border-radius: 10px; box-shadow: 0 4px 8px rgba(0, 0, 0, 0.1);">
-            <h2 style="color: #333;">MASt3R and 3DGS Pipeline Demo</h2>
-            <p style="font-size: 16px; color: #555;">This pipeline is designed for 3D reconstruction using MASt3R and 3DGS.</p>
-            <p style="font-size: 16px; color: #555;">The process is divided into two stages:</p>
-            <ol style="text-align: left; display: inline-block; margin: 0 auto; color: #555;">
+            <h2>MASt3R and 3DGS Pipeline Demo</h2>
+            <p style="font-size: 16px;">This pipeline is designed for 3D reconstruction using MASt3R and 3DGS.</p>
+            <p style="font-size: 16px;">The process is divided into two stages:</p>
+            <ol style="text-align: left; display: inline-block; margin: 0 auto;">
                 <li>MASt3R is used to obtain the initial point cloud and camera parameters.</li>
                 <li>3DGS is then trained on the results from MASt3R to refine the 3D scene representation.</li>
             </ol>
-            <p style="font-size: 16px; color: #555;">For a full version of this pipeline, please visit the repository at:</p>
-            <a href="https://github.com/nerlfield/wild-gaussian-splatting" target="_blank" style="font-size: 16px; color: #007bff; text-decoration: none;">nerlfield/wild-gaussian-splatting</a>
+            <p style="font-size: 16px;">For a full version of this pipeline, please visit the repository at:</p>
+            <a href="https://github.com/nerlfield/wild-gaussian-splatting" target="_blank" style="font-size: 16px; text-decoration: none;">nerlfield/wild-gaussian-splatting</a>
         </div>
         ''')
 

demo/gs_demo.py

@@ -6,7 +6,8 @@ from demo_globals import CACHE_PATH, MODEL, DEVICE, SILENT, DATASET_DIR
 
 def get_dataset_folders(datasets_path):
     try:
-        return [f for f in os.listdir(datasets_path) if os.path.isdir(os.path.join(datasets_path, f))]
+        folders = [f for f in os.listdir(datasets_path) if os.path.isdir(os.path.join(datasets_path, f))]
+        return sorted(folders, key=lambda x: int(x.split('_')[-1]) if x.split('_')[-1].isdigit() else float('inf'))
     except FileNotFoundError:
         return []
 
@@ -38,7 +39,7 @@ def gs_demo_tab():
 
         # Instructions
         gr.Markdown('''
-        <div style="padding: 10px; background-color: #e9f7ef; border-radius: 5px; margin-bottom: 10px;">
+        <div style="padding: 10px; border-radius: 5px; margin-bottom: 10px;">
             <h3>Instructions for 3DGS Demo</h3>
             <ul style="text-align: left; color: #333;">
                 <li>Make sure to press "Refresh Datasets" to obtain an updated list of datasets from Stage 1. They are in the format run_0, run_1, run_...</li>
@@ -58,10 +59,9 @@ def gs_demo_tab():
             print("update_dataset_dropdown, cache_path", CACHE_PATH)
             # Update the dataset folders list
             dataset_folders = get_dataset_folders(dataset_path)
-            # dataset_folders = "/app/data/scenes/"
             print("dataset_folders", dataset_folders)
-            # Only set a default value if there are folders available
-            default_value = dataset_folders[0] if dataset_folders else None
+            # Set the default value to the last run if there are folders available
+            default_value = dataset_folders[-1] if dataset_folders else None
             return gr.Dropdown(label="Select Dataset", choices=dataset_folders, value=default_value)
         
         # Set the update function to be called when the refresh button is clicked
@@ -88,7 +88,7 @@ def gs_demo_tab():
                     densify_from_iter = gr.Number(label="Densify From Iter", value=500)
                     densify_until_iter = gr.Number(label="Densify Until Iter", value=15000)
                     densify_grad_threshold = gr.Number(label="Densify Grad Threshold", value=0.0002)
-        iterations = gr.Slider(label="Iterations", value=7000, minimum=1, maximum=15000, step=5)
+        iterations = gr.Slider(label="Iterations", value=9000, minimum=1, maximum=15000, step=5)
         
         start_button = gr.Button("Start Training")
         

demo/gs_train.py

@@ -51,7 +51,7 @@ class ModelParams:
     data_device: str = "cuda"
     eval: bool = False
 
-@spaces.GPU(duration=20)
+@spaces.GPU(duration=30)
 def train(
     data_source_path, iterations, position_lr_init, position_lr_final, position_lr_delay_mult,
     position_lr_max_steps, feature_lr, opacity_lr, scaling_lr, rotation_lr,

demo/mast3r_demo.py

@@ -25,13 +25,15 @@ from mast3r.cloud_opt.tsdf_optimizer import TSDFPostProcess
 from mast3r.model import AsymmetricMASt3R
 from dust3r.image_pairs import make_pairs
 from dust3r.utils.image import load_images
-from dust3r.utils.device import to_numpy
 from dust3r.viz import add_scene_cam, CAM_COLORS, OPENGL, pts3d_to_trimesh, cat_meshes
 from dust3r.demo import get_args_parser as dust3r_get_args_parser
+from copy import deepcopy
 
 import matplotlib.pyplot as pl
 
 import torch
+import os.path as path
+HERE_PATH = path.normpath(path.dirname(__file__))  # noqa
 
 from demo_globals import CACHE_PATH, MODEL, DEVICE, SILENT, DATASET_DIR
 
@@ -138,7 +140,7 @@ def get_3D_model_from_scene(scene, scene_state, min_conf_thr=2, as_pointcloud=Fa
     return _convert_scene_output_to_glb(outfile, rgbimg, pts3d, msk, focals, cams2world, as_pointcloud=as_pointcloud,
                                         transparent_cams=transparent_cams, cam_size=cam_size, silent=SILENT)
 
-def save_colmap_scene(scene, save_dir, min_conf_thr=2, clean_depth=False):
+def save_colmap_scene(scene, save_dir, min_conf_thr=2, clean_depth=False, mask_images=True):
     if 'save_pointcloud_with_normals' not in globals():
         sys.path.append(os.path.join(os.path.dirname(__file__), '../wild-gaussian-splatting/gaussian-splatting'))
         sys.path.append(os.path.join(os.path.dirname(__file__), '../wild-gaussian-splatting/src'))
@@ -162,10 +164,7 @@ def save_colmap_scene(scene, save_dir, min_conf_thr=2, clean_depth=False):
     pts3d = [i.detach().reshape(imgs[0].shape) for i in pts3d] #
 
     masks = to_numpy([c > min_conf_thr for c in to_numpy(confs)])
-
     # move
-    mask_images = True
-
     save_path, images_path, masks_path, sparse_path = init_filestructure(save_dir)
     save_images_masks(imgs, masks, images_path, masks_path, mask_images)
     save_cameras(focals, principal_points, sparse_path, imgs_shape=imgs.shape)
@@ -173,10 +172,10 @@ def save_colmap_scene(scene, save_dir, min_conf_thr=2, clean_depth=False):
     save_pointcloud_with_normals(imgs, pts3d, masks, sparse_path)
     return save_path
 
-@spaces.GPU(duration=10)
-def get_reconstructed_scene(current_scene_state,
-                            filelist, min_conf_thr, matching_conf_thr,
-                            as_pointcloud, cam_size, shared_intrinsics, **kw):
+@spaces.GPU(duration=30)
+def get_reconstructed_scene(snapshot, current_scene_state,
+                            min_conf_thr, matching_conf_thr,
+                            as_pointcloud, cam_size, shared_intrinsics, clean_depth, filelist, **kw):
     """
     from a list of images, run mast3r inference, sparse global aligner.
     then run get_3D_model_from_scene
@@ -189,11 +188,11 @@ def get_reconstructed_scene(current_scene_state,
         filelist = [filelist[0], filelist[0] + '_2']
 
     lr1 = 0.07
-    niter1 = 500
+    niter1 = 700
     lr2 = 0.014
-    niter2 = 200
-    optim_level = 'refine'
-    mask_sky, clean_depth, transparent_cams = False, True, False
+    niter2 = 300
+    optim_level = 'refine+depth'
+    mask_sky, transparent_cams = False, False
     if len(filelist) < 13:
         scenegraph_type = 'complete'
         winsize = 1
@@ -272,9 +271,8 @@ def mast3r_demo_tab():
         # Title for the MASt3R demo
         gradio.HTML('<h2 style="text-align: center;">MASt3R Demo</h2>')
         
-        # Add instructions for the MASt3R demo
         gradio.HTML('''
-        <div style="padding: 10px; background-color: #e9f7ef; border-radius: 5px; margin-bottom: 10px;">
+        <div style="padding: 10px; border-radius: 5px; margin-bottom: 10px;">
             <h3>Instructions for MASt3R Demo</h3>
             <ul style="text-align: left; color: #333;">
                 <li>Upload images. It is recommended to use no more than 10-12 images to avoid exceeding the 3-minute runtime limit for zeroGPU dynamic resources.</li>
@@ -285,7 +283,7 @@ def mast3r_demo_tab():
         ''')
         
         inputfiles = gradio.File(file_count="multiple")
-
+        snapshot = gradio.Image(None, visible=False)
         run_btn = gradio.Button("Run")
 
         with gradio.Row():
@@ -298,14 +296,64 @@ def mast3r_demo_tab():
             as_pointcloud = gradio.Checkbox(value=True, label="As pointcloud")
             shared_intrinsics = gradio.Checkbox(value=False, label="Shared intrinsics",
                                 info="Only optimize one set of intrinsics for all views")
-
+            clean_depth = gradio.Checkbox(value=False, label="Clean depth")
+        
         outmodel = gradio.Model3D()
         run_btn.click(
             fn=get_reconstructed_scene,
-            inputs=[scene, inputfiles, min_conf_thr, matching_conf_thr,
-                    as_pointcloud, cam_size, shared_intrinsics],
+            inputs=[snapshot, scene, min_conf_thr, matching_conf_thr,
+                    as_pointcloud, cam_size, shared_intrinsics, clean_depth, inputfiles],
             outputs=[scene, outmodel]
         )
-            
+
+        tower_folder = os.path.join(HERE_PATH, '../wild-gaussian-splatting/mast3r/assets/NLE_tower/')
+        turtle_folder = os.path.join(HERE_PATH, '../wild-gaussian-splatting/data/images/turtle_imgs/')
+        puma_folder = os.path.join(HERE_PATH, '../wild-gaussian-splatting/data/images/puma_imgs/')
+        tower_images = [os.path.join(tower_folder, file) for file in os.listdir(tower_folder) if file.endswith('.jpg') and not file.startswith('2679C386-1DC0-4443-81B5-93D7EDE4AB37-83120-000041DADB2EA917')] # my code not addpted to different size input
+        turtle_images = [os.path.join(turtle_folder, file) for file in os.listdir(turtle_folder) if file.endswith('.jpg')]
+        puma_images = [os.path.join(puma_folder, file) for file in os.listdir(puma_folder)[:12] if file.endswith('.jpg')]
+        
+        
+        examples = gradio.Examples(
+            examples=[
+                [
+                    puma_images[0],
+                    None,
+                    1.5, 0.0, 0.2, True, True, False,
+                    puma_images,
+                ]
+            ],
+            inputs=[snapshot, scene, min_conf_thr, matching_conf_thr, cam_size, as_pointcloud, shared_intrinsics, clean_depth, inputfiles],
+            outputs=[scene, outmodel],
+            fn=get_reconstructed_scene,
+        )
+        examples = gradio.Examples(
+            examples=[
+                [
+                    turtle_images[0],
+                    None,
+                    1.5, 0.0, 0.2, True, True, False,
+                    turtle_images,
+                ]
+            ],
+            inputs=[snapshot, scene, min_conf_thr, matching_conf_thr, cam_size, as_pointcloud, shared_intrinsics, clean_depth, inputfiles],
+            outputs=[scene, outmodel],
+            fn=get_reconstructed_scene,
+        )
+        examples = gradio.Examples(
+            examples=[
+                [
+                    tower_images[0],
+                    None,
+                    1.5, 0.0, 0.2, True, True, False,
+                    tower_images,
+                ]
+            ],
+            inputs=[snapshot, scene, min_conf_thr, matching_conf_thr, cam_size, as_pointcloud, shared_intrinsics, clean_depth, inputfiles],
+            outputs=[scene, outmodel],
+            fn=get_reconstructed_scene,
+        )
+
+
     return demo
     

requirements.txt

@@ -12,7 +12,6 @@ scipy
 einops
 trimesh
 
-https://huggingface.co/spaces/JeffreyXiang/TRELLIS/resolve/main/wheels/nvdiffrast-0.3.3-cp310-cp310-linux_x86_64.whl?download=true
 https://huggingface.co/spaces/ostapagon/mast3r-3dgs/resolve/main/wheels/diff_gaussian_rasterization-0.0.0-cp310-cp310-linux_x86_64.whl?download=true
 https://huggingface.co/spaces/ostapagon/mast3r-3dgs/resolve/main/wheels/simple_knn-0.0.0-cp310-cp310-linux_x86_64.whl?download=true
 

wheels/simple_knn-0.0.0-cp310-cp310-linux_x86_64.whl

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:4fca844adde511fb22934188e1311f3d7f43e8a1f22d9c7a21957d5dbd10c2bb
-size 2987149
+oid sha256:03f3f1b4b1e534649efb34ad8f1d8fe2dacf0e11eb9ca83632833e386b884927
+size 3043174