更新接口

yolopart/detector.py

@@ -2,6 +2,7 @@ import cv2
 import numpy as np
 from ultralytics import YOLO
 import os
+from PIL import Image, ImageDraw, ImageFont
 
 class LicensePlateYOLO:
     """
@@ -113,19 +114,38 @@ class LicensePlateYOLO:
             print(f"检测过程中出错: {e}")
             return []
     
-    def draw_detections(self, image, detections):
+    def draw_detections(self, image, detections, plate_numbers=None):
         """
         在图像上绘制检测结果
         
         参数:
             image: 输入图像
             detections: 检测结果列表
+            plate_numbers: 车牌号列表，与detections对应
         
         返回:
             numpy.ndarray: 绘制了检测结果的图像
         """
         draw_image = image.copy()
         
+        # 转换为PIL图像以支持中文字符
+        pil_image = Image.fromarray(cv2.cvtColor(draw_image, cv2.COLOR_BGR2RGB))
+        draw = ImageDraw.Draw(pil_image)
+        
+        # 尝试加载中文字体
+        try:
+            # Windows系统常见的中文字体
+            font_path = "C:/Windows/Fonts/simhei.ttf"  # 黑体
+            if not os.path.exists(font_path):
+                font_path = "C:/Windows/Fonts/msyh.ttc"  # 微软雅黑
+            if not os.path.exists(font_path):
+                font_path = "C:/Windows/Fonts/simsun.ttc"  # 宋体
+            
+            font = ImageFont.truetype(font_path, 20)
+        except:
+            # 如果无法加载字体，使用默认字体
+            font = ImageFont.load_default()
+        
         for i, detection in enumerate(detections):
             box = detection['box']
             keypoints = detection['keypoints']
@@ -133,6 +153,11 @@ class LicensePlateYOLO:
             confidence = detection['confidence']
             incomplete = detection.get('incomplete', False)
             
+            # 获取对应的车牌号
+            plate_number = ""
+            if plate_numbers and i < len(plate_numbers):
+                plate_number = plate_numbers[i]
+            
             # 绘制边界框
             x1, y1, x2, y2 = map(int, box)
             
@@ -140,30 +165,53 @@ class LicensePlateYOLO:
             if class_name == '绿牌':
                 box_color = (0, 255, 0)  # 绿色
             elif class_name == '蓝牌':
-                box_color = (255, 0, 0)  # 蓝色
+                box_color = (0, 0, 255)  # 蓝色
             else:
                 box_color = (128, 128, 128)  # 灰色
             
-            cv2.rectangle(draw_image, (x1, y1), (x2, y2), box_color, 2)
+            # 在PIL图像上绘制边界框
+            draw.rectangle([(x1, y1), (x2, y2)], outline=box_color, width=2)
+            
+            # 构建标签文本
+            if plate_number:
+                label = f"{class_name} {plate_number} {confidence:.2f}"
+            else:
+                label = f"{class_name} {confidence:.2f}"
             
-            # 绘制标签
-            label = f"{class_name} {confidence:.2f}"
             if incomplete:
                 label += " (不完整)"
             
-            # 计算文本大小和位置
-            font = cv2.FONT_HERSHEY_SIMPLEX
-            font_scale = 0.6
-            thickness = 2
-            (text_width, text_height), _ = cv2.getTextSize(label, font, font_scale, thickness)
+            # 计算文本大小
+            bbox = draw.textbbox((0, 0), label, font=font)
+            text_width = bbox[2] - bbox[0]
+            text_height = bbox[3] - bbox[1]
             
             # 绘制文本背景
-            cv2.rectangle(draw_image, (x1, y1 - text_height - 10), 
-                         (x1 + text_width, y1), box_color, -1)
+            draw.rectangle([(x1, y1 - text_height - 10), (x1 + text_width, y1)], 
+                         fill=box_color)
             
             # 绘制文本
-            cv2.putText(draw_image, label, (x1, y1 - 5), 
-                       font, font_scale, (255, 255, 255), thickness)
+            draw.text((x1, y1 - text_height - 5), label, fill=(255, 255, 255), font=font)
+        
+        # 转换回OpenCV格式
+        draw_image = cv2.cvtColor(np.array(pil_image), cv2.COLOR_RGB2BGR)
+        
+        # 绘制关键点和连线（使用OpenCV）
+        for i, detection in enumerate(detections):
+            box = detection['box']
+            keypoints = detection['keypoints']
+            incomplete = detection.get('incomplete', False)
+            
+            x1, y1, x2, y2 = map(int, box)
+            
+            # 根据车牌类型选择颜色
+            class_name = detection['class_name']
+            if class_name == '绿牌':
+                box_color = (0, 255, 0)  # 绿色
+            elif class_name == '蓝牌':
+                box_color = (0, 0, 255)  # 蓝色
+            else:
+                box_color = (128, 128, 128)  # 灰色
             
             # 绘制关键点和连线
             if len(keypoints) >= 4 and not incomplete: