import cv2
import numpy as np
from ultralytics import YOLO
import supervision as sv
from supervision.detection.core import Detections

class PlayerDetector:
    def __init__(self, model_path=None, confidence=0.5, iou_threshold=0.45):
        """
        初始化球员检测器
        
        Args:
            model_path: YOLO模型路径，如果为None则使用预训练模型
            confidence: 检测置信度阈值
            iou_threshold: NMS IOU阈值
        """
        self.confidence_threshold = confidence
        self.iou_threshold = iou_threshold
        
        if model_path:
            self.model = YOLO(model_path)
        else:
            # 使用预训练的YOLOv8模型
            print("加载YOLOv8预训练模型")
            self.model = YOLO("yolov8n.pt")
        
        # 人员检测类别ID (COCO数据集中的person类别为0)
        self.person_class_id = 0
        
        # 初始化监督库的标注工具
        self.box_annotator = sv.BoxAnnotator(
            thickness=2,
            text_thickness=2,
            text_scale=0.5
        )
    
    def detect(self, frame):
        """
        检测帧中的球员
        
        Args:
            frame: 输入的视频帧
            
        Returns:
            检测到的人物边界框列表 [x1, y1, x2, y2, confidence]
        """
        # 使用YOLO模型进行检测，可以调整参数来提高精度
        results = self.model(
            frame, 
            verbose=False, 
            conf=self.confidence_threshold,
            iou=self.iou_threshold,
            classes=[self.person_class_id]  # 只检测人员
        )[0]
        
        # 提取人员检测结果
        detections = []
        
        for r in results.boxes.data.cpu().numpy():
            x1, y1, x2, y2, conf, class_id = r
            
            # 已经在YOLO调用时过滤了类别，这里再次确认是否为人员
            if int(class_id) == self.person_class_id and conf > self.confidence_threshold:
                # 添加一些额外的过滤，比如尺寸过滤（可选）
                box_width = x2 - x1
                box_height = y2 - y1
                box_area = box_width * box_height
                img_area = frame.shape[0] * frame.shape[1]
                
                # 过滤掉太小或太大的检测框
                area_ratio = box_area / img_area
                if 0.0005 < area_ratio < 0.2:  # 根据实际场景调整这些值
                    detections.append([x1, y1, x2, y2, conf])
        
        return np.array(detections) if detections else np.empty((0, 5))
    
    def visualize(self, frame, detections):
        """
        可视化检测结果
        
        Args:
            frame: 原始视频帧
            detections: 检测结果
            
        Returns:
            标注后的视频帧
        """
        # 复制原始帧以避免修改
        annotated_frame = frame.copy()
        
        # 使用supervision库进行可视化（更好的可视化效果）
        if len(detections) > 0:
            # 转换为supervision的Detections格式
            sv_detections = Detections(
                xyxy=detections[:, :4],
                confidence=detections[:, 4],
                class_id=np.zeros(len(detections), dtype=int)  # 全部是人员类别
            )
            
            # 添加标签
            labels = [f"{conf:.2f}" for conf in detections[:, 4]]
            
            # 绘制检测框和标签
            annotated_frame = self.box_annotator.annotate(
                scene=annotated_frame,
                detections=sv_detections,
                labels=labels
            )
            
            return annotated_frame
        
        return annotated_frame
    
    def detect_with_supervision(self, frame):
        """
        使用Supervision库进行检测和结果格式化，提供更标准的输出
        
        Args:
            frame: 输入的视频帧
            
        Returns:
            Supervision库的Detections对象
        """
        # 使用YOLO模型进行检测
        results = self.model(frame, verbose=False, conf=self.confidence_threshold, classes=[self.person_class_id])[0]
        
        # 转换为supervision的Detections格式
        detections = sv.Detections.from_ultralytics(results)
        
        # 过滤出人员类别
        mask = np.array([int(cls) == self.person_class_id for cls in detections.class_id], dtype=bool)
        detections = detections[mask]
        
        return detections