MASt3R_jupyter.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github"
      },
      "source": [
        "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/camenduru/MASt3R-jupyter/blob/main/MASt3R_jupyter.ipynb)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VjYy0F2gZIPR"
      },
      "outputs": [],
      "source": [
        "%cd /content\n",
        "!git clone -b dev --recursive https://github.com/camenduru/mast3r\n",
        "%cd /content/mast3r\n",
        "\n",
        "%cd dust3r/croco/models/curope\n",
        "!python setup.py build_ext --inplace\n",
        "\n",
        "!mkdir -p /content/mast3r/checkpoints\n",
        "!wget https://huggingface.co/camenduru/mast3r/resolve/main/MASt3R_ViTLarge_BaseDecoder_512_catmlpdpt_metric.pth -O /content/mast3r/checkpoints/MASt3R_ViTLarge_BaseDecoder_512_catmlpdpt_metric.pth\n",
        "\n",
        "!pip install roma einops trimesh\n",
        "!pip install /content/mast3r/dust3r/croco/models/curope\n",
        "\n",
        "!wget https://raw.githubusercontent.com/naver/mast3r/refs/heads/main/assets/NLE_tower/01D90321-69C8-439F-B0B0-E87E7634741C-83120-000041DAE419D7AE.jpg -O /content/1.jpg\n",
        "!wget https://raw.githubusercontent.com/naver/mast3r/refs/heads/main/assets/NLE_tower/1AD85EF5-B651-4291-A5C0-7BDB7D966384-83120-000041DADF639E09.jpg -O /content/2.jpg\n",
        "!wget https://raw.githubusercontent.com/naver/mast3r/refs/heads/main/assets/NLE_tower/2679C386-1DC0-4443-81B5-93D7EDE4AB37-83120-000041DADB2EA917.jpg -O /content/3.jpg\n",
        "!wget https://raw.githubusercontent.com/naver/mast3r/refs/heads/main/assets/NLE_tower/28EDBB63-B9F9-42FB-AC86-4852A33ED71B-83120-000041DAF22407A1.jpg -O /content/4.jpg\n",
        "!wget https://raw.githubusercontent.com/naver/mast3r/refs/heads/main/assets/NLE_tower/91E9B685-7A7D-42D7-B933-23A800EE4129-83120-000041DAE12C8176.jpg -O /content/5.jpg\n",
        "!wget https://raw.githubusercontent.com/naver/mast3r/refs/heads/main/assets/NLE_tower/CDBBD885-54C3-4EB4-9181-226059A60EE0-83120-000041DAE0C3D612.jpg -O /content/6.jpg\n",
        "!wget https://raw.githubusercontent.com/naver/mast3r/refs/heads/main/assets/NLE_tower/FF5599FD-768B-431A-AB83-BDA5FB44CB9D-83120-000041DADDE35483.jpg -O /content/7.jpg"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "%cd /content/mast3r\n",
        "\n",
        "import os\n",
        "import torch\n",
        "import math\n",
        "\n",
        "import numpy as np\n",
        "import trimesh\n",
        "import copy, shutil, tempfile\n",
        "from scipy.spatial.transform import Rotation\n",
        "\n",
        "from dust3r.utils.image import load_images\n",
        "from dust3r.image_pairs import make_pairs\n",
        "from dust3r.utils.device import to_numpy\n",
        "from dust3r.viz import add_scene_cam, CAM_COLORS, OPENGL, pts3d_to_trimesh, cat_meshes\n",
        "\n",
        "from mast3r.cloud_opt.sparse_ga import sparse_global_alignment\n",
        "from mast3r.cloud_opt.tsdf_optimizer import TSDFPostProcess\n",
        "\n",
        "class SparseGAState():\n",
        "    def __init__(self, sparse_ga, should_delete=False, cache_dir=None, outfile_name=None):\n",
        "        self.sparse_ga = sparse_ga\n",
        "        self.cache_dir = cache_dir\n",
        "        self.outfile_name = outfile_name\n",
        "        self.should_delete = should_delete\n",
        "\n",
        "    def __del__(self):\n",
        "        if not self.should_delete:\n",
        "            return\n",
        "        if self.cache_dir is not None and os.path.isdir(self.cache_dir):\n",
        "            shutil.rmtree(self.cache_dir)\n",
        "        self.cache_dir = None\n",
        "        if self.outfile_name is not None and os.path.isfile(self.outfile_name):\n",
        "            os.remove(self.outfile_name)\n",
        "        self.outfile_name = None\n",
        "\n",
        "def _convert_scene_output_to_glb(outfile, imgs, pts3d, mask, focals, cams2world, cam_size=0.05,\n",
        "                                 cam_color=None, as_pointcloud=False,\n",
        "                                 transparent_cams=False, silent=False):\n",
        "    assert len(pts3d) == len(mask) <= len(imgs) <= len(cams2world) == len(focals)\n",
        "    pts3d = to_numpy(pts3d)\n",
        "    imgs = to_numpy(imgs)\n",
        "    focals = to_numpy(focals)\n",
        "    cams2world = to_numpy(cams2world)\n",
        "\n",
        "    scene = trimesh.Scene()\n",
        "\n",
        "    # full pointcloud\n",
        "    if as_pointcloud:\n",
        "        pts = np.concatenate([p[m.ravel()] for p, m in zip(pts3d, mask)]).reshape(-1, 3)\n",
        "        col = np.concatenate([p[m] for p, m in zip(imgs, mask)]).reshape(-1, 3)\n",
        "        valid_msk = np.isfinite(pts.sum(axis=1))\n",
        "        pct = trimesh.PointCloud(pts[valid_msk], colors=col[valid_msk])\n",
        "        scene.add_geometry(pct)\n",
        "    else:\n",
        "        meshes = []\n",
        "        for i in range(len(imgs)):\n",
        "            pts3d_i = pts3d[i].reshape(imgs[i].shape)\n",
        "            msk_i = mask[i] & np.isfinite(pts3d_i.sum(axis=-1))\n",
        "            meshes.append(pts3d_to_trimesh(imgs[i], pts3d_i, msk_i))\n",
        "        mesh = trimesh.Trimesh(**cat_meshes(meshes))\n",
        "        scene.add_geometry(mesh)\n",
        "\n",
        "    # add each camera\n",
        "    for i, pose_c2w in enumerate(cams2world):\n",
        "        if isinstance(cam_color, list):\n",
        "            camera_edge_color = cam_color[i]\n",
        "        else:\n",
        "            camera_edge_color = cam_color or CAM_COLORS[i % len(CAM_COLORS)]\n",
        "        add_scene_cam(scene, pose_c2w, camera_edge_color,\n",
        "                      None if transparent_cams else imgs[i], focals[i],\n",
        "                      imsize=imgs[i].shape[1::-1], screen_width=cam_size)\n",
        "\n",
        "    rot = np.eye(4)\n",
        "    rot[:3, :3] = Rotation.from_euler('y', np.deg2rad(180)).as_matrix()\n",
        "    scene.apply_transform(np.linalg.inv(cams2world[0] @ OPENGL @ rot))\n",
        "    if not silent:\n",
        "        print('(exporting 3D scene to', outfile, ')')\n",
        "    scene.export(file_obj=outfile)\n",
        "    return outfile\n",
        "\n",
        "def get_3D_model_from_scene(silent, scene_state, min_conf_thr=2, as_pointcloud=False, mask_sky=False,\n",
        "                            clean_depth=False, transparent_cams=False, cam_size=0.05, TSDF_thresh=0):\n",
        "    \"\"\"\n",
        "    extract 3D_model (glb file) from a reconstructed scene\n",
        "    \"\"\"\n",
        "    if scene_state is None:\n",
        "        return None\n",
        "    outfile = scene_state.outfile_name\n",
        "    if outfile is None:\n",
        "        return None\n",
        "\n",
        "    # get optimized values from scene\n",
        "    scene = scene_state.sparse_ga\n",
        "    rgbimg = scene.imgs\n",
        "    focals = scene.get_focals().cpu()\n",
        "    cams2world = scene.get_im_poses().cpu()\n",
        "\n",
        "    # 3D pointcloud from depthmap, poses and intrinsics\n",
        "    if TSDF_thresh > 0:\n",
        "        tsdf = TSDFPostProcess(scene, TSDF_thresh=TSDF_thresh)\n",
        "        pts3d, _, confs = to_numpy(tsdf.get_dense_pts3d(clean_depth=clean_depth))\n",
        "    else:\n",
        "        pts3d, _, confs = to_numpy(scene.get_dense_pts3d(clean_depth=clean_depth))\n",
        "    msk = to_numpy([c > min_conf_thr for c in confs])\n",
        "    return _convert_scene_output_to_glb(outfile, rgbimg, pts3d, msk, focals, cams2world, as_pointcloud=as_pointcloud,\n",
        "                                        transparent_cams=transparent_cams, cam_size=cam_size, silent=silent)\n",
        "\n",
        "def get_reconstructed_scene(outdir, gradio_delete_cache, model, device, silent, image_size, current_scene_state,\n",
        "                            filelist, optim_level, lr1, niter1, lr2, niter2, min_conf_thr, matching_conf_thr,\n",
        "                            as_pointcloud, mask_sky, clean_depth, transparent_cams, cam_size, scenegraph_type, winsize,\n",
        "                            win_cyclic, refid, TSDF_thresh, shared_intrinsics, **kw):\n",
        "    \"\"\"\n",
        "    from a list of images, run mast3r inference, sparse global aligner.\n",
        "    then run get_3D_model_from_scene\n",
        "    \"\"\"\n",
        "    imgs = load_images(filelist, size=image_size, verbose=not silent)\n",
        "    if len(imgs) == 1:\n",
        "        imgs = [imgs[0], copy.deepcopy(imgs[0])]\n",
        "        imgs[1]['idx'] = 1\n",
        "        filelist = [filelist[0], filelist[0] + '_2']\n",
        "\n",
        "    scene_graph_params = [scenegraph_type]\n",
        "    if scenegraph_type in [\"swin\", \"logwin\"]:\n",
        "        scene_graph_params.append(str(winsize))\n",
        "    elif scenegraph_type == \"oneref\":\n",
        "        scene_graph_params.append(str(refid))\n",
        "    if scenegraph_type in [\"swin\", \"logwin\"] and not win_cyclic:\n",
        "        scene_graph_params.append('noncyclic')\n",
        "    scene_graph = '-'.join(scene_graph_params)\n",
        "    pairs = make_pairs(imgs, scene_graph=scene_graph, prefilter=None, symmetrize=True)\n",
        "    if optim_level == 'coarse':\n",
        "        niter2 = 0\n",
        "    # Sparse GA (forward mast3r -> matching -> 3D optim -> 2D refinement -> triangulation)\n",
        "    if current_scene_state is not None and \\\n",
        "        not current_scene_state.should_delete and \\\n",
        "            current_scene_state.cache_dir is not None:\n",
        "        cache_dir = current_scene_state.cache_dir\n",
        "    elif gradio_delete_cache:\n",
        "        cache_dir = tempfile.mkdtemp(suffix='_cache', dir=outdir)\n",
        "    else:\n",
        "        cache_dir = os.path.join(outdir, 'cache')\n",
        "    os.makedirs(cache_dir, exist_ok=True)\n",
        "    scene = sparse_global_alignment(filelist, pairs, cache_dir,\n",
        "                                    model, lr1=lr1, niter1=niter1, lr2=lr2, niter2=niter2, device=device,\n",
        "                                    opt_depth='depth' in optim_level, shared_intrinsics=shared_intrinsics,\n",
        "                                    matching_conf_thr=matching_conf_thr, **kw)\n",
        "    if current_scene_state is not None and \\\n",
        "        not current_scene_state.should_delete and \\\n",
        "            current_scene_state.outfile_name is not None:\n",
        "        outfile_name = current_scene_state.outfile_name\n",
        "    else:\n",
        "        outfile_name = tempfile.mktemp(suffix='_scene.glb', dir=outdir)\n",
        "\n",
        "    scene_state = SparseGAState(scene, gradio_delete_cache, cache_dir, outfile_name)\n",
        "    outfile = get_3D_model_from_scene(silent, scene_state, min_conf_thr, as_pointcloud, mask_sky,\n",
        "                                      clean_depth, transparent_cams, cam_size, TSDF_thresh)\n",
        "    return scene_state, outfile\n",
        "\n",
        "# from mast3r.demo import get_reconstructed_scene\n",
        "from mast3r.model import AsymmetricMASt3R\n",
        "from mast3r.utils.misc import hash_md5\n",
        "\n",
        "# for gpu >= Ampere and pytorch >= 1.12\n",
        "torch.backends.cuda.matmul.allow_tf32 = True\n",
        "batch_size = 1\n",
        "\n",
        "weights_path = \"/content/mast3r/checkpoints/MASt3R_ViTLarge_BaseDecoder_512_catmlpdpt_metric.pth\"\n",
        "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
        "model = AsymmetricMASt3R.from_pretrained(weights_path).to(device)\n",
        "chkpt_tag = hash_md5(weights_path)\n",
        "\n",
        "# Set the temporary directory\n",
        "temp_dir = '/content/temp'\n",
        "os.makedirs(temp_dir, exist_ok=True)  # Ensure the temp directory exists\n",
        "\n",
        "image_size = 512\n",
        "silent = True\n",
        "gradio_delete_cache = 7200\n",
        "\n",
        "class FileState:\n",
        "    def __init__(self, outfile_name=None):\n",
        "        self.outfile_name = outfile_name\n",
        "\n",
        "    def __del__(self):\n",
        "        if self.outfile_name is not None and os.path.isfile(self.outfile_name):\n",
        "            os.remove(self.outfile_name)\n",
        "        self.outfile_name = None\n",
        "\n",
        "def local_get_reconstructed_scene(filelist, min_conf_thr, matching_conf_thr,\n",
        "                                  as_pointcloud, cam_size,\n",
        "                                  shared_intrinsics, **kw):\n",
        "    lr1 = 0.07\n",
        "    niter1 = 500\n",
        "    lr2 = 0.014\n",
        "    niter2 = 200\n",
        "    optim_level = 'refine'\n",
        "    mask_sky = False\n",
        "    clean_depth = True\n",
        "    transparent_cams = False\n",
        "\n",
        "    if len(filelist) < 5:\n",
        "        scenegraph_type = 'complete'\n",
        "        winsize = 1\n",
        "    else:\n",
        "        scenegraph_type = 'logwin'\n",
        "        half_size = math.ceil((len(filelist) - 1) / 2)\n",
        "        max_winsize = max(1, math.ceil(math.log(half_size, 2)))\n",
        "        winsize = min(5, max_winsize)\n",
        "        \n",
        "    refid = 0\n",
        "    win_cyclic = False\n",
        "    scene_state, outfile = get_reconstructed_scene(temp_dir, gradio_delete_cache, model, device, silent, image_size, None,\n",
        "                                                   filelist, optim_level, lr1, niter1, lr2, niter2, min_conf_thr, matching_conf_thr,\n",
        "                                                   as_pointcloud, mask_sky, clean_depth, transparent_cams, cam_size, scenegraph_type, winsize,\n",
        "                                                   win_cyclic, refid, TSDF_thresh=0, shared_intrinsics=shared_intrinsics, **kw)\n",
        "    \n",
        "    filestate = FileState(scene_state.outfile_name)\n",
        "    scene_state.outfile_name = None\n",
        "    del scene_state\n",
        "    return filestate, outfile\n",
        "\n",
        "def run_example(snapshot, matching_conf_thr, min_conf_thr, cam_size, as_pointcloud, shared_intrinsics, filelist, **kw):\n",
        "    return local_get_reconstructed_scene(filelist, min_conf_thr, matching_conf_thr, as_pointcloud, cam_size, shared_intrinsics, **kw)\n",
        "\n",
        "filelist = [\n",
        "    \"/content/1.jpg\",\n",
        "    \"/content/2.jpg\",\n",
        "    \"/content/3.jpg\",\n",
        "    \"/content/4.jpg\",\n",
        "    \"/content/5.jpg\",\n",
        "    \"/content/6.jpg\",\n",
        "    \"/content/7.jpg\"\n",
        "]\n",
        "matching_conf_thr = 2.0\n",
        "min_conf_thr = 1.5\n",
        "cam_size = 0.2\n",
        "as_pointcloud = False\n",
        "shared_intrinsics = True\n",
        "\n",
        "filestate, outfile = run_example(None, matching_conf_thr, min_conf_thr, cam_size, as_pointcloud, shared_intrinsics, filelist)\n",
        "print(f\"Reconstructed scene saved to: {outfile}\")\n",
        "!mv {outfile} /content/scene.glb"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "T4",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}