import io
from PIL import Image
from ultralytics import YOLO


class ModelHandler:
    def __init__(self):
        """加载 YOLOv11 模型"""
        self.model = YOLO("/opt/nuclio/best.pt")  # 确保路径正确

    def infer(self, image_data, threshold=0.3):
        """
        执行推理
        :param image_data: 图片的二进制数据
        :param threshold: 置信度阈值（默认0.3）
        :return: 符合阈值的检测结果
        """
        image = Image.open(io.BytesIO(image_data))
        results = self.model(image)

        detections = []
        for result in results:
            for box in result.boxes.data.tolist():
                x1, y1, x2, y2, score, class_id = box
                if score >= threshold:  # 过滤低置信度目标
                    detections.append({
                        "confidence": score,
                        "label": self.model.names[int(class_id)],
                        "points": [x1, y1, x2, y2],
                        "type": "rectangle",
                    })

        return detections