Spaces:

blanchon
/

gaussian-splatting-kit

Build error

App Files Files Community

blanchon commited on Oct 5, 2023

Commit

76c1d81

1 Parent(s): cb2de0f

Update code

Browse files

Files changed (3) hide show

server.py +514 -317
services/gaussian_splatting_cuda.py +19 -3
services/rerun.py +92 -74

server.py CHANGED Viewed

@@ -1,303 +1,255 @@
 from pathlib import Path
 import shutil
 import tempfile
 import gradio as gr
 import uuid
 from typing_extensions import TypedDict, Tuple
 from fastapi import FastAPI
-from fastapi.staticfiles import StaticFiles
-import uvicorn
 app = FastAPI()
 # create a static directory to store the static files
-gs_dir = Path(str(tempfile.gettempdir())) / "gaussian_splatting_gradio"
-gs_dir.mkdir(parents=True, exist_ok=True)
-# mount FastAPI StaticFiles server
-app.mount("/static", StaticFiles(directory=gs_dir), name="static")
 StateDict = TypedDict("StateDict", {
     "uuid": str,
 })
-def getHTML():
-    html_body = """
-<body>
-    <div id="progress"></div>
-    <div id="message"></div>
-    <div class="scene" id="spinner">
-        <div class="cube-wrapper">
-            <div class="cube">
-                <div class="cube-faces">
-                    <div class="cube-face bottom"></div>
-                    <div class="cube-face top"></div>
-                    <div class="cube-face left"></div>
-                    <div class="cube-face right"></div>
-                    <div class="cube-face back"></div>
-                    <div class="cube-face front"></div>
-                </div>
-            </div>
-        </div>
-    </div>
-    <canvas id="canvas"></canvas>
-    <div id="quality">
-        <span id="fps"></span>
-    </div>
-    <style>
-        .cube-wrapper {
-            transform-style: preserve-3d;
-        }
-        .cube {
-            transform-style: preserve-3d;
-            transform: rotateX(45deg) rotateZ(45deg);
-            animation: rotation 2s infinite;
-        }
-        .cube-faces {
-            transform-style: preserve-3d;
-            height: 80px;
-            width: 80px;
-            position: relative;
-            transform-origin: 0 0;
-            transform: translateX(0) translateY(0) translateZ(-40px);
-        }
-        .cube-face {
-            position: absolute;
-            inset: 0;
-            background: #0017ff;
-            border: solid 1px #ffffff;
-        }
-        .cube-face.top {
-            transform: translateZ(80px);
-        }
-        .cube-face.front {
-            transform-origin: 0 50%;
-            transform: rotateY(-90deg);
-        }
-        .cube-face.back {
-            transform-origin: 0 50%;
-            transform: rotateY(-90deg) translateZ(-80px);
-        }
-        .cube-face.right {
-            transform-origin: 50% 0;
-            transform: rotateX(-90deg) translateY(-80px);
-        }
-        .cube-face.left {
-            transform-origin: 50% 0;
-            transform: rotateX(-90deg) translateY(-80px) translateZ(80px);
-        }
-        @keyframes rotation {
-            0% {
-                transform: rotateX(45deg) rotateY(0) rotateZ(45deg);
-                animation-timing-function: cubic-bezier(
-                    0.17,
-                    0.84,
-                    0.44,
-                    1
-                );
-            }
-            50% {
-                transform: rotateX(45deg) rotateY(0) rotateZ(225deg);
-                animation-timing-function: cubic-bezier(
-                    0.76,
-                    0.05,
-                    0.86,
-                    0.06
-                );
-            }
-            100% {
-                transform: rotateX(45deg) rotateY(0) rotateZ(405deg);
-                animation-timing-function: cubic-bezier(
-                    0.17,
-                    0.84,
-                    0.44,
-                    1
-                );
-            }
-        }
-        .scene,
-        #message {
-            position: absolute;
-            display: flex;
-            top: 0;
-            right: 0;
-            left: 0;
-            bottom: 0;
-            z-index: 2;
-            height: 100%;
-            width: 100%;
-            align-items: center;
-            justify-content: center;
-        }
-        #message {
-            font-weight: bold;
-            font-size: large;
-            color: red;
-            pointer-events: none;
-        }
-        #progress {
-            position: absolute;
-            top: 0;
-            height: 5px;
-            background: blue;
-            z-index: 99;
-            transition: width 0.1s ease-in-out;
-        }
-        #quality {
-            position: absolute;
-            bottom: 10px;
-            z-index: 999;
-            right: 10px;
-        }
-        #canvas {
-            display: block;
-            position: absolute;
-            top: 0;
-            left: 0;
-            width: 100%;
-            height: 100%;
-            touch-action: none;
-        }
-        #instructions {
-            background: rgba(0,0,0,0.6);
-            white-space: pre-wrap;
-            padding: 10px;
-            border-radius: 10px;
-            font-size: x-small;
-        }
-    </style>
-</body>
 """
-    html = f"""
-<head>
-  <title>3D Gaussian Splatting Viewer</title>
-  <script src="http://zeus.blanchon.cc/dropshare/main.js"></script>
-</head>
-{html_body}
 """
-    return f"""<iframe style="width: 100%; height: 900px" srcdoc='{html}'></iframe>"""
-def createStateSession() -> StateDict:
-    # Create new session
-    session_uuid = str(uuid.uuid4())
-    print("createStateSession")
-    print(session_uuid)
-    return StateDict(
-        uuid=session_uuid,
-    )
 def removeStateSession(session_state_value: StateDict):
     # Clean up previous session
     return StateDict(
         uuid=None,
     )
