iam-ram's picture
Upload 2 files
f2b8349 verified
Raw History Blame Contribute Delete
1.05 kB
from fastapi import FastAPI, File, UploadFile
from fastapi.responses import JSONResponse
from PIL import Image
import io
from ultralytics import YOLO
app = FastAPI()
model = YOLO("best.pt") # Ensure this file is in your project folder
@app.post("/detect")
async def detect(file: UploadFile = File(...)):
contents = await file.read()
image = Image.open(io.BytesIO(contents)).convert("RGB")
results = model(image)
detections = []
for r in results:
for box in r.boxes:
cls_id = int(box.cls)
conf = float(box.conf)
label = model.names[cls_id]
x1, y1, x2, y2 = box.xyxy[0]
detections.append({
"label": label,
"confidence": round(conf, 2),
"box": [
round(float(x1), 2),
round(float(y1), 2),
round(float(x2 - x1), 2),
round(float(y2 - y1), 2)
]
})
return JSONResponse(content={"detections": detections})