Yolov/detection_render.py

import cv2
import torch
import numpy as np
import concurrent.futures
from typing import List, Dict, Tuple, Any, Optional, Union
from dataclasses import dataclass
from collections import defaultdict

from PIL import Image, ImageDraw, ImageFont

# 尝试使用PIL绘制中文
font = None
FONT_PATHS = [
    "simhei.ttf",  # 黑体
    "/usr/share/fonts/truetype/wqy/wqy-microhei.ttc",  # 文泉驿微米黑
    "C:/Windows/Fonts/simhei.ttf",  # Windows黑体
    "/System/Library/Fonts/PingFang.ttc"  # macOS苹方
]

for path in FONT_PATHS:
    try:
        font = ImageFont.truetype(path, 20)
        break
    except:
        continue


@dataclass
class ModelConfig:
    """模型配置数据类"""
    model: Any
    config: Any
    tags: Optional[Dict[str, Dict]] = None
    name: str = "未命名"
    id: Any = None
    device: str = "cuda:0"
    conf_thres: float = 0.25
    iou_thres: float = 0.45
    half: bool = False
    key_valid: Any = None
    model_hash: Any = None
    imgsz: Any = 640

    def __post_init__(self):
        """初始化后处理"""
        if self.tags is None:
            self.tags = {}


@dataclass
class Detection:
    """检测结果数据类"""
    model_idx: int
    model_name: str
    boxes: np.ndarray
    confidences: np.ndarray
    class_ids: np.ndarray
    class_names: List[str]
    tags: Dict[str, Dict]
    raw_result: Any = None  # 新增：保存原始结果对象


class DetectionVisualizer:
    """检测结果可视化器"""

    # 预定义模型颜色
    MODEL_COLORS = [
        (0, 255, 0),  # 绿色
        (255, 0, 0),  # 蓝色
        (0, 0, 255),  # 红色
        (255, 255, 0),  # 青色
        (255, 0, 255),  # 紫色
        (0, 255, 255),  # 黄色
        (128, 0, 128),  # 深紫色
        (0, 128, 128),  # 橄榄色
    ]

    def __init__(self, use_pil: bool = True):
        """
        初始化可视化器

        Args:
            use_pil: 是否使用PIL绘制（支持中文）
        """
        self.use_pil = use_pil
        self.font = font

    def should_draw_detection(self, class_id: int, confidence: float,
                              tags: Dict[str, Dict]) -> Tuple[bool, Optional[Tuple[int, int, int]]]:
        """
        判断是否应该绘制检测框

        Args:
            class_id: 类别ID
            confidence: 置信度
            tags: 标签配置字典

        Returns:
            (是否绘制, 颜色)
        """
        class_id_str = str(class_id)

        # 如果标签配置为空，默认不绘制任何标签（因为需求是只绘制配置的标签）
        if not tags:
            return False, None

        # 如果标签不在配置中，不绘制
        tag_config = tags.get(class_id_str)
        if not tag_config:
            return False, None

        # 检查select标记
        if not tag_config.get('select', True):
            return False, None

        # 检查置信度阈值
        reliability = tag_config.get('reliability', 0)
        if confidence < reliability:
            return False, None

        # 获取自定义颜色
        color = tag_config.get('color')
        if color and isinstance(color, (list, tuple)) and len(color) >= 3:
            return True, tuple(color[:3])

        return True, None

    def draw_with_pil(self, frame: np.ndarray, detections: List[Detection],
                      confidence_threshold: float) -> np.ndarray:
        """使用PIL绘制检测结果（支持中文）"""
        # 转换到PIL格式
        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        pil_image = Image.fromarray(frame_rgb)
        draw = ImageDraw.Draw(pil_image)

        for det in detections:
            model_idx = det.model_idx
            boxes = det.boxes
            confidences = det.confidences
            class_names = det.class_names
            class_ids = det.class_ids
            tags = det.tags

            for i, (box, conf, cls_id, cls_name) in enumerate(
                    zip(boxes, confidences, class_ids, class_names)):

                if conf < confidence_threshold:
                    continue

                should_draw, custom_color = self.should_draw_detection(
                    cls_id, conf, tags)
                if not should_draw:
                    continue

                # 使用自定义颜色或模型颜色
                color = custom_color or self.MODEL_COLORS[model_idx % len(self.MODEL_COLORS)]

                x1, y1, x2, y2 = map(int, box[:4])

                # 标签文本
                label = f"{cls_name} {conf:.2f}"

                # 绘制矩形框
                draw.rectangle([x1, y1, x2, y2], outline=color, width=2)

                # 绘制标签背景
                if self.font:
                    try:
                        text_bbox = draw.textbbox((x1, y1 - 25), label, font=self.font)
                        draw.rectangle(text_bbox, fill=color)
                        draw.text((x1, y1 - 25), label, fill=(255, 255, 255), font=self.font)
                    except:
                        # 字体失败，回退到OpenCV
                        cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)
                        cv2.putText(frame, label, (x1, y1 - 10),
                                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
                else:
                    cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)
                    cv2.putText(frame, label, (x1, y1 - 10),
                                cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)

        # 转换回OpenCV格式
        return cv2.cvtColor(np.array(pil_image), cv2.COLOR_RGB2BGR)

    def draw_with_opencv(self, frame: np.ndarray, detections: List[Detection],
                         confidence_threshold: float) -> np.ndarray:
        """使用OpenCV绘制检测结果"""
        frame_drawn = frame.copy()

        for det in detections:
            model_idx = det.model_idx
            model_name = det.model_name
            boxes = det.boxes
            confidences = det.confidences
            class_names = det.class_names
            class_ids = det.class_ids
            tags = det.tags

            for i, (box, conf, cls_id, cls_name) in enumerate(
                    zip(boxes, confidences, class_ids, class_names)):

                if conf < confidence_threshold:
                    continue

                should_draw, custom_color = self.should_draw_detection(
                    cls_id, conf, tags)
                if not should_draw:
                    continue

                # 使用自定义颜色或模型颜色
                color = custom_color or self.MODEL_COLORS[model_idx % len(self.MODEL_COLORS)]

                x1, y1, x2, y2 = map(int, box[:4])

                # 绘制矩形框
                cv2.rectangle(frame_drawn, (x1, y1), (x2, y2), color, 2)

                # 绘制标签
                label = f"{model_name}: {cls_name} {conf:.2f}"

                # 计算文本大小
                (text_width, text_height), baseline = cv2.getTextSize(
                    label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 2
                )

                # 绘制标签背景
                cv2.rectangle(frame_drawn,
                              (x1, y1 - text_height - 10),
                              (x1 + text_width, y1),
                              color, -1)

                # 绘制文本
                cv2.putText(frame_drawn, label, (x1, y1 - 5),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 2)

        return frame_drawn

    def draw(self, frame: np.ndarray, detections: List[Detection],
             confidence_threshold: float) -> np.ndarray:
        """绘制检测结果"""
        if self.use_pil:
            try:
                return self.draw_with_pil(frame, detections, confidence_threshold)
            except Exception as e:
                print(f"PIL绘制失败，使用OpenCV: {e}")
                return self.draw_with_opencv(frame, detections, confidence_threshold)
        else:
            return self.draw_with_opencv(frame, detections, confidence_threshold)


