Update api.py

api.py

@@ -4,14 +4,10 @@ from fastapi.responses import JSONResponse, StreamingResponse
 import uvicorn
 import logging
 import io
-import os
-from typing import Tuple, Optional
 import time
 import numpy as np
 from PIL import Image
 import cv2
-
-# ML
 from ultralytics import YOLO
 import mediapipe as mp
 
@@ -29,10 +25,7 @@ def verify_api_key(api_key: str = Security(api_key_header)):
 # ==========================
 # 📝 Logger
 # ==========================
-logging.basicConfig(
-    level=logging.INFO,
-    format="%(asctime)s - %(levelname)s - %(message)s"
-)
+logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
 logger = logging.getLogger("stroke-api")
 
 # ==========================
@@ -41,11 +34,7 @@ logger = logging.getLogger("stroke-api")
 app = FastAPI(
     title="Stroke Detection API",
     version="1.2.0",
-    description="""
-    🚑 Stroke Detection API using YOLOv8 + Face Detection (MediaPipe)
-
-    ⚠️ **Disclaimer**: Research/demo only — not a medical device.
-    """
+    description="🚑 Stroke Detection API using YOLOv8 + Face Detection (MediaPipe). Research/demo only."
 )
 
 # ==========================
@@ -66,18 +55,13 @@ mp_face_detection = mp.solutions.face_detection
 ALLOWED_EXT = (".png", ".jpg", ".jpeg")
 ALLOWED_MIME = {"image/png", "image/jpeg"}
 MAX_BYTES = 8 * 1024 * 1024  # 8 MB
-CROP_ON_FACE = True          # recadrer sur le visage détecté
+CROP_ON_FACE = True
 
 def _validate_file(file: UploadFile, raw: bytes):
-    # extension
     if not file.filename.lower().endswith(ALLOWED_EXT):
-        raise HTTPException(status_code=400, detail="Invalid file extension. Use .png/.jpg/.jpeg")
-    # MIME
-    if (file.content_type or "").lower() not in ALLOWED_MIME:
-        # On continue si extension OK mais content_type vide côté client
-        if file.content_type:
-            raise HTTPException(status_code=400, detail="Invalid content-type. Use image/png or image/jpeg")
-    # taille
+        raise HTTPException(status_code=400, detail="Invalid file extension")
+    if (file.content_type or "").lower() not in ALLOWED_MIME and file.content_type:
+        raise HTTPException(status_code=400, detail="Invalid content-type")
     if len(raw) > MAX_BYTES:
         raise HTTPException(status_code=413, detail=f"Image too large. Max {MAX_BYTES//(1024*1024)} MB")
 
@@ -88,10 +72,7 @@ def _read_image_to_numpy(raw: bytes) -> np.ndarray:
     except Exception:
         raise HTTPException(status_code=400, detail="Unreadable image file")
 
-def _largest_face_bbox(np_img: np.ndarray, min_conf: float = 0.6) -> Optional[Tuple[int,int,int,int]]:
-    """
-    Retourne (x1,y1,x2,y2) du plus grand visage détecté, ou None.
-    """
+def _largest_face_bbox(np_img: np.ndarray, min_conf: float = 0.6):
     h, w = np_img.shape[:2]
     with mp_face_detection.FaceDetection(min_detection_confidence=min_conf) as fd:
         results = fd.process(cv2.cvtColor(np_img, cv2.COLOR_RGB2BGR))
