Spaces:
Runtime error
Runtime error
| """ | |
| Face Recognition Backend API on Hugging Face Spaces (Gradio SDK + FastAPI). | |
| Serves MTCNN + ArcFace vector inference with full CORS support for Vercel. | |
| """ | |
| import base64 | |
| import os | |
| import secrets | |
| import sys | |
| import threading | |
| import time | |
| import traceback | |
| from typing import Optional, List, Dict, Any | |
| from fastapi import FastAPI, Request, Response, HTTPException, status | |
| from fastapi.responses import JSONResponse, Response as RawResponse | |
| from fastapi.middleware.cors import CORSMiddleware | |
| from starlette.middleware.base import BaseHTTPMiddleware | |
| from pydantic import BaseModel | |
| import gradio as gr | |
| import pipeline as P | |
| # -------------------------------------------------------------------------- | |
| # State & Sessions | |
| # -------------------------------------------------------------------------- | |
| SESSION_TTL = 4 * 3600 | |
| SESSIONS = {} | |
| _SESS_LOCK = threading.Lock() | |
| _status = {"ready": False, "error": None, "seconds": None} | |
| def _warm(): | |
| t0 = time.time() | |
| try: | |
| P.warmup() | |
| _status["ready"] = True | |
| except Exception as exc: | |
| _status["error"] = f"{type(exc).__name__}: {exc}" | |
| traceback.print_exc() | |
| finally: | |
| _status["seconds"] = round(time.time() - t0, 1) | |
| print(f"[warmup] ready={_status['ready']} in {_status['seconds']}s") | |
| threading.Thread(target=_warm, daemon=True).start() | |
| def _get_session(sid: Optional[str]): | |
| now = time.time() | |
| with _SESS_LOCK: | |
| for dead in [s for s, v in SESSIONS.items() if now - v["touched"] > SESSION_TTL]: | |
| SESSIONS.pop(dead, None) | |
| if not sid or sid not in SESSIONS: | |
| sid = secrets.token_urlsafe(18) | |
| SESSIONS[sid] = {"gallery": P.Gallery(), "dismissed": set(), "touched": now} | |
| SESSIONS[sid]["touched"] = now | |
| return sid, SESSIONS[sid]["gallery"], SESSIONS[sid]["dismissed"] | |
| def _session_payload(gallery, dismissed): | |
| return { | |
| "people": gallery.summary(), | |
| "total_embeddings": sum(len(v) for v in gallery.data.values()), | |
| "duplicates": P.find_duplicates(gallery.data, dismissed), | |
| } | |
| # -------------------------------------------------------------------------- | |
| # FastAPI Application | |
| # -------------------------------------------------------------------------- | |
| api_app = FastAPI(title="Face Recognition API", docs_url="/docs") | |
| # CORS middleware for cross-origin Vercel requests | |
| api_app.add_middleware( | |
| CORSMiddleware, | |
| allow_origins=["*"], | |
| allow_credentials=True, | |
| allow_methods=["*"], | |
| allow_headers=["*"], | |
| expose_headers=["Content-Disposition", "X-Session-Id"], | |
| ) | |
| class SessionMiddleware(BaseHTTPMiddleware): | |
| async def dispatch(self, request: Request, call_next): | |
| sid_in = request.cookies.get("sid") or request.headers.get("X-Session-Id") | |
| sid, gallery, dismissed = _get_session(sid_in) | |
| request.state.sid = sid | |
| request.state.gallery = gallery | |
| request.state.dismissed = dismissed | |
| response = await call_next(request) | |
| # Attach session id to cookie & response header | |
| response.set_cookie( | |
| "sid", sid, httponly=True, samesite="none", secure=True, max_age=SESSION_TTL | |
| ) | |
| response.headers["X-Session-Id"] = sid | |
| return response | |
| api_app.add_middleware(SessionMiddleware) | |
| # -------------------------------------------------------------------------- | |
| # Pydantic Request Models | |
| # -------------------------------------------------------------------------- | |
| class EnrollRequest(BaseModel): | |
| person_id: str | |
| images: List[str] | |
| class RecognizeRequest(BaseModel): | |
| image: str | |
| class MergeRequest(BaseModel): | |
| source: str | |
| target: str | |
| class DismissRequest(BaseModel): | |
| a: str | |
| b: str | |
| # -------------------------------------------------------------------------- | |
| # REST API Endpoints | |
| # -------------------------------------------------------------------------- | |
| async def get_status(request: Request): | |
| people_count = len(request.state.gallery.data) if hasattr(request.state, "gallery") else 0 | |
| return {**_status, "people": people_count} | |
| async def get_session_info(request: Request): | |
| return _session_payload(request.state.gallery, request.state.dismissed) | |
| async def reset_session(request: Request): | |
| with _SESS_LOCK: | |
| SESSIONS[request.state.sid]["gallery"] = P.Gallery() | |
| SESSIONS[request.state.sid]["dismissed"] = set() | |
| return {"ok": True} | |
| async def enroll(req: EnrollRequest, request: Request): | |
| person_id = req.person_id.strip() | |
| images = req.images | |
| if not person_id: | |
| raise HTTPException(status_code=400, detail="A name is required") | |
| if not images: | |
| raise HTTPException(status_code=400, detail="No images supplied") | |
| if not _status["ready"]: | |
| raise HTTPException(status_code=503, detail="Models are still loading, please wait") | |
| results = [] | |
| added = 0 | |
| for idx, data_url in enumerate(images): | |
| try: | |
| bgr = P.decode_data_url(data_url) | |
| faces = P.extract_faces(bgr) | |
| if not faces: | |
| results.append({"index": idx, "ok": False, "reason": "No face detected"}) | |
| continue | |
| if len(faces) > 1: | |
| results.append({ | |
| "index": idx, "ok": False, | |
| "reason": f"{len(faces)} faces in frame - register one face at a time" | |
| }) | |
| continue | |
| face = faces[0] | |
| request.state.gallery.add(person_id, P.embed(face["crop_bgr"])) | |
| added += 1 | |
| results.append({ | |
| "index": idx, | |
| "ok": True, | |
| "confidence": round(face["confidence"], 4), | |
| "crop": P.encode_jpeg_data_url(face["crop_bgr"], quality=75), | |
| }) | |
| except Exception as exc: | |
| results.append({"index": idx, "ok": False, "reason": str(exc)}) | |
| return { | |
| "person_id": person_id, | |
| "added": added, | |
| "total": len(request.state.gallery.data.get(person_id, [])), | |
| "cohesion": P.cohesion(request.state.gallery.data.get(person_id, [])), | |
| "results": results, | |
| } | |
| async def recognize(req: RecognizeRequest, request: Request): | |
| data_url = req.image | |
| if not data_url: | |
| raise HTTPException(status_code=400, detail="Image is required") | |
| if not _status["ready"]: | |
| raise HTTPException(status_code=503, detail="Models are still loading") | |
| if not request.state.gallery.data: | |
| raise HTTPException(status_code=400, detail="Nobody registered yet") | |
| t0 = time.time() | |
| try: | |
| bgr = P.decode_data_url(data_url) | |
| h, w = bgr.shape[:2] | |
| faces = P.extract_faces(bgr) | |
| except Exception as exc: | |
| raise HTTPException(status_code=400, detail=str(exc)) | |
| out = [] | |
| for face in faces: | |
| ranked = request.state.gallery.rank(P.embed(face["crop_bgr"])) | |
| top = ranked[0] if ranked else None | |
| runner = ranked[1] if len(ranked) > 1 else None | |
| out.append({ | |
| "box": face["box"], | |
| "confidence": round(face["confidence"], 4), | |
| "crop": P.encode_jpeg_data_url(face["crop_bgr"], quality=70), | |
| "ranked": [{"id": r["id"], "score": round(r["score"], 4)} for r in ranked], | |
| "top": {"id": top["id"], "score": round(top["score"], 4)} if top else None, | |
| "margin": round(top["score"] - runner["score"], 4) if runner else None, | |
| }) | |
| return { | |
| "faces": out, | |
| "frame": {"width": w, "height": h}, | |
| "elapsed_ms": int((time.time() - t0) * 1000), | |
| } | |
| async def delete_person(person_id: str, request: Request): | |
| ok = request.state.gallery.remove(person_id) | |
| for pair in [p for p in request.state.dismissed if person_id in p]: | |
| request.state.dismissed.discard(pair) | |
| return {"ok": ok} | |
| async def merge_people(req: MergeRequest, request: Request): | |
| source, target = req.source.strip(), req.target.strip() | |
| try: | |
| request.state.gallery.merge(source, target) | |
| except ValueError as exc: | |
| raise HTTPException(status_code=400, detail=str(exc)) | |
| for pair in [p for p in request.state.dismissed if source in p]: | |
| request.state.dismissed.discard(pair) | |
| return {"ok": True, **_session_payload(request.state.gallery, request.state.dismissed)} | |
| async def dismiss_duplicate(req: DismissRequest, request: Request): | |
| a, b = req.a.strip(), req.b.strip() | |
| gal = request.state.gallery | |
| if not a or not b or a == b or a not in gal.data or b not in gal.data: | |
| raise HTTPException(status_code=400, detail="Pick two different registered people") | |
| request.state.dismissed.add(frozenset((a, b))) | |
| return {"ok": True, **_session_payload(request.state.gallery, request.state.dismissed)} | |
| async def export_pkl(request: Request): | |
| if not request.state.gallery.data: | |
| raise HTTPException(status_code=400, detail="Nothing to export") | |
| data = request.state.gallery.to_pickle_bytes() | |
| return RawResponse( | |
| content=data, | |
| media_type="application/octet-stream", | |
| headers={"Content-Disposition": 'attachment; filename="student_embeddings.pkl"'} | |
| ) | |
| async def export_json(request: Request): | |
| return { | |
| "embeddings": request.state.gallery.to_json(), | |
| "dismissed": [sorted(p) for p in request.state.dismissed], | |
| } | |
| async def import_gallery(request: Request): | |
| body = await request.json() | |
| try: | |
| if body.get("json") is not None: | |
| gal = P.Gallery.from_json(body["json"]) | |
| elif body.get("pkl_b64"): | |
| gal = P.Gallery.from_pickle_bytes(base64.b64decode(body["pkl_b64"])) | |
| else: | |
| raise HTTPException(status_code=400, detail="Send either 'json' or 'pkl_b64'") | |
| except Exception as exc: | |
| raise HTTPException(status_code=400, detail=f"Could not load: {exc}") | |
| dismissed = set() | |
| for pair in body.get("dismissed") or []: | |
| if (isinstance(pair, (list, tuple)) and len(pair) == 2 | |
| and all(isinstance(x, str) and x in gal.data for x in pair) | |
| and pair[0] != pair[1]): | |
| dismissed.add(frozenset(pair)) | |
| with _SESS_LOCK: | |
| SESSIONS[request.state.sid]["gallery"] = gal | |
| SESSIONS[request.state.sid]["dismissed"] = dismissed | |
| return {"ok": True, **_session_payload(gal, dismissed)} | |
| # -------------------------------------------------------------------------- | |
| # Gradio Interface (Mounted at /) | |
| # -------------------------------------------------------------------------- | |
| with gr.Blocks(title="Face Recognition Backend API") as demo: | |
| gr.Markdown( | |
| """ | |
| # 🚀 Face Recognition Backend API | |
| **Powered by MTCNN + ArcFace (Retrieval-Augmented Recognition)** | |
| This Hugging Face Space hosts the ML inference API for the attendance application. | |
| - **Status:** Running | |
| - **Documentation:** Interactive OpenAPI docs available at [`/docs`](/docs) | |
| - **Frontend:** Connected to your live Vercel web application | |
| """ | |
| ) | |
| with gr.Row(): | |
| btn = gr.Button("Check Model Readiness") | |
| out = gr.JSON(label="Readiness Info") | |
| btn.click(fn=lambda: {**_status, "status": "online"}, outputs=out) | |
| app = gr.mount_gradio_app(api_app, demo, path="/") | |
| if __name__ == "__main__": | |
| print("[startup] Starting Gradio server", flush=True) | |
| demo.launch( | |
| server_name="0.0.0.0", | |
| server_port=1008, | |
| show_error=True | |
| ) |