File size: 3,722 Bytes
4f37f46 e2af0c7 4f37f46 e2af0c7 4f37f46 18a82fb e2af0c7 18a82fb e2af0c7 4f37f46 e2af0c7 4f37f46 18a82fb e2af0c7 13e7af1 4f37f46 e2af0c7 4f37f46 e2af0c7 4f37f46 2ef860a 4f37f46 e2af0c7 4f37f46 e2af0c7 4f37f46 e2af0c7 4f37f46 13e7af1 18a82fb 4f37f46 13e7af1 4f37f46 13e7af1 4f37f46 e2af0c7 13e7af1 e2af0c7 4f37f46 13e7af1 e2af0c7 13e7af1 e2af0c7 13e7af1 2ef860a 13e7af1 18a82fb 13e7af1 | 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 | # api/app.py
import io
import json
import traceback
from pathlib import Path
from fastapi import FastAPI, File, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse, JSONResponse, Response
from PIL import Image
import numpy as np
import sys
# make sure project root is importable (scripts/inference.py lives at repo root)
ROOT = Path(__file__).resolve().parent.parent
sys.path.append(str(ROOT))
# ---- YOUR ORIGINAL MODEL LOADER (unchanged in spirit) ----
from scripts.inference import ModelInference
# create model ONCE (like your original code)
# if your ModelInference supports map_location/device, set to CPU for Spaces
model = ModelInference(
checkpoint_path=str(ROOT / "checkpoints" / "best_checkpoint.pth"),
multi_task=True
)
# -------- FastAPI app & CORS ----------
app = FastAPI(title="VizRef API")
app.add_middleware(
CORSMiddleware,
allow_origins=["*"], allow_credentials=True,
allow_methods=["*"], allow_headers=["*"],
)
# -------- Serve your UI ----------
@app.get("/")
def root():
return FileResponse(str(ROOT / "ui" / "index.html"))
# -------- API: health ----------
@app.get("/api/health")
def health_check():
return {
"status": "healthy",
"model_loaded": True,
"model_info": {"model_name": "Efficientnet - B0", "multi_task": True},
}
# -------- API: training data ----------
@app.get("/api/training_data")
def get_training_data():
with open(ROOT / "data" / "splits" / "train.json", "r") as f:
data = json.load(f)
return {"status": "success", "data": data}
# -------- API: proxy image ----------
@app.get("/api/proxy_image")
def proxy_image(url: str):
import requests
try:
if url.startswith("http://"):
https_url = url.replace("http://", "https://")
try:
r = requests.get(https_url, timeout=5)
if r.status_code == 200:
return Response(content=r.content, media_type="image/jpeg")
except:
pass
r = requests.get(url, timeout=5)
return Response(content=r.content, media_type="image/jpeg")
except:
return Response(status_code=404)
# --- helper to make numpy types JSON-safe ---
def _pythonify(obj):
if isinstance(obj, dict):
return {k: _pythonify(v) for k, v in obj.items()}
if isinstance(obj, list):
return [_pythonify(v) for v in obj]
if isinstance(obj, (np.integer,)):
return int(obj)
if isinstance(obj, (np.floating,)):
return float(obj)
if isinstance(obj, (np.bool_)):
return bool(obj)
return obj
# -------- API: predict (your original logic, adapted to JSON) ----------
@app.post("/api/predict")
async def predict(file: UploadFile = File(...)):
try:
raw = await file.read()
try:
img = Image.open(io.BytesIO(raw)).convert("RGB")
except Exception as e:
return JSONResponse({"status": "error", "message": f"Cannot read image: {e}"}, status_code=400)
tmp = Path("/tmp/vizref_upload.jpg")
img.save(tmp)
# your original call
result = model.predict_image(str(tmp))
return JSONResponse({
"status": "success",
"predictions": _pythonify(result),
"model_info": {"model_name": "Efficientnet - B0", "multi_task": True},
})
except Exception as e:
print("PREDICT ERROR\n", traceback.format_exc())
return JSONResponse({"status": "error", "message": str(e)}, status_code=500)
finally:
try:
if 'tmp' in locals() and tmp.exists():
tmp.unlink()
except:
pass
|