Remove 3dgs for now

app.py

@@ -8,7 +8,7 @@ import gradio as gr
 from mast3r.demo import get_args_parser
 from mast3r.utils.misc import hash_md5
 from mast3r_demo import mast3r_demo_tab
-from gs_demo import gs_demo_tab
+# from gs_demo import gs_demo_tab
 
 if __name__ == '__main__':
     parser = get_args_parser()
@@ -20,6 +20,7 @@ if __name__ == '__main__':
         server_name = '0.0.0.0' if args.local_network else '127.0.0.1'
 
     weights_path = args.weights if args.weights is not None else "naver/" + args.model_name
+    device = device = 'cuda' if torch.cuda.is_available() else 'cpu'
     chkpt_tag = hash_md5(weights_path)
 
     with tempfile.TemporaryDirectory(suffix='demo') as tmpdirname:
@@ -29,9 +30,9 @@ if __name__ == '__main__':
         with gr.Blocks() as demo:
             with gr.Tabs():
                 with gr.Tab("MASt3R Demo"):
-                    mast3r_demo_tab(cache_path, weights_path, args.device)
-                with gr.Tab("Gaussian Splatting Demo"):
-                    gs_demo_tab(cache_path)
+                    mast3r_demo_tab(cache_path, weights_path, device)
+                # with gr.Tab("Gaussian Splatting Demo"):
+                #     gs_demo_tab(cache_path)
 
             demo.launch(server_name=server_name, server_port=args.server_port) 
 

demo/gs_train.py

@@ -13,21 +13,19 @@ sys.path.append(gaussian_splatting_path)
 
 # Import necessary modules from the gaussian-splatting directory
 from utils.loss_utils import l1_loss, ssim
-# from gaussian_renderer import render, network_gui
 from diff_gaussian_rasterization import GaussianRasterizationSettings, GaussianRasterizer
-
-# from scene import Scene, GaussianModel
+from scene import Scene, GaussianModel
 from utils.general_utils import safe_state
 from utils.image_utils import psnr
 
-# # Dynamically import the train module from the gaussian-splatting directory
-# train_spec = importlib.util.spec_from_file_location("gaussian_splatting_train", os.path.join(gaussian_splatting_path, "train.py"))
-# gaussian_splatting_train = importlib.util.module_from_spec(train_spec)
-# train_spec.loader.exec_module(gaussian_splatting_train)
+# Dynamically import the train module from the gaussian-splatting directory
+train_spec = importlib.util.spec_from_file_location("gaussian_splatting_train", os.path.join(gaussian_splatting_path, "train.py"))
+gaussian_splatting_train = importlib.util.module_from_spec(train_spec)
+train_spec.loader.exec_module(gaussian_splatting_train)
 
-# # Import the necessary functions from the dynamically loaded module
-# prepare_output_and_logger = gaussian_splatting_train.prepare_output_and_logger
-# training_report = gaussian_splatting_train.training_report
+# Import the necessary functions from the dynamically loaded module
+prepare_output_and_logger = gaussian_splatting_train.prepare_output_and_logger
+training_report = gaussian_splatting_train.training_report
 
 from dataclasses import dataclass, field
 
@@ -89,204 +87,203 @@ def train(
     densify_from_iter, densify_until_iter, densify_grad_threshold, random_background
 ):
     print(data_source_path)
-    # # Create instances of the parameter dataclasses
-    # dataset = ModelParams(
-    #     sh_degree=sh_degree,
-    #     source_path=data_source_path,
-    #     model_path=model_path,
-    #     images=images,
-    #     resolution=resolution,
-    #     white_background=white_background,
-    #     data_device=data_device,
-    #     eval=eval
-    # )
+    # Create instances of the parameter dataclasses
+    dataset = ModelParams(
+        sh_degree=sh_degree,
+        source_path=data_source_path,
+        model_path=model_path,
+        images=images,
+        resolution=resolution,
+        white_background=white_background,
+        data_device=data_device,
+        eval=eval
+    )
     
-    # pipe = PipelineParams(
-    #     convert_SHs_python=convert_SHs_python,
-    #     compute_cov3D_python=compute_cov3D_python,
-    #     debug=debug
-    # )
+    pipe = PipelineParams(
+        convert_SHs_python=convert_SHs_python,
+        compute_cov3D_python=compute_cov3D_python,
+        debug=debug
+    )
     
