Update api.py

api.py

@@ -1,3 +1,5 @@
+
+
 import os
 import cv2
 import time
@@ -12,100 +14,196 @@ from PIL import Image
 # -----------------------------
 # 1. Config & Model
 # -----------------------------
-MODEL_STROKE_PATH = "stroke.pt"
+MODEL_IRM_PATH = "best_seg.pt"   # <- place ton modèle ici
 OUTPUT_DIR = "/tmp/outputs"
 os.makedirs(OUTPUT_DIR, exist_ok=True)
 
-# Charger YOLO une seule fois
-model_stroke = YOLO(MODEL_STROKE_PATH)
+# Charger YOLO (segmentation) une seule fois
+model_irm = YOLO(MODEL_IRM_PATH)
 
-BASE_URL = "https://stroke-ia-avc-detect.hf.space"  # ⚠️ à adapter selon ton déploiement
+# ⚠️ Adapte ce BASE_URL selon ton déploiement (ex : https://tondomaine.tld)
+BASE_URL = "https://mediscan.caba31.com"  
 
-# Mapping des classes vers un rapport médical
-CLASS_LABELS = {
-    0: "Hémorragie intracrânienne",
-    1: "Suspicion de zone ischémique",
-    2: "Normale Brain",
-    # 👉 adapte en fonction des classes de ton modèle
+# Mapping optionnel des classes (si ton modèle prédit des classes).
+# Si ton modèle ne prédit pas de classes, on utilisera "Lésion suspecte".
+CLASS_LABELS_IRM = {
+    1: "Lésion suspecte",
+    0: "Anomalie secondaire",
+    # Ajoute/édite selon les classes réelles de ton modèle
 }
 
 # -----------------------------
 # 2. Génération de rapport
 # -----------------------------
-def generate_report(results) -> str:
-    boxes = results[0].boxes
-    if len(boxes) == 0:
-        return "=== RAPPORT AUTOMATIQUE ===\n\nAucune anomalie détectée.\n"
-
-    rapport = "=== RAPPORT AUTOMATIQUE AVC ===\n\n"
-    rapport += f"Nombre de lésions détectées : {len(boxes)}\n\n"
+def generate_report_irm(results, image_shape=None) -> str:
+    """
+    Génère un rapport texte simple à partir des résultats YOLO segmentation.
+    On compte ici le nombre de masques et, si des boxes existent, on donne
+    une estimation de la surface relative (en % de l'image).
+    """
+    # Sécurité si pas de masques
+    if results[0].masks is None or len(results[0].masks.data) == 0:
+        return "=== RAPPORT AUTOMATIQUE IRM ===\n\nAucun masque détecté.\n"
 
-    detected_classes = boxes.cls.cpu().numpy().astype(int)
+    rapport = "=== RAPPORT AUTOMATIQUE IRM ===\n\n"
+    # Nombre de masques détectés
+    n_masks = len(results[0].masks.data)
+    rapport += f"Nombre de masques détectés : {n_masks}\n\n"
 
-    for i, cls_id in enumerate(detected_classes, 1):
-        label = CLASS_LABELS.get(cls_id, f"Classe inconnue {cls_id}")
-        rapport += f"- Lésion {i}: {label}\n"
+    # Récupérer boxes si présentes
+    boxes = results[0].boxes
+    has_boxes = len(boxes) > 0
+
+    # Taille image (hauteur, largeur) si fournie
+    img_h, img_w = (None, None)
+    if image_shape is not None:
+        img_h, img_w = image_shape[0], image_shape[1]
+
+    # Récupérer classes si présentes
+    detected_classes = None
+    try:
+        if has_boxes:
+            detected_classes = boxes.cls.cpu().numpy().astype(int)
+    except Exception:
+        detected_classes = None
+
+    for i in range(n_masks):
+        # Classe si disponible (sinon texte générique)
+        if detected_classes is not None and i < len(detected_classes):
+            cls_id = int(detected_classes[i])
+            label = CLASS_LABELS_IRM.get(cls_id, f"Classe {cls_id}")
+        else:
+            label = "Lésion suspecte"
+
+        rapport += f"- Masque {i+1}: {label}\n"
+
+        # Si on a une boîte, calculer surface approximative
+        if has_boxes and i < len(boxes):
+            try:
+                xyxy = boxes.xyxyn[i]  # valeurs normalisées (0..1) si disponible
+                # parfois boxes.xyxyn peut exister ; sinon fallback sur xyxy
+                if xyxy is not None and len(xyxy) == 4:
+                    # xyxyn = (x1_norm, y1_norm, x2_norm, y2_norm)
+                    x1n, y1n, x2n, y2n = [float(x) for x in xyxy]
+                    if img_h and img_w:
+                        width = (x2n - x1n) * img_w
+                        height = (y2n - y1n) * img_h
+                        area = max(width * height, 0.0)
+                        percent = (area / (img_w * img_h)) * 100 if (img_w and img_h) else None
+                        if percent is not None:
+                            rapport += f"  - Surface estimée: {percent:.2f}% de l'image\n"
+                else:
+                    # fallback using absolute xyxy
+                    xyxy_abs = boxes.xyxy[i].cpu().numpy()
+                    x1, y1, x2, y2 = xyxy_abs
+                    width = max(x2 - x1, 0.0)
+                    height = max(y2 - y1, 0.0)
+                    if img_h and img_w:
+                        area = width * height
+                        percent = (area / (img_w * img_h)) * 100
+                        rapport += f"  - Surface estimée: {percent:.2f}% de l'image\n"
+            except Exception:
+                # si erreur, on ignore l'estimation de surface
+                pass
 
     rapport += "\nRecommandations :\n"
-    rapport += "- Vérifier la concordance clinique.\n"
-    rapport += "- Considérer un suivi neurologique urgent.\n"
+    rapport += "- Corréler ces résultats avec le tableau clinique et un radiologue.\n"
+    rapport += "- Si la lésion est significative, envisager un suivi/consultation spécialisée.\n"
 
     return rapport
 
 # -----------------------------
 # 3. FastAPI
 # -----------------------------
-app = FastAPI(title="Stroke Detection API")
+app = FastAPI(title="IRM Segmentation API")
 app.mount("/files", StaticFiles(directory=OUTPUT_DIR), name="files")
 
-@app.post("/predict/")
-async def predict_stroke(image_file: UploadFile = File(...), conf: float = 0.8):
+@app.post("/predict_irm/")
+async def predict_irm(image_file: UploadFile = File(...), conf: float = 0.8, show_labels: bool = True):
     """
-    Endpoint qui reçoit une image IRM et renvoie une image annotée + rapport texte
+    Endpoint qui reçoit une image IRM (upload) et renvoie :
+    - une image annotée (masques/boxes) accessible via URL
+    - un fichier texte de rapport accessible via URL
     """
     # Sauvegarde temporaire
-    tmp_path = f"/tmp/{image_file.filename}"
+    tmp_path = f"/tmp/{datetime.now().strftime('%Y%m%d_%H%M%S')}_{image_file.filename}"
     with open(tmp_path, "wb") as f:
         f.write(await image_file.read())
 
-    # Charger image
-    image = Image.open(tmp_path).convert("RGB")
-    np_img = np.array(image)
-
-    # Conversion en BGR pour OpenCV
-    np_img = cv2.cvtColor(np_img, cv2.COLOR_RGB2BGR)
-
-    # Prédiction
-    results = model_stroke.predict(source=np_img, conf=conf, verbose=False)
-
-    if len(results[0].boxes) == 0:
+    try:
+        # Charger image avec PIL -> RGB
+        image = Image.open(tmp_path).convert("RGB")
+        np_img = np.array(image)
+
+        # Conversion en BGR pour OpenCV (et pour ultralytics si nécessaire)
+        if np_img.shape[2] == 4:
+            np_img = cv2.cvtColor(np_img, cv2.COLOR_RGBA2BGR)
+        else:
+            np_img = cv2.cvtColor(np_img, cv2.COLOR_RGB2BGR)
+
+        # Prédiction
+        results = model_irm.predict(source=np_img, conf=conf, verbose=False)
+
+        # Vérifier présence de masques
+        if results[0].masks is None or len(results[0].masks.data) == 0:
+            # Nettoyage input
+            os.remove(tmp_path)
+            return {"message": "⚠️ Aucun masque détecté par le modèle IRM."}
+
+        # Annoter image (affiche masques + boxes)
+        annotated_image = results[0].plot(labels=show_labels)
+        # Définir timestamp une seule fois
+        timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
+        
+		# 2. Mask image
+        out_mask_name = f"mask_{timestamp}.png"
+        out_mask_path = os.path.join(OUTPUT_DIR, out_mask_name)
+        mask = np.zeros(np_img.shape[:2], dtype=np.uint8)
+        cv2.rectangle(mask, (50, 50), (200, 200), 255, -1)
+        cv2.imwrite(out_mask_path, mask)
+    
+        # Sauvegarder sortie image
+        
+        out_img_name = f"irm_result_{timestamp}.png"
+        out_img_path = os.path.join(OUTPUT_DIR, out_img_name)
+        
+    
+        # Si annotated_image est PIL Image, convertir ; sinon sauver directement si numpy
+        if hasattr(annotated_image, "save"):
+            # PIL Image
+            annotated_image.save(out_img_path)
+        else:
+            # numpy array (probablement BGR)
+            cv2.imwrite(out_img_path, annotated_image)
+            cv2.imwrite(out_mask_path, mask)
+        # Générer & sauvegarder rapport
+        img_shape = np_img.shape[:2]  # (h, w)
+       
+	   rapport_text = generate_report_irm(results)
+        out_txt_name = f"rapport_irm_{timestamp}.txt"
+        out_txt_path = os.path.join(OUTPUT_DIR, out_txt_name)
+        with open(out_txt_path, "w", encoding="utf-8") as f:
+            f.write(rapport_text)
+
+        # Nettoyage input
         os.remove(tmp_path)
-        return {"message": "⚠️ Aucun AVC détecté."}
-
-    # Annoter l’image
-    annotated_image = results[0].plot(labels=True)
-
-    # Sauvegarder sortie image
-    timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
-    out_img_name = f"stroke_result_{timestamp}.png"
-    out_img_path = os.path.join(OUTPUT_DIR, out_img_name)
-    cv2.imwrite(out_img_path, annotated_image)
-
-    # Sauvegarder rapport
-    rapport_text = generate_report(results)
-    out_txt_name = f"rapport_{timestamp}.txt"
-    out_txt_path = os.path.join(OUTPUT_DIR, out_txt_name)
-    with open(out_txt_path, "w", encoding="utf-8") as f:
-        f.write(rapport_text)
-
-    # Nettoyage input
-    os.remove(tmp_path)
-
-    return {
-        "annotated_result_url": f"{BASE_URL}/files/{out_img_name}",
-        "rapport_url": f"{BASE_URL}/files/{out_txt_name}",
-        "message": "✅ Prédiction réussie avec rapport"
-    }
+
+        return {
+            "annotated_result_url": f"{BASE_URL}/files/{out_img_name}",
+            "rapport_url": f"{BASE_URL}/files/{out_txt_name}",
+            "message": "✅ Prédiction réussie avec rapport"
+            
+            "message": "✅ Prédiction réussie avec rapport"
+        }
+
+    except Exception as e:
+        # Nettoyage fichier temporaire si existant
+        try:
+            if os.path.exists(tmp_path):
+                os.remove(tmp_path)
+        except Exception:
+            pass
+        return {"error": str(e)}
 
 # -----------------------------
 # 4. Auto-cleanup toutes les 10 min