Update handler.py

handler.py

@@ -1,14 +1,20 @@
 from huggingface_hub import hf_hub_download
-hf_hub_download(repo_id="Drazcat-AI/galletas", filename="yolov8_galletas/runs/detect/train/weights/best.pt")
+from typing import Dict, List, Any
 from ultralytics import YOLO
 import cv2
 import torch
 import numpy as np
-model = YOLO('/root/.cache/huggingface/hub/models--Drazcat-AI--galletas/snapshots/d960a6784ad838e30f01690ce813fb4dc7460bc0/yolov8_galletas/runs/detect/train/weights/best.pt')
-results = 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))
+class EndpointHandler():
+    def __init__(self, path=""):
+        hf_hub_download(repo_id="Drazcat-AI/galletas", filename="yolov8_galletas/runs/detect/train/weights/best.pt")
 
-for result in results:
-      cv2.imshow("result", result)
-      cv2.waitKey(0)
+    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')
+        results = 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)