Create handler.py

handler.py

@@ -0,0 +1,52 @@
+from huggingface_hub import hf_hub_download
+from typing import Dict, List, Any
+from ultralytics import YOLO
+import json
+
+class EndpointHandler():
+    def __init__(self, path=""):
+        hf_hub_download(repo_id="Drazcat-AI/sku-100k", filename="yolov8_sku-100k/runs/detect/train/weights/best.pt")
+        self.model = YOLO(hf_hub_download(repo_id="Drazcat-AI/sku-100k", filename="yolov8_sku-100k/runs/detect/train/weights/best.pt", local_files_only=True))
+    
+    def predict_objects(self, image_path):
+        results = self.model(image_path, imgsz=800)
+        predictions = []
+        for box in results[0].boxes:
+            class_id = results[0].names[box.cls[0].item()]
+            cords = box.xywh[0].tolist()
+            cords = [round(x) for x in cords]
+            conf = round(box.conf[0].item(), 2)
+            prediction = {
+            "label": class_id,
+            "score": conf,
+            "box":{
+                "x": cords[0],
+                "y": cords[1],
+                "width": cords[2],
+                "height": cords[3]}
+            }
+            predictions.append(prediction)
+        predictions_array = {"predictions": predictions}   
+
+        return predictions_array
+
+    def __call__(self, event):
+        if "inputs" not in event:
+            return {
+                "statusCode": 400,
+                "body": json.dumps("Error: Please provide an 'inputs' parameter."),
+            }
+
+        image_path = event["inputs"]
+
+        try:
+            predictions = self.predict_objects(image_path)
+            return {
+                "statusCode": 200,
+                "body": json.dumps(predictions),
+            }
+        except Exception as e:
+            return {
+                "statusCode": 500,
+                "body": json.dumps(f"Error: {str(e)}"),
+            }