Update handler.py

handler.py

@@ -1,7 +1,6 @@
 from ultralyticsplus import YOLO
 from typing import List, Dict, Any
 from sahi import ObjectPrediction
-import torch, torchvision
 
 
 DEFAULT_CONFIG = {'conf': 0.25, 'iou': 0.45, 'agnostic_nms': False, 'max_det': 1000}
@@ -22,32 +21,31 @@ class EndpointHandler():
       Return:
             object_predictions
         """
-        # 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')   
+        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')   
         
-        # # perform inference
-        # inputs = data.pop("inputs", data)
-        # result = self.model.predict(inputs['image'])[0]
+        # perform inference
+        inputs = data.pop("inputs", data)
+        result = self.model.predict(inputs['image'])[0]
 
-        # names = self.model.model.names
-        # boxes = result.boxes
+        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
-        return torch.__version__, torchvision.__version__
+        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