Update app.py

app.py

@@ -1,249 +1,232 @@
-# Run-your-script (dynamic) — HF Spaces wrapper for main.py with user inputs
-import os, sys, io, time, glob, json, shlex, subprocess
+import os
+os.environ.setdefault("OMP_NUM_THREADS", "1")  # silence libgomp spam on HF
+os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
+
 from pathlib import Path
-from typing import Optional, Tuple
+import io
+import numpy as np
+from PIL import Image
+
+import torch
+from transformers import GLPNForDepthEstimation, GLPNImageProcessor
 
+import open3d as o3d
 import gradio as gr
-from PIL import Image
 
-# Keep CPU runtimes stable
-os.environ.pop("OMP_NUM_THREADS", None)
-os.environ.setdefault("OMP_NUM_THREADS", "1")
-os.environ.setdefault("OPENBLAS_NUM_THREADS", "1")
-os.environ.setdefault("MKL_NUM_THREADS", "1")
-os.environ.setdefault("NUMEXPR_NUM_THREADS", "1")
-os.environ.setdefault("MPLBACKEND", "Agg")
-
-REPO_ROOT = Path(".").resolve()
-
-# ---- Defaults: adjust to match your script, or override in the UI ----
-DEFAULT_SCRIPT      = str(REPO_ROOT / "main.py")
-DEFAULT_INPUT_PATH  = str(REPO_ROOT / "ROOM.jpg")   # where we'll save the uploaded image
-DEFAULT_WORKDIR     = str(REPO_ROOT)
-DEFAULT_OUTPUT_DIR  = str(REPO_ROOT / "outputs")    # where your script writes results
-
-# ---------- helpers ----------
-def _save_image(img: Image.Image, path: str) -> str:
-    p = Path(path); p.parent.mkdir(parents=True, exist_ok=True)
-    img.convert("RGB").save(p, format="JPEG", quality=95)
-    return str(p)
-
-def _pick_latest(patterns):
-    newest = None; mt = -1
-    for pat in patterns:
-        for fp in glob.glob(pat):
-            try:
-                sz = os.path.getsize(fp)
-                if sz <= 0: continue
-                m = os.path.getmtime(fp)
-                if m > mt:
-                    newest, mt = fp, m
-            except Exception:
-                pass
-    return newest
-
-def _scan_outputs(output_dir: str):
-    od = Path(output_dir)
-    depth = _pick_latest([
-        str(od / "depth_preview.*"),
-        str(od / "*depth*.png"),
-        str(od / "*depth*.jpg"),
-        str(REPO_ROOT / "depth_preview.*"),
-    ])
-    pcd = _pick_latest([
-        str(od / "point_cloud.ply"),
-        str(od / "*.ply"),
-    ])
-    mesh = _pick_latest([
-        str(od / "mesh.obj"),
-        str(od / "*.obj"),
-        str(od / "mesh.ply"),
-        str(od / "*mesh*.ply"),
-        str(od / "*.glb"),
-        str(od / "*.gltf"),
-    ])
-    return depth, pcd, mesh
-
-def _compose_cli(script_path: str, base_args: str, kv_pairs: str):
+
+# ----------------------------
+# Device & model (load once)
+# ----------------------------
+DEVICE = torch.device(
+    "cuda" if torch.cuda.is_available()
+    else ("mps" if getattr(torch.backends, "mps", None) and torch.backends.mps.is_available() else "cpu")
+)
+PROCESSOR = GLPNImageProcessor.from_pretrained("vinvino02/glpn-nyu")
+MODEL = GLPNForDepthEstimation.from_pretrained("vinvino02/glpn-nyu").to(DEVICE).eval()
+
+
+# ----------------------------
+# Helpers (faithful to main.py logic)
+# ----------------------------
+def _resize_like_main(pil_img: Image.Image, cap_h: int = 480):
+    """Mirror your main.py: cap height at 480, then round down to multiple of 32, preserve aspect."""
+    new_h = min(pil_img.height, cap_h)
+    new_h -= (new_h % 32)
+    if new_h < 32:
+        new_h = 32
+    new_w = int(new_h * pil_img.width / pil_img.height)
+    return pil_img.resize((new_w, new_h), Image.BILINEAR), (pil_img.width, pil_img.height)
+
+
+@torch.inference_mode()
+def estimate_depth_glpn(pil_img: Image.Image) -> np.ndarray:
+    """
+    GLPN forward that DOES NOT rely on .post_process_depth()
+    (fix for your AttributeError). We upsample back to the original size manually.
+    Returns depth in float32 (larger = farther).
+    """
+    resized, (orig_w, orig_h) = _resize_like_main(pil_img)
+    inputs = PROCESSOR(images=resized, return_tensors="pt")
+    for k in inputs:
+        inputs[k] = inputs[k].to(DEVICE)
+
+    outputs = MODEL(**inputs)
+    pred = outputs.predicted_depth  # [B, 1, h, w]
+    depth = pred[0, 0].float().detach().cpu().numpy()  # resized size
+
+    # Resize depth back to original image size for downstream Open3D steps
+    depth_img = Image.fromarray(depth)
+    depth_full = depth_img.resize((orig_w, orig_h), Image.BILINEAR)
+    depth_full = np.array(depth_full).astype(np.float32)
+
+    return depth_full
+
+
+def depth_vis(depth: np.ndarray) -> Image.Image:
+    """Normalize depth to 0..255 for a PNG preview (like your matplotlib preview)."""
+    d = depth.copy()
+    d = d - np.nanmin(d)
+    maxv = np.nanmax(d)
+    if maxv <= 0:
+        maxv = 1.0
+    d = (255.0 * d / maxv).astype(np.uint8)
+    return Image.fromarray(d)
+
+
+def rgbd_from_rgb_depth(rgb: Image.Image, depth_f32: np.ndarray) -> o3d.geometry.RGBDImage:
+    """
+    Create Open3D RGBD using an 8-bit depth *preview* for visualization consistency
+    (same as your main.py normalization step).
+    """
+    rgb_np = np.array(rgb)
+    # match your main.py: depth to 0..255 uint8 before feeding create_from_color_and_depth
+    d8 = (depth_f32 * 255.0 / (depth_f32.max() + 1e-8)).astype(np.uint8)
+    depth_o3d = o3d.geometry.Image(d8)
+    color_o3d = o3d.geometry.Image(rgb_np)
+    rgbd = o3d.geometry.RGBDImage.create_from_color_and_depth(
+        color_o3d, depth_o3d, convert_rgb_to_intensity=False
+    )
+    return rgbd
+
+
+def pointcloud_from_rgbd(rgbd: o3d.geometry.RGBDImage, w: int, h: int) -> o3d.geometry.PointCloud:
     """
-    base_args: free-form CLI string (e.g., "--poisson_depth 10 --out outputs")
-    kv_pairs:  JSON or 'key=value key2=value2' → becomes '--key value --key2 value2'
+    Reproduce your simple pinhole intrinsics (fx=fy=500, cx=w/2, cy=h/2) and back-project.
     """
