Update app.py

app.py

@@ -1,52 +1,61 @@
-from fastapi import FastAPI, UploadFile, File
+from fastapi import FastAPI, File, UploadFile
 from fastapi.responses import JSONResponse
 from ultralytics import YOLO
-import os
-import shutil
+from PIL import Image
+from io import BytesIO
+import base64
+import cv2
+import numpy as np
 
 app = FastAPI()
 
-# ✅ Load both YOLO models
-FIELD_MODEL_PATH = "model/field/best.pt"
-PLAYER_MODEL_PATH = "model/player/last.pt"
+# Load models once at startup
+player_model = YOLO("models/player/best.pt")
+field_model = YOLO("models/field/best.pt")
 
-field_model = YOLO(FIELD_MODEL_PATH)
-player_model = YOLO(PLAYER_MODEL_PATH)
 
+@app.get("/")
+def home():
+    return {"message": "Server running ✅ Use /predict/player or /predict/field"}
 
-def run_detection(model, file_path):
-    results = model(file_path)
-    detections = []
-    for r in results:
-        for box in r.boxes:
-            detections.append({
-                "class": int(box.cls),
-                "confidence": float(box.conf),
-                "bbox": box.xyxy[0].tolist()
-            })
-    return detections
 
+def process_image(file, model):
+    # Load uploaded image
+    image = Image.open(file.file).convert("RGB")
+    image_np = np.array(image)
 
-@app.post("/predict/field")
-async def predict_field(file: UploadFile = File(...)):
-    temp_path = f"/tmp/{file.filename}"
-    with open(temp_path, "wb") as buffer:
-        shutil.copyfileobj(file.file, buffer)
-    detections = run_detection(field_model, temp_path)
-    os.remove(temp_path)
-    return JSONResponse({"model": "field", "detections": detections})
+    # Run inference
+    results = model(image_np)
+
+    # Draw detections
+    annotated_frame = results[0].plot()
+
+    # Convert annotated image to bytes
+    _, buffer = cv2.imencode(".jpg", annotated_frame)
+    img_bytes = buffer.tobytes()
+
+    # Encode image as base64 to include in JSON response
+    img_base64 = base64.b64encode(img_bytes).decode("utf-8")
+
+    # Prepare JSON result
+    detections = []
+    for box in results[0].boxes:
+        detections.append({
+            "class": int(box.cls[0]),
+            "confidence": float(box.conf[0]),
+            "bbox": [float(x) for x in box.xyxy[0].tolist()]
+        })
+
+    return {"detections": detections, "image_base64": img_base64}
 
 
 @app.post("/predict/player")
 async def predict_player(file: UploadFile = File(...)):
-    temp_path = f"/tmp/{file.filename}"
-    with open(temp_path, "wb") as buffer:
-        shutil.copyfileobj(file.file, buffer)
-    detections = run_detection(player_model, temp_path)
-    os.remove(temp_path)
-    return JSONResponse({"model": "player", "detections": detections})
+    result = process_image(file, player_model)
+    return JSONResponse(result)
 
 
-@app.get("/")
-def home():
-    return JSONResponse({"message": "Server running. Use /predict/player or /predict/field endpoints."})
+@app.post("/predict/field")
+async def predict_field(file: UploadFile = File(...)):
+    result = process_image(file, field_model)
+    return JSONResponse(result)