-def makeButtonVisible() -> Tuple[gr.Button, gr.Button]:
-    process_button = gr.Button(visible=True)
-    reset_button = gr.Button(visible=False) #TODO: I will bring this back when I figure out how to stop the process
-    return process_button, reset_button
-def resetSession(state: StateDict) -> Tuple[StateDict, gr.Button, gr.Button]:
-    print("resetSession")
-    new_state = removeStateSession(state)
-    process_button = gr.Button(visible=False)
-    reset_button = gr.Button(visible=False)
-    return new_state, process_button, reset_button
-def process(
-        # *args, **kwargs
         session_state_value: StateDict,
-        filepath: str,
         ffmpeg_fps: int,
         ffmpeg_qscale: int,
-        colmap_camera: str,
-    ):
     if session_state_value["uuid"] is None:
         return
-    print("process")
-    # print(args)
-    # print(kwargs)
-    # return
-    print(session_state_value)
-    print(f"Processing {filepath}")
-    try:
-        session_tmpdirname = gs_dir / str(session_state_value['uuid'])
-        session_tmpdirname.mkdir(parents=True, exist_ok=True)
-        print('Created temporary directory', session_tmpdirname)
-        gs_dir_path = Path(session_tmpdirname)
-        logfile_path = Path(session_tmpdirname) / "log.txt"
-        logfile_path.touch()
         with logfile_path.open("w") as log_file:
-            # Create log file
-            logfile_path.touch()
-            from services.ffmpeg import ffmpeg_run
             ffmpeg_run(
-                video_path = Path(filepath),
-                output_path = gs_dir_path,
                 fps = int(ffmpeg_fps),
                 qscale = int(ffmpeg_qscale),
                 stream_file=log_file
             )
-            from services.colmap import colmap
             colmap(
-                source_path=gs_dir_path,
                 camera=str(colmap_camera),
                 stream_file=log_file
             )
-            print("Done with colmap")
-            # Create a zip of the gs_dir_path folder
-            print(gs_dir, gs_dir_path)
-            print(gs_dir_path.name)
-            archive = shutil.make_archive("result", 'zip', gs_dir, gs_dir_path)
-            print('Created zip file', archive)
-            # Move the zip file to the gs_dir_path folder
-            shutil.move(archive, gs_dir_path)
-            from services.gaussian_splatting_cuda import gaussian_splatting_cuda
             gaussian_splatting_cuda(
-                data_path = gs_dir_path,
-                output_path = gs_dir_path / "output",
                 gs_command = str(Path(__file__).parent.absolute() / "build" / 'gaussian_splatting_cuda'),
-                iterations = 100,
-                convergence_rate = 0.01,
-                resolution = 512,
                 enable_cr_monitoring = False,
                 force = False,
                 empty_gpu_cache = False,
                 stream_file = log_file
             )
-    except Exception:
-        pass
-        # print('Error - Removing temporary directory', session_tmpdirname)
-        # shutil.rmtree(session_tmpdirname)
-def updateLog(session_state_value: StateDict) -> str:
     if session_state_value["uuid"] is None:
         return ""
-    log_file = gs_dir / str(session_state_value['uuid']) / "log.txt"
     if not log_file.exists():
         return ""
@@ -306,117 +258,362 @@ def updateLog(session_state_value: StateDict) -> str:
     return logs
 with gr.Blocks() as demo:
     session_state = gr.State({
         "uuid": None,
     })
-    with gr.Row():
-        with gr.Column():
-            video_input = gr.PlayableVideo(
-                format="mp4",
-                source="upload",
-                label="Upload a video",
-                include_audio=False
-            )
-            with gr.Row(variant="panel"):
-                ffmpeg_fps = gr.Number(
-                    label="FFMPEG FPE",
-                    value=1,
-                    minimum=1,
-                    maximum=5,
-                    step=0.10,
                 )
-                ffmpeg_qscale = gr.Number(
-                    label="FFMPEG QSCALE",
-                    value=1,
-                    minimum=1,
-                    maximum=5,
-                    step=1,
                 )
-                colmap_camera = gr.Dropdown(
-                    label="COLMAP Camera",
-                    value="OPENCV",
-                    choices=["OPENCV", "SIMPLE_PINHOLE", "PINHOLE", "SIMPLE_RADIAL", "RADIAL"],
-                )
-        text_log = gr.Textbox(
-            label="Logs",
-            info="Logs",
             interactive=False,
-            show_copy_button=True
         )
-        # text_log = gr.Code(
-        #     label="Logs",
-        #     language=None,
-        #     interactive=False,
-        # )
-    process_button = gr.Button("Process", visible=False)
-    reset_button = gr.ClearButton(
-        components=[video_input, text_log, ffmpeg_fps, ffmpeg_qscale, colmap_camera],
-        label="Reset",
-        visible=False,
     )
-    process_event = process_button.click(
-        fn=process,
-        inputs=[session_state, video_input, ffmpeg_fps, ffmpeg_qscale, colmap_camera],
-        outputs=[],
     )
-    upload_event = video_input.upload(
-        fn=makeButtonVisible,
-        inputs=[],
-        outputs=[process_button, reset_button]
-    ).then(
         fn=createStateSession,
-        inputs=[],
-        outputs=[session_state],
     ).then(
-        fn=updateLog,
         inputs=[session_state],
-        outputs=[text_log],
         every=2,
     )