-    # opt = OptimizationParams(
-    #     iterations=iterations,
-    #     position_lr_init=position_lr_init,
-    #     position_lr_final=position_lr_final,
-    #     position_lr_delay_mult=position_lr_delay_mult,
-    #     position_lr_max_steps=position_lr_max_steps,
-    #     feature_lr=feature_lr,
-    #     opacity_lr=opacity_lr,
-    #     scaling_lr=scaling_lr,
-    #     rotation_lr=rotation_lr,
-    #     percent_dense=percent_dense,
-    #     lambda_dssim=lambda_dssim,
-    #     densification_interval=densification_interval,
-    #     opacity_reset_interval=opacity_reset_interval,
-    #     densify_from_iter=densify_from_iter,
-    #     densify_until_iter=densify_until_iter,
-    #     densify_grad_threshold=densify_grad_threshold,
-    #     random_background=random_background
-    # )
-
-    # args = TrainingArgs()
-
-    # testing_iterations = args.test_iterations
-    # saving_iterations = args.save_iterations 
-    # checkpoint_iterations = args.checkpoint_iterations 
-    # debug_from = args.debug_from
-
-    # tb_writer = prepare_output_and_logger(dataset)
+    opt = OptimizationParams(
+        iterations=iterations,
+        position_lr_init=position_lr_init,
+        position_lr_final=position_lr_final,
+        position_lr_delay_mult=position_lr_delay_mult,
+        position_lr_max_steps=position_lr_max_steps,
+        feature_lr=feature_lr,
+        opacity_lr=opacity_lr,
+        scaling_lr=scaling_lr,
+        rotation_lr=rotation_lr,
+        percent_dense=percent_dense,
+        lambda_dssim=lambda_dssim,
+        densification_interval=densification_interval,
+        opacity_reset_interval=opacity_reset_interval,
+        densify_from_iter=densify_from_iter,
+        densify_until_iter=densify_until_iter,
+        densify_grad_threshold=densify_grad_threshold,
+        random_background=random_background
+    )
+
+    args = TrainingArgs()
+
+    testing_iterations = args.test_iterations
+    saving_iterations = args.save_iterations 
+    checkpoint_iterations = args.checkpoint_iterations 
+    debug_from = args.debug_from
+
+    tb_writer = prepare_output_and_logger(dataset)
     
-    # gaussians = GaussianModel(dataset.sh_degree)
-    # scene = Scene(dataset, gaussians)
-    # gaussians.training_setup(opt)
-
-    # bg_color = [1, 1, 1] if dataset.white_background else [0, 0, 0]
-    # background = torch.tensor(bg_color, dtype=torch.float32, device="cuda")
-
-    # iter_start = torch.cuda.Event(enable_timing = True)
-    # iter_end = torch.cuda.Event(enable_timing = True)
-
-    # viewpoint_stack = None
-    # ema_loss_for_log = 0.0
-    # first_iter = 0
-    # progress_bar = tqdm(range(first_iter, opt.iterations), desc="Training progress")
-    # first_iter += 1
-
-    # point_cloud_path = ""
-    # progress = gr.Progress()  # Initialize the progress bar
-    # for iteration in range(first_iter, opt.iterations + 1):
-    #     iter_start.record()
-    #     gaussians.update_learning_rate(iteration)
+    gaussians = GaussianModel(dataset.sh_degree)
+    scene = Scene(dataset, gaussians)
+    gaussians.training_setup(opt)
+
+    bg_color = [1, 1, 1] if dataset.white_background else [0, 0, 0]
+    background = torch.tensor(bg_color, dtype=torch.float32, device="cuda")
+
+    iter_start = torch.cuda.Event(enable_timing = True)
+    iter_end = torch.cuda.Event(enable_timing = True)
+
+    viewpoint_stack = None
+    ema_loss_for_log = 0.0
+    first_iter = 0
+    progress_bar = tqdm(range(first_iter, opt.iterations), desc="Training progress")
+    first_iter += 1
+
+    point_cloud_path = ""
+    progress = gr.Progress()  # Initialize the progress bar
+    for iteration in range(first_iter, opt.iterations + 1):
+        iter_start.record()
+        gaussians.update_learning_rate(iteration)
         
-    #     # Every 1000 its we increase the levels of SH up to a maximum degree
-    #     if iteration % 1000 == 0:
-    #         gaussians.oneupSHdegree()
+        # Every 1000 its we increase the levels of SH up to a maximum degree
+        if iteration % 1000 == 0:
+            gaussians.oneupSHdegree()
             
