from fastapi import FastAPI, UploadFile, File
from transformers import YolosForObjectDetection, YolosImageProcessor
from PIL import Image
import torch
import io

app = FastAPI()

# Load the YOLOv5 model and processor from Hugging Face
model = YolosForObjectDetection.from_pretrained('ultralytics/yolov5')
processor = YolosImageProcessor.from_pretrained('ultralytics/yolov5')
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)
model.eval()

@app.post("/predict/")
async def predict(file: UploadFile = File(...)):
    contents = await file.read()
    image = Image.open(io.BytesIO(contents))
    inputs = processor(images=image, return_tensors="pt").to(device)
    outputs = model(**inputs)
    return outputs