-    reset_button.click(
-        fn=resetSession,
         inputs=[session_state],
-        outputs=[session_state, process_button, reset_button],
-        cancels=[process_event]
     )
-    video_input.clear(
-        fn=resetSession,
         inputs=[session_state],
-        outputs=[session_state, process_button, reset_button],
-        cancels=[process_event]
     )
-    demo.close
     # gr.LoginButton, gr.LogoutButton
     # gr.HuggingFaceDatasetSaver
     # gr.OAuthProfile
     # with gr.Tab("jsdn"):
-    #     input_mic = gr.HTML(getHTML())
 demo.queue()
-# demo.launch()
 # mount Gradio app to FastAPI app
-app = gr.mount_gradio_app(app, demo, path="/")
-if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=7860)

+import os
 from pathlib import Path
 import shutil
 import tempfile
+from typing import List
 import gradio as gr
 import uuid
 from typing_extensions import TypedDict, Tuple
 from fastapi import FastAPI
 app = FastAPI()
 # create a static directory to store the static files
+GS_DIR = Path(str(tempfile.gettempdir())) / "gaussian_splatting_gradio"
+GS_DIR.mkdir(parents=True, exist_ok=True)
 StateDict = TypedDict("StateDict", {
     "uuid": str,
 })
+# http://localhost:7860/file=/tmp/gradio/c2110a7de804b39754d229de426dc9307bc03aea/page.svelte
+HOST = "localhost"
+PORT = 7860
+home_markdown = """
+...
 """
+step1_markdown = """
+# Step 1 - Split Video into Frames
+In the journey of transforming a video into a 3D Gaussian Splatting, the initial step is the conversion of the video into individual frames. You can here provide a **video file** and specify how much image you want to extract per second (*fps*). The application will then automatically extract the frames from the video and prepare them for the next step in the process.
+However, you can also do this step manually and upload the frames directory by yourself in the next step. In this case, you can skip this step and go directly to the next step.
+Please not that blurry frames will mostlikely result in a bad 3D model. So, make sure that the video is clear enough.
 """
+step2_markdown = """
+# Step 2 - SfM using Colmap
+In this step we use Colmap (https://github.com/colmap/colmap). This process utilizes the frames extracted from the uploaded video to generate camera parameters and a point cloud, which are essential components for the 3D Gaussian Splatting process.
+This step could take a while depending on the number of frames and the resolution. So, please be patient.
+You might want to do this step manually and upload the frames directory by yourself in the next step. In this case, you can skip this step and go directly to the next step.
+"""
+step3_markdown = """
+# Step 3 - 3D Gaussian Splatting
+In this final step we use the 3D Gaussian Splatting Cuda implementation by MrNeRF (https://twitter.com/janusch_patas): https://github.com/MrNeRF/gaussian-splatting-cuda.
+As it's quite rapid to train, you can easily use a high number of iterations.
+"""
+def getPlyFile(session_state_value: StateDict) -> str:
+    return f"/tmp/gaussian_splatting_gradio/{session_state_value['uuid']}/output/final_point_cloud.ply"
+def getCamerasFile(session_state_value: StateDict) -> str:
+    return f"/tmp/gaussian_splatting_gradio/{session_state_value['uuid']}/output/cameras.json"
+def getZipFile(session_state_value: StateDict) -> str:
+    return f"/tmp/gaussian_splatting_gradio/{session_state_value['uuid']}/result.zip"
+def makeResult(session_state_value: StateDict) -> tuple[str, str, str]:
+    ply_file = getPlyFile(session_state_value)
+    cameras_file = getCamerasFile(session_state_value)
+    zip_file = getZipFile(session_state_value)
+    return [ply_file, cameras_file, zip_file]
+# Utility functions
+def createStateSession(previous_session: StateDict) -> StateDict:
+    if previous_session["uuid"] is None:
+        # Create new session
+        session_uuid = str(uuid.uuid4())
+        print("Creating new session: ", session_uuid)
+        session_tmpdirname = GS_DIR / str(session_uuid)
+        session_tmpdirname.mkdir(parents=True, exist_ok=True)
+        print('Created temporary directory: ', session_tmpdirname)
+        session = StateDict(
+            uuid=session_uuid,
+        )
+    else:
+        # Use previous session
+        session = previous_session
+    return session
 def removeStateSession(session_state_value: StateDict):
     # Clean up previous session
+    session_uuid = session_state_value["uuid"]
+    session_tmpdirname = GS_DIR / str(session_uuid)
+    print('Removing temporary directory: ', session_tmpdirname)
+    shutil.rmtree(session_tmpdirname)
     return StateDict(
         uuid=None,
     )
+def makeButtonVisible(btn_value: str) -> gr.Button:
+    return gr.Button(btn_value, visible=True)
+#  Process functions
+def process_ffmpeg(
         session_state_value: StateDict,
+        ffmpeg_input: str,
         ffmpeg_fps: int,
         ffmpeg_qscale: int,
+    ) -> list[str]:
+    # Ensure that a session is active
     if session_state_value["uuid"] is None:
         return