@@ -105,26 +86,22 @@ def _largest_face_bbox(np_img: np.ndarray, min_conf: float = 0.6) -> Optional[Tu
             x2 = int(min(1.0, rel.xmin + rel.width) * w)
             y2 = int(min(1.0, rel.ymin + rel.height) * h)
             boxes.append((x1, y1, x2, y2))
-        # choisir le plus grand
         boxes.sort(key=lambda b: (b[2]-b[0])*(b[3]-b[1]), reverse=True)
         return boxes[0] if boxes else None
 
-def _crop_to_bbox(np_img: np.ndarray, bbox: Tuple[int,int,int,int], margin: float = 0.15) -> np.ndarray:
+def _crop_to_bbox(np_img: np.ndarray, bbox, margin: float = 0.15) -> np.ndarray:
     h, w = np_img.shape[:2]
     x1, y1, x2, y2 = bbox
     bw, bh = x2 - x1, y2 - y1
-    # marge autour du visage
-    dx, dy = int(bw * margin), int(bh * margin)
-    X1 = max(0, x1 - dx)
-    Y1 = max(0, y1 - dy)
-    X2 = min(w, x2 + dx)
-    Y2 = min(h, y2 + dy)
+    dx, dy = int(bw*margin), int(bh*margin)
+    X1, Y1 = max(0,x1-dx), max(0,y1-dy)
+    X2, Y2 = min(w,x2+dx), min(h,y2+dy)
     return np_img[Y1:Y2, X1:X2].copy()
 
-def _annotate_face_box(np_img: np.ndarray, bbox: Tuple[int,int,int,int]) -> np.ndarray:
+def _annotate_face_box(np_img: np.ndarray, bbox) -> np.ndarray:
     annotated = cv2.cvtColor(np_img, cv2.COLOR_RGB2BGR).copy()
     x1, y1, x2, y2 = bbox
-    cv2.rectangle(annotated, (x1, y1), (x2, y2), (0, 255, 0), 2)  # couleur par défaut
+    cv2.rectangle(annotated, (x1, y1), (x2, y2), (0,255,0), 2)
     return annotated
 
 # ==========================
@@ -138,121 +115,88 @@ async def health():
 # 📦 Endpoint JSON
 # ==========================
 @app.post("/v1/predict/")
-async def predict(
-    file: UploadFile = File(...),
-    api_key: str = Depends(verify_api_key)
-):
+async def predict(file: UploadFile = File(...), api_key: str = Depends(verify_api_key)):
     raw = await file.read()
     _validate_file(file, raw)
-
     try:
         np_img = _read_image_to_numpy(raw)
-
-        # 1) Détection visage obligatoire
         face_bbox = _largest_face_bbox(np_img)
         if face_bbox is None:
-            return JSONResponse(
-                status_code=422,
-                content={"status": "error", "message": "Aucun visage humain détecté. Veuillez centrer le visage."}
-            )
-
-        # 2) Option : recadrer sur le visage pour améliorer la détection
+            return JSONResponse(status_code=422, content={"status":"error","message":"Aucun visage détecté"})
         input_img = _crop_to_bbox(np_img, face_bbox) if CROP_ON_FACE else np_img
 
-        # 3) YOLO inference (en mémoire)
         start_time = time.time()
         results = model.predict(source=input_img, verbose=False)
         elapsed = time.time() - start_time
-        
-        
-        # 4) Format des prédictions
-        output = []
-        for r in results:
-            for box in r.boxes:
-                output.append({
-                    "class": r.names[int(box.cls[0].item())],
-                    "confidence": float(box.conf[0].item()),
-                    "bbox": box.xyxy[0].tolist()
-                })
 
-        logger.info(f"/predict {file.filename} -> {len(output)} detections (face ok)")
+        output = [{"class": r.names[int(box.cls[0].item())],
+                   "confidence": float(box.conf[0].item()),
+                   "bbox": box.xyxy[0].tolist()} 
+                  for r in results for box in r.boxes]
+
         return JSONResponse(content={
             "status": "ok",
             "face_detected": True,
             "face_bbox": list(map(int, face_bbox)),
-            "predictions": output
+            "predictions": output,
+            "inference_time_sec": round(elapsed,3)
         })
 
-    except HTTPException:
-        raise
     except Exception as e:
         logger.exception("Error in /v1/predict")
         raise HTTPException(status_code=500, detail=str(e))
 
 # ==========================
-# 🖼️ Endpoint Image (annotée)
+# 🖼️ Endpoint Image annotée
 # ==========================
 @app.post("/v1/predict_image/")
-async def predict_image(
-    file: UploadFile = File(...),
-    api_key: str = Depends(verify_api_key)
-):
+async def predict_image(file: UploadFile = File(...), api_key: str = Depends(verify_api_key)):
     raw = await file.read()
     _validate_file(file, raw)
-
     try:
         np_img = _read_image_to_numpy(raw)
-
-        # 1) Détection visage
         face_bbox = _largest_face_bbox(np_img)
         if face_bbox is None:
-            return JSONResponse(
-                status_code=422,
-                content={"status": "error", "message": "Aucun visage humain détecté. Veuillez centrer le visage."}
-            )
-
-        # 2) Recadrer sur le visage (optionnel)
+            return JSONResponse(status_code=422, content={"status":"error","message":"Aucun visage détecté"})
         input_img = _crop_to_bbox(np_img, face_bbox) if CROP_ON_FACE else np_img
 
-        # 3) YOLO
         start_time = time.time()
         results = model.predict(source=input_img, verbose=False)
         elapsed = time.time() - start_time
 
-        # 4) Annotations YOLO
-        yolo_annot = results[0].plot()  # BGR
-        yolo_annot = cv2.cvtColor(yolo_annot, cv2.COLOR_BGR2RGB)
+        yolo_annot = cv2.cvtColor(results[0].plot(), cv2.COLOR_BGR2RGB)
+        out_rgb = yolo_annot if CROP_ON_FACE else _annotate_face_box(np_img, face_bbox)
 
-        # 5) Si on n’a pas recadré, on dessine aussi le cadre visage sur l’image d’origine
-        if not CROP_ON_FACE:
-            annotated = _annotate_face_box(np_img, face_bbox)
-            # fusion simple : ici on retourne juste l’annot YOLO (non redimensionnée)
-            out_rgb = annotated
-        else:
-            # On retourne l’image annotée sur le crop visage
-            out_rgb = yolo_annot
-
-        # 6) Retour en PNG (stream)
         pil_img = Image.fromarray(out_rgb)
         buf = io.BytesIO()
         pil_img.save(buf, format="PNG")
         buf.seek(0)
-        
-        # 7) Ajouter temps d'inférence dans header
-        headers = {"X-Inference-Time": str(round(elapsed,3))}
 
-        logger.info(f"/predict_image {file.filename} -> face ok + image annotée")
-        return StreamingResponse(buf, media_type="image/png")
+        headers = {"X-Inference-Time": str(round(elapsed,3))}
+        return StreamingResponse(buf, media_type="image/png", headers=headers)
 
-    except HTTPException:
-        raise
     except Exception as e:
         logger.exception("Error in /v1/predict_image")
         raise HTTPException(status_code=500, detail=str(e))
 
+# ==========================
+# 🧪 Test automatique
+# ==========================
+@app.get("/test_upload/")
+async def test_upload():
+    try:
+        file_path = "test.jpg"
+        np_img = _read_image_to_numpy(open(file_path,"rb").read())
+        face_bbox = _largest_face_bbox(np_img)
+        if not face_bbox:
+            return {"status":"error","message":"Aucun visage détecté"}
+        results = model.predict(source=np_img, verbose=False)
+        return {"status":"ok","face_detected":True,"num_detections":len(results[0].boxes)}
+    except Exception as e:
+        return {"status":"error","message": str(e)}
+
 # ==========================
 # 🚀 Lancement local
 # ==========================
 if __name__ == "__main__":
-    # Sur HF Spaces, c’est Gradio/Space qui lance; localement :
     uvicorn.run(app, host="0.0.0.0", port=7860)