Guanc27/check_fonts

1"""2Font Detection API Server.3 4FastAPI application that accepts an image upload and returns the top-K5matching Google Fonts using an OpenCLIP embedding model and FAISS index.6 7Run locally:8    uvicorn api_server:app --reload --port 80009 10Test:11    curl -X POST http://localhost:8000/detect -F "file=@test.png"12"""13 14import io15import json16import os17import time18from contextlib import asynccontextmanager19from pathlib import Path20 21import faiss22import numpy as np23import torch24from fastapi import FastAPI, File, Query, UploadFile25from fastapi.middleware.cors import CORSMiddleware26from fastapi.responses import JSONResponse27from PIL import Image28 29import open_clip30from preprocessing import FontImagePreprocessor31 32# ---------------------------------------------------------------------------33# Configuration (environment variables with sensible local defaults)34# ---------------------------------------------------------------------------35CHECKPOINT_PATH = os.environ.get("CHECKPOINT_PATH", "models/best_model2.pt")36VECTOR_DB_DIR = os.environ.get("VECTOR_DB_DIR", "vector_db")37MODEL_NAME = os.environ.get("MODEL_NAME", "ViT-B-32")38PRETRAINED = os.environ.get("PRETRAINED", "openai")39DEFAULT_TOP_K = int(os.environ.get("DEFAULT_TOP_K", "5"))40 41# ---------------------------------------------------------------------------42# Shared state populated once at startup43# ---------------------------------------------------------------------------44state: dict = {}45 46 47def _google_fonts_url(font_name: str) -> str:48    """Build a Google Fonts specimen URL from a font name."""49    slug = font_name.replace(" ", "+")50    return f"https://fonts.google.com/specimen/{slug}"51 52 53def _load_resources():54    """Load model, FAISS index, metadata, and preprocessor."""55    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")56 57    # -- OpenCLIP model --58    model, _, _preprocess = open_clip.create_model_and_transforms(59        MODEL_NAME, pretrained=PRETRAINED, device=device,60    )61 62    checkpoint_path = Path(CHECKPOINT_PATH)63    if checkpoint_path.exists():64        checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False)65        model.load_state_dict(checkpoint["model_state_dict"])66    else:67        raise FileNotFoundError(f"Checkpoint not found: {checkpoint_path}")68 69    model = model.to(device)70    model.eval()71 72    # -- FAISS index + metadata --73    db_dir = Path(VECTOR_DB_DIR)74    index_path = db_dir / "faiss.index"75    metadata_path = db_dir / "metadata.json"76 77    if not index_path.exists():78        raise FileNotFoundError(f"FAISS index not found: {index_path}")79    if not metadata_path.exists():80        raise FileNotFoundError(f"Metadata not found: {metadata_path}")81 82    index = faiss.read_index(str(index_path))83    with open(metadata_path, "r", encoding="utf-8") as f:84        metadata = json.load(f)85 86    samples = metadata.get("samples", [])87    normalize = metadata.get("normalize", True)88 89    # -- Preprocessor --90    preprocessor = FontImagePreprocessor(target_size=224)91 92    return {93        "model": model,94        "device": device,95        "index": index,96        "samples": samples,97        "metadata": metadata,98        "normalize": normalize,99        "preprocessor": preprocessor,100    }101 102 103@asynccontextmanager104async def lifespan(app: FastAPI):105    """Load heavy resources once at startup; release on shutdown."""106    state.update(_load_resources())107    yield108    state.clear()109 110 111# ---------------------------------------------------------------------------112# FastAPI app113# ---------------------------------------------------------------------------114app = FastAPI(115    title="Font Detection API",116    version="1.0.0",117    lifespan=lifespan,118)119 120app.add_middleware(121    CORSMiddleware,122    allow_origins=["*"],  # Chrome extensions use chrome-extension:// origin123    allow_credentials=True,124    allow_methods=["*"],125    allow_headers=["*"],126)127 128 129# ---------------------------------------------------------------------------130# Endpoints131# ---------------------------------------------------------------------------132 133 134@app.get("/health")135async def health():136    """Return model / index status."""137    if not state:138        return JSONResponse({"status": "loading"}, status_code=503)139 140    metadata = state["metadata"]141    return {142        "status": "ok",143        "model": MODEL_NAME,144        "checkpoint": CHECKPOINT_PATH,145        "device": str(state["device"]),146        "index_vectors": state["index"].ntotal,147        "font_count": len(metadata.get("font_to_id", {})),148    }149 150 151@app.post("/detect")152async def detect(153    file: UploadFile = File(...),154    top_k: int = Query(DEFAULT_TOP_K, ge=1, le=50),155    mode: str = Query("single", pattern="^(single|multi)$"),156):157    """Identify the font in an uploaded image.158 159    Parameters160    ----------161    file : uploaded image (PNG / JPEG / WebP / etc.)162    top_k : number of results to return (default 5)163    mode : ``"single"`` (pre-cropped region, default) or ``"multi"``164           (full photo — runs EasyOCR text detection first)165    """166    start = time.perf_counter()167 168    # -- Read & validate image --169    contents = await file.read()170    try:171        image = Image.open(io.BytesIO(contents))172        image.verify()173        # Re-open after verify (verify consumes the stream)174        image = Image.open(io.BytesIO(contents)).convert("RGB")175    except Exception:176        return JSONResponse(177            {"error": "Invalid image file."}, status_code=400,178        )179 180    preprocessor: FontImagePreprocessor = state["preprocessor"]181    model = state["model"]182    device = state["device"]183    index = state["index"]184    samples = state["samples"]185    normalize = state["normalize"]186 187    # -- Preprocess --188    if mode == "multi":189        tensor = preprocessor.preprocess_for_model(image)190    else:191        tensor = preprocessor.preprocess_single(image)192 193    regions_detected = tensor.shape[0]194 195    # -- Embed --196    tensor = tensor.to(device)197    with torch.no_grad():198        features = model.encode_image(tensor)199        if normalize:200            features = features / features.norm(dim=-1, keepdim=True)201 202    # Average embeddings when multiple regions are detected203    if features.shape[0] > 1:204        features = features.mean(dim=0, keepdim=True)205 206    query = features.cpu().numpy().astype("float32")207 208    # -- Search FAISS --209    fetch_k = min(top_k * 5, index.ntotal)210    scores, indices = index.search(query, fetch_k)211 212    # Aggregate best score per font213    font_best: dict[str, tuple[float, dict]] = {}214    for idx, score in zip(indices[0], scores[0]):215        if idx < 0 or idx >= len(samples):216            continue217        sample = samples[int(idx)]218        name = sample["font_name"]219        if name not in font_best or score > font_best[name][0]:220            font_best[name] = (float(score), sample)221 222    ranked = sorted(font_best.values(), key=lambda x: x[0], reverse=True)223 224    matches = []225    for rank, (score, sample) in enumerate(ranked[:top_k], start=1):226        matches.append({227            "rank": rank,228            "font_name": sample["font_name"],229            "score": round(score, 4),230            "google_fonts_url": _google_fonts_url(sample["font_name"]),231        })232 233    elapsed_ms = (time.perf_counter() - start) * 1000234 235    return {236        "matches": matches,237        "processing_time_ms": round(elapsed_ms, 1),238        "regions_detected": regions_detected,239    }240