skumar1998 commited on
Commit
f851167
·
1 Parent(s): b9ddb2d

Create handler.py

Browse files
Files changed (1) hide show
  1. handler.py +30 -0
handler.py ADDED
@@ -0,0 +1,30 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import Dict
2
+ from ultralytics import YOLO
3
+ import requests
4
+ from io import BytesIO
5
+ from PIL import Image
6
+ import time
7
+
8
+
9
+ class EndpointHandler():
10
+ def __init__(self, path=""):
11
+ # load the model
12
+ print('Loading the yolo model')
13
+ model_path = "yolov8m_detect_usdl.pt"
14
+ self.model = YOLO(model_path)
15
+ print('Yolo model loaded successfully')
16
+
17
+
18
+ def __call__(self, data: Dict) -> Dict:
19
+
20
+ image_url = data.get('image_url')
21
+ print(f'Image url is : {image_url}')
22
+ response = requests.get(image_url)
23
+ pil_image = Image.open(BytesIO(response.content))
24
+
25
+ print('Model inference started....')
26
+ t1 = time.time()
27
+ results = self.model(pil_image)
28
+ print(f'TIME Model inference: {time.time() - t1}')
29
+ # postprocess the prediction
30
+ return {"bbox": results[0].boxes.data.tolist()}