Update app.py

app.py

@@ -1,194 +1,249 @@
-# ===== Minimal app wrapper around your script logic =====
-# Steps: GLPN depth -> RGBD -> point cloud -> Poisson mesh -> save to outputs/
-# No extra options, no video. Matches a standard single-image pipeline.
+# Run-your-script (dynamic) — HF Spaces wrapper for main.py with user inputs
+import os, sys, io, time, glob, json, shlex, subprocess
+from pathlib import Path
+from typing import Optional, Tuple
 
-# ---- keep libgomp quiet (optional) ----
-import os, multiprocessing
+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")
-
-# ---- imports ----
-import io, time, traceback
-from pathlib import Path
-from typing import Tuple, Optional
-
-import numpy as np
-from PIL import Image
-import torch
-import gradio as gr
-import open3d as o3d
-from matplotlib import pyplot as plt
-from matplotlib.figure import Figure
-from transformers import GLPNForDepthEstimation, GLPNImageProcessor
-
-# ---- device & model ----
-DEVICE = torch.device(
-    "mps" if torch.backends.mps.is_available()
-    else ("cuda" if torch.cuda.is_available() else "cpu")
-)
-MODEL_ID = "vinvino02/glpn-nyu"
-_PROCESSOR: Optional[GLPNImageProcessor] = None
-_MODEL: Optional[GLPNForDepthEstimation] = None
-
-def _load_model_once() -> Tuple[GLPNImageProcessor, GLPNForDepthEstimation]:
-    global _PROCESSOR, _MODEL
-    if _PROCESSOR is None or _MODEL is None:
-        _PROCESSOR = GLPNImageProcessor.from_pretrained(MODEL_ID)
-        _MODEL = GLPNForDepthEstimation.from_pretrained(MODEL_ID).to(DEVICE)
-        _MODEL.eval()
-    return _PROCESSOR, _MODEL
-
-# ---- helpers mirroring script behavior ----
-def _depth_from_glpn(pil_img: Image.Image) -> np.ndarray:
-    """Return depth normalized to [0,1], upsampled to original size (matches typical GLPN usage)."""
-    processor, model = _load_model_once()
-
-    # your script-style resize: height <= 480 and multiple of 32
-    new_h = min(pil_img.height, 480)
-    new_h -= (new_h % 32)
-    new_w = int(new_h * pil_img.width / pil_img.height)
-    resized = pil_img.resize((new_w, new_h))
-
-    inputs = processor(images=resized, return_tensors="pt")
-    inputs = {k: v.to(DEVICE) for k, v in inputs.items()}
-
-    with torch.no_grad():
-        outputs = model(**inputs)           # (1, h, w)
-        pred = outputs.predicted_depth
-        # upsample logits back to original image size
-        pred = torch.nn.functional.interpolate(
-            pred.unsqueeze(1), size=pil_img.size[::-1],
-            mode="bicubic", align_corners=False
-        ).squeeze(1)
-    depth = pred[0].detach().float().cpu().numpy()
-    # normalize like a typical script
-    depth -= depth.min()
-    if depth.max() > 0:
-        depth /= depth.max()
-    return depth.astype(np.float32)
-
-def _depth_preview_img(depth01: np.ndarray) -> Image.Image:
-    fig: Figure = plt.figure(figsize=(5, 5), dpi=120)
-    ax = fig.add_subplot(111)
-    ax.axis("off")
-    ax.imshow(depth01, cmap="viridis")
-    buf = io.BytesIO()
-    plt.savefig(buf, format="png", bbox_inches="tight", pad_inches=0)
-    plt.close(fig)
-    buf.seek(0)
-    return Image.open(buf)
-
-def _make_rgbd(pil_img: Image.Image, depth_norm01: np.ndarray) -> o3d.geometry.RGBDImage:
-    # map normalized depth to a metric-like range (like a simple fixed scale in scripts)
-    depth_mm = (depth_norm01 * 3000.0).astype(np.float32)  # ~0–3 m span; matches common simple pipelines
-    depth_o3d = o3d.geometry.Image(depth_mm)
-    color_o3d = o3d.geometry.Image(np.array(pil_img.convert("RGB")))
-    return o3d.geometry.RGBDImage.create_from_color_and_depth(
-        color_o3d, depth_o3d,
-        convert_rgb_to_intensity=False,
-        depth_scale=1000.0,   # 1000 units == 1 meter
-        depth_trunc=10.0,     # ignore far depths
+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):
+    """
+    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'
+    """
+    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 = []
+
+    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"
     )
 
