Update handler.py

handler.py

@@ -1,28 +1,59 @@
-from subprocess import run
+#from subprocess import run
 #run("pip install ultralytics", shell=True, check=True)
 #run("pip install opencv-python", shell=True, check=True)
 from huggingface_hub import hf_hub_download
 from typing import Dict, List, Any
 from ultralytics import YOLO
-import cv2
-import torch
-import numpy as np
+#import cv2
+#import torch
+#import numpy as np
+
 class EndpointHandler():
     def __init__(self, path=""):
         hf_hub_download(repo_id="Drazcat-AI/galletas", filename="yolov8_galletas/runs/detect/train/weights/best.pt")
-        #self.model = YOLO(hf_hub_download(repo_id="Drazcat-AI/galletas", filename="yolov8_galletas/runs/detect/train/weights/best.pt", local_files_only=True))
+        self.model = YOLO(hf_hub_download(repo_id="Drazcat-AI/galletas", filename="yolov8_galletas/runs/detect/train/weights/best.pt", local_files_only=True))
+    
+    def predict_objects(self, image_path):
+        results = self.model(image_path, imgsz=800)
+        
+        for result in results:
+            print(result)
 
+        return result
+
+    def __call__(self, event, context):
+        if "image_path" not in event:
+            return {
+                "statusCode": 400,
+                "body": json.dumps("Error: Please provide an 'image_path' parameter."),
+            }
+
+        image_path = event["image_path"]
+
+        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)}"),
+            }
+            
+    """
     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:    
 
-        #model = YOLO('/root/.cache/huggingface/hub/models--Drazcat-AI--galletas/snapshots/d960a6784ad838e30f01690ce813fb4dc7460bc0/yolov8_galletas/runs/detect/train/weights/best.pt')
-        model = YOLO(hf_hub_download(repo_id="Drazcat-AI/galletas", filename="yolov8_galletas/runs/detect/train/weights/best.pt", local_files_only=True))
-        results = model('https://rocketpin-ml-images.s3.amazonaws.com/smu/visits/1120/IMG_20230609_173208_902.jpg', imgsz=800)
+        #model = YOLO(hf_hub_download(repo_id="Drazcat-AI/galletas", filename="yolov8_galletas/runs/detect/train/weights/best.pt", local_files_only=True))
+        results = self.model('https://rocketpin-ml-images.s3.amazonaws.com/smu/visits/1120/IMG_20230609_173208_902.jpg', imgsz=800)
         img = cv2.imread('IMG_20230609_173208_902.jpg')
         img = cv2.resize(img, (640, 800))
         
         for result in results:
               cv2.imshow("result", result)
               cv2.waitKey(0)
+     """         
 """
 class ObjectDetectionHandler:
     def __init__(self, model_name):