+    # Set up session directory
+    session_path = GS_DIR / str(session_state_value['uuid'])
+    logfile_path = Path(session_path) / "ffmpeg_log.txt"
+    logfile_path.touch()
+    try:
+        from services.ffmpeg import ffmpeg_run
         with logfile_path.open("w") as log_file:
             ffmpeg_run(
+                video_path = Path(ffmpeg_input),
+                output_path = session_path,
                 fps = int(ffmpeg_fps),
                 qscale = int(ffmpeg_qscale),
                 stream_file=log_file
             )
+        print("Done with ffmpeg")
+    except Exception as e:
+        print(f"Error - {e}")
+        # print('Error - Removing temporary directory', session_path)
+        # shutil.rmtree(session_path)
+    # Get the list of all the file of (session_path / "input")
+    list_of_jpgs = [str(f) for f in (session_path / "input").glob("*.jpg")]
+    return list_of_jpgs
+def processColmap(
+        session_state_value: StateDict,
+        colmap_inputs: List[tempfile.NamedTemporaryFile],
+        colmap_camera: str,
+        enable_rerun: bool
+    ) -> Tuple[str, str]:
+    # Ensure that a session is active
+    if session_state_value["uuid"] is None:
+        return "", ""
+    # Set up session directory
+    session_path = GS_DIR / str(session_state_value['uuid'])
+    logfile_path = Path(session_path) / "colmap_log.txt"
+    logfile_path.touch()
+    rerunfile_path = Path(session_path) / "rerun_page.html"
+    rerunfile_path.touch()
+    (session_path / "input").mkdir(parents=True, exist_ok=True)
+    for file in colmap_inputs:
+        print("copying", file.name, "to", session_path / "input")
+        shutil.copy(file.name, session_path / "input")
+    try:
+        from services.colmap import colmap
+        with logfile_path.open("w") as log_file:
             colmap(
+                source_path=session_path,
                 camera=str(colmap_camera),
                 stream_file=log_file
             )
+        print("Done with colmap")
+        if enable_rerun:
+            from services.rerun import read_and_log_sparse_reconstruction
+            html = read_and_log_sparse_reconstruction(
+                exp_name = str(session_state_value['uuid']),
+                dataset_path = session_path,
+            )
+            print("Done with rerun")
+        else:
+            html = "Rerun was disable !"
+        with rerunfile_path.open("w") as rerunfile:
+            rerunfile.write(html)
+    except Exception as e:
+        print(f"Error - {e}")
+        # print('Error - Removing temporary directory', session_path)
+        # shutil.rmtree(session_path)
+    # zip the session_path folder
+    archive = shutil.make_archive("result", 'zip', GS_DIR, session_path)
+    print('Created zip file', archive)
+    return archive, rerunfile_path
+def processGaussianSplattingCuda(
+        session_state_value: StateDict,
+        gs_input: tempfile.NamedTemporaryFile,
+        gs_iterations: int,
+        gs_convergence_rate: float,
+        gs_resolution: int,
+    ) -> Tuple[str, str]:
+    # Ensure that a session is active
+    if session_state_value["uuid"] is None:
+        return
+    # Set up session directory
+    session_path = GS_DIR / str(session_state_value['uuid'])
+    logfile_path = Path(session_path) / "gaussian_splatting_cuda_log.txt"
+    logfile_path.touch()
+    # Unzip the gs_input file to the session_path
+    shutil.unpack_archive(gs_input.name, session_path)
+    # Copy the gs_input directory to the session_path
+    # shutil.copytree(gs_input, session_path)
+    try:
+        from services.gaussian_splatting_cuda import gaussian_splatting_cuda
+        with logfile_path.open("w") as log_file:
             gaussian_splatting_cuda(
+                data_path = session_path,
+                output_path = session_path / "output",
                 gs_command = str(Path(__file__).parent.absolute() / "build" / 'gaussian_splatting_cuda'),
+                iterations = int(gs_iterations),
+                convergence_rate = float(gs_convergence_rate),
+                resolution = int(gs_resolution),
                 enable_cr_monitoring = False,
                 force = False,
                 empty_gpu_cache = False,
                 stream_file = log_file
             )
+        print("Done with gaussian_splatting_cuda")
+        # Create a zip of the session_path folder
+        archive = shutil.make_archive("result", 'zip', GS_DIR, session_path)
+        print('Created zip file', archive)
+        # Move the zip file to the session_path folder
+        shutil.move(archive, session_path)
+    except Exception as e:
+        print(f"Error - {e}")
+        # print('Error - Removing temporary directory', session_path)
+        # shutil.rmtree(session_path)
+    return (
+        session_path / "output" / "final_point_cloud.ply",
+        session_path / "output" / "cameras.json",
+    )
+def updateLog(logname:str, session_state_value: StateDict) -> str:
     if session_state_value["uuid"] is None:
         return ""
+    log_file = GS_DIR / str(session_state_value['uuid']) / f"{logname}.txt"
     if not log_file.exists():
         return ""
     return logs
+def bindStep1Step2(step1_output: list[tempfile.NamedTemporaryFile]) -> list[str]:
+    return [file.name for file in step1_output]
+def bindStep2Step3(step2_output: tempfile.NamedTemporaryFile) -> str:
+    return step2_output.name
+def makeRerunIframe(rerun_html : tempfile.NamedTemporaryFile) -> str:
+    # If rerun_html is bigger than 300MB, then we don't show it
+    print(f"Rerun file size: {os.stat(rerun_html.name).st_size}")
+    if os.stat(rerun_html.name).st_size > 100_000_000:
+        print("Rerun file is too big, not showing it")
+        return ""
+    filepath = rerun_html.name
+    print("filepath", filepath)
+    return f"""<iframe src="/file={filepath}" width="100%"; height="1080px"></iframe>"""
 with gr.Blocks() as demo:
