from fastapi import FastAPI, UploadFile, File
from fastapi.responses import HTMLResponse, Response
import numpy as np
import cv2
from ultralytics import YOLO

app = FastAPI()

model = YOLO("detection.pt")

@app.get("/", response_class=HTMLResponse)
def home():
    with open("index.html", "r") as f:
        return f.read()

@app.post("/detect")
async def detect(file: UploadFile = File(...)):
    contents = await file.read()

    nparr = np.frombuffer(contents, np.uint8)
    frame = cv2.imdecode(nparr, cv2.IMREAD_COLOR)

    results = model(frame)
    annotated_frame = results[0].plot()

    _, buffer = cv2.imencode(".jpg", annotated_frame)

    return Response(content=buffer.tobytes(), media_type="image/jpeg")