-    args = [sys.executable, script_path]
-
-    # Add free-form args (if provided)
-    if base_args and base_args.strip():
-        args.extend(shlex.split(base_args.strip()))
-
-    # Add key=value pairs
-    if kv_pairs and kv_pairs.strip():
-        # try JSON first
-        as_json = None
-        try:
-            as_json = json.loads(kv_pairs)
-        except Exception:
-            pass
-        if isinstance(as_json, dict):
-            for k, v in as_json.items():
-                if k.startswith("--"): args.append(k)
-                else: args.append(f"--{k}")
-                if v is not True and v is not None:
-                    args.append(str(v))
-        else:
-            # fallback: split by spaces, accept k=v tokens
-            for token in shlex.split(kv_pairs.strip()):
-                if "=" in token:
-                    k, v = token.split("=", 1)
-                    if k.startswith("--"): args.append(k)
-                    else: args.append(f"--{k}")
-                    args.append(v)
-                else:
-                    # allow plain flags like --use_poisson
-                    args.append(token)
-
-    return args
-
-# ---------- streaming runner ----------
-def _run_streaming(
-    image,
-    script_path,
-    input_path,
-    workdir,
-    output_dir,
-    freeform_args,      # raw CLI string
-    kv_args,            # k=v pairs or JSON
-    extra_env_json      # ENV as JSON (optional)
-):
-    depth_path = None; pcd_path = None; mesh_path = None
-    viewer_path = None
-    log_buf = []
+    K = o3d.camera.PinholeCameraIntrinsic()
+    K.set_intrinsics(w, h, 500.0, 500.0, w / 2.0, h / 2.0)
+    pcd = o3d.geometry.PointCloud.create_from_rgbd_image(rgbd, K)
+    return pcd
 
