detect / app.py
jefhgofk's picture
Update app.py
a926bbe verified
raw
history blame contribute delete
719 Bytes
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")