main.py

from fastapi import FastAPI, File, UploadFile, Response
from PIL import Image
import torch
import io

app = FastAPI()

# Load the YOLOv5 model
model = torch.hub.load('ultralytics/yolov5', 'custom', path='best.pt')

@app.post("/detect/")
async def detect(file: UploadFile = File(...)):
    # Read image file
    image_data = await file.read()
    image = Image.open(io.BytesIO(image_data))

    # Perform inference
    results = model(image)

    # Render the results on the image
    results.render()

    # Convert the image to bytes
    img_bytes = io.BytesIO()
    image_with_boxes = Image.fromarray(results.ims[0])
    image_with_boxes.save(img_bytes, format='JPEG')
    img_bytes.seek(0)

    # Create a response with the image
    return Response(content=img_bytes.getvalue(), media_type="image/jpeg")