od-test-1 / app.py
kedimestan's picture
Update app.py
0d857ca verified
raw
history blame contribute delete
No virus
1.18 kB
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