Add initial project files including assets, styles, and utility functions

.gitignore

@@ -0,0 +1,5 @@
+# Byte-compiled / optimized / DLL files
+venv/
+dust3r/checkpoints/
+dust3r/dust3r.egg-info/
+__pycache__/

assets/fullscreen.js

@@ -0,0 +1,12 @@
+() => {
+    const el = document.getElementById("model-container");
+    if (!el) return;
+
+    if (el.requestFullscreen) {
+        el.requestFullscreen();
+    } else if (el.webkitRequestFullscreen) {
+        el.webkitRequestFullscreen();
+    } else if (el.msRequestFullscreen) {
+        el.msRequestFullscreen();
+    }
+}

assets/style.css

@@ -0,0 +1,53 @@
+footer {display: none !important;}
+#gradio-menu, .built-with, .api-link, #settings-button {display: none !important;}
+
+:root {
+    --primary-500: #FFFFFF !important;
+    --body-background-fill: #000000 !important;
+    --block-background-fill: #000000 !important;
+    --input-background-fill: #000000 !important;
+    --border-color-primary: #333333 !important;
+    --background-fill-secondary: #000000 !important;
+}
+
+.gradio-container {
+    background-color: #000000 !important;
+    color: #FFFFFF !important;
+    font-family: 'Inter', system-ui, sans-serif !important;
+}
+
+button.primary {
+    background-color: #FFFFFF !important;
+    color: #000000 !important;
+    border-radius: 0px !important;
+    font-weight: 600 !important;
+    text-transform: uppercase;
+    letter-spacing: 1px;
+}
+
+button.primary:hover {
+    background-color: #B2B2B2 !important;
+}
+
+.label {
+    color: #808080 !important;
+    text-transform: uppercase;
+    font-size: 11px !important;
+}
+
+#model-container:fullscreen {
+    background-color: black;
+    width: 100vw;
+    height: 100vh;
+}
+
+.generating::after {
+    color: #808080;
+    font-size: 10px;
+    letter-spacing: 2px;
+    animation: blink 1.2s infinite;
+}
+
+@keyframes blink {
+    50% { opacity: 0; }
+}

gradio_ui.py

@@ -0,0 +1,57 @@
+import gradio as gr
+import functools
+from pathlib import Path
+from pipeline import run_pipeline
+
+ASSETS = Path(__file__).parent / "assets"
+CSS = (ASSETS / "style.css").read_text()
+FULLSCREEN_JS = (ASSETS / "fullscreen.js").read_text()
+
+
+def build_ui(outdir, model, device):
+
+    pipeline = functools.partial(run_pipeline, outdir, model, device, 512)
+
+    with gr.Blocks(
+        title="3D Object Reconstruction",
+        css=CSS,
+        theme=gr.themes.Base(),
+        fill_width=True,
+    ) as app:
+
+        gr.Markdown("# 3D Object Reconstruction")
+
+        with gr.Row():
+
+            with gr.Column(scale=1):
+                files = gr.File(file_count="multiple", label="Images")
+                run_btn = gr.Button("Run Inference", variant="primary")
+
+                with gr.Accordion("Settings", open=False):
+                    iters = gr.Slider(100, 1000, 300, step=50, label="Alignment Iteration")
+                    as_pc = gr.Checkbox(True, label="Render as Point Cloud")
+                    refine = gr.Checkbox(True, label="Filter Background Points")
+                    clean = gr.Checkbox(True, label="Clean-up depthmaps")
+
+            with gr.Column(scale=2):
+                model3d = gr.Model3D(
+                    label="3D Output",
+                    height=600,
+                    elem_id="model-container",
+                )
+                fs_btn = gr.Button("Toggle Full Screen ⛶", size="sm")
+
+        gr.Markdown("---")
+        gallery = gr.Gallery(columns=3, label="RGB | DEPTH | CONFIDENCE")
+
+        fs_btn.click(None, None, None, js=FULLSCREEN_JS)
+
+        state = gr.State()
+        run_btn.click(
+            fn=pipeline,
+            inputs=[files, iters, as_pc, refine, clean],
+            outputs=[state, model3d, gallery],
+            show_progress="minimal",
+        )
+
+    return app

