Update app.py

app.py

@@ -1,54 +1,47 @@
 from fastapi import FastAPI, UploadFile, File
 from fastapi.responses import JSONResponse
-from fastapi.responses import StreamingResponse
 from ultralytics import YOLO
-from PIL import Image
-import io
+import os
+import shutil
 
-app = FastAPI(title="YOLO API - Football Model")
+app = FastAPI()
 
-# Load model
-try:
-    model = YOLO("best.pt")
-    print("✅ Loaded best.pt")
-except:
-    model = YOLO("last.pt")
-    print("⚠️ Loaded last.pt")
+# ✅ Load both YOLO models
+FIELD_MODEL_PATH = "models/field/best.pt"
+PLAYER_MODEL_PATH = "models/player/last.pt"
 
-@app.get("/")
-def home():
-    return {"message": "YOLO Football Model API is running!"}
+field_model = YOLO(FIELD_MODEL_PATH)
+player_model = YOLO(PLAYER_MODEL_PATH)
 
 
-@app.post("/predict")
-async def predict(file: UploadFile = File(...)):
-    contents = await file.read()
-    image = Image.open(io.BytesIO(contents)).convert("RGB")
-
-    # Run inference
-    results = model(image)
-
-    # Get JSON detections
+def run_detection(model, file_path):
+    results = model(file_path)
     detections = []
-    for box in results[0].boxes:
-        detections.append({
-            "class": int(box.cls),
-            "confidence": float(box.conf),
-            "bbox": box.xyxy[0].tolist()
-        })
-
-    # Save image with bounding boxes drawn by YOLO
-    annotated_image = results[0].plot()  # numpy array (BGR)
-    annotated_image = Image.fromarray(annotated_image[..., ::-1])  # convert BGR→RGB
-
-    # Convert to bytes for response
-    img_bytes = io.BytesIO()
-    annotated_image.save(img_bytes, format="JPEG")
-    img_bytes.seek(0)
-
-    # Return both JSON and image
-    return StreamingResponse(
-        img_bytes,
-        media_type="image/jpeg",
-        headers={"detections": str(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
+
+
+@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})
+
+
+@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})