Update handler.py

handler.py

@@ -15,22 +15,23 @@ class EndpointHandler():
         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'))
-        print('Yolo model loaded successfully')
+        logging.info(f'TIME: loading the model {time.time() - t1}')
 
 
     def __call__(self, data: Dict) -> Dict:
 
-        logging.info(f'Value of data is : {data}')
+        # logging.info(f'Value of data is : {data}')
 
-        # 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))
+        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}')
+        print('Model inference started....')
+        t1 = time.time()
+        results = self.model(pil_image)
+        print(f'TIME Model inference: {time.time() - t1}')
         # postprocess the prediction -> results[0].boxes.data.tolist()
-        return {"bbox": data}
+        return {"bbox": results[0].boxes.data.tolist()}