Spaces:
Sleeping
Sleeping
File size: 4,485 Bytes
bc4c433 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 | """Gradio occupancy loop: fill an OBJ AABB, then classify with ``best.pt``.
Reuses the viewer helper (``obj_fill``, ``infer_job``, ``model_access``).
Does not import Gradio and does not change the Three.js viewer.
"""
from __future__ import annotations
import base64
from pathlib import Path
from typing import Any
import numpy as np
# Occupancy is the installed ``scatteringnet`` package (``pip install -e .``).
# This folder is not a package so it cannot shadow pip ``gradio``.
_GRADIO_DIR = Path(__file__).resolve().parent
_REPO = _GRADIO_DIR.parents[1]
from scatteringnet.config import load_config
from scatteringnet.viewer.infer_job import infer_uploaded_obj, pred_from_probs # noqa: E402
from scatteringnet.viewer.model_access import inspect_run_id, list_viewer_models # noqa: E402
from scatteringnet.viewer.obj_fill import ( # noqa: E402
fill_aabb_lattice,
spacing_from_slider,
triangles_from_obj_text,
)
# INSPECT alias from docs/inspect_checkpoint.yaml (not the newest best.pt).
INSPECT_RUN_ID = inspect_run_id()
# Viewer allows 200k; Plotly + Spaces need a tighter lattice.
MAX_FILL_POINTS = 80_000
MAX_OBJ_BYTES = 32 * 1024 * 1024
# Slider default matches the Three.js Density control (spacing ~0.15).
DEFAULT_DENSITY = 71
DEFAULT_CUT = 0.50
def models_root() -> Path:
"""Repo ``models/`` (``models/<run_id>/best.pt``)."""
return _REPO / "models"
def list_run_ids(root: Path | None = None) -> list[str]:
"""Newest-first run folder names that have a ``best.pt``."""
rows = list_viewer_models(root or models_root())
return [str(row["id"]) for row in rows]
def default_run_id(root: Path | None = None) -> str:
"""INSPECT pointer if that ``best.pt`` exists, else newest, else empty."""
ids = list_run_ids(root)
wanted = inspect_run_id()
if wanted and wanted in ids:
return wanted
return ids[0] if ids else ""
def runtime_cfg(cfg: Any | None = None):
"""
Device / batch from YAML. ``data_dir`` is not required (uploaded OBJ).
"""
if cfg is not None:
return cfg
return load_config(require_existing_data_dir=False)
def apply_cut(probs: np.ndarray, threshold: float) -> tuple[np.ndarray, int, int]:
"""Hard labels from stored sigmoid probs (no extra GPU pass)."""
pred = pred_from_probs(probs, threshold)
n_in = int(np.count_nonzero(pred > 0))
return pred, n_in, int(pred.shape[0]) - n_in
def fill_and_infer(
obj_text: str,
*,
obj_name: str,
run_id: str,
density: float,
models: Path | None = None,
cfg: Any | None = None,
) -> dict[str, Any]:
"""
Job B for Gradio: AABB lattice + occupancy forward.
Returns vertices/faces (for the Plotly mesh), query XYZ, sigmoid
probs, timings, and the run id. The UI recuts ``probs`` locally.
"""
raw = str(obj_text or "")
if not raw.strip():
raise ValueError("OBJ is empty")
if len(raw.encode("utf-8")) > MAX_OBJ_BYTES:
raise ValueError(f"OBJ too large (max {MAX_OBJ_BYTES} bytes)")
name = str(run_id or "").strip()
if not name:
raise ValueError("select a checkpoint under models/<run_id>/best.pt")
vertices, faces = triangles_from_obj_text(raw)
spacing = spacing_from_slider(density)
points, used, grid = fill_aabb_lattice(
vertices, spacing, max_points=MAX_FILL_POINTS
)
ckpt_root = Path(models) if models is not None else models_root()
out = infer_uploaded_obj(
run_id=name,
models_root=ckpt_root,
data_dir=None,
obj_name=str(obj_name or "upload.obj"),
obj_text=raw,
points=points,
cfg=runtime_cfg(cfg),
)
probs = np.frombuffer(base64.b64decode(out["prob_b64"]), dtype=np.float32).reshape(
-1
)
pred, n_in, n_out = apply_cut(probs, DEFAULT_CUT)
return {
"vertices": np.ascontiguousarray(vertices, dtype=np.float32),
"faces": np.ascontiguousarray(faces, dtype=np.int32),
"points": np.ascontiguousarray(points, dtype=np.float32),
"probs": np.ascontiguousarray(probs, dtype=np.float32),
"pred": pred,
"n": int(out["n"]),
"n_inside": n_in,
"n_outside": n_out,
"used_spacing": float(used),
"grid": [int(grid[0]), int(grid[1]), int(grid[2])],
"run_id": str(out["run_id"]),
"shape_encoder": str(out["shape_encoder"]),
"timings": dict(out["timings"]),
"obj_name": str(obj_name or "upload.obj"),
}
|