Update app.py

app.py

@@ -1,37 +1,250 @@
+# app.py
+from __future__ import annotations
+
+import datetime as dt
+import io
+import os
+import shutil
+import subprocess
+import textwrap
+import uuid
+from pathlib import Path
+from typing import List, Optional, Tuple
+
 import gradio as gr
 from PIL import Image
 
-def run_pipeline(gallery, max_px, match_mode, use_gpu_sift, voxel, depth, tris):
+# Optional: Open3D for meshing
+try:
+    import open3d as o3d
+except Exception:
+    o3d = None  # We’ll still run COLMAP and return the fused point cloud if meshing libs aren’t present.
+
+# Be gentle with HF CPU boxes that choke on many threads
+os.environ.setdefault("OMP_NUM_THREADS", "4")
+
+def _run(cmd: List[str], cwd: Optional[Path] = None, env: Optional[dict] = None) -> Tuple[int, str]:
+    """Run a subprocess and capture merged stdout/stderr as text."""
+    p = subprocess.run(
+        cmd,
+        cwd=str(cwd) if cwd else None,
+        env=env,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.STDOUT,
+        text=True,
+    )
+    return p.returncode, p.stdout
+
+def _ensure_tool(tool: str) -> bool:
+    return shutil.which(tool) is not None
+
+def _save_images(files: List[gr.File], out_dir: Path, max_px: int) -> None:
+    out_dir.mkdir(parents=True, exist_ok=True)
+    for f in files:
+        # gr.File returns a path str in .name on Spaces
+        src = Path(f.name)
+        with Image.open(src) as im:
+            im = im.convert("RGB")
+            w, h = im.size
+            scale = min(max_px / max(w, h), 1.0)
+            if scale < 1.0:
+                im = im.resize((int(w * scale), int(h * scale)))
+            im.save(out_dir / src.name, quality=92)
+
+def run_pipeline(
+    files: List[gr.File],
+    max_px: int,
+    match_mode: str,
+    use_gpu_sift: bool,
+    voxel: float,
+    depth: int,
+    tris: int,
+):
+    logs: List[str] = []
     try:
