""" AegisRoad v3.0 — YOLOv11 Inference Server Hugging Face Spaces Deployment """ import os import io import base64 import time import random from pathlib import Path import numpy as np from PIL import Image, ImageDraw import cv2 import uvicorn from fastapi import FastAPI, File, UploadFile, HTTPException from fastapi.middleware.cors import CORSMiddleware from fastapi.responses import JSONResponse, FileResponse from fastapi.staticfiles import StaticFiles # ── App setup ────────────────────────────────────────────────────────────── app = FastAPI( title="AegisRoad Inference API", description="YOLOv11 road damage detection endpoint", version="3.0.0", ) app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"], ) # Mount static files BEFORE defining routes if Path("static").exists(): app.mount("/static", StaticFiles(directory="static"), name="static") # ── Constants ────────────────────────────────────────────────────────────── MODEL_PATH = Path("models/best.pt") # Fallback class map if model doesn't embed names FALLBACK_CLASS_NAMES = {0: "D00", 1: "D10", 2: "D20", 3: "D40"} CLASS_LABELS = { "D00": "Longitudinal Crack", "D10": "Transverse Crack", "D20": "Alligator Cracking", "D40": "Pothole", } SEVERITY_MAP = { "D00": "low", "D10": "medium", "D20": "high", "D40": "critical", } CLASS_COLORS_RGB = { "D00": (255, 220, 0), "D10": (255, 140, 0), "D20": (255, 60, 0), "D40": (220, 0, 0), } # ── Model loading ────────────────────────────────────────────────────────── model = None model_class_names = FALLBACK_CLASS_NAMES.copy() def load_model(): global model, model_class_names if MODEL_PATH.exists(): try: from ultralytics import YOLO model = YOLO(str(MODEL_PATH)) # ✅ Read class names FROM the model itself — fixes "unknown" class bug raw_names = model.names # e.g. {0: 'D00', 1: 'D10', ...} model_class_names = {k: str(v).upper() for k, v in raw_names.items()} print(f"✅ Model loaded from {MODEL_PATH}") print(f"✅ Model classes: {model_class_names}") except Exception as e: print(f"⚠️ Could not load model: {e}. Running in demo mode.") model = None else: print(f"⚠️ {MODEL_PATH} not found. Running in demo mode.") model = None @app.on_event("startup") async def startup_event(): load_model() # ── Helper: draw bounding boxes ──────────────────────────────────────────── def draw_boxes(image_pil: Image.Image, detections: list) -> Image.Image: draw = ImageDraw.Draw(image_pil) for det in detections: cls_name = det["class"] conf = det["confidence"] x1, y1, x2, y2 = det["bbox"] color = CLASS_COLORS_RGB.get(cls_name, (255, 255, 255)) color_hex = "#{:02x}{:02x}{:02x}".format(*color) # Box draw.rectangle([x1, y1, x2, y2], outline=color_hex, width=3) # Label background + text label = f"{cls_name} {conf:.0%}" text_y = max(y1 - 22, 0) draw.rectangle( [x1, text_y, x1 + len(label) * 8 + 6, text_y + 20], fill=(0, 0, 0) ) draw.text((x1 + 3, text_y + 2), label, fill=color_hex) return image_pil # ── Helper: PIL → base64 ─────────────────────────────────────────────────── def pil_to_b64(img: Image.Image, fmt: str = "JPEG") -> str: buf = io.BytesIO() img.save(buf, format=fmt, quality=88) return base64.b64encode(buf.getvalue()).decode() # ── Demo mode fallback ───────────────────────────────────────────────────── def demo_detections(w: int, h: int) -> list: candidates = [ {"class": "D40", "confidence": round(random.uniform(0.72, 0.94), 2), "bbox": [int(w*0.15), int(h*0.35), int(w*0.42), int(h*0.65)]}, {"class": "D00", "confidence": round(random.uniform(0.55, 0.78), 2), "bbox": [int(w*0.50), int(h*0.20), int(w*0.85), int(h*0.45)]}, {"class": "D20", "confidence": round(random.uniform(0.60, 0.80), 2), "bbox": [int(w*0.05), int(h*0.60), int(w*0.30), int(h*0.90)]}, ] return random.sample(candidates, random.randint(1, 3)) # ── Routes ───────────────────────────────────────────────────────────────── # ✅ FIX: serve the HTML frontend at root instead of JSON @app.get("/") async def root(): return FileResponse("static/index.html") @app.get("/health") async def health(): return { "status": "ok", "model_loaded": model is not None, "mode": "live" if model is not None else "demo", } # ✅ NEW: debug endpoint — tells you exactly what class names your model uses @app.get("/debug") async def debug(): if model is None: return {"model_loaded": False, "mode": "demo"} return { "model_loaded": True, "mode": "live", "model_path": str(MODEL_PATH), "model_classes": model_class_names, "num_classes": len(model_class_names), "model_task": model.task, } @app.get("/classes") async def get_classes(): return { "classes": [ { "id": k, "code": v, "label": CLASS_LABELS.get(v, v), "severity": SEVERITY_MAP.get(v, "unknown"), } for k, v in model_class_names.items() ] } @app.post("/predict") async def predict(file: UploadFile = File(...)): if not file.content_type.startswith("image/"): raise HTTPException(status_code=400, detail="File must be an image.") start = time.time() contents = await file.read() try: image = Image.open(io.BytesIO(contents)).convert("RGB") except Exception: raise HTTPException(status_code=400, detail="Could not read image file.") w, h = image.size # ── Live inference ──────────────────────────────────────────────────── if model is not None: # ✅ Tuned thresholds: conf=0.30 reduces noise, iou=0.4 reduces duplicate boxes results = model(image, conf=0.30, iou=0.4)[0] detections = [] for box in results.boxes: cls_id = int(box.cls[0]) # ✅ Use model's own embedded class names — fixes "unknown" bug cls_name = model_class_names.get(cls_id, f"cls_{cls_id}") conf = float(box.conf[0]) x1, y1, x2, y2 = [round(v) for v in box.xyxy[0].tolist()] detections.append({ "class": cls_name, "label": CLASS_LABELS.get(cls_name, cls_name), "confidence": round(conf, 4), "severity": SEVERITY_MAP.get(cls_name, "low"), "bbox": [x1, y1, x2, y2], }) demo = False # ── Demo fallback ───────────────────────────────────────────────────── else: detections = demo_detections(w, h) for d in detections: d["label"] = CLASS_LABELS.get(d["class"], d["class"]) d["severity"] = SEVERITY_MAP.get(d["class"], "unknown") demo = True # ── Annotate image ──────────────────────────────────────────────────── annotated = draw_boxes(image.copy(), detections) annotated_b64 = pil_to_b64(annotated) original_b64 = pil_to_b64(image) elapsed = round(time.time() - start, 3) # ── Summary stats ───────────────────────────────────────────────────── severity_counts = {"low": 0, "medium": 0, "high": 0, "critical": 0} for d in detections: sev = d["severity"] if sev in severity_counts: severity_counts[sev] += 1 road_score = max(0, 100 - ( severity_counts["critical"] * 30 + severity_counts["high"] * 20 + severity_counts["medium"] * 10 + severity_counts["low"] * 5 )) return JSONResponse({ "success": True, "demo_mode": demo, "inference_ms": int(elapsed * 1000), "image_size": {"width": w, "height": h}, "detections": detections, "detection_count": len(detections), "severity_summary": severity_counts, "road_health_score": road_score, "annotated_image": annotated_b64, "original_image": original_b64, })