import numpy as np from typing import List, Tuple import cv2 from concurrent.futures import ThreadPoolExecutor, as_completed def multi_model_inference( 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]]: """ 多模型并行推理 Args: models: 已加载的模型列表，每个元素是包含'model'键的字典 frame: 视频帧 (BGR格式) confidence_threshold: 全局置信度阈值 parallel: 是否并行推理 use_plot_for_single: 当只有一个模型时，是否使用result.plot()绘制 Returns: (绘制完成的帧, 检测结果列表) """ if len(models) == 0: return frame, [] original_frame = frame.copy() def inference_single_model(model_info: dict, img: np.ndarray) -> dict: """单模型推理""" try: model = model_info['model'] # 使用模型配置中的参数，如果不存在则使用默认值 conf_thres = model_info.get('conf_thres', confidence_threshold) iou_thres = model_info.get('iou_thres', 0.45) imgsz = model_info.get('imgsz', 640) device = model_info.get('device', 'cpu') # 执行推理 results = model.predict( source=img, imgsz=imgsz, conf=conf_thres, iou=iou_thres, verbose=False, device=device, max_det=300 ) result = results[0] if len(results) > 0 else None model_name = model_info.get('name', f"model_{model_info.get('id', 'unknown')}") model_id = model_info.get('id', 0) if result and result.boxes is not None and len(result.boxes): # 提取检测信息 boxes_xyxy = result.boxes.xyxy.cpu().numpy() confidences = result.boxes.conf.cpu().numpy() class_ids = result.boxes.cls.cpu().numpy().astype(int) class_names = [result.names[int(cls_id)] for cls_id in class_ids] return { 'model_idx': model_id, 'model_name': model_name, 'boxes': boxes_xyxy, 'confidences': confidences, 'class_ids': class_ids, 'class_names': class_names, 'tags': model_info.get('tags', []), 'raw_result': result, 'success': True } else: return { 'model_idx': model_id, 'model_name': model_name, 'boxes': np.array([]), 'confidences': np.array([]), 'class_ids': np.array([]), 'class_names': [], 'tags': model_info.get('tags', []), 'raw_result': result, 'success': False } except Exception as e: print(f"模型 {model_info.get('name', 'unknown')} 推理失败: {str(e)}") model_id = model_info.get('id', 0) model_name = model_info.get('name', f"model_{model_id}") return { 'model_idx': model_id, 'model_name': model_name, 'boxes': np.array([]), 'confidences': np.array([]), 'class_ids': np.array([]), 'class_names': [], 'tags': model_info.get('tags', []), 'raw_result': None, 'success': False } # 执行推理 detections = [] if parallel and len(models) > 1: # 并行推理 max_workers = min(len(models), 4) with ThreadPoolExecutor(max_workers=max_workers) as executor: # 提交推理任务 future_to_model = {} for model_info in models: future = executor.submit(inference_single_model, model_info, original_frame) future_to_model[future] = model_info # 收集结果 for future in as_completed(future_to_model.keys()): try: result = future.result(timeout=30.0) detections.append(result) except Exception as e: model_info = future_to_model[future] model_name = model_info.get('name', 'unknown') print(f"并行推理任务失败 - 模型 {model_name}: {str(e)}") else: # 串行推理 for model_info in models: result = inference_single_model(model_info, original_frame) detections.append(result) # 结果绘制 plotted_frame = original_frame.copy() if len(detections) == 1 and use_plot_for_single: # 单个模型直接使用YOLO的plot()方法 detection = detections[0] if detection['success'] and detection['raw_result'] is not None: try: # 使用YOLO自带的plot方法 plotted_frame_rgb = detection['raw_result'].plot( img=original_frame, conf=True, labels=True, boxes=True, line_width=2 ) # 确保返回的是BGR格式 if plotted_frame_rgb.shape[-1] == 3: plotted_frame = cv2.cvtColor(plotted_frame_rgb, cv2.COLOR_RGB2BGR) else: plotted_frame = plotted_frame_rgb except Exception as e: print(f"使用YOLO plot()绘制失败: {str(e)}") # 单个模型时不使用自定义绘制，直接返回原图 plotted_frame = original_frame.copy() else: plotted_frame = original_frame.copy() elif len(detections) > 0: # 多个模型使用自定义绘制 plotted_frame = plot_custom_results(original_frame, detections) else: plotted_frame = original_frame.copy() return plotted_frame, detections def plot_custom_results(frame: np.ndarray, detections: List[dict]) -> np.ndarray: """自定义绘制多个模型的结果""" result_frame = frame.copy() # 预定义的颜色列表 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), # 橄榄色 ] for idx, detection in enumerate(detections): if not detection['success'] or len(detection['boxes']) == 0: continue color = colors[idx % len(colors)] model_name = detection['model_name'] # 绘制每个检测框 for box, conf, cls_id, cls_name in zip( detection['boxes'], detection['confidences'], detection['class_ids'], detection['class_names'] ): x1, y1, x2, y2 = map(int, box) # 绘制边界框 cv2.rectangle(result_frame, (x1, y1), (x2, y2), color, 2) # 准备标签 label = f"{model_name}:{cls_name} {conf:.2f}" # 计算文本尺寸 font = cv2.FONT_HERSHEY_SIMPLEX font_scale = 0.5 thickness = 2 (text_width, text_height), baseline = cv2.getTextSize( label, font, font_scale, thickness ) # 绘制标签背景 label_y1 = max(0, y1 - text_height - baseline - 5) label_y2 = max(0, y1) label_x2 = x1 + text_width if label_y1 < label_y2 and label_x2 > x1: cv2.rectangle( result_frame, (x1, label_y1), (label_x2, label_y2), color, -1 ) # 绘制标签文本 text_y = max(baseline + 5, label_y1 + text_height + baseline) cv2.putText( result_frame, label, (x1, text_y), font, font_scale, (255, 255, 255), thickness ) return result_frame # 简化版本，用于快速集成 def multi_model_inference_simple( models: List[dict], frame: np.ndarray, confidence_threshold: float = 0.25 ) -> Tuple[np.ndarray, List[dict]]: """ 简化的多模型推理（串行，使用plot绘制） Args: models: 已加载的模型列表 frame: 视频帧 confidence_threshold: 置信度阈值 Returns: (绘制完成的帧, 检测结果列表) """ if len(models) == 0: return frame, [] original_frame = frame.copy() detections = [] # 串行推理 for model_info in models: try: model = model_info['model'] conf_thres = model_info.get('conf_thres', confidence_threshold) iou_thres = model_info.get('iou_thres', 0.45) imgsz = model_info.get('imgsz', 640) device = model_info.get('device', 'cpu') # 执行推理 results = model.predict( source=original_frame, imgsz=imgsz, conf=conf_thres, iou=iou_thres, verbose=False, device=device, max_det=300 ) result = results[0] if len(results) > 0 else None model_name = model_info.get('name', f"model_{model_info.get('id', 'unknown')}") model_id = model_info.get('id', 0) detection = { 'model_idx': model_id, 'model_name': model_name, 'raw_result': result, 'success': result is not None and result.boxes is not None and len(result.boxes) > 0 } detections.append(detection) except Exception as e: print(f"模型推理失败: {str(e)}") detections.append({ 'model_idx': model_info.get('id', 0), 'model_name': model_info.get('name', 'unknown'), 'raw_result': None, 'success': False }) # 绘制结果 plotted_frame = original_frame.copy() if len(models) == 1: # 单个模型使用plot() detection = detections[0] if detection['success'] and detection['raw_result'] is not None: try: plotted_frame_rgb = detection['raw_result'].plot( img=original_frame, conf=True, labels=True, boxes=True ) if plotted_frame_rgb.shape[-1] == 3: plotted_frame = cv2.cvtColor(plotted_frame_rgb, cv2.COLOR_RGB2BGR) except: plotted_frame = original_frame.copy() else: # 多个模型使用自定义绘制 plotted_frame = plot_custom_results(original_frame, detections) return plotted_frame, detections