import numpy as np import cv2 from PIL import Image, ImageDraw, ImageFont import os from collections import defaultdict class OptimizedDetectionRenderer: def __init__(self, font_path=None): self.font_cache = {} self.font_path = font_path self.drawn_labels = set() # 记录已经绘制过的标签 self.detection_cache = defaultdict(list) # 缓存同一位置的检测结果 def _get_font(self, size): """获取字体对象，带缓存""" if size in self.font_cache: return self.font_cache[size] font_paths = [] if self.font_path and os.path.exists(self.font_path): font_paths.append(self.font_path) # 添加常用字体路径 font_paths.extend([ "simhei.ttf", "msyh.ttc", "C:/Windows/Fonts/simhei.ttf", "C:/Windows/Fonts/msyh.ttc", "C:/Windows/Fonts/Deng.ttf", "C:/Windows/Fonts/simsun.ttc", "/System/Library/Fonts/PingFang.ttc", # macOS "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf" # Linux ]) font = None for path in font_paths: if os.path.exists(path): try: font = ImageFont.truetype(path, size, encoding="utf-8") print(f"加载字体: {path}") break except: continue if font is None: try: font = ImageFont.load_default() except: font = ImageFont.load_default() self.font_cache[size] = font return font def compute_iou(self, box1, box2): """计算两个边界框的IoU""" x1 = max(box1[0], box2[0]) y1 = max(box1[1], box2[1]) x2 = min(box1[2], box2[2]) y2 = min(box1[3], box2[3]) if x1 >= x2 or y1 >= y2: return 0.0 intersection = (x2 - x1) * (y2 - y1) area1 = (box1[2] - box1[0]) * (box1[3] - box1[1]) area2 = (box2[2] - box2[0]) * (box2[3] - box2[1]) return intersection / (area1 + area2 - intersection) def filter_duplicate_detections(self, detections, iou_threshold=0.3): """过滤重复检测结果，保留置信度最高的""" filtered_detections = [] # 按置信度降序排序 print(detections) sorted_detections = sorted(detections, key=lambda x: x['confidence'], reverse=True) for det in sorted_detections: is_duplicate = False for kept_det in filtered_detections: iou = self.compute_iou(det['box'], kept_det['box']) # 如果IoU超过阈值，认为是同一目标 if iou > iou_threshold: is_duplicate = True # 如果是相同类别，用置信度高的 if det['class_name'] == kept_det['class_name']: # 如果当前检测置信度更高，替换 if det['confidence'] > kept_det['confidence']: filtered_detections.remove(kept_det) filtered_detections.append(det) break if not is_duplicate: filtered_detections.append(det) return filtered_detections def draw_text_with_background(self, img, text, position, font_size=20, text_color=(0, 255, 0), bg_color=None, padding=5): """绘制带背景的文本（自动调整背景大小）""" try: # 转换为PIL图像 img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) pil_img = Image.fromarray(img_rgb) draw = ImageDraw.Draw(pil_img) # 获取字体 font = self._get_font(font_size) # 计算文本尺寸 if hasattr(font, 'getbbox'): # 新版本PIL bbox = font.getbbox(text) text_width = bbox[2] - bbox[0] text_height = bbox[3] - bbox[1] else: # 旧版本PIL text_width, text_height = font.getsize(text) # 计算背景位置 x, y = position bg_x1 = x - padding bg_y1 = y - padding bg_x2 = x + text_width + padding bg_y2 = y + text_height + padding # 如果提供了背景颜色，绘制背景 if bg_color: # 将BGR转换为RGB rgb_bg_color = bg_color[::-1] draw.rectangle([bg_x1, bg_y1, bg_x2, bg_y2], fill=rgb_bg_color) # 绘制文本 rgb_text_color = text_color[::-1] draw.text(position, text, font=font, fill=rgb_text_color) # 转换回OpenCV格式 return cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR) except Exception as e: print(f"文本渲染失败: {e}，使用后备方案") # 后备方案 if bg_color: cv2.rectangle(img, (position[0] - padding, position[1] - padding), (position[0] + len(text) * 10 + padding, position[1] + font_size + padding), bg_color, cv2.FILLED) cv2.putText(img, text, position, cv2.FONT_HERSHEY_SIMPLEX, font_size / 30, text_color, 2) return img def draw_detection(self, frame, detection_info, model_config, show_model_name=True, show_confidence=True): """在帧上绘制检测结果""" x1, y1, x2, y2 = map(int, detection_info['box']) color = tuple(detection_info['color']) line_width = model_config.get('line_width', 2) # 绘制边界框 cv2.rectangle(frame, (x1, y1), (x2, y2), color, line_width) # 准备标签文本 label_parts = [] if show_model_name and 'model_name' in detection_info: label_parts.append(detection_info['model_name']) if 'class_name' in detection_info: label_parts.append(detection_info['class_name']) if show_confidence and 'confidence' in detection_info: label_parts.append(f"{detection_info['confidence']:.2f}") label = ":".join(label_parts) # 生成标签的唯一标识 label_key = f"{detection_info['class_name']}_{x1}_{y1}" # 检查是否已经绘制过相同类别的标签（在一定区域内） label_drawn = False for drawn_label in self.drawn_labels: drawn_class, drawn_x, drawn_y = drawn_label.split('_') drawn_x, drawn_y = int(drawn_x), int(drawn_y) # 计算距离，如果很近且类别相同，认为已经绘制过 distance = np.sqrt((x1 - drawn_x) ** 2 + (y1 - drawn_y) ** 2) if distance < 50 and detection_info['class_name'] == drawn_class: label_drawn = True break if not label_drawn: # 计算标签位置（放在框的上方，如果上方空间不够则放在下方） label_y = y1 - 20 if label_y < 20: # 如果上方空间不够 label_y = y2 + 20 # 绘制带背景的标签 frame = self.draw_text_with_background( frame, label, (x1, label_y), font_size=model_config.get('font_size', 20), text_color=(255, 255, 255), bg_color=color, padding=3 ) # 记录已绘制的标签 self.drawn_labels.add(label_key) return frame def draw_all_detections(self, frame, all_detections, model_configs, enable_nms=True, show_model_name=True): """绘制所有检测结果（主入口函数）""" # 重置已绘制标签记录 self.drawn_labels.clear() if not all_detections: return frame # 如果需要，过滤重复检测 if enable_nms: filtered_detections = self.filter_duplicate_detections(all_detections) else: filtered_detections = all_detections # 按置信度排序，先绘制置信度低的，再绘制置信度高的 sorted_detections = sorted(filtered_detections, key=lambda x: x.get('confidence', 0)) # 绘制每个检测结果 for detection in sorted_detections: model_id = detection.get('model_id') model_config = model_configs.get(model_id, {}) frame = self.draw_detection(frame, detection, model_config, show_model_name) return frame def put_text_simple(self, img, text, position, font_size=20, color=(0, 255, 0)): """简化版文本绘制函数""" return self.draw_text_with_background( img, text, position, font_size, color, None, 0 ) # 使用示例 def main(): # 初始化渲染器 renderer = OptimizedDetectionRenderer() # 模拟多模型检测结果 detections = [ { 'model_id': 'yolov8n', 'model_name': 'YOLOv8', 'class_id': 0, 'class_name': 'person', 'confidence': 0.85, 'box': [100, 100, 200, 300], 'reliability': 0.9, 'color': (0, 255, 0) }, { 'model_id': 'yolov8s', 'model_name': 'YOLOv8s', 'class_id': 0, 'class_name': 'person', 'confidence': 0.75, 'box': [110, 110, 210, 310], # 与第一个重叠 'reliability': 0.8, 'color': (0, 0, 255) }, { 'model_id': 'yolov8n', 'model_name': 'YOLOv8', 'class_id': 2, 'class_name': 'car', 'confidence': 0.95, 'box': [300, 150, 450, 250], 'reliability': 0.95, 'color': (255, 0, 0) } ] # 模型配置 model_configs = { 'yolov8n': {'line_width': 2, 'font_size': 20}, 'yolov8s': {'line_width': 2, 'font_size': 18} } # 读取测试图像 frame = np.zeros((480, 640, 3), dtype=np.uint8) # 绘制所有检测结果 frame = renderer.draw_all_detections( frame, detections, model_configs, enable_nms=True, # 启用NMS去重 show_model_name=True ) # 显示结果 cv2.imshow('Detections', frame) cv2.waitKey(0) cv2.destroyAllWindows() if __name__ == "__main__": main()