Update handler.py

handler.py

@@ -1,16 +1,15 @@
 from ultralyticsplus import YOLO
-from typing import List, Dict, Any
-from sahi import ObjectPrediction
+from typing import Dict, Any, List
 
 
 DEFAULT_CONFIG = {'conf': 0.25, 'iou': 0.45, 'agnostic_nms': False, 'max_det': 1000}
-
+BOX_KEYS = ['xmin', 'ymin', 'xmax', 'ymax']
 
 class EndpointHandler():
-    def __init__(self, path=""):
+    def __init__(self):
         self.model = YOLO('ultralyticsplus/yolov8s')
 
-    def __call__(self, data: str) -> List[ObjectPrediction]:
+    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
         """
        data args:
             image: image path to segment
@@ -19,33 +18,37 @@ class EndpointHandler():
                      agnostic_nms - NMS class-agnostic: True / False,
                      max_det - maximum number of detections per image)
       Return:
-            object_predictions
+            A :obj: `dict` | `dict`: {scores, labels, boxes}
         """
-        config = DEFAULT_CONFIG
+        inputs = data.pop("inputs", data)
+        input_config = inputs.pop("config", DEFAULT_CONFIG)
+        config = {**DEFAULT_CONFIG, **input_config}
+        
+        if config is None:
+            config = DEFAULT_CONFIG
         # Set model parameters
         self.model.overrides['conf'] = config.get('conf')
         self.model.overrides['iou'] = config.get('iou')
         self.model.overrides['agnostic_nms'] = config.get('agnostic_nms')
         self.model.overrides['max_det'] = config.get('max_det')   
         
+        # Get label idx-to-name
+        names = model.model.names
+        
         # perform inference
-        inputs = data.pop("inputs", data)
         result = self.model.predict(inputs['image'])[0]
-
-        names = self.model.model.names
-        boxes = result.boxes
-
-        object_predictions = []
-        if boxes is not None:
-            det_ind = 0
-            for xyxy, conf, cls in zip(boxes.xyxy, boxes.conf, boxes.cls):
-                object_prediction = ObjectPrediction(
-                    bbox=xyxy.tolist(),
-                    category_name=names[int(cls)],
-                    category_id=int(cls),
-                    score=conf,
-                )
-                object_predictions.append(object_prediction)
-                det_ind += 1
-        return object_predictions
-    
+        prediction = []
+        for score, label, box in zip(result.boxes.conf, result.boxes.cls, result.boxes.xyxy):
+            item_score = score.item()
+            item_label = names[int(label)]
+            item_box = box.to(dtype=int).tolist()
+            
+            item_prediction = {
+                'score': item_score,
+                'label': item_label,
+                'box': dict(zip(BOX_KEYS, item_box))
+            }
+            
+            prediction.append(item_prediction)
+        
+        return prediction