class YOLOModelWrapper:
    """YOLO模型包装器"""

    @staticmethod
    def infer_yolov8(model: Any, frame: np.ndarray, conf_thres: float,
                     iou_thres: float, imgsz: int, half: bool, device: str):
        """YOLOv8推理"""
        with torch.no_grad():
            results = model(frame, conf=conf_thres, verbose=False,
                            iou=iou_thres, imgsz=imgsz, half=half, device=device)

            if not results or len(results) == 0:
                return np.array([]), np.array([]), np.array([]), [], None

            result = results[0]
            if not hasattr(result, 'boxes'):
                return np.array([]), np.array([]), np.array([]), [], None

            boxes = result.boxes.xyxy.cpu().numpy()
            confidences = result.boxes.conf.cpu().numpy()
            class_ids = result.boxes.cls.cpu().numpy().astype(int)
            class_names = [result.names.get(cid, str(cid)) for cid in class_ids]

            return boxes, confidences, class_ids, class_names, result

    @staticmethod
    def infer_yolov5(model: Any, frame: np.ndarray):
        """YOLOv5推理"""
        results = model(frame)

        if not hasattr(results, 'xyxy'):
            return np.array([]), np.array([]), np.array([]), [], None

        detections = results.xyxy[0].cpu().numpy()
        if len(detections) == 0:
            return np.array([]), np.array([]), np.array([]), [], None

        boxes = detections[:, :4]
        confidences = detections[:, 4]
        class_ids = detections[:, 5].astype(int)

        if hasattr(results, 'names'):
            class_names = [results.names.get(cid, str(cid)) for cid in class_ids]
        else:
            class_names = [str(cid) for cid in class_ids]

        return boxes, confidences, class_ids, class_names, results


