|
|
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() |
|
|
|
|
|
|
|
|
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): |
|
|
|
|
|
image = Image.open(file.file).convert("RGB") |
|
|
image_np = np.array(image) |
|
|
|
|
|
|
|
|
results = model(image_np) |
|
|
|
|
|
|
|
|
annotated_frame = results[0].plot() |
|
|
|
|
|
|
|
|
_, buffer = cv2.imencode(".jpg", annotated_frame) |
|
|
img_bytes = buffer.tobytes() |
|
|
|
|
|
|
|
|
img_base64 = base64.b64encode(img_bytes).decode("utf-8") |
|
|
|
|
|
|
|
|
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) |
|
|
|