Create app.py

app.py

@@ -0,0 +1,101 @@
+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")
+CONF_THRES = float(os.getenv("CONF_THRES", "0.25"))
+
+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()
+    results = model.predict(image, conf=CONF_THRES, 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,
+    }