handler.py

from huggingface_hub import hf_hub_download
from typing import Dict, List, Any
from ultralytics import YOLO

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))
    
    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)}"),
            }