+    #############################
+    ########## State ############
+    #############################
     session_state = gr.State({
         "uuid": None,
     })
+    #############################
+    ###### UI Components ########
+    #############################
+    gr.Markdown("# Gaussian Splatting Kit")
+    gr.Markdown("Click on the **Duplicate** button to create a new instance of this app.")
+    duplicate_button = gr.DuplicateButton()
+    gr.Markdown(value=home_markdown)
+    with gr.Tab("Slit Video into Frames"):
+        step1_description = gr.Markdown(step1_markdown)
+        # Video Frames
+        with gr.Row():
+            # Video Frames - Inputs
+            with gr.Column():
+                # Video Frames - Inputs - Video File
+                step1_input = gr.PlayableVideo(
+                    format="mp4",
+                    source="upload",
+                    label="Upload a video",
+                    include_audio=False
+                )
+                # Video Frames - Inputs - Parameters
+                with gr.Row(variant="panel"):
+                    # Video Frames - Inputs - Parameters - FFMPEG FPS
+                    step1_fps = gr.Number(
+                        label="FFMPEG Fps",
+                        value=1,
+                        minimum=1,
+                        maximum=5,
+                        step=0.10,
+                    )
+                    # Video Frames - Inputs - Parameters - FFMPEG Qscale
+                    step1_qscale = gr.Number(
+                        label="FFMPEG Qscale",
+                        value=1,
+                        minimum=1,
+                        maximum=5,
+                        step=1,
+                    )
+            # Video Frames - Outputs
+            with gr.Column():
+                # Video Frames - Outputs - Video File
+                step1_output = gr.File(
+                    label="Frames",
+                    file_count="directory",
+                    type="file",
+                    interactive=False,
+                )
+                # Video Frames - Outputs - Logs
+                step1_logs = gr.Textbox(
+                    label="Videos Logs",
+                    interactive=False,
+                    show_copy_button=True
                 )
+        # Video Frames - Process Button
+        step1_processbtn = gr.Button("Process", visible=True)
+        # Video Frames - Visualize
+        # Video Frames - Visualize -
+        # step1_visualize_gallery = gr.Gallery()
+    with gr.Tab("Colmap"):
+        step2_description = gr.Markdown(step2_markdown)
+        # Colmap
+        with gr.Row():
+            # Colmap - Inputs
+            with gr.Column():
+                # Colmap - Inputs - Frames Directory
+                step2_input = gr.File(
+                    label="Upload a frames directory",
+                    file_count="directory",
+                    type="file",
+                    interactive=True,
                 )
+                # Colmap - Inputs - Parameters
+                with gr.Row(variant="panel"):
+                    # Colmap - Inputs - Parameters - Colmap Camera
+                    step2_camera = gr.Dropdown(
+                        label="COLMAP Camera",
+                        value="OPENCV",
+                        choices=["OPENCV", "SIMPLE_PINHOLE", "PINHOLE", "SIMPLE_RADIAL", "RADIAL"],
+                    )
+                    # Colmap - Inputs - Parameters - Enable Rerun
+                    step2_rerun = gr.Checkbox(
+                        value=True,
+                        label="Enable Rerun",
+                    )
+            # Colmap - Outputs
+            with gr.Column():
+                # Colmap - Outputs - Video File
+                step2_output = gr.File(
+                    label="Colmap",
+                    file_count="single",
+                    file_types=[".zip"],
+                    type="file",
+                    interactive=False,
+                )
+                # Colmap - Outputs - Logs
+                step2_logs = gr.Textbox(
+                    label="Colmap Logs",
+                    interactive=False,
+                    show_copy_button=True
+                )
+        # Colmap - Process Button
+        step2_processbtn = gr.Button("Process", visible=True)
+        # Colmap - Visualize
+        # Colmap - Visualize - Rerun HTML File
+        step_2_visualize_html = gr.File(
+            label="Rerun HTML",
+            file_count="single",
+            file_types=[".html"],
+            type="file",
             interactive=False,
+            visible=False
         )
