app.py

import os
import time
from typing import Any, Dict, List

import cv2
import numpy as np
from fastapi import FastAPI, File, UploadFile, HTTPException
from huggingface_hub import hf_hub_download
from ultralytics import YOLO


MODEL_REPO_ID = os.getenv("MODEL_REPO_ID", "HudatersU/road_maintanance")
MODEL_FILENAME = os.getenv("MODEL_FILENAME", "pothole_best2.pt") #pothole_best2.pt
CONF_THRES = float(os.getenv("CONF_THRES", "0.45"))

app = FastAPI(title="Road Maintenance Detection API")

model = None


@app.on_event("startup")
def load_model():
    global model

    model_path = hf_hub_download(
        repo_id=MODEL_REPO_ID,
        filename=MODEL_FILENAME,
    )

    model = YOLO(model_path)


@app.get("/")
def root():
    return {
        "status": "ok",
        "model_repo": MODEL_REPO_ID,
        "model_file": MODEL_FILENAME,
        "endpoints": ["/health", "/predict"],
    }


@app.get("/health")
def health():
    return {
        "status": "ok",
        "model_loaded": model is not None,
    }


@app.post("/predict")
async def predict(file: UploadFile = File(...)) -> Dict[str, Any]:
    if model is None:
        raise HTTPException(status_code=503, detail="Model is not loaded yet")

    image_bytes = await file.read()

    np_arr = np.frombuffer(image_bytes, np.uint8)
    image = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)

    if image is None:
        raise HTTPException(status_code=400, detail="Invalid image")

    h, w = image.shape[:2]

    start = time.time()
    infer_image = cv2.resize(image, (640, 480))
    results = model.predict(infer_image, conf=CONF_THRES, imgsz=640, verbose=False)
    elapsed_ms = round((time.time() - start) * 1000, 2)

    detections: List[Dict[str, Any]] = []

    for r in results:
        if r.boxes is None:
            continue

        for i, box in enumerate(r.boxes):
            x1, y1, x2, y2 = box.xyxy[0].tolist()
            conf = float(box.conf[0])
            cls_id = int(box.cls[0])
            cls_name = model.names.get(cls_id, str(cls_id))

            det = {
                "class_id": cls_id,
                "class_name": cls_name,
                "confidence": conf,
                "box": [x1, y1, x2, y2],
            }

            if r.masks is not None and r.masks.xy is not None and i < len(r.masks.xy):
                det["polygon"] = r.masks.xy[i].tolist()

            detections.append(det)

    return {
        "image": {
            "width": w,
            "height": h,
        },
        "inference_ms": elapsed_ms,
        "detections": detections,
    }