File size: 673 Bytes
a83d16c
 
 
 
 
 
 
170f4f4
a83d16c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
from fastapi import FastAPI, UploadFile, File
from ultralytics import YOLO
import uvicorn
from PIL import Image
import io

app = FastAPI()
model = YOLO("model.pt")  # Load the trained YOLOv11/12 model

@app.get("/")
def read_root():
    return {"message": "Welcome to the Garbage Detection API"}

@app.post("/predict/")
async def predict(file: UploadFile = File(...)):
    image_bytes = await file.read()
    image = Image.open(io.BytesIO(image_bytes)).convert("RGB")

    results = model(image)
    boxes = results[0].boxes.xyxy.tolist()  # List of [x1, y1, x2, y2]
    confidences = results[0].boxes.conf.tolist()

    return {"boxes": boxes, "confidences": confidences}