Update api.py

api.py

@@ -1,15 +1,19 @@
 from fastapi import FastAPI, UploadFile, File, HTTPException, Security, Depends
 from fastapi.security.api_key import APIKeyHeader
-from fastapi.responses import FileResponse, JSONResponse
+from fastapi.responses import JSONResponse, StreamingResponse
 import uvicorn
-import shutil
-import os
+import io
+import numpy as np
+from PIL import Image
+import cv2
 from ultralytics import YOLO
+import requests
+import os
 
 # ==========================
 # 🔑 Sécurité : API Key
 # ==========================
-API_KEY = "12345"  # <-- Change ici avant de partager
+API_KEY = "1234"  # <-- Change ici avant de partager
 api_key_header = APIKeyHeader(name="X-API-Key")
 
 def verify_api_key(api_key: str = Security(api_key_header)):
@@ -43,13 +47,13 @@ async def predict(
     api_key: str = Depends(verify_api_key)
 ):
     try:
-        # Sauvegarde temporaire
-        temp_file = f"temp_{file.filename}"
-        with open(temp_file, "wb") as buffer:
-            shutil.copyfileobj(file.file, buffer)
+        # Lire directement en mémoire
+        contents = await file.read()
+        image = Image.open(io.BytesIO(contents)).convert("RGB")
+        np_image = np.array(image)
 
         # Prédiction YOLO
-        results = model.predict(temp_file)
+        results = model.predict(np_image, conf=0.5, verbose=False)
 
         output = []
         for r in results:
@@ -57,16 +61,13 @@ async def predict(
                 output.append({
                     "class": r.names[int(box.cls[0].item())],
                     "confidence": float(box.conf[0].item()),
-                    "bbox": box.xyxy[0].tolist()  # [x1, y1, x2, y2]
+                    "bbox": box.xyxy[0].tolist()
                 })
 
         return JSONResponse(content={"predictions": output})
 
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
-    finally:
-        if os.path.exists(temp_file):
-            os.remove(temp_file)
 
 # ==========================
 # 📦 Endpoint Image
@@ -77,25 +78,65 @@ async def predict_image(
     api_key: str = Depends(verify_api_key)
 ):
     try:
-        # Sauvegarde temporaire
-        temp_file = f"temp_{file.filename}"
-        with open(temp_file, "wb") as buffer:
-            shutil.copyfileobj(file.file, buffer)
+        # Lire directement en mémoire
+        contents = await file.read()
+        image = Image.open(io.BytesIO(contents)).convert("RGB")
+        np_image = np.array(image)
 
-        # Prédiction YOLO avec sauvegarde image annotée
-        results = model.predict(temp_file, save=True, project="runs", name="detect", exist_ok=True)
+        # Prédiction YOLO + image annotée
+        results = model.predict(np_image, conf=0.5, verbose=False)
+        annotated = results[0].plot()
 
-        # Récupérer la dernière image annotée
-        output_dir = results[0].save_dir
-        output_path = os.path.join(output_dir, os.path.basename(temp_file))
+        # Convertir en bytes pour StreamingResponse
+        annotated_pil = Image.fromarray(cv2.cvtColor(annotated, cv2.COLOR_BGR2RGB))
+        img_byte_arr = io.BytesIO()
+        annotated_pil.save(img_byte_arr, format="PNG")
+        img_byte_arr.seek(0)
 
-        return FileResponse(output_path, media_type="image/png")
+        return StreamingResponse(img_byte_arr, media_type="image/png")
 
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
-    finally:
-        if os.path.exists(temp_file):
-            os.remove(temp_file)
+
+# ==========================
+# 🧪 Endpoint Test interne
+# ==========================
+@app.get("/test_request/")
+async def test_request():
+    """
+    Test interne de l'API déployée sur Hugging Face.
+    Utilise une image locale 'test.jpg' (⚠️ à placer dans ton repo Space).
+    """
+    try:
+        file_path = "test.jpg"  # ⚠️ Mets une image dans ton Space
+        base_url = "https://stroke-ia-api.hf.space"  # ⚠️ adapte au nom exact de ton Space
+
+        if not os.path.exists(file_path):
+            return {"error": f"{file_path} introuvable dans le Space."}
+
+        # Test JSON
+        url_predict = f"{base_url}/v1/predict/"
+        files = {"file": open(file_path, "rb")}
+        headers = {"X-API-Key": API_KEY}
+        response = requests.post(url_predict, files=files, headers=headers)
+        json_result = response.json()
+
+        # Test image annotée
+        url_img = f"{base_url}/v1/predict_image/"
+        files = {"file": open(file_path, "rb")}
+        response_img = requests.post(url_img, files=files, headers=headers)
+
+        with open("result.png", "wb") as f:
+            f.write(response_img.content)
+
+        return {
+            "message": "✅ Test request exécuté sur Hugging Face API. Résultats sauvegardés.",
+            "json_result": json_result,
+            "saved_image": "result.png"
+        }
+
+    except Exception as e:
+        return {"error": str(e)}
 
 # ==========================
 # 🚀 Lancement local