+
+def filter_pointcloud(pcd: o3d.geometry.PointCloud):
+    """
+    Statistical outlier removal ~ your 'noise removal' step. Tuned conservatively.
+    """
+    if len(pcd.points) == 0:
+        return pcd
+    cl, ind = pcd.remove_statistical_outlier(nb_neighbors=20, std_ratio=2.0)
+    pcd_f = pcd.select_by_index(ind)
+    pcd_f.estimate_normals(
+        search_param=o3d.geometry.KDTreeSearchParamHybrid(radius=0.05, max_nn=30)
+    )
+    return pcd_f
+
+
+def poisson_mesh(pcd: o3d.geometry.PointCloud, rotate_up=True) -> o3d.geometry.TriangleMesh:
+    if len(pcd.points) == 0:
+        return o3d.geometry.TriangleMesh()
+    mesh, _ = o3d.geometry.TriangleMesh.create_from_point_cloud_poisson(
+        pcd, depth=10, n_threads=1
+    )
+    # Flip like your main.py (rotate X by pi)
+    if rotate_up:
+        R = mesh.get_rotation_matrix_from_xyz((np.pi, 0.0, 0.0))
+        mesh.rotate(R, center=(0, 0, 0))
+    mesh.compute_vertex_normals()
+    return mesh
+
+
+def o3d_to_ply_bytes(geom: o3d.geometry.Geometry) -> bytes:
+    """Serialize an Open3D geometry to .ply bytes (so Gradio can offer a download)."""
+    tmp = Path("tmp_out.ply")
+    if isinstance(geom, o3d.geometry.PointCloud):
+        o3d.io.write_point_cloud(str(tmp), geom)
+    else:
+        o3d.io.write_triangle_mesh(str(tmp), geom)
+    data = tmp.read_bytes()
+    tmp.unlink(missing_ok=True)
+    return data
+
+
+def render_point_count(pcd: o3d.geometry.PointCloud) -> str:
+    return f"Points: {len(pcd.points):,}"
+
+
+def render_face_count(mesh: o3d.geometry.TriangleMesh) -> str:
+    return f"Vertices: {len(mesh.vertices):,}  |  Triangles: {len(mesh.triangles):,}"
+
+
+# ----------------------------
+# Gradio pipeline
+# ----------------------------
+def pipeline(image: Image.Image):
+    logs = []
     if image is None:
