Update handler.py

handler.py

@@ -51,15 +51,13 @@ class EndpointHandler:
         self.annotator = BoxAnnotator()
 
     def __call__(self, data: Dict[str, Any]) -> Any:
-        # inputs => {"inputs": data}
-        #   data => {
-        #     "image": url/base64,
-        #     "image_size": tuple(int, int) / list(int),
-        #     "bbox_threshold": float,
-        #     "iou_threshold": float,
-        #     [disabled] "use_paddleocr": bool,
-        #     "draw_bboxes": bool,
-        #   }
+        # data should contain the following:
+        #  "inputs": {
+        #    "image": url/base64,
+        #    (optional) "image_size": tuple(int, int) / list(int),
+        #    (optional) "bbox_threshold": float,
+        #    (optional) "iou_threshold": float,
+        #  }
         data = data.pop("inputs")
 
         # read image from either url or base64 encoding
@@ -198,7 +196,7 @@ class EndpointHandler:
     def get_som_labeled_img(
         self,
         image: ImageType,
-        image_size: Optional[List[int]] = None,
+        image_size: Optional[Dict[Literal["w", "h"], int]] = None,
         ocr_texts: Optional[List[str]] = None,
         ocr_bboxes: Optional[List[List[int]]] = None,
         bbox_threshold: float = 0.01,
@@ -211,11 +209,13 @@ class EndpointHandler:
 
         w, h = image.size
         if image_size is None:
-            image_size = [h, w]
+            imgsz = {"h": h, "w": w}
+        else:
+            imgsz = [image_size.get("h", h), image_size.get("w", w)]
 
         out = self.yolo.predict(
             image,
-            imgsz=image_size,
+            imgsz=imgsz,
             conf=bbox_threshold,
             iou=iou_threshold or 0.7,
             verbose=False,
@@ -652,4 +652,4 @@ class BoxAnnotator:
             text_background_y1,
             text_background_x2,
             text_background_y2,
-        )
+        )