File size: 9,228 Bytes
e7d52aa dc56a72 e7d52aa dc56a72 e7d52aa |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 |
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'
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
# 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)
|