+        # Colmap - Visualize - Rerun HTML
+        step_2_visualize = gr.HTML("Rerun", visible=True)
+    with gr.Tab("Gaussian Splatting"):
+        step3_description = gr.Markdown(step3_markdown)
+        # Gaussian Splatting
+        with gr.Row():
+            # Gaussian Splatting - Inputs
+            with gr.Column():
+                # Gaussian Splatting - Inputs - Colmap + Frames
+                step3_input = gr.File(
+                    label="Upload a colmap + frames directory",
+                    file_count="single",
+                    file_types=[".zip"],
+                    type="file",
+                    interactive=True,
+                )
+                # Gaussian Splatting - Inputs - Parameters
+                with gr.Row(variant="panel"):
+                    # Gaussian Splatting - Inputs - Parameters - GS Iterations
+                    step3_iterations = gr.Number(
+                        label="GS Iterations",
+                        value=10_000,
+                        minimum=1_000,
+                        maximum=50_000,
+                        step=1_000,
+                    )
+                    # Gaussian Splatting - Inputs - Parameters - GS Convergence Rate
+                    step3_convergence_rate = gr.Number(
+                        label="GS Convergence Rate",
+                        value=0.01,
+                        minimum=0.01,
+                        maximum=1,
+                        step=0.01,
+                    )
+                    # Gaussian Splatting - Inputs - Parameters - GS Resolution
+                    step3_resolution = gr.Number(
+                        label="GS Resolution",
+                        value=512,
+                        minimum=128,
+                        maximum=1024,
+                        step=128,
+                    )
+            # Gaussian Splatting - Outputs
+            with gr.Column():
+                with gr.Row():
+                    # Gaussian Splatting - Outputs - PLY File
+                    step3_output1 = gr.File(
+                        label="PLY File",
+                        file_count="single",
+                        type="file",
+                        interactive=False,
+                    )
+                    # Gaussian Splatting - Outputs - Cameras File
+                    step3_output2 = gr.File(
+                        label="Cameras File",
+                        file_count="single",
+                        type="file",
+                        interactive=False,
+                    )
+                # Gaussian Splatting - Outputs - Logs
+                step3_logs = gr.Textbox(
+                    label="Gaussian Splatting Logs",
+                    interactive=False,
+                    show_copy_button=True
+                )
+        # Gaussian Splatting - Process Button
+        step3_processbtn = gr.Button("Process", visible=True)
+        # Gaussian Splatting - Visualize
+        # Gaussian Splatting - Visualize - Antimatter15 HTML
+        # step_3_visualize = gr.HTML(getAntimatter15HTML(), visible=True)
+        step_3_visualize = gr.Button("Visualize", visible=True, link="https://antimatter15.com/splat/")
+    #############################
+    ########## Events ###########
+    #############################
+    ### Step 1
+    # Make the process button visible when a video is uploaded
+    step1_upload_event = step1_input.upload(
+        fn=createStateSession,
+        inputs=[session_state],
+        outputs=[session_state]
+    ).success(
+        fn=makeButtonVisible,
+        inputs=[step1_processbtn],
+        outputs=[step1_processbtn],
     )
+    # Do the processing when the process button is clicked
+    step1_processevent = step1_processbtn.click(
+        fn=process_ffmpeg,
+        inputs=[session_state, step1_input, step1_fps, step1_qscale],
+        outputs=[step1_output],
+    ).success(
+        fn=bindStep1Step2,
+        inputs=[step1_output],
+        outputs=[step2_input],
+    ).success(
+        fn=makeButtonVisible,
+        inputs=[step2_processbtn],
+        outputs=[step2_processbtn],
     )
+    # Update the logs every 2 seconds
+    step1_logsevent = step1_processbtn.click(
+        fn=lambda session: updateLog("ffmpeg_log", session),
+        inputs=[session_state],
+        outputs=[step1_logs],
+        every=2,
+    )
+    ## Step 2
+    # Make the process button visible when a video is uploaded
+    step2_upload_event = step2_input.upload(
         fn=createStateSession,
+        inputs=[session_state],
+        outputs=[session_state]
+    ).success(
+        fn=makeButtonVisible,
+        inputs=[step2_processbtn],
+        outputs=[step2_processbtn],
+    )
+    # Do the processing when the process button is clicked
+    step2_processevent = step2_processbtn.click(
+        fn=processColmap,
+        inputs=[session_state, step2_input, step2_camera, step2_rerun],
+        outputs=[step2_output, step_2_visualize_html]
+    ).success(
+        fn=bindStep2Step3,
+        inputs=[step2_output],
+        outputs=[step3_input],
+    ).success(
+        fn=makeButtonVisible,
+        inputs=[step3_processbtn],
+        outputs=[step3_processbtn],
     ).then(
+        fn=makeRerunIframe,
+        inputs=[step_2_visualize_html],
+        outputs=[step_2_visualize],
+    )
+    # Update the logs every 2 seconds
+    step2_logsevent = step2_processbtn.click(
+        fn=lambda session: updateLog("colmap_log", session),
         inputs=[session_state],
+        outputs=[step2_logs],
         every=2,
     )
+    ## Step 3
+    # Make the process button visible when a video is uploaded
+    step3_upload_event = step3_input.upload(
+        fn=createStateSession,
         inputs=[session_state],
+        outputs=[session_state]
+    ).success(
+        fn=makeButtonVisible,
+        inputs=[step3_processbtn],
+        outputs=[step3_processbtn],
     )
+    # Do the processing when the process button is clicked
+    step3_processevent = step3_processbtn.click(
+        fn=processGaussianSplattingCuda,
+        inputs=[session_state, step3_input, step3_iterations, step3_convergence_rate, step3_resolution],
+        outputs=[step3_output1, step3_output2]
+    )
+    # .success(
+    #     fn=lambda x: x,
+    #     inputs=[step3_output1, step3_output2],
+    #     outputs=[],
+    # )
+    # Update the logs every 2 seconds
+    step3_logsevent = step3_processbtn.click(
+        fn=lambda session: updateLog("gaussian_splatting_cuda_log", session),
         inputs=[session_state],
+        outputs=[step3_logs],
+        every=2,
     )