-def _pcd_from_rgbd(rgbd: o3d.geometry.RGBDImage) -> o3d.geometry.PointCloud:
-    h = np.asarray(rgbd.depth).shape[0]
-    w = np.asarray(rgbd.depth).shape[1]
-    # simple pinhole intrinsics (script-style constants)
-    fx = fy = 575.0
-    cx, cy = w / 2.0, h / 2.0
-    intr = o3d.camera.PinholeCameraIntrinsic(w, h, fx, fy, cx, cy)
-    pcd = o3d.geometry.PointCloud.create_from_rgbd_image(rgbd, intr)
-    # flip to a more natural orientation (Open3D default is +Z forward)
-    pcd.transform([[1,0,0,0],[0,-1,0,0],[0,0,-1,0],[0,0,0,1]])
-    return pcd
-
-def _pcd_clean_and_normals(pcd: o3d.geometry.PointCloud) -> o3d.geometry.PointCloud:
-    if len(pcd.points) == 0:
-        return pcd
-    # lightweight clean like a typical script
-    pcd, _ = pcd.remove_statistical_outlier(nb_neighbors=20, std_ratio=2.0)
-    pcd.estimate_normals(search_param=o3d.geometry.KDTreeSearchParamHybrid(radius=0.05, max_nn=30))
-    pcd.orient_normals_consistent_tangent_plane(10)
-    return pcd
-
-def _poisson_mesh(pcd: o3d.geometry.PointCloud, depth: int = 8) -> o3d.geometry.TriangleMesh:
-    mesh, _ = o3d.geometry.TriangleMesh.create_from_point_cloud_poisson(pcd, depth=depth)
-    mesh.remove_duplicated_vertices()
-    mesh.remove_degenerate_triangles()
-    mesh.remove_non_manifold_edges()
-    mesh.remove_unreferenced_vertices()
-    mesh.compute_vertex_normals()
-    return mesh
-
-# ---- core pipeline (1:1 with a standard single-image script) ----
-def pipeline(image: Image.Image):
-    t0 = time.time()
-    out_dir = Path("outputs")
-    out_dir.mkdir(parents=True, exist_ok=True)
-
-    # 1) depth
-    depth01 = _depth_from_glpn(image)
-    depth_png = _depth_preview_img(depth01)
-    depth_png_path = out_dir / "depth_preview.png"
-    depth_png.save(depth_png_path)
-
-    # 2) rgbd
-    rgbd = _make_rgbd(image, depth01)
-
-    # 3) point cloud
-    pcd = _pcd_from_rgbd(rgbd)
-    pcd = _pcd_clean_and_normals(pcd)
-    pcd_path = out_dir / "point_cloud.ply"
-    o3d.io.write_point_cloud(str(pcd_path), pcd)
-
-    # 4) mesh (Poisson)
-    mesh = _poisson_mesh(pcd, depth=8)   # matches common default in script examples
-    # try OBJ first for viewer compatibility; fallback to PLY
-    mesh_obj = out_dir / "mesh.obj"
-    if not o3d.io.write_triangle_mesh(str(mesh_obj), mesh):
-        mesh_obj = out_dir / "mesh.ply"
-        o3d.io.write_triangle_mesh(str(mesh_obj), mesh)
-
-    elapsed = time.time() - t0
-    log_text = f"Done in {elapsed:.1f}s\nSaved:\n- {depth_png_path}\n- {pcd_path}\n- {mesh_obj}"
-    return depth_png, str(mesh_obj), str(pcd_path), str(mesh_obj), log_text
-
-# ---- very small UI, just like running the script ----
-with gr.Blocks(title="2D → 3D (Script-Exact)") as demo:
-    gr.Markdown("### 2D → 3D Reconstruction — matches your Python script\nUpload an image → Depth • Point Cloud • Poisson Mesh (saved in `outputs/`).")
     with gr.Row():
         with gr.Column(scale=1):
-            img = gr.Image(type="pil", label="Input image", interactive=True)
-            run_btn = gr.Button("Run", variant="primary")
+            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="pil", label="Depth (preview)")
-                with gr.Tab("3D Mesh"):
-                    viewer = gr.Model3D(label="Interactive Mesh (OBJ/PLY)")
+                    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)")
+                    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="Log", lines=8)
+                    logs = gr.Textbox(label="Live logs", lines=20)
 
-    def _on_run(image):
-        if image is None:
-            return None, None, None, "Please upload an image."
-        try:
-            depth_im, model3d_path, pcd_path, mesh_path, log = pipeline(image)
-            return depth_im, model3d_path, pcd_path, mesh_path, log
-        except Exception as e:
-            tb = traceback.format_exc()
-            return None, None, None, None, f"ERROR: {e}\n\n{tb}"
-
-    run_btn.click(_on_run, inputs=[img], outputs=[depth_img, viewer, pcd_file, mesh_file, logs])
+    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]
+    )
 
-# Optional queue to keep the session stable on small CPUs
-demo.queue(concurrency_count=1, max_size=4, status_update_rate=1.0)
+# 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)
+    demo.launch(show_error=True, server_keepalive_timeout=180)