def multi_model_inference(
        models_config: List[Union[Dict, ModelConfig]],
        frame: np.ndarray,
        confidence_threshold: float = 0.25,
        parallel: bool = True,
        use_pil: bool = True,
        use_plot_for_single: bool = True  # 新增参数：是否对单个模型使用plot绘制
) -> Tuple[np.ndarray, List[Detection]]:
    """
    多模型并行推理

    Args:
        models_config: 模型配置列表
        frame: 视频帧 (BGR格式)
        confidence_threshold: 全局置信度阈值
        parallel: 是否并行推理
        use_pil: 是否使用PIL绘制
        use_plot_for_single: 当只有一个模型时，是否使用result.plot()绘制

    Returns:
        (绘制完成的帧, 检测结果列表)
    """
    # 转换为ModelConfig对象
    configs = []
    for cfg in models_config:
        if isinstance(cfg, dict):
            configs.append(ModelConfig(**cfg))
        else:
            configs.append(cfg)

    def single_model_inference(model_cfg: ModelConfig, model_idx: int) -> Detection:
        """单个模型的推理函数"""
        try:
            model = model_cfg.model

            # 根据模型类型进行推理
            if hasattr(model, 'predict'):  # YOLOv8
                boxes, confidences, class_ids, class_names, raw_result = YOLOModelWrapper.infer_yolov8(
                    model, frame, model_cfg.conf_thres, model_cfg.iou_thres,
                    model_cfg.imgsz, model_cfg.half, model_cfg.device
                )
            elif hasattr(model, '__call__'):  # YOLOv5
                boxes, confidences, class_ids, class_names, raw_result = YOLOModelWrapper.infer_yolov5(model, frame)
            else:
                raise ValueError(f"不支持的模型类型: {type(model)}")

            return Detection(
                model_idx=model_idx,
                model_name=model_cfg.name,
                boxes=boxes,
                confidences=confidences,
                class_ids=class_ids,
                class_names=class_names,
                tags=model_cfg.tags,
                raw_result=raw_result
            )

        except Exception as e:
            print(f"模型 {model_idx} ({model_cfg.name}) 推理失败: {e}")
            return Detection(
                model_idx=model_idx,
                model_name=model_cfg.name,
                boxes=np.array([]),
                confidences=np.array([]),
                class_ids=np.array([]),
                class_names=[],
                tags=model_cfg.tags,
                raw_result=None
            )

    # 并行推理
    if parallel and len(configs) > 1:
        detections = []
        with concurrent.futures.ThreadPoolExecutor(max_workers=len(configs)) as executor:
            # 提交任务
            future_to_idx = {
                executor.submit(single_model_inference, cfg, idx): idx
                for idx, cfg in enumerate(configs)
            }

            # 收集结果
            for future in concurrent.futures.as_completed(future_to_idx):
                try:
                    detection = future.result(timeout=2.0)
                    detections.append(detection)
                except concurrent.futures.TimeoutError:
                    idx = future_to_idx[future]
                    print(f"模型 {idx} 推理超时")
                    # 创建空的检测结果
                    detections.append(Detection(
                        model_idx=idx,
                        model_name=configs[idx].name,
                        boxes=np.array([]),
                        confidences=np.array([]),
                        class_ids=np.array([]),
                        class_names=[],
                        tags=configs[idx].tags,
                        raw_result=None
                    ))
                except Exception as e:
                    print(f"模型推理异常: {e}")
    else:
        # 顺序推理
        detections = [single_model_inference(cfg, idx)
                      for idx, cfg in enumerate(configs)]

    # 按照模型索引排序
    detections.sort(key=lambda x: x.model_idx)

    # 新增：如果只有一个模型且配置使用plot绘制，则使用result.plot()
    if len(configs) == 1 and use_plot_for_single:
        single_detection = detections[0]
        if single_detection.raw_result is not None:
            try:
                # 检查是否有plot方法
                if hasattr(single_detection.raw_result, 'plot'):
                    # 使用plot方法绘制结果
                    frame_drawn = single_detection.raw_result.plot()
                    # 确保返回的是numpy数组
                    if not isinstance(frame_drawn, np.ndarray):
                        frame_drawn = np.array(frame_drawn)
                    # plot方法通常返回RGB图像，转换为BGR
                    if len(frame_drawn.shape) == 3 and frame_drawn.shape[2] == 3:
                        frame_drawn = cv2.cvtColor(frame_drawn, cv2.COLOR_RGB2BGR)
                    print(f"使用 {single_detection.model_name} 的 plot() 方法绘制结果")
                    return frame_drawn, detections
                else:
                    print(f"模型 {single_detection.model_name} 的结果对象没有 plot() 方法，使用自定义绘制")
            except Exception as e:
                print(f"使用 plot() 方法绘制失败: {e}，回退到自定义绘制")

    # 绘制结果
    visualizer = DetectionVisualizer(use_pil=use_pil)
    frame_drawn = visualizer.draw(frame.copy(), detections, confidence_threshold)

    return frame_drawn, detections