+    # reset_button = gr.ClearButton(
+    #     components=[video_input, text_log, ffmpeg_fps, ffmpeg_qscale, colmap_camera],
+    #     label="Reset",
+    #     visible=False,
+    # )
+    # print(f"async (x) => {{ {getJS(url='http://0.0.0.0:7860/output/37c7ae54-7752-4e7b-8ba9-bab32c86b316/output/point_cloud/iteration_100/point_cloud.ply')} }}")
+    # show_button.click(
+    #     fn=None,
+    #     inputs=[],
+    #     outputs=[],
+    #     _js=f"async (x) => {{ {getJS(url='http://0.0.0.0:7860/output/37c7ae54-7752-4e7b-8ba9-bab32c86b316/output/point_cloud/iteration_100/point_cloud.ply')} }}"
+    # ).then(
+    #     fn=None,
+    #     inputs=[],
+    #     outputs=[],
+    #     _js=f"async (x) => {{ {getJS(url='http://0.0.0.0:7860/output/37c7ae54-7752-4e7b-8ba9-bab32c86b316/output/point_cloud/iteration_100/point_cloud.ply')} }}"
+    # )
     # gr.LoginButton, gr.LogoutButton
     # gr.HuggingFaceDatasetSaver
     # gr.OAuthProfile
     # with gr.Tab("jsdn"):
+    #     input_mic = gr.HTML(getRerunHTML())
 demo.queue()
+demo.launch()
 # mount Gradio app to FastAPI app
+# app = gr.mount_gradio_app(app, demo, path="/")
+# if __name__ == "__main__":
+#     uvicorn.run(app, host="0.0.0.0", port=7860, ws_max_size=16777216*1000)

services/gaussian_splatting_cuda.py CHANGED Viewed

@@ -3,6 +3,8 @@ from pathlib import Path
 import subprocess
 from typing import Optional
 from rich.console import Console
 console = Console()
