from ultralytics import YOLO
import cv2
import numpy as np
from PIL import Image, ImageDraw, ImageFont
import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit

class YOLODetector:
    def __init__(self, model_path='models/best.engine'):
        # 加载 TensorRT 模型
        self.model = YOLO(model_path, task="detect")
        # 英文类别名称到中文的映射
        self.class_name_mapping = {
            'pedestrian': '行人',
            'people': '人群',
            'bicycle': '自行车',
            'car': '轿车',
            'van': '面包车',
            'truck': '卡车',
            'tricycle': '三轮车',
            'awning-tricycle': '篷式三轮车',
            'bus': '公交车',
            'motor': '摩托车'
        }
        # 为每个类别设置固定的RGB颜色
        self.color_mapping = {
            'pedestrian': (71, 0, 36),      # 勃艮第红
            'people': (0, 255, 0),          # 绿色
            'bicycle': (0, 49, 83),         # 普鲁士蓝
            'car': (0, 47, 167),            # 克莱茵蓝
            'van': (128, 0, 128),           # 紫色
            'truck': (212, 72, 72),         # 缇香红
            'tricycle': (0, 49, 83),        # 橙色
            'awning-tricycle': (251, 220, 106), # 申布伦黄
            'bus': (73, 45, 34),            # 凡戴克棕
            'motor': (1, 132, 127)          # 马尔斯绿
        }
        # 初始化类别计数器
        self.class_counts = {cls_name: 0 for cls_name in self.class_name_mapping.keys()}
        # 初始化字体
        try:
            self.font = ImageFont.truetype("simhei.ttf", 20)
        except IOError:
            self.font = ImageFont.load_default()

    def detect_and_draw_English(self, frame, conf=0.3, iou=0.5):
        """
        对输入帧进行目标检测并返回绘制结果
        
        Args:
            frame: 输入的图像帧（BGR格式）
            conf: 置信度阈值
            iou: IOU阈值
        
        Returns:
            annotated_frame: 绘制了检测结果的图像帧
        """
        try:
            # 进行 YOLO 目标检测
            results = self.model(
                frame,
                conf=conf,
                iou=iou,
                half=True,
            )
            result = results[0]
            
            # 使用YOLO自带的绘制功能
            annotated_frame = result.plot()
            
            return annotated_frame
            
        except Exception as e:
            print(f"Detection error: {e}")
            return frame

    def detect_and_draw_Chinese(self, frame, conf=0.2, iou=0.3):
        """
        对输入帧进行目标检测并绘制中文标注
        
        Args:
            frame: 输入的图像帧（BGR格式）
            conf: 置信度阈值
            iou: IOU阈值
        
        Returns:
            annotated_frame: 绘制了检测结果的图像帧
        """
        try:
            # 进行 YOLO 目标检测
            results = self.model(
                frame,
                conf=conf,
                iou=iou,
                # half=True,
            )
            result = results[0]
            
            # 获取原始帧的副本
            img = frame.copy()
            
            # 转换为PIL图像以绘制中文
            pil_img = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
            draw = ImageDraw.Draw(pil_img)
            
            # 绘制检测结果
            for box in result.boxes:
                # 获取边框坐标
                x1, y1, x2, y2 = box.xyxy[0].cpu().numpy().astype(int)
                
                # 获取类别ID和置信度
                cls_id = int(box.cls[0].item())
                conf = box.conf[0].item()
                
                # 获取类别名称并转换为中文
                cls_name = result.names[cls_id]
                chinese_name = self.class_name_mapping.get(cls_name, cls_name)
                
                # 获取该类别的颜色
                color = self.color_mapping.get(cls_name, (255, 255, 255))
                
                # 绘制边框
                draw.rectangle([(x1, y1), (x2, y2)], outline=color, width=3)
                
                # 准备标签文本
                text = f"{chinese_name} {conf:.2f}"
                text_size = draw.textbbox((0, 0), text, font=self.font)
                text_width = text_size[2] - text_size[0]
                text_height = text_size[3] - text_size[1]
                
                # 绘制标签背景（使用与边框相同的颜色）
                draw.rectangle(
                    [(x1, y1 - text_height - 4), (x1 + text_width, y1)],
                    fill=color
                )
                
                # 绘制白色文本
                draw.text(
                    (x1, y1 - text_height - 2),
                    text,
                    fill=(255, 255, 255),  # 白色文本
                    font=self.font
                )
            
            # 转换回OpenCV格式
            return cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR)
            
        except Exception as e:
            print(f"Detection error: {e}")
            return frame