Update main.py

main.py

@@ -1,12 +1,10 @@
 from fastapi import FastAPI, UploadFile, File
 from fastapi.middleware.cors import CORSMiddleware
 from ultralytics import YOLO
-from PIL import Image
-import io, requests, os, tempfile
-import cv2
+import os, tempfile
 
 app = FastAPI()
-#
+
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -25,12 +23,10 @@ app.add_middleware(
 #             f.write(r.content)
 #         print(" Model ready")
 
-# download_model()
 print("Loading model...")
 model = YOLO("best.pt")
 print("Model ready")
 
-
 CLASS_NAMES = {
     0: 'crack',
     1: 'other',
@@ -52,62 +48,24 @@ def root():
 @app.post("/detect")
 async def detect(file: UploadFile = File(...)):
     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)
-        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
-
-        # نحلل 5 frames موزعة على الفيديو
-        sample_points  = [int(total_frames * i / 5) for i in range(1, 6)]
-        all_detections = []
-
-        for frame_num in sample_points:
-            cap.set(cv2.CAP_PROP_POS_FRAMES, frame_num)
-            ret, frame = cap.read()
-            if not ret:
-                continue
 
-            img     = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
-            results = model(img)
+    suffix = '.' + file.filename.split('.')[-1]
+    with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as tmp:
+        tmp.write(contents)
+        tmp_path = tmp.name
 
-            for box in results[0].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
-                })
+    results = model.predict(
+        source  = tmp_path,
+        conf    = 0.25,
+        verbose = False,
+        stream  = True
+    )
 
-        cap.release()
-        os.unlink(tmp_path)
+    all_detections = []
+    frame_num      = 0
 
-        return {
-            "total"     : len(all_detections),
-            "detections": all_detections
-        }
-
-    # ── صورة ──
-    else:
-        try:
-            img = Image.open(io.BytesIO(contents))
-        except Exception as e:
-            return {"error": str(e), "total": 0, "detections": []}
-
-        results        = model(img)
-        all_detections = []
-
-        for box in results[0].boxes:
+    for result in results:
+        for box in result.boxes:
             cls   = int(box.cls)
             conf  = float(box.conf)
             xywhn = box.xywhn[0].tolist()
@@ -117,10 +75,13 @@ async def detect(file: UploadFile = File(...)):
                 "confidence" : round(conf, 3),
                 "severity"   : get_severity(conf, area),
                 "bbox"       : xywhn,
-                "frame"      : 0
+                "frame"      : frame_num
             })
+        frame_num += 1
+
+    os.unlink(tmp_path)
 
-        return {
-            "total"     : len(all_detections),
-            "detections": all_detections
-        }
+    return {
+        "total"     : len(all_detections),
+        "detections": all_detections
+    }