File size: 1,430 Bytes
f851167 c8ae6b5 4504977 ac15725 f851167 c8ae6b5 4504977 c8ae6b5 c9cf887 77955ed 83d24c3 77955ed ac15725 f851167 565ea1f 77955ed 8ab6d35 77955ed f851167 8ab6d35 77955ed 8ab6d35 4504977 77955ed |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 |
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()} |