-    #     # Pick a random Camera
-    #     if not viewpoint_stack:
-    #         viewpoint_stack = scene.getTrainCameras().copy()
-    #     viewpoint_cam = viewpoint_stack.pop(randint(0, len(viewpoint_stack)-1))
+        # Pick a random Camera
+        if not viewpoint_stack:
+            viewpoint_stack = scene.getTrainCameras().copy()
+        viewpoint_cam = viewpoint_stack.pop(randint(0, len(viewpoint_stack)-1))
         
-    #     # Render
-    #     if (iteration - 1) == debug_from:
-    #         pipe.debug = True
-    #     bg = torch.rand((3), device="cuda") if opt.random_background else background
+        # Render
+        if (iteration - 1) == debug_from:
+            pipe.debug = True
+        bg = torch.rand((3), device="cuda") if opt.random_background else background
         
-    #     render_pkg = render(viewpoint_cam, gaussians, pipe, bg)
-    #     image, viewspace_point_tensor, visibility_filter, radii = render_pkg["render"], render_pkg["viewspace_points"], render_pkg["visibility_filter"], render_pkg["radii"]
-
-    #     # Loss
-    #     gt_image = viewpoint_cam.original_image.cuda()
-    #     Ll1 = l1_loss(image, gt_image)
-    #     loss = (1.0 - opt.lambda_dssim) * Ll1 + opt.lambda_dssim * (1.0 - ssim(image, gt_image))
-    #     loss.backward()
-    #     iter_end.record()
+        render_pkg = render(viewpoint_cam, gaussians, pipe, bg)
+        image, viewspace_point_tensor, visibility_filter, radii = render_pkg["render"], render_pkg["viewspace_points"], render_pkg["visibility_filter"], render_pkg["radii"]
+
+        # Loss
+        gt_image = viewpoint_cam.original_image.cuda()
+        Ll1 = l1_loss(image, gt_image)
+        loss = (1.0 - opt.lambda_dssim) * Ll1 + opt.lambda_dssim * (1.0 - ssim(image, gt_image))
+        loss.backward()
+        iter_end.record()
         
