| import torch | |
| from ultralytics import YOLO | |
| from PIL import Image | |
| import io | |
| # Load YOLOv8 model | |
| model = YOLO("model.pt") | |
| def predict(image_bytes): | |
| img = Image.open(io.BytesIO(image_bytes)) | |
| results = model.predict(img) | |
| output = [] | |
| for result in results: | |
| for i in range(len(result.boxes)): | |
| output.append({ | |
| "bbox": result.boxes.xyxy[i].tolist(), | |
| "class": int(result.boxes.cls[i].item()), | |
| "confidence": float(result.boxes.conf[i].item()) | |
| }) | |
| return output | |