Upload 3 files

cervezas_v2-12/runs/detect/train/weights/best.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2c3ac5e75e25800ba22b48c1174c3ecffb62404f7b93c58c232fb291f2910e9c
+size 87836713

handler.py

@@ -1,52 +1,58 @@
-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/cervezas", filename="yolov8_cervezas/runs/detect/train/weights/best.pt")
-        self.model = YOLO(hf_hub_download(repo_id="Drazcat-AI/cervezas", filename="yolov8_cervezas/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)}"),
+from huggingface_hub import hf_hub_download
+from typing import Dict, List, Any
+from ultralytics import YOLO
+import json
+import urllib.request
+import cv2
+from io import BytesIO
+import numpy as np
+class EndpointHandler():
+    def __init__(self, path=""):
+        hf_hub_download(repo_id="Drazcat-AI/cervezas", filename="cervezas_v2-12/runs/detect/train/weights/best.pt")
+        self.model = YOLO(hf_hub_download(repo_id="Drazcat-AI/cervezas", filename="cervezas_v2-12/runs/detect/train/weights/best.pt", local_files_only=True))
+    
+    def predict_objects(self, image_path, image_size_m):
+        results = self.model(image_path, imgsz=[1280, 960])
+        predictions = []
+        for box in results[0].boxes:
+            class_id = results[0].names[box.cls[0].item()]
+            cords = box.xywh[0].tolist()
+            
+            conf = box.conf[0].item()
+            prediction = {
+            "x": round(cords[0]*image_size_m[0]),
+            "y": round(cords[1]*image_size_m[1]),
+            "width": round(cords[2]*image_size_m[0]),
+            "height": round(cords[3]*image_size_m[1]),
+            "confidence": conf,
+            "class": class_id
+            }
+            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:
+            with urllib.request.urlopen(image_path) as response:
+                image_content = np.asarray(bytearray(response.read()), dtype=np.uint8)
+            image = cv2.imdecode(image_content, cv2.IMREAD_COLOR)
+            
+            predictions = self.predict_objects(image, (1,1))
+            return {
+                "statusCode": 200,
+                "body": json.dumps(predictions),
+            }
+        except Exception as e:
+            return {
+                "statusCode": 500,
+                "body": json.dumps(f"Error: {str(e)}"),
             }

requirements.txt

@@ -1,2 +1,3 @@
-ultralytics
-torchvision==0.12.0
+ultralytics==8.3.61
+opencv-python
+numpy