Spaces:
Sleeping
Sleeping
File size: 1,612 Bytes
992bcd9 5c049ba ff32071 992bcd9 ff32071 86a5c6a ff32071 992bcd9 8ee3652 ff32071 992bcd9 5c049ba 86a5c6a 992bcd9 86a5c6a 992bcd9 86a5c6a 992bcd9 9cd40e7 992bcd9 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 |
from fastapi import FastAPI, File, UploadFile
from fastapi.responses import JSONResponse
from ultralytics import YOLO
from PIL import Image
from io import BytesIO
import base64
import cv2
import numpy as np
app = FastAPI()
# Load models once at startup
player_model = YOLO("model/player/best.pt")
field_model = YOLO("model/field/best.pt")
@app.get("/")
def home():
return {"message": "Server running ✅ Use /predict/player or /predict/field"}
def process_image(file, model):
# Load uploaded image
image = Image.open(file.file).convert("RGB")
image_np = np.array(image)
# 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(...)):
result = process_image(file, player_model)
return JSONResponse(result)
@app.post("/predict/field")
async def predict_field(file: UploadFile = File(...)):
result = process_image(file, field_model)
return JSONResponse(result)
|