main.py

from fastapi import FastAPI, UploadFile, File
from fastapi.middleware.cors import CORSMiddleware
from ultralytics import YOLO
import os, tempfile

app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)
#test hi
# هذا لو كنت ابغى يعتمد على مودل موجود بالهقنق فيس
# MODEL_URL = "https://huggingface.co/Rahaf2001/sabiq-road-detection/resolve/main/best.pt"

# def download_model():
#     if not os.path.exists("best.pt"):
#         print("Downloading model...")
#         r = requests.get(MODEL_URL)
#         with open("best.pt", "wb") as f:
#             f.write(r.content)
#         print(" Model ready")

print("Loading model...")
model = YOLO("best.pt")
print("Model ready")

CLASS_NAMES = {
    0: 'crack',
    1: 'other',
    2: 'pothole'
}

def get_severity(conf, area):
    if conf > 0.85 and area > 0.05:
        return 'high'
    elif conf > 0.65:
        return 'medium'
    else:
        return 'low'

@app.get("/")
def root():
    return {"status": "SABIQ API running"}

@app.post("/detect")
async def detect(file: UploadFile = File(...)):
    contents = await file.read()

    suffix = '.' + file.filename.split('.')[-1]
    with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as tmp:
        tmp.write(contents)
        tmp_path = tmp.name

    results = model.predict(
        source  = tmp_path,
        conf    = 0.25,
        verbose = False,
        stream  = True
    )

    all_detections = []
    frame_num      = 0

    for result in results:
        for box in result.boxes:
            cls   = int(box.cls)
            conf  = float(box.conf)
            xywhn = box.xywhn[0].tolist()
            area  = xywhn[2] * xywhn[3]
            all_detections.append({
                "damage_type": CLASS_NAMES.get(cls, 'other'),
                "confidence" : round(conf, 3),
                "severity"   : get_severity(conf, area),
                "bbox"       : xywhn,
                "frame"      : frame_num
            })
        frame_num += 1

    os.unlink(tmp_path)

    return {
        "total"     : len(all_detections),
        "detections": all_detections
    }