Spaces:
Runtime error
Runtime error
File size: 5,307 Bytes
e7f7cda | 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 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 | import os
import io
import json
import base64
import asyncio
import numpy as np
from PIL import Image
from contextlib import asynccontextmanager
from fastapi import FastAPI, HTTPException, status, Header, Depends
from fastapi.responses import JSONResponse
from pydantic import BaseModel
import onnxruntime as ort
# Add parent dir to path so we can import registry module
import sys
from pathlib import Path
sys.path.append(str(Path(__file__).parent.parent))
from registry.download_model import download_artifacts, TARGET_DIR, VERSION
# Global state
is_model_ready = False
ort_session = None
class_labels = []
temperature = 1.0
INTERNAL_SECRET = os.environ.get("INTERNAL_SECRET")
if not INTERNAL_SECRET:
raise ValueError("CRITICAL: INTERNAL_SECRET environment variable is missing.")
@asynccontextmanager
async def lifespan(app: FastAPI):
global is_model_ready, ort_session, class_labels, temperature
# 1. Download artifacts
# NOTE: In production on HF Spaces, HF_TOKEN must be set if repo is private
# We run this in a thread to not block the event loop
print("Initializing Hugging Face Space Lifespan...")
try:
await asyncio.to_thread(download_artifacts)
except Exception as e:
print(f"Failed to download artifacts: {e}")
# 2. Load the ONNX model
onnx_path = TARGET_DIR / f"cropguard_{VERSION}.onnx"
if not onnx_path.exists():
print(f"Warning: Model not found at {onnx_path}. Space will fail health checks.")
else:
# Load ONNX model
providers = ['CPUExecutionProvider']
ort_session = ort.InferenceSession(str(onnx_path), providers=providers)
# Load labels
json_path = TARGET_DIR / f"cropguard_{VERSION}.json"
if json_path.exists():
with open(json_path, 'r') as f:
metadata = json.load(f)
class_labels = metadata.get("classes", [])
# Load temperature (optional)
temp_path = TARGET_DIR / "temperature.json"
if temp_path.exists():
with open(temp_path, 'r') as f:
temperature = json.load(f).get("temperature", 1.0)
is_model_ready = True
print("Model loaded successfully. Engine ready.")
yield
app = FastAPI(lifespan=lifespan)
class InferenceRequest(BaseModel):
image_b64: str
def verify_internal_secret(x_internal_secret: str = Header(None)):
if x_internal_secret != INTERNAL_SECRET:
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
detail="Forbidden: Invalid or missing X-Internal-Secret header"
)
return x_internal_secret
@app.get("/health")
async def health_check():
if not is_model_ready:
return JSONResponse(status_code=503, content={"status": "loading"})
return {"status": "ready"}
def preprocess_image(image_b64: str):
# Decode base64
if "," in image_b64:
image_b64 = image_b64.split(",")[1]
image_bytes = base64.b64decode(image_b64)
image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
# Resize to 224x224 (ConvNeXt standard)
image = image.resize((224, 224), Image.Resampling.BILINEAR)
# Convert to numpy array and normalize
# ConvNeXt ImageNet Mean: [0.485, 0.456, 0.406]
# ConvNeXt ImageNet Std: [0.229, 0.224, 0.225]
img_arr = np.array(image, dtype=np.float32) / 255.0
mean = np.array([0.485, 0.456, 0.406], dtype=np.float32)
std = np.array([0.229, 0.224, 0.225], dtype=np.float32)
img_arr = (img_arr - mean) / std
# HWC to CHW format
img_arr = np.transpose(img_arr, (2, 0, 1))
# Add batch dimension
img_arr = np.expand_dims(img_arr, axis=0)
return img_arr
def run_inference(img_arr: np.ndarray):
input_name = ort_session.get_inputs()[0].name
outputs = ort_session.run(None, {input_name: img_arr})
logits = outputs[0][0]
# Apply temperature scaling
scaled_logits = logits / temperature
# Softmax
exp_logits = np.exp(scaled_logits - np.max(scaled_logits))
probs = exp_logits / exp_logits.sum()
pred_idx = int(np.argmax(probs))
confidence = float(probs[pred_idx])
label = class_labels[pred_idx] if pred_idx < len(class_labels) else f"class_{pred_idx}"
return label, confidence
@app.post("/infer", dependencies=[Depends(verify_internal_secret)])
async def infer(payload: InferenceRequest):
if not is_model_ready:
return JSONResponse(status_code=503, content={"status": "loading"})
try:
# 1. Preprocess
img_arr = await asyncio.to_thread(preprocess_image, payload.image_b64)
# 2. Run Inference
label, confidence = await asyncio.to_thread(run_inference, img_arr)
# 3. Generate Heatmap (Mock for now, as True CAM requires PyTorch hooks)
mock_heatmap = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8BQDwAEhQGAhKmMIQAAAABJRU5ErkJggg=="
return {
"raw_label": label,
"confidence": confidence,
"heatmap_b64": mock_heatmap
}
except Exception as e:
print(f"Inference error: {e}")
raise HTTPException(status_code=500, detail="Inference failed")
|