152 lines
5.2 KiB
Python
152 lines
5.2 KiB
Python
|
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
|