import cv2
from ultralytics import YOLO

# YOLOv8Detector类：用于加载YOLO模型，进行目标检测并在图像上绘制信息
class YOLOv8Detector:
    def __init__(self):
        """
        初始化YOLOv8检测器
        :param classes: 自定义类别名称列表（暂未使用）
        """   
        self.model = YOLO(r"model/yolov8n.pt")  # 使用YOLOv8官方预训练模型
        # self.model = YOLO(r"model/best.pt")  # 使用自训练模型
           
        self.class_names = self.model.names  # 获取类别名称
        self._frame_counter = 0  # 帧计数器（用于调试打印）

    # +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
    # 执行YOLO目标检测，并在图像中绘制检测框、中心坐标和尺寸信息（不再过滤类别）
    # +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
    def detect_and_draw(self, image):
        height, width = image.shape[:2]

        # 使用YOLO模型进行检测
        results = self.model(image, imgsz=640)
        frame_id = self._frame_counter
        self._frame_counter += 1

        result_info = []  # 存储每个目标的检测信息
        # print(f"\n[FRAME {frame_id}] Detection Result:")

        for result in results:
            boxes = result.boxes

            # 遍历所有检测框（不再按类别筛选）
            for i in range(len(boxes)):
                cls_id = int(boxes.cls[i].item())  # 类别索引
                cls_name = self.class_names[cls_id] if cls_id < len(self.class_names) else f"class_{cls_id}"
                score = float(boxes.conf[i].item())  # 置信度
                
                # （1）boxes.xyxy[i]提取第i个种类边界框坐标[x1, y1, x2, y2]；（2）tolist()实现将PyTorch 的张量（Tensor）[x1, y1, x2, y2]转换为 普通的 Python 列表；（3）使用 map() 函数将列表中的 浮点数坐标转换为整数坐标，因为像素坐标用于图像绘制时必须是整数。
                x1_box, y1_box, x2_box, y2_box = map(int, boxes.xyxy[i].tolist())  

                # 计算中心点与宽高
                x_center = (x1_box + x2_box) / 2
                y_center = (y1_box + y2_box) / 2
                bbox_width = x2_box - x1_box
                bbox_height = y2_box - y1_box

                # 打印调试信息
                # print(f"[{cls_name}] Score: {score:.2f}")
                # print(f"    Center: ({x_center:.1f}, {y_center:.1f})")
                # print(f"    BBox: Width={bbox_width}px, Height={bbox_height}px")
                # print(f"    左上角坐标: ({x1_box}px, {y1_box}px)")

                # 绘制检测框（绿色），参数2表示线宽
                cv2.rectangle(image, (x1_box, y1_box), (x2_box, y2_box), (0, 255, 0), 2)

                # 绘制框右上角顶部标签（类别名+得分）
                label = f"{cls_name} {score:.2f}"
                cv2.putText(image, label, (x2_box, y1_box - 60),
                            cv2.FONT_HERSHEY_SIMPLEX, 2.5, (0, 255, 0), 3)
                # 绘制框右上角顶部标签（类别名+得分）
                cv2.putText(image, f"L-corner-coor:{x1_box, y1_box}", (x1_box, y1_box - 20),
                            cv2.FONT_HERSHEY_SIMPLEX, 2.5, (0, 255, 0), 3)
                
                # 在框底部左下角显示中心坐标和尺寸信息
                cv2.putText(image, f"Center: ({x_center:.1f}, {y_center:.1f})",
                            (x1_box, y2_box + 70),
                            cv2.FONT_HERSHEY_SIMPLEX, 2.5, (0, 255, 255), 3)
                cv2.putText(image, f"Width: {bbox_width}px, Height: {bbox_height}px",
                            (x1_box, y2_box + 150),
                            cv2.FONT_HERSHEY_SIMPLEX, 2.5, (0, 255, 255), 3)

                # 保存检测信息字典."L-corner-coor":(x1_box, y1_box)左上角坐标点
                result_dict = {
                    "class": cls_name,
                    "score": score,
                    "center": (x_center, y_center),
                    "bbox": (bbox_width, bbox_height),
                    "L-corner-coor":(x1_box, y1_box)
                }
                result_info.append(result_dict)

        return image, result_info