Update handler.py

handler.py

@@ -11,21 +11,19 @@ class EndpointHandler():
         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=800)
-        #results = self.model(image_path)
+    def predict_objects(self, image_path, image_size_m):
+        results = self.model(image_path)
         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]
+            #cords = [round(x) for x in cords]
             conf = box.conf[0].item()
             prediction = {
-            "x": cords[0],
-            "y": cords[1],
-            "width": cords[2],
-            "height": cords[3],
+            "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
             }
@@ -44,12 +42,12 @@ class EndpointHandler():
         image_path = event["inputs"]
 
         try:
-            #url image
-            #data=urlopen(image_path).read()
-            #image = Image.open(BytesIO(data))
-            #bytes image
-            image=image_path
-            predictions = self.predict_objects(image)
+            data=urlopen(image_path).read()
+            image = Image.open(BytesIO(data))
+            image_size = image.size
+            image = image.resize([1280,960])
+            predictions = self.predict_objects(image, [image_size[0]/1280,image_size[1]/960])
+            
             return {
                 "statusCode": 200,
                 "body": json.dumps(predictions),