src/mast3r_src/demo.py

#!/usr/bin/env python3
# Copyright (C) 2024-present Naver Corporation. All rights reserved.
# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
#
# --------------------------------------------------------
# gradio demo
# --------------------------------------------------------
import math
import gradio
import os
import torch
import numpy as np
import tempfile
import functools
import trimesh
import copy
from scipy.spatial.transform import Rotation

from mast3r.cloud_opt.sparse_ga import sparse_global_alignment
from mast3r.cloud_opt.tsdf_optimizer import TSDFPostProcess

from mast3r.model import AsymmetricMASt3R
from mast3r.utils.misc import hash_md5
import mast3r.utils.path_to_dust3r # noqa
from dust3r.image_pairs import make_pairs
from dust3r.utils.image import load_images
from dust3r.utils.device import to_numpy
from dust3r.viz import add_scene_cam, CAM_COLORS, OPENGL, pts3d_to_trimesh, cat_meshes
from dust3r.demo import get_args_parser as dust3r_get_args_parser

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 get_args_parser():
 parser = dust3r_get_args_parser()
 parser.add_argument('--share', action='store_true')

 actions = parser._actions
 for action in actions:
 if action.dest == 'model_name':
 action.choices = ["MASt3R_ViTLarge_BaseDecoder_512_catmlpdpt_metric"]
 # change defaults
 parser.prog = 'mast3r demo'
 return parser


def _convert_scene_output_to_glb(outdir, imgs, pts3d, mask, focals, cams2world, cam_size=0.05,
 cam_color=None, as_pointcloud=False,
 transparent_cams=False, silent=False):
 assert len(pts3d) == len(mask) <= len(imgs) <= len(cams2world) == len(focals)
 pts3d = to_numpy(pts3d)
 imgs = to_numpy(imgs)
 focals = to_numpy(focals)
 cams2world = to_numpy(cams2world)

 scene = trimesh.Scene()

 # full pointcloud
 if as_pointcloud:
 pts = np.concatenate([p[m.ravel()] for p, m in zip(pts3d, mask)])
 col = np.concatenate([p[m] for p, m in zip(imgs, mask)])
 pct = trimesh.PointCloud(pts.reshape(-1, 3), colors=col.reshape(-1, 3))
 scene.add_geometry(pct)
 else:
 meshes = []
 for i in range(len(imgs)):
 meshes.append(pts3d_to_trimesh(imgs[i], pts3d[i].reshape(imgs[i].shape), mask[i]))
 mesh = trimesh.Trimesh(**cat_meshes(meshes))
 scene.add_geometry(mesh)

 # add each camera
 for i, pose_c2w in enumerate(cams2world):
 if isinstance(cam_color, list):
 camera_edge_color = cam_color[i]
 else:
 camera_edge_color = cam_color or CAM_COLORS[i % len(CAM_COLORS)]
 add_scene_cam(scene, pose_c2w, camera_edge_color,
 None if transparent_cams else imgs[i], focals[i],
 imsize=imgs[i].shape[1::-1], screen_width=cam_size)

 rot = np.eye(4)
 rot[:3, :3] = Rotation.from_euler('y', np.deg2rad(180)).as_matrix()
 scene.apply_transform(np.linalg.inv(cams2world[0] @ OPENGL @ rot))
 outfile = os.path.join(outdir, 'scene.glb')
 if not silent:
 print('(exporting 3D scene to', outfile, ')')
 scene.export(file_obj=outfile)
 return outfile


def get_3D_model_from_scene(outdir, silent, scene, min_conf_thr=2, as_pointcloud=False, mask_sky=False,
 clean_depth=False, transparent_cams=False, cam_size=0.05, TSDF_thresh=0):
 """
 extract 3D_model (glb file) from a reconstructed scene
 """
 if scene is None:
 return None

 # get optimized values from scene
 rgbimg = scene.imgs
 focals = scene.get_focals().cpu()
 cams2world = scene.get_im_poses().cpu()

 # 3D pointcloud from depthmap, poses and intrinsics
 if TSDF_thresh > 0:
 tsdf = TSDFPostProcess(scene, TSDF_thresh=TSDF_thresh)
 pts3d, _, confs = to_numpy(tsdf.get_dense_pts3d(clean_depth=clean_depth))
 else:
 pts3d, _, confs = to_numpy(scene.get_dense_pts3d(clean_depth=clean_depth))
 msk = to_numpy([c > min_conf_thr for c in confs])
 return _convert_scene_output_to_glb(outdir, rgbimg, pts3d, msk, focals, cams2world, as_pointcloud=as_pointcloud,
 transparent_cams=transparent_cams, cam_size=cam_size, silent=silent)


