File size: 569 Bytes
9ef0c02
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
from fastapi import FastAPI, File, UploadFile
from fastapi.responses import JSONResponse
from ultralytics import YOLO
from PIL import Image
import io

app = FastAPI()
model = YOLO("best.pt")  # make sure this file exists

@app.post("/detect")
async def detect_boxes(file: UploadFile = File(...)):
    contents = await file.read()
    img = Image.open(io.BytesIO(contents)).convert("RGB")

    # Inference
    results = model.predict(img, conf=0.3, iou=0.5, verbose=False)
    box_count = len(results[0].boxes)

    return JSONResponse(content={"box_count": box_count})