Create app.py

app.py

@@ -1,28 +1,169 @@
-from dust3r.demo import *
-
-if True:
-    # parser = get_args_parser()
-    # args = parser.parse_args()
-    # args = {
-    #     "--weights": "./dust3r/checkpoints/DUSt3R_ViTLarge_BaseDecoder_512_dpt.pth",
-    #     "tmp_dir": None,
-    #     "server_name": None,
-    #     "device": "cuda",
-        
-    # }
-
-    # if args.tmp_dir is not None:
-    #     tmp_path = args.tmp_dir
-    #     os.makedirs(tmp_path, exist_ok=True)
-    #     tempfile.tempdir = tmp_path
-
-    if False: #args.server_name is not None:
+import os
+import tempfile
+import gradio
+import argparse
+import math
+import torch
+import numpy as np
+import trimesh
+import copy
+import functools
+from scipy.spatial.transform import Rotation
+
+from dust3r.inference import inference
+from dust3r.model import AsymmetricCroCo3DStereo
+from dust3r.image_pairs import make_pairs
+from dust3r.utils.image import load_images, rgb
+from dust3r.utils.device import to_numpy
+from dust3r.viz import add_scene_cam, CAM_COLORS, OPENGL, pts3d_to_trimesh, cat_meshes
+from dust3r.cloud_opt import global_aligner, GlobalAlignerMode
+
+import matplotlib.pyplot as pl
+pl.ion()
+
+torch.backends.cuda.matmul.allow_tf32 = True  # for gpu >= Ampere and pytorch >= 1.12
+batch_size = 1
+
+def run_dust3r_inference(args, filelist, schedule, niter, min_conf_thr, as_pointcloud,
+                          mask_sky, clean_depth, transparent_cams, cam_size,
+                          scenegraph_type, winsize, refid):
+    tmpdirname = tempfile.mkdtemp(suffix='dust3r_gradio_demo')
+    if args.tmp_dir is not None:
+        tmp_path = args.tmp_dir
+        os.makedirs(tmp_path, exist_ok=True)
+        tempfile.tempdir = tmp_path
+
+    if args.server_name is not None:
         server_name = args.server_name
     else:
