# Copyright (C) 2024-present Naver Corporation. All rights reserved. # Licensed under CC BY-NC-SA 4.0 (non-commercial use only). # # -------------------------------------------------------- # masst3r demo # -------------------------------------------------------- import spaces import os import sys import os.path as path import torch import tempfile import gradio HERE_PATH = path.normpath(path.dirname(__file__)) # noqa MASt3R_REPO_PATH = path.normpath(path.join(HERE_PATH, './mast3r')) # noqa sys.path.insert(0, MASt3R_REPO_PATH) # noqa from mast3r.demo import get_reconstructed_scene, get_3D_model_from_scene, set_scenegraph_options from mast3r.model import AsymmetricMASt3R from mast3r.utils.misc import hash_md5 import matplotlib.pyplot as pl pl.ion() # for gpu >= Ampere and pytorch >= 1.12 torch.backends.cuda.matmul.allow_tf32 = True batch_size = 1 weights_path = "naver/" + 'MASt3R_ViTLarge_BaseDecoder_512_catmlpdpt_metric' device = 'cuda' if torch.cuda.is_available() else 'cpu' model = AsymmetricMASt3R.from_pretrained(weights_path).to(device) chkpt_tag = hash_md5(weights_path) tmpdirname = "tmp/gradio" image_size = 512 silent = True gradio_delete_cache = 7200 @spaces.GPU() def local_get_reconstructed_scene(current_scene_state, filelist, optim_level, lr1, niter1, lr2, niter2, min_conf_thr, matching_conf_thr, as_pointcloud, mask_sky, clean_depth, transparent_cams, cam_size, scenegraph_type, winsize, win_cyclic, refid, TSDF_thresh, shared_intrinsics, **kw): return get_reconstructed_scene(tmpdirname, gradio_delete_cache, model, device, silent, image_size, current_scene_state, filelist, optim_level, lr1, niter1, lr2, niter2, min_conf_thr, matching_conf_thr, as_pointcloud, mask_sky, clean_depth, transparent_cams, cam_size, scenegraph_type, winsize, win_cyclic, refid, TSDF_thresh, shared_intrinsics, **kw) @spaces.GPU() def local_get_3D_model_from_scene(scene_state, min_conf_thr=2, as_pointcloud=False, mask_sky=False, clean_depth=False, transparent_cams=False, cam_size=0.05, TSDF_thresh=0): return get_3D_model_from_scene(silent, scene_state, min_conf_thr, as_pointcloud, mask_sky, clean_depth, transparent_cams, cam_size, TSDF_thresh) recon_fun = local_get_reconstructed_scene model_from_scene_fun = local_get_3D_model_from_scene def get_context(delete_cache): css = """.gradio-container {margin: 0 !important; min-width: 100%};""" title = "MASt3R Demo" if delete_cache: return gradio.Blocks(css=css, title=title, delete_cache=(delete_cache, delete_cache)) else: return gradio.Blocks(css=css, title="MASt3R Demo") # for compatibility with older versions with get_context(gradio_delete_cache) 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('

MASt3R Demo

') with gradio.Column(): inputfiles = gradio.File(file_count="multiple") with gradio.Row(): with gradio.Column(): with gradio.Row(): lr1 = gradio.Slider(label="Coarse LR", value=0.07, minimum=0.01, maximum=0.2, step=0.01) niter1 = gradio.Number(value=500, precision=0, minimum=0, maximum=10_000, label="num_iterations", info="For coarse alignment!") lr2 = gradio.Slider(label="Fine LR", value=0.014, minimum=0.005, maximum=0.05, step=0.001) niter2 = gradio.Number(value=200, precision=0, minimum=0, maximum=100_000, label="num_iterations", info="For refinement!") optim_level = gradio.Dropdown(["coarse", "refine", "refine+depth"], value='refine', label="OptLevel", info="Optimization level") with gradio.Row(): matching_conf_thr = gradio.Slider(label="Matching Confidence Thr", value=5., minimum=0., maximum=30., step=0.1, info="Before Fallback to Regr3D!") shared_intrinsics = gradio.Checkbox(value=False, label="Shared intrinsics", info="Only optimize one set of intrinsics for all views") scenegraph_type = gradio.Dropdown([("complete: all possible image pairs", "complete"), ("swin: sliding window", "swin"), ("logwin: sliding window with long range", "logwin"), ("oneref: match one image with all", "oneref")], value='complete', label="Scenegraph", info="Define how to make pairs", interactive=True) with gradio.Column(visible=False) as win_col: winsize = gradio.Slider(label="Scene Graph: Window Size", value=1, minimum=1, maximum=1, step=1) win_cyclic = gradio.Checkbox(value=False, label="Cyclic sequence") 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=1.5, minimum=0.0, maximum=10, step=0.1) # adjust the camera size in the output pointcloud cam_size = gradio.Slider(label="cam_size", value=0.2, minimum=0.001, maximum=1.0, step=0.001) TSDF_thresh = gradio.Slider(label="TSDF Threshold", value=0., minimum=0., maximum=1., step=0.01) with gradio.Row(): as_pointcloud = gradio.Checkbox(value=True, 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() # events scenegraph_type.change(set_scenegraph_options, inputs=[inputfiles, win_cyclic, refid, scenegraph_type], outputs=[win_col, winsize, win_cyclic, refid]) inputfiles.change(set_scenegraph_options, inputs=[inputfiles, win_cyclic, refid, scenegraph_type], outputs=[win_col, winsize, win_cyclic, refid]) win_cyclic.change(set_scenegraph_options, inputs=[inputfiles, win_cyclic, refid, scenegraph_type], outputs=[win_col, winsize, win_cyclic, refid]) run_btn.click(fn=recon_fun, inputs=[scene, inputfiles, optim_level, lr1, niter1, lr2, niter2, min_conf_thr, matching_conf_thr, as_pointcloud, mask_sky, clean_depth, transparent_cams, cam_size, scenegraph_type, winsize, win_cyclic, refid, TSDF_thresh, shared_intrinsics], outputs=[scene, outmodel]) min_conf_thr.release(fn=model_from_scene_fun, inputs=[scene, min_conf_thr, as_pointcloud, mask_sky, clean_depth, transparent_cams, cam_size, TSDF_thresh], 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, TSDF_thresh], outputs=outmodel) TSDF_thresh.change(fn=model_from_scene_fun, inputs=[scene, min_conf_thr, as_pointcloud, mask_sky, clean_depth, transparent_cams, cam_size, TSDF_thresh], 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, TSDF_thresh], 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, TSDF_thresh], 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, TSDF_thresh], outputs=outmodel) transparent_cams.change(model_from_scene_fun, inputs=[scene, min_conf_thr, as_pointcloud, mask_sky, clean_depth, transparent_cams, cam_size, TSDF_thresh], outputs=outmodel) demo.launch(share=None, server_name=None, server_port=None)