@@ -45,10 +47,16 @@ def gaussian_splatting_cuda_training(
     Set custom average onvergence rate for the training process. Requires the flag --enable-cr-monitoring to be set.
     """
     cmd = [
         gs_command,
-        f"--data-path={data_path.as_posix()}"
-        f"--output-path={output_path.as_posix()}"
         f"--iter={iterations}",
         # TODO: Enable these options and put the right defaults in the function signature
         # f"--convergence-rate={convergence_rate}",
@@ -105,4 +113,12 @@ def gaussian_splatting_cuda(
         force,
         empty_gpu_cache,
         stream_file
-    )

 import subprocess
 from typing import Optional
 from rich.console import Console
+import os
+import shutil
 console = Console()
     Set custom average onvergence rate for the training process. Requires the flag --enable-cr-monitoring to be set.
     """
+    # export LC_ALL=C
+    # export LANG=C
+    os.environ["LC_ALL"] = "C"
+    os.environ["LANG"] = "C"
     cmd = [
         gs_command,
+        f"--data-path={data_path.as_posix()}",
+        f"--output-path={output_path.as_posix()}",
         f"--iter={iterations}",
         # TODO: Enable these options and put the right defaults in the function signature
         # f"--convergence-rate={convergence_rate}",
         force,
         empty_gpu_cache,
         stream_file
+    )
+    # Copy the /output/point_cloud/iteration_{iteration}/point_cloud.ply to the output_path
+    shutil.copyfile(
+        src=output_path / "point_cloud" / f"iteration_{iterations}" / "point_cloud.ply",
+        dst=output_path / "final_point_cloud.ply"
+    )
+    console.log(f"📄 Final point cloud saved to {output_path / 'final_point_cloud.ply'}")

services/rerun.py CHANGED Viewed

@@ -1,85 +1,103 @@
 import os
 import re
 from pathlib import Path
 import numpy as np
 import rerun as rr  # pip install rerun-sdk
-from utils.read_write_model import read_model
 # From https://github.com/rerun-io/rerun/tree/main/examples/python/structure_from_motion
 def read_and_log_sparse_reconstruction(
         exp_name: str,
         dataset_path: Path,
-        output_path: Path,
-        filter_output: bool = False,
-        filter_min_visible: int = 2_000
-    ) -> None:
-    rr.init(exp_name)
-    cameras, images, points3D = read_model(dataset_path / "sparse", ext=".bin")
-    if filter_output:
-        # Filter out noisy points
-        points3D = {id: point for id, point in points3D.items() if point.rgb.any() and len(point.image_ids) > 4}
-    rr.log_view_coordinates("/", up="-Y", timeless=True)
-    # Iterate through images (video frames) logging data related to each frame.
-    for image in sorted(images.values(), key=lambda im: im.name):  # type: ignore[no-any-return]
-        image_file = dataset_path / "images" / image.name
-        if not os.path.exists(image_file):
-            continue
-        # COLMAP sets image ids that don't match the original video frame
-        idx_match = re.search(r"\d+", image.name)
-        assert idx_match is not None
-        frame_idx = int(idx_match.group(0))
-        quat_xyzw = image.qvec[[1, 2, 3, 0]]  # COLMAP uses wxyz quaternions
-        camera = cameras[image.camera_id]
-        np.array([1.0, 1.0])
-        visible = [id != -1 and points3D.get(id) is not None for id in image.point3D_ids]
-        visible_ids = image.point3D_ids[visible]
-        if filter_output and len(visible_ids) < filter_min_visible:
-            continue
-        visible_xyzs = [points3D[id] for id in visible_ids]
-        visible_xys = image.xys[visible]
-        rr.set_time_sequence("frame", frame_idx)
-        points = [point.xyz for point in visible_xyzs]
-        point_colors = [point.rgb for point in visible_xyzs]
-        point_errors = [point.error for point in visible_xyzs]
-        rr.log_scalar("plot/avg_reproj_err", np.mean(point_errors), color=[240, 45, 58])
-        rr.log_points("points", points, colors=point_colors, ext={"error": point_errors})
-        # COLMAP's camera transform is "camera from world"
-        rr.log_transform3d(
-            "camera", rr.TranslationRotationScale3D(image.tvec, rr.Quaternion(xyzw=quat_xyzw)), from_parent=True
-        )
-        rr.log_view_coordinates("camera", xyz="RDF")  # X=Right, Y=Down, Z=Forward
-        # Log camera intrinsics
-        assert camera.model == "PINHOLE"
-        rr.log_pinhole(
-            "camera/image",
-            width=camera.width,
-            height=camera.height,
-            focal_length_px=camera.params[:2],
-            principal_point_px=camera.params[2:],
-        )
-        rr.log_image_file("camera/image", img_path=dataset_path / "images" / image.name)
-        rr.log_points("camera/image/keypoints", visible_xys, colors=[34, 138, 167])
-    rerun_output_directory = output_path / "rerun"
-    rerun_output_directory.mkdir(parents=True, exist_ok=True)
-    rerun_output_file = rerun_output_directory / "recording.rrd"
-    rr.save(rerun_output_file.as_posix())

 import os
 import re
 from pathlib import Path
+from typing import Optional
 import numpy as np
 import rerun as rr  # pip install rerun-sdk
+from services.utils.read_write_model import read_model
 # From https://github.com/rerun-io/rerun/tree/main/examples/python/structure_from_motion
 def read_and_log_sparse_reconstruction(
         exp_name: str,
         dataset_path: Path,
+        max_image_number: Optional[int] = 15,
+        filter_output: bool = True,
+        filter_min_visible: int = 50,
+        filter_max_visible: int = 500
+    ) -> str:
+    try:
+        rr.init(exp_name)
+        rec = rr.memory_recording()
+        cameras, images, points3D = read_model(dataset_path / "sparse", ext=".bin")
+        print(f"Loaded {len(cameras)} cameras, {len(images)} images, and {len(points3D)} points3D")
+        if filter_output:
+            # Filter out noisy points
+            points3D = {id: point for id, point in points3D.items() if point.rgb.any() and len(point.image_ids) > 4}
+        rr.log_view_coordinates("/", up="-Y", timeless=True)
+        print(f"Number of image frames: {len(images)}")
+        if max_image_number is not None:
+            # Sample the image sequence to reduce output size
+            image_ids = sorted(images.keys())
+            image_ids = np.random.permutation(image_ids)[:max_image_number]
+            images = {id: images[id] for id in image_ids}
+        print(f"Number of image frames: {len(images)}")
+        # Iterate through images (video frames) logging data related to each frame.
+        for image in sorted(images.values(), key=lambda im: im.name):  # type: ignore[no-any-return]
+            image_file = dataset_path / "images" / image.name
+            if not os.path.exists(image_file):
+                continue
+            # COLMAP sets image ids that don't match the original video frame
+            idx_match = re.search(r"\d+", image.name)
+            assert idx_match is not None
+            frame_idx = int(idx_match.group(0))
+            quat_xyzw = image.qvec[[1, 2, 3, 0]]  # COLMAP uses wxyz quaternions
+            camera = cameras[image.camera_id]
+            np.array([1.0, 1.0])
+            visible = [id != -1 and points3D.get(id) is not None for id in image.point3D_ids]
+            print(f"Frame {frame_idx} has {np.sum(visible)} visible points")
+            # Randomly sample points to reduce output size
+            if filter_output and np.sum(visible) > filter_max_visible:
+                visible_indices = np.random.permutation(np.where(visible)[0])[:filter_max_visible]
+                visible = np.zeros_like(visible, dtype=bool)
+                visible[visible_indices] = True
+                print(f"Frame {frame_idx} has {np.sum(visible)} visible points after sampling")
+            visible_ids = image.point3D_ids[visible]
+            if filter_output and len(visible_ids) < filter_min_visible:
+                continue
+            visible_xyzs = [points3D[id] for id in visible_ids]
+            visible_xys = image.xys[visible]
+            rr.set_time_sequence("frame", frame_idx)
+            points = [point.xyz for point in visible_xyzs]
+            point_colors = [point.rgb for point in visible_xyzs]
+            point_errors = [point.error for point in visible_xyzs]
+            rr.log_scalar("plot/avg_reproj_err", np.mean(point_errors), color=[240, 45, 58])
+            rr.log_points("points", points, colors=point_colors, ext={"error": point_errors})
+            # COLMAP's camera transform is "camera from world"
+            rr.log_transform3d(
+                "camera", rr.TranslationRotationScale3D(image.tvec, rr.Quaternion(xyzw=quat_xyzw)), from_parent=True
+            )
+            rr.log_view_coordinates("camera", xyz="RDF")  # X=Right, Y=Down, Z=Forward
+            # Log camera intrinsics
+            assert camera.model == "PINHOLE"
+            rr.log_pinhole(
+                "camera/image",
+                width=camera.width,
+                height=camera.height,
+                focal_length_px=camera.params[:2],
+                principal_point_px=camera.params[2:],
+            )
+            rr.log_image_file("camera/image", img_path=dataset_path / "images" / image.name)
+            rr.log_points("camera/image/keypoints", visible_xys, colors=[34, 138, 167])
+        return rec.as_html()
+    except Exception as e:
+        print(e)