Add more fixes to demo files

app.py

@@ -20,6 +20,7 @@ if __name__ == '__main__':
     # else:
     server_name = '0.0.0.0'# if args.local_network else '127.0.0.1'
 
+    # weights_path = '/app/wild-gaussian-splatting/mast3r/checkpoints/MASt3R_ViTLarge_BaseDecoder_512_catmlpdpt_metric.pth'
     weights_path = "naver/MASt3R_ViTLarge_BaseDecoder_512_catmlpdpt_metric"#args.weights if args.weights is not None else  + MASt3R_ViTLarge_BaseDecoder_512_catmlpdpt_metric
     device = device = 'cuda' if torch.cuda.is_available() else 'cpu'
     chkpt_tag = hash_md5(weights_path)
@@ -36,5 +37,6 @@ if __name__ == '__main__':
                     gs_demo_tab(cache_path)
 
             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=5555) 
 
 # python3 demo.py --weights "/app/mast3r/checkpoints/MASt3R_ViTLarge_BaseDecoder_512_catmlpdpt_metric.pth" --device "cuda" --server_port 3334 --local_network "$@"

demo/gs_demo.py

@@ -38,7 +38,7 @@ def gs_demo_tab(cache_path):
         def update_dataset_dropdown():
             print("update_dataset_dropdown, cache_path", cache_path)
             # Update the dataset folders list
-            dataset_folders = get_dataset_folders(datasets_path)
+            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
@@ -109,7 +109,7 @@ def gs_demo_tab(cache_path):
         
         def handle_training_complete(selected_folder, *args):
             # Construct the full path to the selected dataset
-            selected_data_path = os.path.join(datasets_path, selected_folder)
+            selected_data_path = os.path.join(dataset_path, selected_folder)
             # Call the training function with the full path
             video_path, model_path = train(selected_data_path, *args)
             # Then return all required outputs

demo/gs_train.py

@@ -8,7 +8,7 @@ import gradio as gr
 import importlib.util
 from dataclasses import dataclass, field
 
-import spaces
+# import spaces
 
 
 @dataclass
@@ -60,7 +60,7 @@ class TrainingArgs:
     checkpoint_iterations: list[int] = field(default_factory=lambda: [7_000, 15_000, 30_000])
     start_checkpoint: str = None
 
-@spaces.GPU(duration=20)
+# @spaces.GPU(duration=20)
 def train(
     data_source_path, sh_degree, model_path, images, resolution, white_background, data_device, eval,
     convert_SHs_python, compute_cov3D_python, debug,
@@ -77,7 +77,7 @@ def train(
 
         # Import necessary modules from the gaussian-splatting directory
         from utils.loss_utils import l1_loss, ssim
-        from diff_gaussian_rasterization import GaussianRasterizationSettings, GaussianRasterizer
+        from gaussian_renderer import render
         from scene import Scene, GaussianModel
         from utils.general_utils import safe_state
         from utils.image_utils import psnr

demo/mast3r_demo.py

@@ -139,42 +139,42 @@ def get_3D_model_from_scene(silent, scene_state, min_conf_thr=2, as_pointcloud=F
     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):
-    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'))
-        from colmap_dataset_utils import (
-            inv,
-            init_filestructure,
-            save_images_masks,
-            save_cameras,
-            save_imagestxt,
-            save_pointcloud,
-            save_pointcloud_with_normals
-        )
-
-    cam2world = scene.get_im_poses().detach().cpu().numpy()
-    world2cam = inv(cam2world) #
-    principal_points = scene.get_principal_points().detach().cpu().numpy()
-    focals = scene.get_focals().detach().cpu().numpy()[..., None]
-    imgs = np.array(scene.imgs)
-
-    pts3d, _, confs = scene.get_dense_pts3d(clean_depth=clean_depth)
-    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)
-    save_imagestxt(world2cam, sparse_path)
-    save_pointcloud_with_normals(imgs, pts3d, masks, sparse_path)
-    return save_path
-
-@spaces.GPU(duration=20)
+# def save_colmap_scene(scene, save_dir, min_conf_thr=2, clean_depth=False):
+#     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'))
+#         from colmap_dataset_utils import (
+#             inv,
+#             init_filestructure,
+#             save_images_masks,
+#             save_cameras,
+#             save_imagestxt,
+#             save_pointcloud,
+#             save_pointcloud_with_normals
+#         )
+
+#     cam2world = scene.get_im_poses().detach().cpu().numpy()
+#     world2cam = inv(cam2world) #
+#     principal_points = scene.get_principal_points().detach().cpu().numpy()
+#     focals = scene.get_focals().detach().cpu().numpy()[..., None]
+#     imgs = np.array(scene.imgs)
+
+#     pts3d, _, confs = scene.get_dense_pts3d(clean_depth=clean_depth)
+#     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)
+#     save_imagestxt(world2cam, sparse_path)
+#     save_pointcloud_with_normals(imgs, pts3d, masks, sparse_path)
+#     return save_path
+
+@spaces.GPU(duration=10)
 def get_reconstructed_scene(outdir, model, device, silent, image_size, current_scene_state,
                             filelist, optim_level, lr1, niter1, lr2, niter2, min_conf_thr, matching_conf_thr,
                             as_pointcloud, mask_sky, clean_depth, transparent_cams, cam_size, scenegraph_type, winsize,