aditya0929's picture
Fix model path in main.py
170f4f4
raw
history blame contribute delete
673 Bytes
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}