| 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") |
|
|