Add lazy caching for the examples

app.py

@@ -22,7 +22,7 @@ def end_session(req: gr.Request):
 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);">
+        <div style="text-align: center; padding: 20px; color: #333;">
             <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>
@@ -30,7 +30,7 @@ if __name__ == '__main__':
                 <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;">Note: After a page reload, any generated MASt3R datasets in the 3DGS tab will be deleted.</p>
+            <p style="font-size: 16px; font-weight: bold;">Note: After a page reload, any generated MASt3R datasets in the 3DGS tab will be deleted.</p>
             <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>
@@ -45,4 +45,5 @@ if __name__ == '__main__':
         demo.load(start_session)
         demo.unload(end_session)
 
-        demo.launch(show_error=True, share=None, server_name=None, server_port=None)
+        demo.launch(show_error=True, share=None, server_name='0.0.0.0', server_port=8111)
+        # demo.launch(show_error=True, share=None, server_name=None, server_port=None)

demo/demo_globals.py

@@ -12,6 +12,8 @@ weights_path = "naver/MASt3R_ViTLarge_BaseDecoder_512_catmlpdpt_metric"
 
 CACHE_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'tmp')
 os.makedirs(CACHE_PATH, exist_ok=True)
+EXAMPLE_PATH = os.path.join(CACHE_PATH, 'examples_datasets')
+os.makedirs(EXAMPLE_PATH, exist_ok=True)
 
 DEVICE = device = 'cuda' if torch.cuda.is_available() else 'cpu'
 MODEL = AsymmetricMASt3R.from_pretrained(weights_path).to(DEVICE)

demo/gs_demo.py

@@ -2,14 +2,16 @@ import gradio as gr
 from gs_train import train
 import os
 
-from demo_globals import CACHE_PATH, MODEL, DEVICE, SILENT, DATASET_DIR
+from demo_globals import CACHE_PATH, EXAMPLE_PATH, MODEL, DEVICE, SILENT, DATASET_DIR
 
 def get_dataset_folders(datasets_path):
-    try:
-        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 []
+
+    folder = []
+    if os.path.isdir(datasets_path):
+        folder += [f for f in os.listdir(datasets_path) if os.path.isdir(os.path.join(datasets_path, f))]
+    if os.path.isdir(EXAMPLE_PATH):
+        folder += [f for f in os.listdir(EXAMPLE_PATH) if os.path.isdir(os.path.join(EXAMPLE_PATH, f))]
+    return sorted(folder, key=lambda x: int(x.split('_')[-1]) if x.split('_')[-1].isdigit() else float('inf'))
 
 def gs_demo_tab():
     # datasets_path = "/app/data/scenes/"
@@ -112,11 +114,12 @@ def gs_demo_tab():
         
         def handle_training_complete(selected_folder, req: gr.Request, *args):
             USER_DIR = os.path.join(CACHE_PATH, str(req.session_hash))
-            dataset_path = os.path.join(USER_DIR, DATASET_DIR)
-            # Construct the full path to the selected dataset
-            selected_data_path = os.path.join(dataset_path, selected_folder)
+            if 'run' in selected_folder:
+                dataset_path = os.path.join(USER_DIR, DATASET_DIR, selected_folder)
+            else:
+                dataset_path = os.path.join(EXAMPLE_PATH, selected_folder)
             # Call the training function with the full path