-    #     with torch.no_grad():
-    #         # Progress bar
-    #         ema_loss_for_log = 0.4 * loss.item() + 0.6 * ema_loss_for_log
-    #         if iteration % 10 == 0:
-    #             progress_bar.set_postfix({"Loss": f"{ema_loss_for_log:.{7}f}"})
-    #             progress_bar.update(10)
-    #             progress(iteration / opt.iterations)  # Update Gradio progress bar
-    #         if iteration == opt.iterations:
-    #             progress_bar.close()
-
-    #         # Log and save
-    #         training_report(tb_writer, iteration, Ll1, loss, l1_loss, iter_start.elapsed_time(iter_end), testing_iterations, scene, render, (pipe, background))
-    #         if (iteration == opt.iterations):
-    #             point_cloud_path = os.path.join(os.path.join(dataset.model_path, "point_cloud/iteration_{}".format(iteration)), "point_cloud.ply")
-    #             print("\n[ITER {}] Saving Gaussians to {}".format(iteration, point_cloud_path))
-    #             scene.save(iteration)
-
-    #         # Densification
-    #         if iteration < opt.densify_until_iter:
-    #             # Keep track of max radii in image-space for pruning
-    #             gaussians.max_radii2D[visibility_filter] = torch.max(gaussians.max_radii2D[visibility_filter], radii[visibility_filter])
-    #             gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter)
-
-    #             if iteration > opt.densify_from_iter and iteration % opt.densification_interval == 0:
-    #                 size_threshold = 20 if iteration > opt.opacity_reset_interval else None
-    #                 gaussians.densify_and_prune(opt.densify_grad_threshold, 0.005, scene.cameras_extent, size_threshold)
-
-    #             if iteration % opt.opacity_reset_interval == 0 or (dataset.white_background and iteration == opt.densify_from_iter):
-    #                 gaussians.reset_opacity()
-
-    #             # Optimizer step
-    #             if iteration < opt.iterations:
-    #                 gaussians.optimizer.step()
-    #                 gaussians.optimizer.zero_grad(set_to_none = True)
-
-    #             if (iteration == opt.iterations):
-    #                 print("\n[ITER {}] Saving Checkpoint".format(iteration))
-    #                 torch.save((gaussians.capture(), iteration), scene.model_path + "/chkpnt" + str(iteration) + ".pth")
-
-
-    # from os import makedirs
-    # from utils.graphics_utils import focal2fov, fov2focal, getProjectionMatrix
-    # import torchvision
-    # import subprocess
-
-    # @torch.no_grad()
-    # def render_path(dataset : ModelParams, iteration : int, pipeline : PipelineParams, render_resize_method='crop'):
-    #     """
-    #     render_resize_method: crop, pad
-    #     """
-    #     gaussians = GaussianModel(dataset.sh_degree)
-    #     scene = Scene(dataset, gaussians, load_iteration=iteration, shuffle=False)
-
-    #     iteration = scene.loaded_iter
-
-    #     bg_color = [1,1,1] if dataset.white_background else [0, 0, 0]
-    #     background = torch.tensor(bg_color, dtype=torch.float32, device="cuda")
-
-    #     model_path = dataset.model_path
-    #     name = "render"
-
-    #     views = scene.getRenderCameras()
-
-    #     # print(len(views))
-    #     render_path = os.path.join(model_path, name, "ours_{}".format(iteration), "renders")
-
-    #     makedirs(render_path, exist_ok=True)
-
-    #     for idx, view in enumerate(tqdm(views, desc="Rendering progress")):
-    #         if render_resize_method == 'crop':
-    #             image_size = 256
-    #         elif render_resize_method == 'pad':
-    #             image_size = max(view.image_width, view.image_height)
-    #         else:
-    #             raise NotImplementedError
-    #         view.original_image = torch.zeros((3, image_size, image_size), device=view.original_image.device)
-    #         focal_length_x = fov2focal(view.FoVx, view.image_width)
-    #         focal_length_y = fov2focal(view.FoVy, view.image_height)
-    #         view.image_width = image_size
-    #         view.image_height = image_size
-    #         view.FoVx = focal2fov(focal_length_x, image_size)
-    #         view.FoVy = focal2fov(focal_length_y, image_size)
-    #         view.projection_matrix = getProjectionMatrix(znear=view.znear, zfar=view.zfar, fovX=view.FoVx, fovY=view.FoVy).transpose(0,1).cuda().float()
-    #         view.full_proj_transform = (view.world_view_transform.unsqueeze(0).bmm(view.projection_matrix.unsqueeze(0))).squeeze(0)
-
-    #         render_pkg = render(view, gaussians, pipeline, background)
-    #         rendering = render_pkg["render"]
-    #         torchvision.utils.save_image(rendering, os.path.join(render_path, '{0:05d}'.format(idx) + ".png"))
-
-    #     # Use ffmpeg to output video
-    #     renders_path = os.path.join(model_path, name, "ours_{}".format(iteration), "renders.mp4")
-    #     # Use ffmpeg to output video
-    #     subprocess.run(["ffmpeg", "-y", 
-    #                 "-framerate", "24",
-    #                 "-i", os.path.join(render_path, "%05d.png"), 
-    #                 "-vf", "pad=ceil(iw/2)*2:ceil(ih/2)*2",
-    #                 "-c:v", "libx264", 
-    #                 "-pix_fmt", "yuv420p",
-    #                 "-crf", "23", 
-    #                 # "-pix_fmt", "yuv420p",  # Set pixel format for compatibility
-    #                 renders_path], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL
-    #                 )
-    #     return renders_path
+        with torch.no_grad():
+            # Progress bar
+            ema_loss_for_log = 0.4 * loss.item() + 0.6 * ema_loss_for_log
+            if iteration % 10 == 0:
+                progress_bar.set_postfix({"Loss": f"{ema_loss_for_log:.{7}f}"})
+                progress_bar.update(10)
+                progress(iteration / opt.iterations)  # Update Gradio progress bar
+            if iteration == opt.iterations:
+                progress_bar.close()
+
+            # Log and save
+            training_report(tb_writer, iteration, Ll1, loss, l1_loss, iter_start.elapsed_time(iter_end), testing_iterations, scene, render, (pipe, background))
+            if (iteration == opt.iterations):
+                point_cloud_path = os.path.join(os.path.join(dataset.model_path, "point_cloud/iteration_{}".format(iteration)), "point_cloud.ply")
+                print("\n[ITER {}] Saving Gaussians to {}".format(iteration, point_cloud_path))
+                scene.save(iteration)
+
+            # Densification
+            if iteration < opt.densify_until_iter:
+                # Keep track of max radii in image-space for pruning
+                gaussians.max_radii2D[visibility_filter] = torch.max(gaussians.max_radii2D[visibility_filter], radii[visibility_filter])
+                gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter)
+
+                if iteration > opt.densify_from_iter and iteration % opt.densification_interval == 0:
+                    size_threshold = 20 if iteration > opt.opacity_reset_interval else None
+                    gaussians.densify_and_prune(opt.densify_grad_threshold, 0.005, scene.cameras_extent, size_threshold)
+
+                if iteration % opt.opacity_reset_interval == 0 or (dataset.white_background and iteration == opt.densify_from_iter):
+                    gaussians.reset_opacity()
+
+                # Optimizer step
+                if iteration < opt.iterations:
+                    gaussians.optimizer.step()
+                    gaussians.optimizer.zero_grad(set_to_none = True)
+
+                if (iteration == opt.iterations):
+                    print("\n[ITER {}] Saving Checkpoint".format(iteration))
+                    torch.save((gaussians.capture(), iteration), scene.model_path + "/chkpnt" + str(iteration) + ".pth")
+
+
+    from os import makedirs
+    from utils.graphics_utils import focal2fov, fov2focal, getProjectionMatrix
+    import torchvision
+    import subprocess
+
+    @torch.no_grad()
+    def render_path(dataset : ModelParams, iteration : int, pipeline : PipelineParams, render_resize_method='crop'):
+        """
+        render_resize_method: crop, pad
+        """
+        gaussians = GaussianModel(dataset.sh_degree)
+        scene = Scene(dataset, gaussians, load_iteration=iteration, shuffle=False)
+
+        iteration = scene.loaded_iter
+
+        bg_color = [1,1,1] if dataset.white_background else [0, 0, 0]
+        background = torch.tensor(bg_color, dtype=torch.float32, device="cuda")
+
+        model_path = dataset.model_path
+        name = "render"
+
+        views = scene.getRenderCameras()
+
+        # print(len(views))
+        render_path = os.path.join(model_path, name, "ours_{}".format(iteration), "renders")
+
+        makedirs(render_path, exist_ok=True)
+
+        for idx, view in enumerate(tqdm(views, desc="Rendering progress")):
+            if render_resize_method == 'crop':
+                image_size = 256
+            elif render_resize_method == 'pad':
+                image_size = max(view.image_width, view.image_height)
+            else:
+                raise NotImplementedError
+            view.original_image = torch.zeros((3, image_size, image_size), device=view.original_image.device)
+            focal_length_x = fov2focal(view.FoVx, view.image_width)
+            focal_length_y = fov2focal(view.FoVy, view.image_height)
+            view.image_width = image_size
+            view.image_height = image_size
+            view.FoVx = focal2fov(focal_length_x, image_size)
+            view.FoVy = focal2fov(focal_length_y, image_size)
+            view.projection_matrix = getProjectionMatrix(znear=view.znear, zfar=view.zfar, fovX=view.FoVx, fovY=view.FoVy).transpose(0,1).cuda().float()
+            view.full_proj_transform = (view.world_view_transform.unsqueeze(0).bmm(view.projection_matrix.unsqueeze(0))).squeeze(0)
+
+            render_pkg = render(view, gaussians, pipeline, background)
+            rendering = render_pkg["render"]
+            torchvision.utils.save_image(rendering, os.path.join(render_path, '{0:05d}'.format(idx) + ".png"))
+
+        # Use ffmpeg to output video
+        renders_path = os.path.join(model_path, name, "ours_{}".format(iteration), "renders.mp4")
+        # Use ffmpeg to output video
+        subprocess.run(["ffmpeg", "-y", 
+                    "-framerate", "24",
+                    "-i", os.path.join(render_path, "%05d.png"), 
+                    "-vf", "pad=ceil(iw/2)*2:ceil(ih/2)*2",
+                    "-c:v", "libx264", 
+                    "-pix_fmt", "yuv420p",
+                    "-crf", "23", 
+                    # "-pix_fmt", "yuv420p",  # Set pixel format for compatibility
+                    renders_path], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL
+                    )
+        return renders_path
     
-    # renders_path = render_path(dataset, opt.iterations, pipe, render_resize_method='crop')
+    renders_path = render_path(dataset, opt.iterations, pipe, render_resize_method='crop')
 
-    # return renders_path, point_cloud_path
-    return None, None
+    return renders_path, point_cloud_path