model.py

@@ -0,0 +1,18 @@
+import torch
+
+def initialize(model_path, device):
+    ckpt = torch.load(model_path, map_location="cpu", weights_only=False)
+    args = ckpt["args"].model.replace(
+        "ManyAR_PatchEmbed", "PatchEmbedDust3R"
+    )
+
+    if "landscape_only" not in args:
+        args = args[:-1] + ", landscape_only=False)"
+    else:
+        args = args.replace(" ", "").replace(
+            "landscape_only=True", "landscape_only=False"
+        )
+
+    net = eval(args)
+    net.load_state_dict(ckpt["model"], strict=False)
+    return net.to(device)

pipeline.py

@@ -0,0 +1,49 @@
+import torch, copy, numpy as np, trimesh, matplotlib.pyplot as plt
+from dust3r.image_pairs import make_pairs
+from dust3r.utils.image import load_images, rgb
+from dust3r.utils.device import to_numpy, to_cpu, collate_with_cat as collate
+from dust3r.viz import pts3d_to_trimesh, cat_meshes
+from dust3r.cloud_opt import global_aligner, GlobalAlignerMode
+from utils import preprocess
+
+BATCH_SIZE = 1
+
+def interleave(img1, img2):
+    out = {}
+    for k, v1 in img1.items():
+        v2 = img2[k]
+        if isinstance(v1, torch.Tensor):
+            out[k] = torch.stack((v1, v2), dim=1).flatten(0, 1)
+        else:
+            out[k] = [x for p in zip(v1, v2) for x in p]
+    return out
+
+def inference(pairs, model, device):
+    results = []
+    for i in range(0, len(pairs), BATCH_SIZE):
+        batch = collate(pairs[i:i+BATCH_SIZE])
+        for view in batch:
+            for k in ["img", "pts3d", "valid_mask", "camera_pose", "camera_intrinsics"]:
+                if k in view:
+                    view[k] = view[k].to(device)
+        v1, v2 = batch
+        v1, v2 = interleave(v1, v2), interleave(v2, v1)
+
+        with torch.cuda.amp.autocast():
+            p1, p2 = model(v1, v2)
+        results.append(to_cpu(dict(view1=v1, view2=v2, pred1=p1, pred2=p2)))
+
+    return collate(results, lists=True)
+
+def create_glb(outdir, scene):
+    meshes = [
+        pts3d_to_trimesh(scene.imgs[i], scene.get_pts3d()[i], scene.get_masks()[i])
+        for i in range(len(scene.imgs))
+    ]
+    mesh = trimesh.Trimesh(**cat_meshes(meshes))
+    mesh.apply_translation(-mesh.centroid)
+
+    scene_out = trimesh.Scene(mesh)
+    out = f"{outdir}/object.glb"
+    scene_out.export(out)
+    return out

setup.sh

@@ -10,6 +10,8 @@ fi
 
 pip install -r requirements.txt
 
+pip install https://github.com/camenduru/wheels/releases/download/colab/curope-0.0.0-cp310-cp310-linux_x86_64.whl
+
 mkdir -p dust3r/checkpoints
 
 WEIGHTS=dust3r/checkpoints/DUSt3R_ViTLarge_BaseDecoder_512_dpt.pth

utils.py

@@ -0,0 +1,14 @@
+import os, tempfile
+from PIL import Image, ImageEnhance
+
+def preprocess(image_paths):
+    cleaned = []
+    for i, path in enumerate(image_paths):
+        img = Image.open(path).convert("RGB")
+        img = ImageEnhance.Contrast(img).enhance(1.2)
+        img = ImageEnhance.Sharpness(img).enhance(1.5)
+
+        out = os.path.join(tempfile.gettempdir(), f"input_refined_{i}.png")
+        img.save(out)
+        cleaned.append(out)
+    return cleaned