| from typing import Dict |
| from ultralytics import YOLO |
| import requests |
| from io import BytesIO |
| from PIL import Image |
| import time |
|
|
|
|
| class EndpointHandler(): |
| def __init__(self, path=""): |
| |
| print('Loading the yolo model') |
| model_path = "yolov8m_detect_usdl.pt" |
| self.model = YOLO(model_path) |
| print('Yolo model loaded successfully') |
|
|
|
|
| def __call__(self, data: Dict) -> Dict: |
|
|
| image_url = data.get('image_url') |
| print(f'Image url is : {image_url}') |
| response = requests.get(image_url) |
| pil_image = Image.open(BytesIO(response.content)) |
| |
| print('Model inference started....') |
| t1 = time.time() |
| results = self.model(pil_image) |
| print(f'TIME Model inference: {time.time() - t1}') |
| |
| return {"bbox": results[0].boxes.data.tolist()} |
