Create api.py

api.py

@@ -0,0 +1,249 @@
+from fastapi import FastAPI, UploadFile, File, HTTPException, Security, Depends
+from fastapi.security.api_key import APIKeyHeader
+from fastapi.responses import JSONResponse, StreamingResponse
+import uvicorn
+import logging
+import io
+import os
+from typing import Tuple, Optional
+
+import numpy as np
+from PIL import Image
+import cv2
+
+# ML
+from ultralytics import YOLO
+import mediapipe as mp
+
+# ==========================
+# 🔑 Sécurité : API Key
+# ==========================
+API_KEY = "1234"  # ⚠️ change avant de partager
+api_key_header = APIKeyHeader(name="X-API-Key")
+
+def verify_api_key(api_key: str = Security(api_key_header)):
+    if api_key != API_KEY:
+        raise HTTPException(status_code=403, detail="Forbidden")
+    return api_key
+
+# ==========================
+# 📝 Logger
+# ==========================
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s - %(levelname)s - %(message)s"
+)
+logger = logging.getLogger("stroke-api")
+
+# ==========================
+# 🚀 App
+# ==========================
+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.
+    """
+)
+
+# ==========================
+# 📦 Chargement modèles
+# ==========================
+try:
+    model = YOLO("best.pt")
+    logger.info("✅ YOLO model loaded.")
+except Exception as e:
+    logger.exception("❌ Failed to load YOLO model")
+    raise RuntimeError(f"Model loading failed: {e}")
+
+mp_face_detection = mp.solutions.face_detection
+
+# ==========================
+# 🔧 Utilitaires
+# ==========================
+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é
+
+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
+    if len(raw) > MAX_BYTES:
+        raise HTTPException(status_code=413, detail=f"Image too large. Max {MAX_BYTES//(1024*1024)} MB")
+
+def _read_image_to_numpy(raw: bytes) -> np.ndarray:
+    try:
+        img = Image.open(io.BytesIO(raw)).convert("RGB")
+        return np.array(img)
+    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.
+    """
+    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))
+        if not results.detections:
+            return None
+        boxes = []
+        for det in results.detections:
+            rel = det.location_data.relative_bounding_box
+            x1 = int(max(0, rel.xmin) * w)
+            y1 = int(max(0, rel.ymin) * h)
+            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:
+    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)
+    return np_img[Y1:Y2, X1:X2].copy()
+
+def _annotate_face_box(np_img: np.ndarray, bbox: Tuple[int,int,int,int]) -> 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
+    return annotated
+
+# ==========================
+# 🩺 Healthcheck
+# ==========================
+@app.get("/health")
+async def health():
+    return {"status": "ok", "model_loaded": True}
+
+# ==========================
+# 📦 Endpoint JSON
+# ==========================
+@app.post("/v1/predict/")
+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
+        input_img = _crop_to_bbox(np_img, face_bbox) if CROP_ON_FACE else np_img
+
+        # 3) YOLO inference (en mémoire)
+        results = model.predict(source=input_img, verbose=False)
+
+        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)")
+        return JSONResponse(content={
+            "status": "ok",
+            "face_detected": True,
+            "face_bbox": list(map(int, face_bbox)),
+            "predictions": output
+        })
+
+    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)
+# ==========================
+@app.post("/v1/predict_image/")
+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)
+        input_img = _crop_to_bbox(np_img, face_bbox) if CROP_ON_FACE else np_img
+
+        # 3) YOLO
+        results = model.predict(source=input_img, verbose=False)
+
+        # 4) Annotations YOLO
+        yolo_annot = results[0].plot()  # BGR
+        yolo_annot = cv2.cvtColor(yolo_annot, cv2.COLOR_BGR2RGB)
+
+        # 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)
+
+        logger.info(f"/predict_image {file.filename} -> face ok + image annotée")
+        return StreamingResponse(buf, media_type="image/png")
+
+    except HTTPException:
+        raise
+    except Exception as e:
+        logger.exception("Error in /v1/predict_image")
+        raise HTTPException(status_code=500, detail=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)