-        logs = []
-        # TODO: your pipeline here
-        # return (preview_path, [file_paths...], "\n".join(logs[-6:]), gr.update(visible=True))
-        return None, [], "Ready.", gr.update(visible=False)
+        if not files:
+            return None, [], "Please upload 3–30 images.", gr.update(visible=False)
+
+        if not _ensure_tool("colmap"):
+            return None, [], "COLMAP not found. Make sure `packages.txt` contains `colmap`.", gr.update(visible=False)
+
+        run_id = dt.datetime.now().strftime("run_%Y%m%d_%H%M%S_") + uuid.uuid4().hex[:8]
+        run_dir = Path("runs") / run_id
+        imgs_dir = run_dir / "images"
+        db = run_dir / "db.db"
+        sparse_dir = run_dir / "sparse"
+        dense_dir = run_dir / "dense"
+        run_dir.mkdir(parents=True, exist_ok=True)
+
+        logs.append(f"Workspace: {run_dir}")
+        _save_images(files, imgs_dir, max_px)
+        img_count = len(list(imgs_dir.glob("*")))
+        logs.append(f"Ingested {img_count} image(s). Max side capped at {max_px}px")
+
+        # 1) Features
+        feat_cmd = [
+            "colmap", "feature_extractor",
+            "--database_path", str(db),
+            "--image_path", str(imgs_dir),
+            "--ImageReader.single_camera", "1",
+            "--SiftExtraction.use_gpu", "1" if (use_gpu_sift and _ensure_tool("nvidia-smi")) else "0",
+        ]
+        code, out = _run(feat_cmd, cwd=run_dir)
+        logs.append("[feature_extractor]\n" + out)
+        if code != 0:
+            raise RuntimeError("COLMAP feature extraction failed.")
+
+        # 2) Matching
+        if match_mode == "sequential":
+            match_cmd = ["colmap", "sequential_matcher", "--database_path", str(db)]
+        elif match_mode == "exhaustive":
+            match_cmd = ["colmap", "exhaustive_matcher", "--database_path", str(db)]
+        else:
+            # Spatial matcher needs priors; default to exhaustive if none
+            match_cmd = ["colmap", "exhaustive_matcher", "--database_path", str(db)]
+        code, out = _run(match_cmd, cwd=run_dir)
+        logs.append(f"[{match_mode}_matcher]\n" + out)
+        if code != 0:
+            raise RuntimeError("COLMAP matching failed.")
+
+        # 3) Sparse reconstruction
+        sparse_dir.mkdir(exist_ok=True)
+        code, out = _run(
+            ["colmap", "mapper", "--database_path", str(db), "--image_path", str(imgs_dir), "--output_path", str(sparse_dir)],
+            cwd=run_dir,
+        )
+        logs.append("[mapper]\n" + out)
+        if code != 0 or not any((sparse_dir).glob("*/cameras.txt")):
+            raise RuntimeError("COLMAP mapper failed or produced no model.")
+
+        model_dirs = sorted(sparse_dir.glob("*"))
+        model_dir = model_dirs[0]
+
+        # 4) Undistort & dense
+        code, out = _run(
+            ["colmap", "image_undistorter", "--image_path", str(imgs_dir), "--input_path", str(model_dir), "--output_path", str(dense_dir), "--output_type", "COLMAP"],
+            cwd=run_dir,
+        )
+        logs.append("[image_undistorter]\n" + out)
+        if code != 0:
+            raise RuntimeError("Undistortion failed.")
+
+        code, out = _run(
+            ["colmap", "patch_match_stereo", "--workspace_path", str(dense_dir), "--workspace_format", "COLMAP", "--PatchMatchStereo.geom_consistency", "true"],
+            cwd=run_dir,
+        )
+        logs.append("[patch_match_stereo]\n" + out)
+        if code != 0:
+            raise RuntimeError("PatchMatch failed.")
+
+        fused = run_dir / "fused.ply"
+        code, out = _run(
+            ["colmap", "stereo_fusion", "--workspace_path", str(dense_dir), "--workspace_format", "COLMAP", "--input_type", "geometric", "--output_path", str(fused)],
+            cwd=run_dir,
+        )
+        logs.append("[stereo_fusion]\n" + out)
+        if code != 0 or not fused.exists():
+            raise RuntimeError("Fusion failed.")
+
+        # 5) Meshing (Open3D). If not available, just return fused point cloud.
+        mesh_paths = []
+        preview_path = fused  # default to point cloud preview
+
+        if o3d is not None:
+            pcd = o3d.io.read_point_cloud(str(fused))
+            if voxel and voxel > 0:
+                pcd = pcd.voxel_down_sample(voxel)
+            pcd.estimate_normals(o3d.geometry.KDTreeSearchParamKNN(knn=20))
+
+            # Poisson surface reconstruction
+            mesh, _ = o3d.geometry.TriangleMesh.create_from_point_cloud_poisson(pcd, depth=depth)
+            mesh.remove_degenerate_triangles()
+            mesh.remove_duplicated_triangles()
+            mesh.remove_duplicated_vertices()
+            mesh.remove_non_manifold_edges()
+
+            if tris and tris > 0:
+                mesh = mesh.simplify_quadric_decimation(tris)
+
+            mesh.compute_vertex_normals()
+
+            mesh_ply = run_dir / "mesh.ply"
+            mesh_obj = run_dir / "mesh.obj"
+            o3d.io.write_triangle_mesh(str(mesh_ply), mesh)
+            o3d.io.write_triangle_mesh(str(mesh_obj), mesh)
+            mesh_paths = [mesh_ply, mesh_obj]
+            preview_path = mesh_ply
+
+        files_out = [preview_path] + mesh_paths
+        file_list = [str(p) for p in files_out if Path(p).exists()]
+
+        return str(preview_path), file_list, "\n".join(logs[-80:]), gr.update(visible=True)
+
     except Exception as e:
