Fix bug with tmp folder for multiuser

app.py

@@ -2,11 +2,22 @@ import sys
 sys.path.append('wild-gaussian-splatting/mast3r/')
 sys.path.append('demo/')
 
+import os
 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
+from demo_globals import CACHE_PATH
+import shutil
+
+def start_session(req: gr.Request):
+    user_dir = os.path.join(CACHE_PATH, str(req.session_hash))
+    os.makedirs(user_dir, exist_ok=True)
+    
+def end_session(req: gr.Request):
+    user_dir = os.path.join(CACHE_PATH, str(req.session_hash))
+    shutil.rmtree(user_dir)
 
 if __name__ == '__main__':
     with gr.Blocks() as demo:
@@ -30,4 +41,7 @@ if __name__ == '__main__':
             with gr.Tab("3DGS"):
                 gs_demo_tab()
 
+        demo.load(start_session)
+        demo.unload(end_session)
+
         demo.launch(show_error=True, share=None, server_name=None, server_port=None)

demo/demo_globals.py

@@ -9,10 +9,8 @@ from mast3r.model import AsymmetricMASt3R
 
 DATASET_DIR = "colmap_data"
 weights_path = "naver/MASt3R_ViTLarge_BaseDecoder_512_catmlpdpt_metric"
-# weights_path = '/app/wild-gaussian-splatting/mast3r/checkpoints/MASt3R_ViTLarge_BaseDecoder_512_catmlpdpt_metric.pth'
-tmpdirname = tempfile.TemporaryDirectory(suffix='demo')
-chkpt_tag = hash_md5(weights_path)
-CACHE_PATH = os.path.join(tmpdirname.name, chkpt_tag)
+
+CACHE_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'tmp')
 os.makedirs(CACHE_PATH, exist_ok=True)
 
 DEVICE = device = 'cuda' if torch.cuda.is_available() else 'cpu'

demo/gs_demo.py

@@ -13,7 +13,7 @@ def get_dataset_folders(datasets_path):
 
 def gs_demo_tab():
     # datasets_path = "/app/data/scenes/"
-    dataset_path = os.path.join(CACHE_PATH, DATASET_DIR)
+
     def start_training(selected_folder, *args):
         selected_data_path = os.path.join(datasets_path, selected_folder)
         return train(selected_data_path, *args)
@@ -55,8 +55,10 @@ def gs_demo_tab():
         refresh_button = gr.Button("Refresh Datasets", elem_classes="refresh-button")
         dataset_dropdown = gr.Dropdown(label="Select Dataset", choices=[], value="")
 
-        def update_dataset_dropdown():
-            print("update_dataset_dropdown, cache_path", CACHE_PATH)
+        def update_dataset_dropdown(req: gr.Request):
+            USER_DIR = os.path.join(CACHE_PATH, str(req.session_hash))
+            print("update_dataset_dropdown, user_path", USER_DIR)
+            dataset_path = os.path.join(USER_DIR, DATASET_DIR)
             # Update the dataset folders list
             dataset_folders = get_dataset_folders(dataset_path)
             print("dataset_folders", dataset_folders)
@@ -108,7 +110,9 @@ def gs_demo_tab():
         load_model_button = gr.Button("Load 3D Model", interactive=False)
         output = gr.Model3D(label="3D Model Output", visible=False)
         
-        def handle_training_complete(selected_folder, *args):
+        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)
             # Call the training function with the full path

demo/gs_train.py

@@ -160,7 +160,7 @@ def train(
 
             # Log and save
             if (iteration == opt.iterations):
-                point_cloud_path = os.path.join(os.path.join(dataset.model_path, "point_cloud/iteration_{}".format(iteration)), "point_cloud.ply")
+                point_cloud_path = os.path.join(os.path.join(data_source_path, "point_cloud/iteration_{}".format(iteration)), "point_cloud.ply")
                 print("\n[ITER {}] Saving Gaussians to {}".format(iteration, point_cloud_path))
                 scene.save(iteration)
 

demo/mast3r_demo.py

@@ -176,11 +176,12 @@ def save_colmap_scene(scene, save_dir, min_conf_thr=2, clean_depth=False, mask_i
 @spaces.GPU(duration=20)
 def get_reconstructed_scene(snapshot, current_scene_state,
                             min_conf_thr, matching_conf_thr,
-                            as_pointcloud, cam_size, shared_intrinsics, clean_depth, filelist, **kw):
+                            as_pointcloud, cam_size, shared_intrinsics, clean_depth, filelist, 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))
     image_size = 512
     imgs = load_images(filelist, size=image_size, verbose=not SILENT)
     if len(imgs) == 1:
@@ -216,7 +217,7 @@ def get_reconstructed_scene(snapshot, current_scene_state,
     scene_graph = '-'.join(scene_graph_params)
     pairs = make_pairs(imgs, scene_graph=scene_graph, prefilter=None, symmetrize=True)
 
-    base_cache_dir = os.path.join(CACHE_PATH, 'cache')
+    base_cache_dir = os.path.join(USER_DIR, 'cache')
     os.makedirs(base_cache_dir, exist_ok=True)
     def get_next_dir(base_dir):
         run_counter = 0
@@ -235,7 +236,7 @@ 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(CACHE_PATH, DATASET_DIR)
+    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)
     # 
@@ -245,7 +246,7 @@ def get_reconstructed_scene(snapshot, current_scene_state,
        current_scene_state.outfile_name is not None:
         outfile_name = current_scene_state.outfile_name
     else:
-        outfile_name = tempfile.mktemp(suffix='_scene.glb', dir=CACHE_PATH)
+        outfile_name = tempfile.mktemp(suffix='_scene.glb', dir=USER_DIR)
 
     scene_state = SparseGAState(cache_dir, outfile_name)
     outfile = get_3D_model_from_scene(scene, scene_state, min_conf_thr, as_pointcloud, mask_sky,
@@ -258,9 +259,6 @@ def get_reconstructed_scene(snapshot, current_scene_state,
 
 
 def mast3r_demo_tab():
-    if not SILENT:
-        print('Outputing stuff in', CACHE_PATH)
-
     def get_context():
         css = """.gradio-container {margin: 0 !important; min-width: 100%};"""
         title = "MASt3R Demo"