Update handler.py

handler.py

@@ -8,13 +8,13 @@ from io import BytesIO
 
 class EndpointHandler():
     def __init__(self, path=""):
-        hf_hub_download(repo_id="Drazcat-AI/flejes", filename="yolov8_flejes/runs/detect/train/weights/last.pt")
-        self.model = YOLO(hf_hub_download(repo_id="Drazcat-AI/flejes", filename="yolov8_flejes/runs/detect/train/weights/last.pt", local_files_only=True))
+        hf_hub_download(repo_id="Drazcat-AI/flejes", filename="yolov8_flejes/runs/detect/train/weights/best.pt")
+        self.model = YOLO(hf_hub_download(repo_id="Drazcat-AI/flejes", filename="yolov8_flejes/runs/detect/train/weights/best.pt", local_files_only=True))
     
     def predict_objects(self, image_path):
-        #results = self.model(image_path, imgsz=1280)
+        results = self.model(image_path, imgsz=1280)
         #results = self.model(image_path, imgsz=800)
-        results = self.model(image_path)
+        #results = self.model(image_path)
         predictions = []
         for box in results[0].boxes:
             class_id = results[0].names[box.cls[0].item()]
@@ -44,9 +44,11 @@ class EndpointHandler():
         image_path = event["inputs"]
 
         try:
-            data=urlopen(image_path).read()
-            #image = Image.Image()._new(Image.open(BytesIO(data)))
-            image = Image.open(BytesIO(data))
+            #url image
+            #data=urlopen(image_path).read()
+            #image = Image.open(BytesIO(data))
+            #bytes image
+            image=image_path
             predictions = self.predict_objects(image)
             return {
                 "statusCode": 200,