|
from typing import Dict |
|
from ultralytics import YOLO |
|
import requests |
|
from io import BytesIO |
|
from PIL import Image |
|
import time |
|
import logging |
|
import os |
|
import torch |
|
import torchvision |
|
|
|
|
|
class EndpointHandler(): |
|
def __init__(self, path=""): |
|
|
|
logging.info(f'value of path is : {path}') |
|
current_directory = os.getcwd() |
|
logging.info(f'current dir: {current_directory}') |
|
logging.info(f'all files: {os.listdir(path)}') |
|
t1 = time.time() |
|
self.model = YOLO(os.path.join(path, 'yolov8m_detect_usdl.pt')) |
|
logging.info(f'TIME: loading the model {time.time() - t1}') |
|
logging.info(f'torch version : {torch.__version__}') |
|
logging.info(f'torch vision version: {torchvision.__version__}') |
|
|
|
|
|
def __call__(self, data: Dict) -> Dict: |
|
|
|
logging.info(f'data is : {data}') |
|
logging.info(f'type of data is : {type(data)}') |
|
image_url = data.get('image_url') |
|
logging.info(f'Image url is : {image_url}') |
|
response = requests.get(image_url) |
|
pil_image = Image.open(BytesIO(response.content)) |
|
|
|
logging.info('Model inference started....') |
|
t1 = time.time() |
|
results = self.model(pil_image) |
|
logging.info(f'TIME Model inference: {time.time() - t1}') |
|
|
|
return {"bbox": results[0].boxes.data.tolist()} |