Spaces:
Sleeping
Sleeping
| """Phase 2 — FastAPI inference service for the HF Docker Space (port 7860). | |
| Loads ALL switchable CLIP models + their indexes once at startup, asserts each | |
| index matches its model, then serves POST /guess (multipart, X-API-Key) and | |
| GET /health (no auth). Pick a model per request with ?model=fast|pro. | |
| """ | |
| import io | |
| import json | |
| import os | |
| import sys | |
| import time | |
| from collections import Counter | |
| from contextlib import asynccontextmanager | |
| from pathlib import Path | |
| import numpy as np | |
| from fastapi import FastAPI, File, Header, HTTPException, Query, UploadFile | |
| from fastapi.middleware.cors import CORSMiddleware | |
| from PIL import Image | |
| sys.path.insert(0, str(Path(__file__).resolve().parent / "shared")) | |
| from version import MODELS, DEFAULT_MODEL # noqa: E402 | |
| from embedder import Embedder # noqa: E402 | |
| import guess as guesslib # noqa: E402 | |
| import atlas as atlaslib # noqa: E402 | |
| import region as regionlib # noqa: E402 | |
| DATA = Path(__file__).resolve().parent / "data" | |
| API_KEY = os.environ.get("API_KEY", "") | |
| ALLOWED_ORIGIN = os.environ.get("ALLOWED_ORIGIN", "*") | |
| MAX_BYTES = 4 * 1024 * 1024 | |
| RATE_LIMIT_PER_MIN = 30 | |
| # Region kNN indexes live in a free HF Dataset repo (the 1GB Space can't hold them all). | |
| # Resolve a region file from local data/ if present, else download from the dataset (cached). | |
| INDEX_DATASET = os.environ.get("INDEX_DATASET", "GeoguessrAngular/geobot-indexes") | |
| def region_file_path(fname): | |
| local = DATA / fname | |
| if local.exists(): | |
| return local | |
| from huggingface_hub import hf_hub_download | |
| return Path(hf_hub_download(repo_id=INDEX_DATASET, filename=fname, repo_type="dataset")) | |
| Image.MAX_IMAGE_PIXELS = 50_000_000 | |
| STATE = {} | |
| _req_times = [] # global in-memory rate limiter timestamps | |
| async def lifespan(app: FastAPI): | |
| centroids = json.loads((DATA / "centroids.json").read_text(encoding="utf-8")) | |
| priors_path = DATA / "priors.json" | |
| priors = json.loads(priors_path.read_text(encoding="utf-8")) if priors_path.exists() else None | |
| models = {} | |
| embedder_cache = {} # model_id -> Embedder (heavy weights loaded once) | |
| index_cache = {} # index_file -> (index, rows, count, model_id, index_version) | |
| for key, m in MODELS.items(): | |
| if m["model_id"] not in embedder_cache: | |
| embedder_cache[m["model_id"]] = Embedder(m["model_id"]) | |
| # --- region kNN locator (e.g. Serbia) --- | |
| if m.get("region_file"): | |
| try: | |
| ref_emb, ref_lat, ref_lng = regionlib.load_region(region_file_path(m["region_file"])) | |
| models[key] = { | |
| "type": "region", "embedder": embedder_cache[m["model_id"]], | |
| "ref_emb": ref_emb, "ref_lat": ref_lat, "ref_lng": ref_lng, | |
| "country_slug": m.get("country_slug", key), "country_name": m.get("country_name", key), | |
| "model_id": m["model_id"], "label": m.get("label", key), | |
| } | |
| # Difficulty cuts whole LOCATIONS, not frames: ?frac= (0..1) keeps | |
| # that fraction of the bot's known locations (each pano's frames | |
| # kept together), so frac=0.05 = the bot knows 5% of locations. | |
| # All frames of a pano share lat/lng → group rows by coord, permute | |
| # the locations (seeded), and store each row's location rank. At | |
| # query time keep rows whose location rank < k. Any % works with no | |
| # redeploy; the admin panel tunes the fractions server-side. | |
| n = int(ref_emb.shape[0]) | |
| if n > 60: | |
| keys = np.round(np.stack([ref_lat, ref_lng], axis=1), 5) | |
| uniq, inv = np.unique(keys, axis=0, return_inverse=True) | |
| inv = np.asarray(inv).ravel() | |
| nloc = int(uniq.shape[0]) | |
| loc_order = np.random.default_rng(1234).permutation(nloc) | |
| loc_rank = np.empty(nloc, dtype=np.int64) | |
| loc_rank[loc_order] = np.arange(nloc) | |
| models[key]["row_rank"] = loc_rank[inv].astype(np.int32) | |
| models[key]["nloc"] = nloc | |
| print(f"Loaded model '{key}': region kNN {ref_emb.shape} " | |
| f"({models[key].get('nloc', '?')} locs), {m['model_id']}" | |
| f"{' +tierable' if 'row_rank' in models[key] else ''}") | |
| except Exception as e: | |
| print(f"SKIP model '{key}': failed to load {m['region_file']}: {e}") | |
| continue | |
| # --- learned classifier head (Atlas) --- | |
| if m.get("head_file"): | |
| W, b, classes = atlaslib.load_head(DATA / m["head_file"]) | |
| models[key] = { | |
| "type": "head", "embedder": embedder_cache[m["model_id"]], | |
| "W": W, "b": b, "classes": classes, | |
| "model_id": m["model_id"], "label": m.get("label", key), | |
| } | |
| print(f"Loaded model '{key}': head W{W.shape} {len(classes)} classes, {m['model_id']}") | |
| continue | |
| # --- retrieval (index) model --- | |
| if m["index_file"] not in index_cache: | |
| meta = json.loads((DATA / m["meta_file"]).read_text(encoding="utf-8")) | |
| index = np.load(DATA / m["index_file"]).astype(np.float32) | |
| rows = meta["rows"] | |
| assert index.shape[0] == len(rows), f"[{key}] index/meta row mismatch" | |
| index_cache[m["index_file"]] = ( | |
| index, rows, Counter(r["country"] for r in rows), | |
| meta["model_id"], meta["index_version"]) | |
| index, rows, count, meta_mid, meta_iv = index_cache[m["index_file"]] | |
| if meta_mid != m["model_id"] or meta_iv != m["index_version"]: | |
| raise RuntimeError( | |
| f"[{key}] index/version mismatch: meta has {meta_mid}/{meta_iv}, " | |
| f"expected {m['model_id']}/{m['index_version']}") | |
| text_vecs, text_countries = None, None | |
| if m.get("text_file"): | |
| text_vecs = np.load(DATA / m["text_file"]).astype(np.float32) | |
| text_countries = json.loads((DATA / m["text_countries_file"]).read_text(encoding="utf-8")) | |
| models[key] = { | |
| "type": "retrieval", "embedder": embedder_cache[m["model_id"]], | |
| "index": index, "rows": rows, "count": count, | |
| "model_id": m["model_id"], "index_version": m["index_version"], | |
| "label": m.get("label", key), | |
| "text_vecs": text_vecs, "text_countries": text_countries, | |
| } | |
| print(f"Loaded model '{key}': index {index.shape}, {m['model_id']} {m['index_version']}" | |
| f"{', +zeroshot-text' if text_vecs is not None else ''}") | |
| # Optional script (writing-system) branch for the Atlas head. | |
| script_vecs = script_names = country_scripts = None | |
| if (DATA / "script_text.npy").exists() and (DATA / "country_scripts.json").exists(): | |
| script_vecs = np.load(DATA / "script_text.npy").astype(np.float32) | |
| script_names = json.loads((DATA / "script_names.json").read_text(encoding="utf-8")) | |
| country_scripts = json.loads((DATA / "country_scripts.json").read_text(encoding="utf-8")) | |
| print(f"Loaded script branch: {len(script_names)} scripts, {len(country_scripts)} country maps") | |
| STATE["models"] = models | |
| STATE["centroids"] = centroids | |
| STATE["priors"] = priors | |
| STATE["script_vecs"] = script_vecs | |
| STATE["script_names"] = script_names | |
| STATE["country_scripts"] = country_scripts | |
| STATE["cfg"] = guesslib.Config() | |
| print(f"Ready. models={list(models)} default={DEFAULT_MODEL} " | |
| f"centroids={len(centroids)} priors={len(priors) if priors else 0}") | |
| yield | |
| STATE.clear() | |
| app = FastAPI(title="GeoBot", lifespan=lifespan) | |
| app.add_middleware(CORSMiddleware, allow_origins=[ALLOWED_ORIGIN] if ALLOWED_ORIGIN != "*" else ["*"], | |
| allow_methods=["*"], allow_headers=["*"]) | |
| def _check_key(x_api_key): | |
| if not API_KEY: | |
| return # unset key disables auth (local dev) | |
| if x_api_key != API_KEY: | |
| raise HTTPException(status_code=401, detail="bad or missing API key") | |
| def _rate_limit(): | |
| now = time.time() | |
| cutoff = now - 60 | |
| while _req_times and _req_times[0] < cutoff: | |
| _req_times.pop(0) | |
| if len(_req_times) >= RATE_LIMIT_PER_MIN: | |
| raise HTTPException(status_code=429, detail="rate limited") | |
| _req_times.append(now) | |
| async def _read_image(upload: UploadFile): | |
| raw = await upload.read() | |
| if len(raw) > MAX_BYTES: | |
| raise HTTPException(status_code=413, detail="image too large (>4 MB)") | |
| try: | |
| return Image.open(io.BytesIO(raw)).convert("RGB") | |
| except Exception: | |
| raise HTTPException(status_code=400, detail="undecodable image") | |
| def health(): | |
| models = STATE.get("models", {}) | |
| return { | |
| "status": "ok", | |
| "default_model": DEFAULT_MODEL, | |
| "models": { | |
| k: {"type": v.get("type"), "model_id": v["model_id"], "label": v["label"], | |
| **({"index_size": int(v["index"].shape[0]), "index_version": v["index_version"]} | |
| if v.get("type") == "retrieval" else | |
| {"refs": int(v["ref_emb"].shape[0]), | |
| "locations": int(v.get("nloc") or 0), | |
| "tierable": v.get("row_rank") is not None} | |
| if v.get("type") == "region" else | |
| {"classes": len(v["classes"])})} | |
| for k, v in models.items() | |
| }, | |
| } | |
| async def do_guess(images: list[UploadFile] = File(None), | |
| image0: UploadFile = File(None), image180: UploadFile = File(None), | |
| skill: float = Query(3.0, ge=0.0, le=3.0), # legacy: 1=easy 2=medium 3=full | |
| frac: float = Query(None, ge=0.0, le=1.0), # index density 0..1 (server-computed; wins over skill) | |
| model: str = Query(DEFAULT_MODEL), | |
| x_api_key: str = Header(None, alias="X-API-Key")): | |
| _check_key(x_api_key) | |
| _rate_limit() | |
| # Accept either N frames under repeated field "images", or legacy image0/image180. | |
| uploads = [u for u in (images or []) if u is not None] | |
| if not uploads: | |
| uploads = [u for u in (image0, image180) if u is not None] | |
| if not uploads: | |
| raise HTTPException(status_code=400, detail="no images") | |
| key = model if model in STATE["models"] else DEFAULT_MODEL | |
| M = STATE["models"][key] | |
| t0 = time.time() | |
| pil = [await _read_image(u) for u in uploads] | |
| emb = M["embedder"].embed(pil) | |
| try: | |
| if M.get("type") == "region": | |
| # Difficulty = fraction of LOCATIONS the bot knows. Prefer server-sent | |
| # ?frac=; else map the legacy skill (1->0.05, 2->0.25, 3->full). Keep | |
| # every frame of the first k locations (k = frac * nloc, >=5 locations). | |
| re_, rl_, rn_ = M["ref_emb"], M["ref_lat"], M["ref_lng"] | |
| f = frac if frac is not None else (0.05 if skill <= 1.5 else 0.25 if skill <= 2.5 else 1.0) | |
| rr, nloc = M.get("row_rank"), M.get("nloc") | |
| if rr is not None and nloc and f < 0.999: | |
| k = min(nloc, max(5, int(round(f * nloc)))) | |
| sel = np.where(rr < k)[0] | |
| re_, rl_, rn_ = re_[sel], rl_[sel], rn_[sel] | |
| # "human mistake" spread by level: easy often takes another option, | |
| # medium rarely, hard never (best geo-medoid). | |
| spread = 0.6 if skill <= 1.5 else 0.3 if skill <= 2.5 else 0.0 | |
| result = regionlib.predict(list(emb), re_, rl_, rn_, | |
| M["country_slug"], M["country_name"], spread=spread) | |
| elif M.get("type") == "head": | |
| result = atlaslib.predict(list(emb), M["W"], M["b"], M["classes"], | |
| STATE["centroids"], STATE["cfg"], STATE["priors"], | |
| STATE["script_vecs"], STATE["script_names"], | |
| STATE["country_scripts"]) | |
| else: | |
| result = guesslib.guess(list(emb), M["index"], M["rows"], | |
| STATE["centroids"], M["count"], STATE["cfg"], | |
| STATE["priors"], M["text_vecs"], M["text_countries"]) | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=f"guess failed: {type(e).__name__}") | |
| result["timing_ms"] = int((time.time() - t0) * 1000) | |
| result["model"] = key | |
| result["model_id"] = M["model_id"] | |
| result["index_version"] = M.get("index_version") | |
| print(f"[{key}] guess winner={result['country']} conf={result['confidence']} " | |
| f"ms={result['timing_ms']}") | |
| return result | |