Update handler.py

handler.py

@@ -12,7 +12,7 @@ class EndpointHandler():
         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, image_size_m):
-        results = self.model(image_path)
+        results = self.model(image_path, imgsz=[1280, 960])
         predictions = []
         for box in results[0].boxes:
             class_id = results[0].names[box.cls[0].item()]
@@ -45,13 +45,17 @@ class EndpointHandler():
             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)
+            """
             image_size = image.shape
             if image.shape[0]>image.shape[0]:
                 x, y = 1280, 960
             else:
                 y, x = 1280, 960
             image = cv2.resize(image, (x, y))
+            
             predictions = self.predict_objects(image, [image_size[0]/x,image_size[1]/y])
+            """
+            predictions = self.predict_objects(image, (1,1))
             return {
                 "statusCode": 200,
                 "body": json.dumps(predictions),