from fastapi import FastAPI, UploadFile, File
from transformers import pipeline
from PIL import Image
import io

# Initialize the FastAPI app
app = FastAPI()

# Load the YOLOv5 model for object detection
from transformers import pipeline
from PIL import Image

pipe = pipeline("object-detection", model="facebook/detr-resnet-50")



@app.post("/detect_objects/")
async def detect_objects(file: UploadFile = File(...)):
    # Read image from uploaded file
    image_bytes = await file.read()
    image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
    
    # Perform object detection
    results = pipe(image)
    
    # Extract bounding box coordinates and labels
    bounding_boxes = []
    for result in results:
        box = result['box']
        bbox = {
            'label': result['label'],
            'confidence': result['score'],
            'x_min': box['xmin'],
            'y_min': box['ymin'],
            'x_max': box['xmax'],
            'y_max': box['ymax']
        }
        bounding_boxes.append(bbox)
    
    # Return bounding box coordinates as response
    return {"bounding_boxes": bounding_boxes}

# To run the server, use: uvicorn app:app --reload