Create handler.py

handler.py

@@ -0,0 +1,46 @@
+from typing import Dict, List, Any
+from ultralytics import YOLO
+
+class EndpointHandler():
+    def __init__(self, path=""):
+        # Preload all the elements you are going to need at inference.
+        # pseudo:
+        # self.model= load_model(path)
+        yolov8_model_name = 'yolov8_2023-07-19_yolov8m.pt'
+        self.model = YOLO(yolov8_model_name)    
+
+    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+        """
+        data args:
+            inputs (:obj: `str` | `PIL.Image` | `np.array`)
+            kwargs
+        Return:
+            A :obj:`list` | `dict`: will be serialized and returned
+        """
+        # Get the prediction
+        result = self.model(data['inputs'])
+        # Get the original image with channel shifted
+        img = result[0].orig_img[:,:,::-1]
+        H, W, _ = img.shape
+        annotated = img.copy()
+        # Modify crop so that it is square
+        try:
+            x1, y1, x2, y2 = result[0].boxes.xyxy.numpy().astype('int')[0]
+            if result[0].boxes.conf[0].item() < 0.75: # if low in confidence
+                x1, y1, x2, y2 = 0, 0, W, H
+            else:
+                annotated = result[0].plot(labels=False, conf=False)[:,:,::-1]
+        except: # in case there is no detection
+            x1, y1, x2, y2 = 0, 0, W, H
+
+        h, w = y2-y1, x2-x1
+        offset = abs(h-w) // 2
+        if h > w:
+            x1 = max(x1 - offset, 0)
+            x2 = min(x2 + offset, W)
+        else:
+            y1 = max(y1 - offset, 0)
+            y2 = min(y2 + offset, H)
+        new_image = img[y1:y2, x1:x2]
+        # Return the annotated original image with the square cropped
+        return annotated, new_image