# 兼容旧接口
def multi_model_inference_legacy(_models: List[Dict], frame: np.ndarray,
                                 confidence_threshold: float = 0.25,
                                 parallel: bool = True,
                                 use_plot_for_single: bool = True) -> Tuple[np.ndarray, List[Dict]]:
    """旧接口兼容函数"""
    frame_drawn, detections = multi_model_inference(
        _models, frame, confidence_threshold, parallel, use_plot_for_single=use_plot_for_single
    )

    # 转换为旧格式
    old_detections = []
    for det in detections:
        old_detections.append({
            'model_idx': det.model_idx,
            'model_name': det.model_name,
            'boxes': det.boxes,
            'confidences': det.confidences,
            'class_ids': det.class_ids,
            'class_names': det.class_names,
            'tags': det.tags
        })

    return frame_drawn, old_detections


if __name__ == '__main__':
    from ultralytics import YOLO

    print('加载模型中')
    model_paths = [
        r"F:\PyModelScope\Yolov\models\yolov8m.pt",
        r"F:\PyModelScope\Yolov\models\car.pt"
    ]

    print("预热模型...")
    model_list = []
    for i, path in enumerate(model_paths):
        model = YOLO(path)
        model.to('cuda:0' if torch.cuda.is_available() else 'cpu')
        model.eval()

        # 创建模型配置
        model_config = ModelConfig(
            model=model,
            name=f'model_{i}',
            device='cuda:0',
            conf_thres=0.45,
            iou_thres=0.45,
            half=False,
            imgsz=1920,
            tags={
                "0": {"name": "汽车", "reliability": 0.4, "select": True, "color": [0, 255, 0]},
                "1": {"name": "行人", "reliability": 0.3, "select": True, "color": [255, 0, 0]},
                "2": {"name": "自行车", "reliability": 0.5, "select": False, "color": [0, 0, 255]}
            },
            config=None
        )
        model_list.append(model_config)

    print("模型预热完成")
    image_path = r"F:\PyModelScope\Yolov\images\444.png"
    frame = cv2.imread(image_path)

    # 测试单个模型的情况 - 使用plot绘制
    print("\n=== 测试单个模型 (使用plot绘制) ===")
    single_model_list = [model_list[0]]  # 只使用第一个模型
    frame_drawn_single, detections_single = multi_model_inference(
        single_model_list, frame, use_plot_for_single=True
    )
    print(f"单个模型检测结果: {len(detections_single[0].boxes)} 个目标")
    cv2.imwrite("uploads/result_single_plot.jpg", frame_drawn_single)
    print("结果已保存到 uploads/result_single_plot.jpg")

    # 测试单个模型的情况 - 强制使用自定义绘制
    print("\n=== 测试单个模型 (强制使用自定义绘制) ===")
    frame_drawn_single_custom, detections_single_custom = multi_model_inference(
        single_model_list, frame, use_plot_for_single=False
    )
    print(f"单个模型自定义绘制结果: {len(detections_single_custom[0].boxes)} 个目标")
    cv2.imwrite("uploads/result_single_custom.jpg", frame_drawn_single_custom)
    print("结果已保存到 uploads/result_single_custom.jpg")

    # 测试多个模型的情况
    print("\n=== 测试多个模型 (使用自定义绘制) ===")
    frame_drawn_multi, detections_multi = multi_model_inference(
        model_list, frame, use_plot_for_single=True  # 即使设为True，多个模型也会使用自定义绘制
    )
    print(f"多个模型检测结果:")
    for det in detections_multi:
        print(f"  模型 {det.model_name}: 检测到 {len(det.boxes)} 个目标")
        for box, conf, cls_id, cls_name in zip(det.boxes, det.confidences,
                                               det.class_ids, det.class_names):
            print(f"box:{box},conf:{conf},cls_id:{cls_id},cls_name:{cls_name}")
            if conf >= 0.25:  # 全局阈值
                should_draw, color = DetectionVisualizer().should_draw_detection(
                    cls_id, conf, det.tags)
                status = "绘制" if should_draw else "不绘制"
                print(f"    {cls_name} (置信度: {conf:.2f}): {status}")

    cv2.imwrite("uploads/result_multi.jpg", frame_drawn_multi)
    print("结果已保存到 uploads/result_multi.jpg")