Update api.py

api.py

@@ -1,109 +1,110 @@
-from fastapi import FastAPI, UploadFile, File
-from fastapi.responses import JSONResponse, StreamingResponse
-from fastapi.middleware.cors import CORSMiddleware
-from ultralytics import YOLO
-import numpy as np
-from PIL import Image
-import io
-import cv2
-
-# Load YOLO model
-model = YOLO("best.pt")
-
-# Class labels
-CLASS_NAMES = [
-    "normalEye",
-    "normalMouth",
-    "strokeEyeMid",
-    "strokeEyeSevere",
-    "strokeEyeWeak",
-    "strokeMouthMid",
-    "strokeMouthSevere",
-    "strokeMouthWeak"
-]
-
-# Initialize FastAPI app
-app = FastAPI(
-    title="Stroke-IA Detection API",
-    description="REST API for stroke sign detection (tech demo, not medical advice).",
-    version="1.0"
-)
-
-# ✅ Enable CORS (to avoid fetch issues in Swagger UI or front-end)
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
-
-@app.get("/")
-async def root():
-    return {"message": "Stroke-IA API is running. Use /predict/ or /predict_image/."}
-
-@app.post("/predict/")
-async def predict(file: UploadFile = File(...)):
-    """
-    Returns JSON with detections (no image).
-    """
-    try:
-        contents = await file.read()
-        image = Image.open(io.BytesIO(contents)).convert("RGB")
-        np_image = np.array(image)
-
-        results = model.predict(source=np_image, conf=0.85, verbose=False)
-
-        if len(results[0].boxes) == 0:
-            return {
-                "message": "✅ No stroke signs detected (confidence ≥ 85%)",
-                "detections": [],
-                "summary": "Healthy face detected with no significant asymmetry."
-            }
-
-        detections = []
-        for box, score, cls in zip(results[0].boxes.xyxy.tolist(),
-                                   results[0].boxes.conf.tolist(),
-                                   results[0].boxes.cls.tolist()):
-            label = CLASS_NAMES[int(cls)]
-            detections.append({
-                "box": box,
-                "score": float(score),
-                "class": int(cls),
-                "label": label
-            })
-
-        best_det = max(detections, key=lambda x: x["score"])
-        summary = f"⚠️ {best_det['label']} detected with {best_det['score']*100:.1f}% confidence."
-
-        return {
-            "message": "⚠️ Possible stroke signs detected",
-            "detections": detections,
-            "summary": summary
-        }
-
-    except Exception as e:
-        return JSONResponse({"error": str(e)}, status_code=500)
-
-@app.post("/predict_image/")
-async def predict_image(file: UploadFile = File(...)):
-    """
-    Returns the annotated image only (PNG).
-    """
-    try:
-        contents = await file.read()
-        image = Image.open(io.BytesIO(contents)).convert("RGB")
-        np_image = np.array(image)
-
-        results = model.predict(source=np_image, conf=0.85, verbose=False)
-
-        annotated = results[0].plot()
-        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 StreamingResponse(img_byte_arr, media_type="image/png")
-
-    except Exception as e:
-        return JSONResponse({"error": str(e)}, status_code=500)
+from fastapi import FastAPI, UploadFile, File
+from fastapi.responses import JSONResponse, StreamingResponse
+from fastapi.middleware.cors import CORSMiddleware
+from ultralytics import YOLO
+import numpy as np
+from PIL import Image
+import io
+import cv2
+
+# Load YOLO model
+MODEL_PATH = os.path.join(os.path.dirname(__file__), "best.pt")
+model = YOLO(MODEL_PATH)
+
+# Class labels
+CLASS_NAMES = [
+    "normalEye",
+    "normalMouth",
+    "strokeEyeMid",
+    "strokeEyeSevere",
+    "strokeEyeWeak",
+    "strokeMouthMid",
+    "strokeMouthSevere",
+    "strokeMouthWeak"
+]
+
+# Initialize FastAPI app
+app = FastAPI(
+    title="Stroke-IA Detection API",
+    description="REST API for stroke sign detection (tech demo, not medical advice).",
+    version="1.0"
+)
+
+# ✅ Enable CORS (to avoid fetch issues in Swagger UI or front-end)
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+@app.get("/")
+async def root():
+    return {"message": "Stroke-IA API is running. Use /predict/ or /predict_image/."}
+
+@app.post("/predict/")
+async def predict(file: UploadFile = File(...)):
+    """
+    Returns JSON with detections (no image).
+    """
+    try:
+        contents = await file.read()
+        image = Image.open(io.BytesIO(contents)).convert("RGB")
+        np_image = np.array(image)
+
+        results = model.predict(source=np_image, conf=0.85, verbose=False)
+
+        if len(results[0].boxes) == 0:
+            return {
+                "message": "✅ No stroke signs detected (confidence ≥ 85%)",
+                "detections": [],
+                "summary": "Healthy face detected with no significant asymmetry."
+            }
+
+        detections = []
+        for box, score, cls in zip(results[0].boxes.xyxy.tolist(),
+                                   results[0].boxes.conf.tolist(),
+                                   results[0].boxes.cls.tolist()):
+            label = CLASS_NAMES[int(cls)]
+            detections.append({
+                "box": box,
+                "score": float(score),
+                "class": int(cls),
+                "label": label
+            })
+
+        best_det = max(detections, key=lambda x: x["score"])
+        summary = f"⚠️ {best_det['label']} detected with {best_det['score']*100:.1f}% confidence."
+
+        return {
+            "message": "⚠️ Possible stroke signs detected",
+            "detections": detections,
+            "summary": summary
+        }
+
+    except Exception as e:
+        return JSONResponse({"error": str(e)}, status_code=500)
+
+@app.post("/predict_image/")
+async def predict_image(file: UploadFile = File(...)):
+    """
+    Returns the annotated image only (PNG).
+    """
+    try:
+        contents = await file.read()
+        image = Image.open(io.BytesIO(contents)).convert("RGB")
+        np_image = np.array(image)
+
+        results = model.predict(source=np_image, conf=0.85, verbose=False)
+
+        annotated = results[0].plot()
+        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 StreamingResponse(img_byte_arr, media_type="image/png")
+
+    except Exception as e:
+        return JSONResponse({"error": str(e)}, status_code=500)