def get_reconstructed_scene(outdir, model, device, silent, image_size, 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):
 """
 from a list of images, run mast3r inference, sparse global aligner.
 then run get_3D_model_from_scene
 """
 imgs = load_images(filelist, size=image_size, verbose=not silent)
 if len(imgs) == 1:
 imgs = [imgs[0], copy.deepcopy(imgs[0])]
 imgs[1]['idx'] = 1
 filelist = [filelist[0], filelist[0] + '_2']

 scene_graph_params = [scenegraph_type]
 if scenegraph_type in ["swin", "logwin"]:
 scene_graph_params.append(str(winsize))
 elif scenegraph_type == "oneref":
 scene_graph_params.append(str(refid))
 if scenegraph_type in ["swin", "logwin"] and not win_cyclic:
 scene_graph_params.append('noncyclic')
 scene_graph = '-'.join(scene_graph_params)
 pairs = make_pairs(imgs, scene_graph=scene_graph, prefilter=None, symmetrize=True)
 if optim_level == 'coarse':
 niter2 = 0
 # Sparse GA (forward mast3r -> matching -> 3D optim -> 2D refinement -> triangulation)
 scene = sparse_global_alignment(filelist, pairs, os.path.join(outdir, 'cache'),
 model, lr1=lr1, niter1=niter1, lr2=lr2, niter2=niter2, device=device,
 opt_depth='depth' in optim_level, shared_intrinsics=shared_intrinsics,
 matching_conf_thr=matching_conf_thr, **kw)
 outfile = get_3D_model_from_scene(outdir, silent, scene, min_conf_thr, as_pointcloud, mask_sky,
 clean_depth, transparent_cams, cam_size, TSDF_thresh)
 return scene, outfile


def set_scenegraph_options(inputfiles, win_cyclic, refid, scenegraph_type):
 num_files = len(inputfiles) if inputfiles is not None else 1
 show_win_controls = scenegraph_type in ["swin", "logwin"]
 show_winsize = scenegraph_type in ["swin", "logwin"]
 show_cyclic = scenegraph_type in ["swin", "logwin"]
 max_winsize, min_winsize = 1, 1
 if scenegraph_type == "swin":
 if win_cyclic:
 max_winsize = max(1, math.ceil((num_files - 1) / 2))
 else:
 max_winsize = num_files - 1
 elif scenegraph_type == "logwin":
 if win_cyclic:
 half_size = math.ceil((num_files - 1) / 2)
 max_winsize = max(1, math.ceil(math.log(half_size, 2)))
 else:
 max_winsize = max(1, math.ceil(math.log(num_files, 2)))
 winsize = gradio.Slider(label="Scene Graph: Window Size", value=max_winsize,
 minimum=min_winsize, maximum=max_winsize, step=1, visible=show_winsize)
 win_cyclic = gradio.Checkbox(value=win_cyclic, label="Cyclic sequence", visible=show_cyclic)
 win_col = gradio.Column(visible=show_win_controls)
 refid = gradio.Slider(label="Scene Graph: Id", value=0, minimum=0,
 maximum=num_files - 1, step=1, visible=scenegraph_type == 'oneref')
 return win_col, winsize, win_cyclic, refid


def main_demo(tmpdirname, model, device, image_size, server_name, server_port, silent=False, share=False):
 if not silent:
 print('Outputing stuff in', tmpdirname)

 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="MASt3R 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;">MASt3R Demo</h2>')
 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", "swin", "logwin", "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=[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=True, server_name=server_name, server_port=server_port)


if __name__ == '__main__':
 parser = get_args_parser()
 args = parser.parse_args()

 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 = AsymmetricMASt3R.from_pretrained(weights_path).to(args.device)
 chkpt_tag = hash_md5(weights_path)

 # mast3r will write the 3D model inside tmpdirname/chkpt_tag
 if args.tmp_dir is not None:
 tmpdirname = args.tmp_dir
 cache_path = os.path.join(tmpdirname, chkpt_tag)
 os.makedirs(cache_path, exist_ok=True)
 main_demo(cache_path, model, args.device, args.image_size, server_name, args.server_port, silent=args.silent,
 share=args.share)
 else:
 with tempfile.TemporaryDirectory(suffix='_mast3r_gradio_demo') as tmpdirname:
 cache_path = os.path.join(tmpdirname, chkpt_tag)
 os.makedirs(cache_path, exist_ok=True)
 main_demo(tmpdirname, model, args.device, args.image_size,
 server_name, args.server_port, silent=args.silent,
 share=args.share)