-        server_name = '0.0.0.0' # if args.local_network else '127.0.0.1'
+        server_name = '0.0.0.0' if args.local_network else '127.0.0.1'
+
+    if args.weights is not None:
+        weights_path = args.weights
+    else:
+        weights_path = "naver/" + args.model_name
+    model = AsymmetricCroCo3DStereo.from_pretrained(weights_path).to(args.device)
+
+    recon_fun = functools.partial(get_reconstructed_scene, tmpdirname, model, args.device, args.image_size)
+    model_from_scene_fun = functools.partial(get_3D_model_from_scene, tmpdirname)
+
+    # Run the reconstruction function
+    scene, outfile, imgs = recon_fun(filelist, schedule, niter, min_conf_thr, as_pointcloud,
+                                     mask_sky, clean_depth, transparent_cams, cam_size,
+                                     scenegraph_type, winsize, refid)
+
+    # Return the result
+    return outfile, imgs
+
+def main_demo(tmpdirname, model, device, image_size, server_name, server_port, silent=False):
+    recon_fun = functools.partial(get_reconstructed_scene, tmpdirname, model, device, silent, image_size)
+    model_from_scene_fun = functools.partial(get_3D_model_from_scene, tmpdirname, silent)
+    with gradio.Blocks(css=""".gradio-container {margin: 0 !important; min-width: 100%};""", title="DUSt3R Demo") as demo:
+        # scene state is save so that you can change conf_thr, cam_size... without rerunning the inference
+        scene = gradio.State(None)
+        gradio.HTML('<h2 style="text-align: center;">DUSt3R Demo</h2>')
+        with gradio.Column():
+            inputfiles = gradio.File(file_count="multiple")
+            with gradio.Row():
+                schedule = gradio.Dropdown(["linear", "cosine"],
+                                           value='linear', label="schedule", info="For global alignment!")
+                niter = gradio.Number(value=300, precision=0, minimum=0, maximum=5000,
+                                      label="num_iterations", info="For global alignment!")
+                scenegraph_type = gradio.Dropdown(["complete", "swin", "oneref"],
+                                                  value='complete', label="Scenegraph",
+                                                  info="Define how to make pairs",
+                                                  interactive=True)
+                winsize = gradio.Slider(label="Scene Graph: Window Size", value=1,
+                                        minimum=1, maximum=1, step=1, visible=False)
+                refid = gradio.Slider(label="Scene Graph: Id", value=0, minimum=0, maximum=0, step=1, visible=False)
+
+            run_btn = gradio.Button("Run")
+
+            with gradio.Row():
+                # adjust the confidence threshold
+                min_conf_thr = gradio.Slider(label="min_conf_thr", value=3.0, minimum=1.0, maximum=20, step=0.1)
+                # adjust the camera size in the output pointcloud
+                cam_size = gradio.Slider(label="cam_size", value=0.05, minimum=0.001, maximum=0.1, step=0.001)
+            with gradio.Row():
+                as_pointcloud = gradio.Checkbox(value=False, label="As pointcloud")
+                # two post process implemented
+                mask_sky = gradio.Checkbox(value=False, label="Mask sky")
+                clean_depth = gradio.Checkbox(value=True, label="Clean-up depthmaps")
+                transparent_cams = gradio.Checkbox(value=False, label="Transparent cameras")
+
+            outmodel = gradio.Model3D()
+            outgallery = gradio.Gallery(label='rgb,depth,confidence', columns=3, height="100%")
+
+            # events
+            scenegraph_type.change(set_scenegraph_options,
+                                   inputs=[inputfiles, winsize, refid, scenegraph_type],
+                                   outputs=[winsize, refid])
+            inputfiles.change(set_scenegraph_options,
+                              inputs=[inputfiles, winsize, refid, scenegraph_type],
+                              outputs=[winsize, refid])
+            run_btn.click(fn=recon_fun,
+                          inputs=[inputfiles, schedule, niter, min_conf_thr, as_pointcloud,
+                                  mask_sky, clean_depth, transparent_cams, cam_size,
+                                  scenegraph_type, winsize, refid],
+                          outputs=[scene, outmodel, outgallery])
+            min_conf_thr.release(fn=model_from_scene_fun,
+                                 inputs=[scene, min_conf_thr, as_pointcloud, mask_sky,
+                                         clean_depth, transparent_cams, cam_size],
+                                 outputs=outmodel)
+            cam_size.change(fn=model_from_scene_fun,
+                            inputs=[scene, min_conf_thr, as_pointcloud, mask_sky,
+                                    clean_depth, transparent_cams, cam_size],
+                            outputs=outmodel)
+            as_pointcloud.change(fn=model_from_scene_fun,
+                                 inputs=[scene, min_conf_thr, as_pointcloud, mask_sky,
+                                         clean_depth, transparent_cams, cam_size],
+                                 outputs=outmodel)
+            mask_sky.change(fn=model_from_scene_fun,
+                            inputs=[scene, min_conf_thr, as_pointcloud, mask_sky,
+                                    clean_depth, transparent_cams, cam_size],
+                            outputs=outmodel)
+            clean_depth.change(fn=model_from_scene_fun,
+                               inputs=[scene, min_conf_thr, as_pointcloud, mask_sky,
+                                       clean_depth, transparent_cams, cam_size],
+                               outputs=outmodel)
+            transparent_cams.change(model_from_scene_fun,
+                                    inputs=[scene, min_conf_thr, as_pointcloud, mask_sky,
+                                            clean_depth, transparent_cams, cam_size],
+                                    outputs=outmodel)
+    demo.launch(share=False, server_name=server_name, server_port=server_port)
+
+# Gradio interface components
+inputfiles = gradio.File(label="Input Images", file_count="multiple")
+schedule = gradio.Dropdown(["linear", "cosine"], label="Schedule")
+niter = gradio.Number(label="Number of Iterations", min=0, max=5000, step=1)
+min_conf_thr = gradio.Slider(label="Minimum Confidence Threshold", min=1.0, max=20.0, step=0.1, default=3.0)
+as_pointcloud = gradio.Checkbox(label="As Pointcloud", default=False)
+mask_sky = gradio.Checkbox(label="Mask Sky", default=False)
+clean_depth = gradio.Checkbox(label="Clean-up Depthmaps", default=True)
+transparent_cams = gradio.Checkbox(label="Transparent Cameras", default=False)
+cam_size = gradio.Slider(label="Camera Size", min=0.001, max=0.1, step=0.001, default=0.05)
+scenegraph_type = gradio.Dropdown(["complete", "swin", "oneref"], label="Scene Graph Type")
+winsize = gradio.Slider(label="Window Size", min=1, max=1, step=1, default=1)
+refid = gradio.Slider(label="Reference ID", min=0, max=0, step=1, default=0)
+
+# Function to connect Gradio inputs to your main logic
+def run_inference(filelist, schedule, niter, min_conf_thr, as_pointcloud, mask_sky, clean_depth,
+                  transparent_cams, cam_size, scenegraph_type, winsize, refid):
+    args = None  # You need to define your args here
+    outfile, imgs = run_dust3r_inference(args, filelist, schedule, niter, min_conf_thr, as_pointcloud,
+                                          mask_sky, clean_depth, transparent_cams, cam_size,
+                                          scenegraph_type, winsize, refid)
+    return imgs
 
-    model = load_model("./dust3r/checkpoints/DUSt3R_ViTLarge_BaseDecoder_512_dpt.pth", "cuda")
-    # dust3r will write the 3D model inside tmpdirname
-    with tempfile.TemporaryDirectory(suffix='dust3r_gradio_demo') as tmpdirname:
-        print('Outputing stuff in', tmpdirname)
-        main_demo(tmpdirname, model, "cuda", 512, server_name, 7860)
+# Launch the Gradio interface
+iface = gradio.Interface(
+    fn=run_inference,
+    inputs=[inputfiles, schedule, niter, min_conf_thr, as_pointcloud, mask_sky, clean_depth,
+            transparent_cams, cam_size, scenegraph_type, winsize, refid],
+    outputs=gradio.Gallery(label="Output Images", columns=3),
+    title="DUSt3R Demo",
+    description="Reconstruct 3D scenes from input images using DUSt3R.",
+    server_name='0.0.0.0',
+    server_port=7860
+)
+iface.launch(share=True)