-            video_path, model_path = train(selected_data_path, *args)
+            video_path, model_path = train(dataset_path, *args)
             # Then return all required outputs
             return [
                 video_path,           # video output

demo/mast3r_demo.py

@@ -18,6 +18,7 @@ import copy
 from scipy.spatial.transform import Rotation
 import tempfile
 import shutil
+import typing
 
 from mast3r.cloud_opt.sparse_ga import sparse_global_alignment
 from mast3r.cloud_opt.tsdf_optimizer import TSDFPostProcess
@@ -36,7 +37,7 @@ 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
+from demo_globals import CACHE_PATH, EXAMPLE_PATH, MODEL, DEVICE, SILENT, DATASET_DIR
 
 class SparseGAState():
     def __init__(self, cache_dir=None, outfile_name=None):
@@ -112,16 +113,11 @@ def _convert_scene_output_to_glb(outfile, imgs, pts3d, mask, focals, cams2world,
     return outfile
 
 
-def get_3D_model_from_scene(scene, scene_state, min_conf_thr=2, as_pointcloud=False, mask_sky=False,
+def get_3D_model_from_scene(scene, outfile, min_conf_thr=2, as_pointcloud=False, mask_sky=False,
                             clean_depth=False, transparent_cams=False, cam_size=0.05, TSDF_thresh=0):
     """
     extract 3D_model (glb file) from a reconstructed scene
     """
-    if scene_state is None:
-        return None
-    outfile = scene_state.outfile_name
-    if outfile is None:
-        return None
 
     # # get optimized values from scene
     # scene = scenescene_state.sparse_ga
@@ -174,14 +170,20 @@ def save_colmap_scene(scene, save_dir, min_conf_thr=2, clean_depth=False, mask_i
     return save_path
 
 @spaces.GPU(duration=160)
-def get_reconstructed_scene(snapshot, current_scene_state,
+def get_reconstructed_scene(snapshot,
                             min_conf_thr, matching_conf_thr,
-                            as_pointcloud, cam_size, shared_intrinsics, clean_depth, filelist, req: gradio.Request, **kw):
+                            as_pointcloud, cam_size, shared_intrinsics, clean_depth, filelist, example_name, req: gradio.Request, **kw):
     """
     from a list of images, run mast3r inference, sparse global aligner.
     then run get_3D_model_from_scene
     """
-    USER_DIR = os.path.join(CACHE_PATH, str(req.session_hash))
+
+    if example_name != '':
+        USER_DIR = os.path.join(CACHE_PATH, example_name)
+    else:
+        USER_DIR = os.path.join(CACHE_PATH, str(req.session_hash))
+    os.makedirs(USER_DIR, exist_ok=True)
+    
     image_size = 512
     imgs = load_images(filelist, size=image_size, verbose=not SILENT)
     if len(imgs) == 1:
@@ -236,26 +238,25 @@ def get_reconstructed_scene(snapshot, current_scene_state,
                                     opt_depth='depth' in optim_level, shared_intrinsics=shared_intrinsics,
                                     matching_conf_thr=matching_conf_thr, **kw)
 
-    base_colmapdata_dir = os.path.join(USER_DIR, DATASET_DIR)
-    os.makedirs(base_colmapdata_dir, exist_ok=True)
-    colmap_data_dir = get_next_dir(base_colmapdata_dir)
-    # 
-    save_colmap_scene(scene, colmap_data_dir, min_conf_thr, clean_depth)
 
-    if current_scene_state is not None and \
-       current_scene_state.outfile_name is not None:
-        outfile_name = current_scene_state.outfile_name
+    if example_name:
+        colmap_data_dir = os.path.join(EXAMPLE_PATH, example_name)
     else:
-        outfile_name = tempfile.mktemp(suffix='_scene.glb', dir=USER_DIR)
+        colmap_data_dir = get_next_dir(os.path.join(USER_DIR, DATASET_DIR))
+    os.makedirs(colmap_data_dir, exist_ok=True)
+    
+
+    save_colmap_scene(scene, colmap_data_dir, min_conf_thr, clean_depth)
 
-    scene_state = SparseGAState(cache_dir, outfile_name)
-    outfile = get_3D_model_from_scene(scene, scene_state, min_conf_thr, as_pointcloud, mask_sky,
+    outfile_name = os.path.join(USER_DIR, 'default_scene.glb')
+
+    outfile = get_3D_model_from_scene(scene, outfile_name, min_conf_thr, as_pointcloud, mask_sky,
                                       clean_depth, transparent_cams, cam_size, TSDF_thresh)
     print(f"colmap_data_dir: {colmap_data_dir}")
     print(f"outfile_name: {outfile_name}")
     print(f"cache_dir: {cache_dir}")
     torch.cuda.empty_cache()
-    return scene_state, outfile
+    return outfile
 
 
 def mast3r_demo_tab():
@@ -285,6 +286,12 @@ def mast3r_demo_tab():
         snapshot = gradio.Image(None, visible=False)
         run_btn = gradio.Button("Run")
 
+        dummy_req = gradio.Request()
+        dummy_text = gradio.Textbox(value="", visible=False)
+
+        example_name = gradio.Textbox(value="", visible=False)
+
+
         with gradio.Row():
             matching_conf_thr = gradio.Slider(label="Matching Confidence Thr", value=2.,
                                 minimum=0., maximum=30., step=0.1,
@@ -300,9 +307,9 @@ def mast3r_demo_tab():
         outmodel = gradio.Model3D()
         run_btn.click(
             fn=get_reconstructed_scene,
-            inputs=[snapshot, scene, min_conf_thr, matching_conf_thr,
-                    as_pointcloud, cam_size, shared_intrinsics, clean_depth, inputfiles],
-            outputs=[scene, outmodel]
+            inputs=[snapshot, min_conf_thr, matching_conf_thr,
+                    as_pointcloud, cam_size, shared_intrinsics, clean_depth, inputfiles, dummy_text],
+            outputs=[outmodel]
         )
 
         tower_folder = os.path.join(HERE_PATH, '../wild-gaussian-splatting/mast3r/assets/NLE_tower/')
@@ -315,35 +322,49 @@ def mast3r_demo_tab():
         examples = gradio.Examples(
             examples=[
                 [
-                    turtle_images[0],
-                    None,
+                    puma_images[0],
                     1.5, 0.0, 0.2, True, True, False,
-                    turtle_images,
+                    puma_images,
+                    'puma',
+                    None,
                 ]
             ],
-            inputs=[snapshot, scene, min_conf_thr, matching_conf_thr, cam_size, as_pointcloud, shared_intrinsics, clean_depth, inputfiles],
+            inputs=[snapshot, min_conf_thr, matching_conf_thr, cam_size, as_pointcloud, shared_intrinsics, clean_depth, inputfiles, example_name],
+            fn=get_reconstructed_scene,
+            outputs=[outmodel],
+            run_on_click=True,
+            cache_examples='lazy',
         )
         examples = gradio.Examples(
             examples=[
                 [
-                    puma_images[0],
-                    None,
+                    turtle_images[0],
                     1.5, 0.0, 0.2, True, True, False,
-                    puma_images,
+                    turtle_images,
+                    'turtle',
+                    None
                 ]
             ],
-            inputs=[snapshot, scene, min_conf_thr, matching_conf_thr, cam_size, as_pointcloud, shared_intrinsics, clean_depth, inputfiles],
+            inputs=[snapshot, min_conf_thr, matching_conf_thr, cam_size, as_pointcloud, shared_intrinsics, clean_depth, inputfiles, example_name],
+            fn=get_reconstructed_scene,
+            outputs=[outmodel],
+            run_on_click=True,
+            cache_examples='lazy',
         )
         examples = gradio.Examples(
             examples=[
                 [
                     tower_images[0],
-                    None,
                     1.5, 0.0, 0.2, True, False, False,
                     tower_images,
+                    'tower',
                 ]
             ],
-            inputs=[snapshot, scene, min_conf_thr, matching_conf_thr, cam_size, as_pointcloud, shared_intrinsics, clean_depth, inputfiles],
+            inputs=[snapshot, min_conf_thr, matching_conf_thr, cam_size, as_pointcloud, shared_intrinsics, clean_depth, inputfiles, example_name],
+            fn=get_reconstructed_scene,
+            outputs=[outmodel],
+            run_on_click=True,
+            cache_examples='lazy',
         )