Spaces:
Running on CPU Upgrade
Running on CPU Upgrade
| """Sand3D Explorer: a JSON API over the catalog, and the static page that reads it. | |
| Environment: `HF_TOKEN` (a Space secret, read access to the private datasets), `EXPLORER_DATA` | |
| (where the catalog is written, default /tmp/sand3d-explorer), `EXPLORER_SPLITS` (comma list, to | |
| load only some splits while developing), `EXPLORER_ASSETS_DIR` (a local copy of the assets | |
| instead of the Hub's). | |
| """ | |
| from __future__ import annotations | |
| import json | |
| import os | |
| import threading | |
| from pathlib import Path | |
| from fastapi import FastAPI, HTTPException, Query | |
| from fastapi.responses import FileResponse, JSONResponse, Response | |
| from fastapi.staticfiles import StaticFiles | |
| from . import search | |
| from .catalog import SOURCES, Catalog, Dataset | |
| HERE = Path(__file__).resolve().parent.parent | |
| CONFIG = json.loads((HERE / "config.json").read_text()) | |
| STATIC = HERE / "static" | |
| PAGE = 60 | |
| cat = Catalog([Dataset(**d) for d in CONFIG["datasets"]], | |
| Path(os.environ.get("EXPLORER_DATA", "/tmp/sand3d-explorer"))) | |
| encoder = search.Encoder(None) | |
| cache = search.Cache() | |
| def _load() -> None: | |
| try: | |
| splits = tuple(s for s in os.environ.get("EXPLORER_SPLITS", "").split(",") if s) or None | |
| adir = os.environ.get("EXPLORER_ASSETS_DIR") | |
| cat.load(os.environ.get("HF_TOKEN"), CONFIG.get("assets_repo"), splits, Path(adir) if adir else None) | |
| if cat.emb is not None: | |
| encoder.spec = getattr(cat, "model", None) | |
| cat.status["stage"] = "loading the text encoder" | |
| encoder.load() | |
| cat.status["stage"] = "ready" | |
| cat.status["encoder"] = encoder.model is not None | |
| except Exception as exc: # noqa: BLE001 — say it on the page rather than die silently | |
| cat.status.update(stage="failed", error=f"{type(exc).__name__}: {exc}"[:500]) | |
| raise | |
| app = FastAPI(title="Sand3D Explorer", docs_url=None, redoc_url=None) | |
| RELOAD_S = 600 | |
| def _reload_loop() -> None: | |
| """Every ten minutes, the GLBs published since: the cached results forgotten when they moved.""" | |
| import time # noqa: PLC0415 | |
| while True: | |
| time.sleep(RELOAD_S) | |
| try: | |
| if cat.reload_glbs(): | |
| with cache.lock: | |
| cache.d.clear() | |
| except Exception as exc: # noqa: BLE001 — the next round tries again | |
| cat.status["reload_error"] = f"{type(exc).__name__}: {exc}"[:200] | |
| def _start() -> None: | |
| threading.Thread(target=_load, daemon=True).start() | |
| threading.Thread(target=_reload_loop, daemon=True).start() | |
| def _ready() -> None: | |
| if not cat.status.get("ready"): | |
| raise HTTPException(503, detail=cat.status) | |
| def status(): | |
| return cat.status | |
| def facets(): | |
| _ready() | |
| return {**cat.facets(), "semantic": encoder.model is not None, "library_url": CONFIG["library_url"]} | |
| def do_search(q: str = "", mode: str = "meaning", dataset: str = "", family: str = "", | |
| articulated: bool = False, glb: bool = False, parts: bool = False, offset: int = Query(0, ge=0), | |
| limit: int = Query(PAGE, ge=1, le=200)): | |
| _ready() | |
| q = q.strip()[:200] | |
| used = "meaning" if (q and mode == "meaning" and encoder.model is not None) else ("words" if q else "all") | |
| key = (used, q.lower(), dataset, family, articulated, glb, parts) | |
| def make(): | |
| m = cat.mask(dataset or None, family or None, articulated, glb, parts) | |
| if used == "meaning": | |
| return search.ranked(cat, encoder.encode(q), m) | |
| return search.words(cat, q, m) | |
| hits = cache.get(key, make) | |
| return {"total": len(hits), "offset": offset, "mode": used, | |
| "items": [cat.card(i) for i in hits[offset:offset + limit]]} | |
| def obj(oid: str): | |
| _ready() | |
| i = cat.by_id.get(oid) | |
| if i is None: | |
| raise HTTPException(404, "no such object") | |
| r = cat.rows[i] | |
| ds = next(d for d in cat.datasets if d.key == r["_dataset"]) | |
| split = r["_split_name"] | |
| return {**cat.card(i), "program": r["program"], "entry": r.get("entry"), "source": r["source"], | |
| "source_label": SOURCES.get(r["source"], r["source"]), "family_raw": r["family"], | |
| "split": r["split"], "f_score": r.get("f_score"), "license": r.get("license"), | |
| "source_license": r.get("source_license"), "license_status": r.get("license_status"), | |
| "attribution": r.get("attribution"), "asset_id": r.get("asset_id"), | |
| "sand3d_version": r.get("sand3d_version"), "usda_bytes": r.get("usda_bytes"), | |
| "dataset_label": ds.label, "repo": ds.repo, | |
| "row_url": f"{CONFIG['hub_url']}/{ds.repo}/viewer/default/{split}?row={r['_row']}", | |
| "usda_url": f"{CONFIG['hub_url']}/{ds.repo}/blob/main/{r['usda_path']}", | |
| "glb_url": _glb_url(r), "assembly": _assembly(r), "library_url": CONFIG["library_url"], "has_similar": bool(cat.emb is not None and cat.emb_present[i])} | |
| def _assembly(r: dict) -> dict | None: | |
| """A part's assembly, when its id names one (`<assembly>__<part>`) and the catalog holds it.""" | |
| if not r["_part"] or "__" not in r["id"]: | |
| return None | |
| i = cat.by_id.get(r["id"].split("__", 1)[0]) | |
| return cat.card(i) if i is not None else None | |
| def _glb_url(r: dict) -> str | None: | |
| """The GLB's URL, versioned by the assets' revision so a rebuilt GLB is never served from a cache.""" | |
| if not r["_glb"]: | |
| return None | |
| return f"/glb/{r['id']}.glb?v={(cat.assets_rev or 'local')[:12]}" | |
| def similar(oid: str, limit: int = Query(12, ge=1, le=60), same_dataset: bool = False): | |
| _ready() | |
| i = cat.by_id.get(oid) | |
| if i is None or cat.emb is None or not cat.emb_present[i]: | |
| return {"items": []} | |
| m = cat.mask(cat.rows[i]["_dataset"] if same_dataset else None, None, False, False) | |
| return {"items": [cat.card(j) for j in search.ranked(cat, cat.emb[i], m, exclude=i)[:limit]]} | |
| def render(oid: str): | |
| _ready() | |
| i = cat.by_id.get(oid) | |
| if i is None: | |
| raise HTTPException(404) | |
| return Response(cat.render(i), media_type="image/png", | |
| headers={"Cache-Control": "private, max-age=3600"}) | |
| def glb(oid: str): | |
| _ready() | |
| i = cat.by_id.get(oid) | |
| path = cat.glb_path(i) if i is not None else None | |
| if path is None: | |
| raise HTTPException(404) | |
| return FileResponse(path, media_type="model/gltf-binary", headers={"Cache-Control": "private, max-age=86400"}) | |
| #: the page itself is always revalidated: its assets carry a version (`?v=`), the page does not | |
| PAGE_HEADERS = {"Cache-Control": "no-cache"} | |
| def page_object(oid: str): | |
| return FileResponse(STATIC / "index.html", headers=PAGE_HEADERS) | |
| def page_index(): | |
| return FileResponse(STATIC / "index.html", headers=PAGE_HEADERS) | |
| app.mount("/static", StaticFiles(directory=STATIC), name="static") | |
| def _http(_req, exc: HTTPException): | |
| return JSONResponse({"detail": exc.detail}, status_code=exc.status_code) | |