-        return None, None, f"[ERROR]\n{e}", gr.update(visible=False)
+        logs.append("\n[ERROR]\n" + textwrap.fill(str(e), width=100))
+        return None, [], "\n".join(logs[-120:]), gr.update(visible=False)
+
+def build_ui():
+    with gr.Blocks(title="Sparse Multi-View 3D (Urban Planning)", theme=gr.themes.Soft()) as demo:
+        gr.Markdown(
+            """
+            # 🗺️ Sparse Multi-View 3D for Urban Planning
 
-def ui():
-    with gr.Blocks(title="Sparse Multi-View 3D (Urban Planning)") as demo:
-        gr.Markdown("# Sparse Multi-View 3D for Urban Planning")
+            Upload **3–30 photos** of a scene (streetscape, plaza, façade). We estimate camera poses with **COLMAP**,
+            build a **dense point cloud**, and (optionally) **mesh** it with Open3D.
+
+            **Tips for sparse captures:** overlap ~60–70%, vary viewpoint (walk an arc), avoid moving cars/people.
+            """
+        )
         with gr.Row():
             with gr.Column(scale=2):
-                imgs = gr.UploadButton("Upload images (JPG/PNG)", file_types=["image"], file_count="multiple")
-                gallery = gr.Gallery(label="Selected", columns=6, height=160)
-                imgs.upload(lambda files: [(f.name, Image.open(f).convert('RGB')) for f in files], imgs, gallery)
-                max_px = gr.Slider(1024, 4096, value=2400, step=64, label="Max image size (px)")
-                match_mode = gr.Radio(["exhaustive", "sequential", "spatial"], value="sequential", label="Matching mode")
-                use_gpu_sift = gr.Checkbox(True, label="Use GPU for SIFT (if available)")
-                voxel = gr.Slider(0.0, 0.05, value=0.01, step=0.005, label="Voxel downsample")
-                depth = gr.Slider(6, 12, value=9, step=1, label="Poisson depth")
-                tris = gr.Slider(0, 500_000, value=150_000, step=10_000, label="Target triangles")
-                run = gr.Button("▶ Reconstruct 3D")
+                images = gr.File(label="Upload images (JPG/PNG)", file_types=["image"], file_count="multiple")
+                gallery = gr.Gallery(label="Preview", columns=6, height=160)
+
+                def _show_gallery(files: List[gr.File]):
+                    rows = []
+                    for f in files or []:
+                        try:
+                            with Image.open(f.name) as im:
+                                rows.append((Path(f.name).name, im.convert("RGB")))
+                        except Exception:
+                            pass
+                    return rows
+
+                images.change(_show_gallery, inputs=images, outputs=gallery)
+
+                with gr.Accordion("Reconstruction settings", open=False):
+                    max_px = gr.Slider(1024, 4096, value=2400, step=64, label="Max image size (px, longest side)")
+                    match_mode = gr.Radio(["exhaustive", "sequential", "spatial"], value="sequential", label="Matching mode")
+                    use_gpu_sift = gr.Checkbox(True, label="Use GPU SIFT if available")
+
+                with gr.Accordion("Meshing", open=True):
+                    voxel = gr.Slider(0.0, 0.05, value=0.01, step=0.005, label="Voxel downsample (m, approx units)")
+                    depth = gr.Slider(6, 12, value=9, step=1, label="Poisson depth (higher → more detail)")
+                    tris = gr.Slider(0, 500_000, value=150_000, step=10_000, label="Target triangles (0 = keep)")
+
+                run = gr.Button("▶ Reconstruct 3D", variant="primary")
+
             with gr.Column(scale=1):
-                preview = gr.Model3D(label="Mesh preview", visible=False)
-                files_out = gr.Files(label="Downloads (GLB/OBJ/PLY)")
-                logs = gr.Markdown("Logs will appear here after running…")
-        run.click(run_pipeline, [gallery, max_px, match_mode, use_gpu_sift, voxel, depth, tris],
-                  [preview, files_out, logs, preview], queue=True)
+                preview = gr.Model3D(label="Preview (PLY/OBJ)", visible=False)
+                outputs = gr.Files(label="Downloads")
+                logs = gr.Markdown("Logs will appear here…")
+
+        run.click(
+            run_pipeline,
+            inputs=[images, max_px, match_mode, use_gpu_sift, voxel, depth, tris],
+            outputs=[preview, outputs, logs, preview],
+            queue=True,
+        )
+
+        gr.Markdown(
+            """
+            ### Notes & Scaling
+            - Results are in **arbitrary units** (SfM scale). For metric scale, align in GIS/CAD with known distances.
+            - Outdoor scenes with repetitive textures (glass/trees) can be challenging—add more oblique views if possible.
+            """
+        )
     return demo
 
 if __name__ == "__main__":
-    ui().launch()
+    build_ui().launch()