Spaces:
Running
on
Zero
Running
on
Zero
| # Copyright (C) 2024-present Naver Corporation. All rights reserved. | |
| # Licensed under CC BY-NC-SA 4.0 (non-commercial use only). | |
| # | |
| # -------------------------------------------------------- | |
| # mast3r 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 mast3r.utils.path_to_dust3r # noqa | |
| from dust3r.demo import set_print_with_timestamp | |
| import matplotlib.pyplot as pl | |
| pl.ion() | |
| # for gpu >= Ampere and pytorch >= 1.12 | |
| torch.backends.cuda.matmul.allow_tf32 = True | |
| batch_size = 1 | |
| set_print_with_timestamp() | |
| 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 | |
| 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 | |
| def run_example(snapshot, matching_conf_thr, min_conf_thr, cam_size, as_pointcloud, shared_intrinsics, filelist, **kw): | |
| return local_get_reconstructed_scene(filelist, min_conf_thr, matching_conf_thr, as_pointcloud, cam_size, shared_intrinsics, **kw) | |
| 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('<p>Upload one or multiple images (wait for them to be fully uploaded before hitting the run button). ' | |
| 'We tested with up to 18 images before running into the allocation timeout - set at 3 minutes but your mileage may vary. ' | |
| 'At the very bottom of this page, you will find an example. If you click on it, it will pull the 3D reconstruction from 7 images of the small Naver Labs Europe tower from cache. ' | |
| 'If you want to try larger image collections, you can find the more complete version of this demo that you can run locally ' | |
| 'and more details about the method at <a href="https://github.com/naver/mast3r">github.com/naver/mast3r</a>. ' | |
| 'The checkpoint used in this demo is available at <a href="https://huggingface.co/naver/MASt3R_ViTLarge_BaseDecoder_512_catmlpdpt_metric">huggingface.co/naver/MASt3R_ViTLarge_BaseDecoder_512_catmlpdpt_metric</a>.</p>') | |
| with gradio.Column(): | |
| inputfiles = gradio.File(file_count="multiple") | |
| snapshot = gradio.Image(None, visible=False) | |
| 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/FF5599FD-768B-431A-AB83-BDA5FB44CB9D-83120-000041DADDE35483.jpg'), | |
| 0.0, 1.5, 0.2, True, False, | |
| [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')] | |
| ] | |
| ], | |
| inputs=[snapshot, matching_conf_thr, min_conf_thr, cam_size, as_pointcloud, shared_intrinsics, inputfiles], | |
| outputs=[filestate, outmodel], | |
| fn=run_example, | |
| 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(show_error=True, share=None, server_name=None, server_port=None) | |
| shutil.rmtree(tmpdirname) | |