-        yield None, None, None, None, "Please upload an image."
-        return
-
-    # Save input where the script expects it
-    try:
-        saved = _save_image(image, input_path)
-        log_buf.append(f"[app] Saved input → {saved}")
-    except Exception as e:
-        yield None, None, None, None, f"[Save error] {e}"
-        return
-
-    # Compose CLI
-    try:
-        args = _compose_cli(script_path, freeform_args, kv_args)
-        # If the script expects a positional image path, add it here (uncomment if needed):
-        # args.extend([saved])
-        log_buf.append(f"[app] Running: {' '.join(shlex.quote(a) for a in args)}")
-    except Exception as e:
-        yield None, None, None, None, f"[Args error] {e}"
-        return
-
-    # Build environment
-    env = os.environ.copy()
-    if extra_env_json and extra_env_json.strip():
-        try:
-            env.update(json.loads(extra_env_json))
-        except Exception as e:
-            yield None, None, None, None, f"[ENV JSON parse error] {e}"
-            return
-
-    # Launch process, stream logs
-    try:
-        proc = subprocess.Popen(
-            args, cwd=workdir, env=env,
-            stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
-            text=True, bufsize=1
-        )
-    except Exception as e:
-        yield None, None, None, None, f"[Run error] {e}"
-        return
-
-    last_yield = time.time()
-    for line in iter(proc.stdout.readline, ""):
-        log_buf.append(line.rstrip("\n"))
-        if time.time() - last_yield > 1.0:
-            d, p, m = _scan_outputs(output_dir)
-            depth_path = depth_path or d
-            pcd_path   = pcd_path or p
-            mesh_path  = mesh_path or m
-            viewer_path = mesh_path or pcd_path
-            yield depth_path, viewer_path, pcd_path, mesh_path, "\n".join(log_buf[-800:])
-            last_yield = time.time()
-
-    proc.wait()
-
-    # Final scan
-    d, p, m = _scan_outputs(output_dir)
-    depth_path = depth_path or d
-    pcd_path   = pcd_path or p
-    mesh_path  = mesh_path or m
-    viewer_path = mesh_path or pcd_path
-    log_buf.append(f"[app] Script finished with return code {proc.returncode}")
-
-    yield depth_path, viewer_path, pcd_path, mesh_path, "\n".join(log_buf[-2000:])
-
-# ---------- UI ----------
-with gr.Blocks(title="Run main.py — Dynamic Inputs") as demo:
-    gr.Markdown(
-        "## Run your `main.py` with dynamic user inputs\n"
-        "- Upload an image (we’ll save it to the path your script expects)\n"
-        "- Enter **CLI arguments** and/or **key=value** pairs (auto-converted to `--key value`)\n"
-        "- We stream stdout/stderr live and show any depth/PCD/mesh files your script writes\n"
+        raise gr.Error("Please upload an image of a room.")
+
+    logs.append("Step 1 — Loaded image.")
+    image = image.convert("RGB")
+    w, h = image.size
+
+    # Depth
+    logs.append("Step 2 — Estimating depth with GLPN (vinvino02/glpn-nyu)…")
+    depth = estimate_depth_glpn(image)
+    depth_preview = depth_vis(depth)
+
+    # RGBD
+    logs.append("Step 3 — Creating RGBD image…")
+    rgbd = rgbd_from_rgb_depth(image, depth)
+
+    # Point cloud
+    logs.append("Step 4 — Back-projecting to point cloud…")
+    pcd = pointcloud_from_rgbd(rgbd, w, h)
+
+    logs.append("Step 5 — Filtering noise & estimating normals…")
+    pcd_f = filter_pointcloud(pcd)
+
+    # Mesh
+    logs.append("Step 6 — Poisson surface reconstruction…")
+    mesh = poisson_mesh(pcd_f, rotate_up=True)
+
+    # Prepare downloads
+    logs.append("Step 7 — Preparing downloads…")
+    pcd_bytes = o3d_to_ply_bytes(pcd_f)
+    mesh_bytes = o3d_to_ply_bytes(mesh)
+
+    # Small text stats
+    pcd_stats = render_point_count(pcd_f)
+    mesh_stats = render_face_count(mesh)
+
+    logs.append("Done.")
+
+    return (
+        image,             # RGB preview
+        depth_preview,     # Depth preview
+        pcd_stats,         # point cloud stats
+        mesh_stats,        # mesh stats
+        ("point_cloud.ply", pcd_bytes),
+        ("mesh.ply", mesh_bytes),
+        "\n".join(logs),
     )
 
+
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# 2D → 3D (GLPN → RGBD → Point Cloud → Poisson Mesh)\nUpload a single image to reproduce your main.py workflow.")
+
     with gr.Row():
-        with gr.Column(scale=1):
-            img = gr.Image(type="pil", label="Upload image", interactive=True)
-
-            with gr.Accordion("Script paths", open=False):
-                script_path = gr.Textbox(value=DEFAULT_SCRIPT, label="Script path (e.g., main.py)")
-                input_path  = gr.Textbox(value=DEFAULT_INPUT_PATH, label="Save uploaded image to (path your script reads)")
-                workdir     = gr.Textbox(value=DEFAULT_WORKDIR, label="Working directory")
-                output_dir  = gr.Textbox(value=DEFAULT_OUTPUT_DIR, label="Output directory to scan")
-
-            with gr.Accordion("Arguments", open=True):
-                freeform_args = gr.Textbox(
-                    value="",
-                    placeholder="e.g., --poisson_depth 10 --out outputs",
-                    label="CLI arguments (free-form)"
-                )
-                kv_args = gr.Textbox(
-                    value="",
-                    placeholder='JSON or k=v (space-separated). e.g., {"poisson_depth":10, "out":"outputs"} or poisson_depth=10 out=outputs',
-                    label="Key=Value (auto → --key value)"
-                )
-
-            with gr.Accordion("Environment (optional)", open=False):
-                extra_env = gr.Textbox(
-                    value="{}",
-                    label="ENV as JSON",
-                    placeholder='e.g., {"OMP_NUM_THREADS":"1"}'
-                )
-
-            run_btn = gr.Button("Run script", variant="primary")
-
-        with gr.Column(scale=2):
-            with gr.Tabs():
-                with gr.Tab("Depth"):
-                    depth_img = gr.Image(type="filepath", label="Depth preview (detected)")
-                with gr.Tab("3D Reconstruction"):
-                    model3d = gr.Model3D(label="Mesh / Point Cloud (OBJ/PLY/GLB/GLTF)")
-                with gr.Tab("Downloads"):
-                    pcd_file  = gr.File(label="Point cloud (PLY)")
-                    mesh_file = gr.File(label="Mesh (OBJ/PLY/GLB/GLTF)")
-                with gr.Tab("Logs"):
-                    logs = gr.Textbox(label="Live logs", lines=20)
-
-    run_btn.click(
-        _run_streaming,
-        inputs=[img, script_path, input_path, workdir, output_dir, freeform_args, kv_args, extra_env],
-        outputs=[depth_img, model3d, pcd_file, mesh_file, logs]
+        with gr.Column():
+            inp = gr.Image(type="pil", label="Input Image")
+            run = gr.Button("Reconstruct 3D", variant="primary")
+            log_box = gr.Textbox(label="Log", lines=14, interactive=False)
+
+        with gr.Column():
+            rgb_out = gr.Image(label="RGB Preview", interactive=False)
+            depth_out = gr.Image(label="Depth Preview (8-bit normalized)", interactive=False)
+
+            pc_txt = gr.Markdown()
+            mesh_txt = gr.Markdown()
+
+            pc_file = gr.File(label="Download Point Cloud (.ply)")
+            mesh_file = gr.File(label="Download Mesh (.ply)")
+
+    run.click(
+        fn=pipeline,
+        inputs=[inp],
+        outputs=[rgb_out, depth_out, pc_txt, mesh_txt, pc_file, mesh_file, log_box],
+        api_name="reconstruct",
     )
 
-# Keep long jobs alive & serialized
-demo.queue(concurrency_count=1, max_size=8, status_update_rate=1.0)
-
-if __name__ == "__main__":
-    demo.launch(show_error=True, server_keepalive_timeout=180)
+# IMPORTANT: older Spaces error came from using unsupported args like concurrency_count.
+demo.queue()  # default queue works across Gradio 4.x
+demo.launch()