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 "" with log_file.open("r") as log_file: logs = log_file.read() 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"""""" 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)