Update app.py

app.py

@@ -1,23 +1,25 @@
 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 torch
 
 app = FastAPI(title="YOLO API - Football Model")
 
-# Load model from local file
+# Load model
 try:
     model = YOLO("best.pt")
-    print("✅ Loaded best.pt model")
-except Exception as e:
-    print("⚠️ best.pt not found, trying last.pt:", e)
+    print("✅ Loaded best.pt")
+except:
     model = YOLO("last.pt")
+    print("⚠️ Loaded last.pt")
 
 @app.get("/")
 def home():
     return {"message": "YOLO Football Model API is running!"}
 
+
 @app.post("/predict")
 async def predict(file: UploadFile = File(...)):
     contents = await file.read()
@@ -26,6 +28,7 @@ async def predict(file: UploadFile = File(...)):
     # Run inference
     results = model(image)
 
+    # Get JSON detections
     detections = []
     for box in results[0].boxes:
         detections.append({
@@ -34,4 +37,18 @@ async def predict(file: UploadFile = File(...)):
             "bbox": box.xyxy[0].tolist()
         })
 
-    return {"detections": detections}
+    # 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)}
+    )