File size: 6,200 Bytes
6151e97 2a90110 6151e97 7494687 6d2638a 0c19783 6151e97 0c19783 6151e97 18276d5 6151e97 e83109b 6151e97 6d2638a 7494687 0c19783 7494687 0c19783 7494687 2a90110 0c19783 7494687 0c19783 2f74cb2 7494687 2f74cb2 0c19783 7494687 0c19783 7494687 2f74cb2 7494687 0c19783 7494687 6d2638a |
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 |
# 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
import shutil
import math
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
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 = tempfile.mkdtemp(suffix='_mast3r_gradio_demo')
image_size = 512
silent = True
gradio_delete_cache = 7200
class FileState:
def __init__(self, outfile_name=None):
self.outfile_name = outfile_name
def __del__(self):
if self.outfile_name is not None and os.path.isfile(self.outfile_name):
os.remove(self.outfile_name)
self.outfile_name = None
@spaces.GPU()
def local_get_reconstructed_scene(filelist, min_conf_thr, matching_conf_thr,
as_pointcloud, cam_size,
shared_intrinsics, **kw):
lr1 = 0.07
niter1 = 500
lr2 = 0.014
niter2 = 200
optim_level = 'refine'
mask_sky, clean_depth, transparent_cams = False, True, False
if len(filelist) < 5:
scenegraph_type = 'complete'
winsize = 1
else:
scenegraph_type = 'logwin'
half_size = math.ceil((len(filelist) - 1) / 2)
max_winsize = max(1, math.ceil(math.log(half_size, 2)))
winsize = min(5, max_winsize)
refid = 0
win_cyclic = False
scene_state, outfile = get_reconstructed_scene(tmpdirname, gradio_delete_cache, model, device, silent, image_size, None,
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=0, shared_intrinsics=shared_intrinsics, **kw)
filestate = FileState(scene_state.outfile_name)
scene_state.outfile_name = None
del scene_state
return filestate, outfile
css = """.gradio-container {margin: 0 !important; min-width: 100%};"""
title = "MASt3R Demo"
with gradio.Blocks(css=css, title=title, delete_cache=(gradio_delete_cache, gradio_delete_cache)) as demo:
filestate = gradio.State(None)
gradio.HTML('<h2 style="text-align: center;">3D Reconstruction with MASt3R</h2>')
gradio.HTML('<h3 style="text-align: center;">Upload one or multiple images</h3>')
with gradio.Column():
inputfiles = gradio.File(file_count="multiple")
with gradio.Row():
matching_conf_thr = gradio.Slider(label="Matching Confidence Thr", value=2.,
minimum=0., maximum=30., step=0.1,
info="Before Fallback to Regr3D!")
# 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)
with gradio.Row():
as_pointcloud = gradio.Checkbox(value=True, label="As pointcloud")
shared_intrinsics = gradio.Checkbox(value=False, label="Shared intrinsics",
info="Only optimize one set of intrinsics for all views")
run_btn = gradio.Button("Run")
outmodel = gradio.Model3D()
examples = gradio.Examples(
examples=[
[
[os.path.join(HERE_PATH, 'mast3r/assets/NLE_tower/01D90321-69C8-439F-B0B0-E87E7634741C-83120-000041DAE419D7AE.jpg'),
os.path.join(
HERE_PATH, 'mast3r/assets/NLE_tower/1AD85EF5-B651-4291-A5C0-7BDB7D966384-83120-000041DADF639E09.jpg'),
os.path.join(
HERE_PATH, 'mast3r/assets/NLE_tower/28EDBB63-B9F9-42FB-AC86-4852A33ED71B-83120-000041DAF22407A1.jpg'),
os.path.join(
HERE_PATH, 'mast3r/assets/NLE_tower/91E9B685-7A7D-42D7-B933-23A800EE4129-83120-000041DAE12C8176.jpg'),
os.path.join(
HERE_PATH, 'mast3r/assets/NLE_tower/2679C386-1DC0-4443-81B5-93D7EDE4AB37-83120-000041DADB2EA917.jpg'),
os.path.join(
HERE_PATH, 'mast3r/assets/NLE_tower/CDBBD885-54C3-4EB4-9181-226059A60EE0-83120-000041DAE0C3D612.jpg'),
os.path.join(HERE_PATH, 'mast3r/assets/NLE_tower/FF5599FD-768B-431A-AB83-BDA5FB44CB9D-83120-000041DADDE35483.jpg')],
1.5, 0.0, True, 0.2, False
]
],
inputs=[inputfiles, min_conf_thr, matching_conf_thr, as_pointcloud, cam_size, shared_intrinsics],
outputs=[filestate, outmodel],
fn=local_get_reconstructed_scene,
cache_examples="lazy",
)
# events
run_btn.click(fn=local_get_reconstructed_scene,
inputs=[inputfiles, min_conf_thr, matching_conf_thr,
as_pointcloud,
cam_size, shared_intrinsics],
outputs=[filestate, outmodel])
demo.launch(share=None, server_name=None, server_port=None)
shutil.rmtree(tmpdirname)
|