yolotesting / handler.py
skumar1998's picture
Update handler.py
ac15725
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=""):
# load the model
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}')
# postprocess the prediction -> results[0].boxes.data.tolist()
return {"bbox": results[0].boxes.data.tolist()}