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()}