Update main.py

main.py

@@ -2,7 +2,8 @@ from fastapi import FastAPI, UploadFile, File
 from fastapi.middleware.cors import CORSMiddleware
 from ultralytics import YOLO
 from PIL import Image
-import io, requests, os
+import io, requests, os, tempfile
+import cv2
 
 app = FastAPI()
 
@@ -21,36 +22,74 @@ def download_model():
         r = requests.get(MODEL_URL)
         with open("best.pt", "wb") as f:
             f.write(r.content)
-        print(" Model ready")
+        print("✅ Model ready")
 
 download_model()
 model = YOLO("best.pt")
 
-CLASS_NAMES = { 0: 'crack', 1: 'other',2: 'pothole'}
+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'
+    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 "}
+    return {"status": "SABIQ API running"}
 
 @app.post("/detect")
 async def detect(file: UploadFile = File(...)):
-    img     = Image.open(io.BytesIO(await file.read()))
-    results = model(img)
-    detections = []
+    contents = await file.read()
+    filename = file.filename.lower()
+
+    # لو فيديو
+    if any(filename.endswith(ext) for ext in ['.mp4', '.avi', '.mov', '.mkv']):
+        with tempfile.NamedTemporaryFile(suffix='.mp4', delete=False) as tmp:
+            tmp.write(contents)
+            tmp_path = tmp.name
+
+        cap = cv2.VideoCapture(tmp_path)
+        ret, frame = cap.read()
+        cap.release()
+        os.unlink(tmp_path)
+
+        if not ret:
+            return {"error": "cannot read video", "total": 0, "detections": []}
+
+        img = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
+
+    # لو صورة
+    else:
+        try:
+            img = Image.open(io.BytesIO(contents))
+        except Exception as e:
+            return {"error": str(e), "total": 0, "detections": []}
+
+    results     = model(img)
+    detections  = []
+
     for box in results[0].boxes:
-        cls      = int(box.cls)
-        conf     = float(box.conf)
-        xywhn    = box.xywhn[0].tolist()
-        area     = xywhn[2] * xywhn[3]
+        cls    = int(box.cls)
+        conf   = float(box.conf)
+        xywhn  = box.xywhn[0].tolist()
+        area   = xywhn[2] * xywhn[3]
+
         detections.append({
             "damage_type": CLASS_NAMES.get(cls, 'other'),
             "confidence" : round(conf, 3),
             "severity"   : get_severity(conf, area),
             "bbox"       : xywhn
         })
-    return {"total": len(detections), "detections": detections}
+
+    return {
+        "total"     : len(detections),
+        "detections": detections
+    }