Update app.py

app.py

@@ -5,48 +5,37 @@ import os
 import cv2
 import base64
 from ultralytics import YOLO
-
-# ─── Load Models ──────────────────────────────────────────────────────────────
-person_model = YOLO("yolov8n.pt")  
-
-# ─── App Setup ────────────────────────────────────────────────────────────────
+# ─── Load Models 
+person_model = YOLO("yolov8n.pt")
+statue_model = YOLO("best.pt")
+# ─── App Setup 
 app = FastAPI()
-
 UPLOAD_FOLDER = "/tmp/uploads"
 os.makedirs(UPLOAD_FOLDER, exist_ok=True)
-
 app.mount("/uploads", StaticFiles(directory=UPLOAD_FOLDER), name="uploads")
-
-PERSON_CONF = 0.30   # threshold 
-STATUE_CONF = 0.20   # threshold
-
+CONF_THRESHOLD = 0.30
 router = APIRouter()
-
-
 def draw_label(image, label, x1, y1, x2, y2, box_color, text_color):
+    
     font = cv2.FONT_HERSHEY_SIMPLEX
     font_scale = 0.6
     thickness = 2
     (tw, th), _ = cv2.getTextSize(label, font, font_scale, thickness)
-
+    
     if y1 - th - 8 >= 0:
         label_y1 = y1 - th - 8
         label_y2 = y1
         text_y = y1 - 5
     else:
+        
         label_y1 = y1
         label_y2 = y1 + th + 8
         text_y = y1 + th + 3
-
     cv2.rectangle(image, (x1, label_y1), (x1 + tw + 4, label_y2), box_color, -1)
     cv2.putText(image, label, (x1 + 2, text_y), font, font_scale, text_color, thickness)
-
-
 @app.get("/")
 def root():
     return {"message": "AI API is running 🚀"}
-
-
 @router.post("/predict-image")
 async def predict_image(
     request: Request,
@@ -54,45 +43,27 @@ async def predict_image(
 ):
     safe_filename = file.filename.replace(" ", "_")
     file_path = os.path.join(UPLOAD_FOLDER, safe_filename)
-
     with open(file_path, "wb") as buffer:
         shutil.copyfileobj(file.file, buffer)
-
     image = cv2.imread(file_path)
-
     if image is None:
         return {"error": "Invalid image"}
-
     detections = []
     person_count = 0
     statue_count = 0
-
-    # ── Person Detection ──────────────────────────────────────────────────────
-    # imgsz=1280 يخلي الموديل يشوف الصورة بدقة أعلى داخلياً
-    # فبيلاقط أشخاص صغار/بعيدين كانوا بيضيعوا على الدقة الافتراضية (640)
-    person_results = person_model(
-        file_path,
-        imgsz=1280,
-        conf=PERSON_CONF,
-        iou=0.5,        # نسبة أعلى شوية من 0.45 تمنع لغي أشخاص متقاربين بالغلط
-        max_det=300
-    )
-
+    # ── Person Detection 
+    person_results = person_model(file_path)
     for box in person_results[0].boxes:
         cls_id = int(box.cls)
         if cls_id != 0:
             continue
-
         conf = float(box.conf)
-        if conf < PERSON_CONF:
+        if conf < CONF_THRESHOLD:
             continue
-
         x1, y1, x2, y2 = map(int, box.xyxy[0])
-
         cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)
         draw_label(image, f"Person {conf:.2f}", x1, y1, x2, y2,
                    box_color=(0, 255, 0), text_color=(0, 0, 0))
-
         detections.append({
             "type": "person",
             "name": "Person",
@@ -100,28 +71,18 @@ async def predict_image(
             "bbox": [x1, y1, x2, y2]
         })
         person_count += 1
-
-    # ── Statue Detection ──────────────────────────────────────────────────────
-    # سيبناها كما كانت بالضبط (imgsz عادي + conf منخفض) عشان تفضل تكشف التماثيل صح
-    statue_results = statue_model(
-        file_path,
-        conf=STATUE_CONF
-    )
-
+    # ── Statue Detection 
+    statue_results = statue_model(file_path)
     for box in statue_results[0].boxes:
         conf = float(box.conf)
-        if conf < STATUE_CONF:
+        if conf < CONF_THRESHOLD:
             continue
-
         cls_id = int(box.cls)
         statue_name = statue_results[0].names[cls_id]
-
         x1, y1, x2, y2 = map(int, box.xyxy[0])
-
         cv2.rectangle(image, (x1, y1), (x2, y2), (0, 0, 255), 2)
         draw_label(image, f"{statue_name} {conf:.2f}", x1, y1, x2, y2,
                    box_color=(0, 0, 255), text_color=(255, 255, 255))
-
         detections.append({
             "type": "statue",
             "name": statue_name,
@@ -129,17 +90,13 @@ async def predict_image(
             "bbox": [x1, y1, x2, y2]
         })
         statue_count += 1
-
-    # ── Save Output Image ─────────────────────────────────────────────────────
+    # ── Save Output Image 
     output_filename = f"output_{safe_filename}"
     output_path = os.path.join(UPLOAD_FOLDER, output_filename)
     cv2.imwrite(output_path, image)
-
     with open(output_path, "rb") as img_file:
         image_base64 = base64.b64encode(img_file.read()).decode("utf-8")
-
     image_url = f"{request.base_url}uploads/{output_filename}"
-
     return {
         "total_count": len(detections),
         "persons": person_count,
@@ -148,6 +105,5 @@ async def predict_image(
         "output_image_base64": f"data:image/jpeg;base64,{image_base64}",
         "detections": detections
     }
+app.include_router(router)
 
-
-app.include_router(router)