128 lines
4.7 KiB
Python
128 lines
4.7 KiB
Python
from ultralytics import YOLO
|
||
import cv2
|
||
import torch
|
||
import numpy as np
|
||
from PIL import ImageFont, ImageDraw, Image
|
||
|
||
|
||
class YOLODetector:
|
||
def __init__(self, model_path='best_weights/best.pt'):
|
||
try:
|
||
# 初始化YOLO模型
|
||
self.model = YOLO(model_path)
|
||
print(f"成功加载模型：{model_path}")
|
||
|
||
self.font = ImageFont.truetype("msyh.ttc", 21)
|
||
# 检测配置
|
||
self.predict_config = {
|
||
'conf_thres': 0.25,
|
||
'iou_thres': 0.30,
|
||
'imgsz': 640,
|
||
'line_width': 2
|
||
}
|
||
|
||
# 设置设备
|
||
self.device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
|
||
if self.device == 'cuda:0':
|
||
print(f"使用GPU: {torch.cuda.get_device_name(0)}")
|
||
torch.backends.cudnn.benchmark = True
|
||
else:
|
||
print("警告：检测，将回退到CPU模式")
|
||
|
||
self.model.to(self.device)
|
||
|
||
except Exception as e:
|
||
print(f"模型初始化失败: {str(e)}")
|
||
raise
|
||
|
||
def process_frame(self, frame, roi=None):
|
||
"""处理单帧图像"""
|
||
try:
|
||
if frame is None or frame.size == 0:
|
||
print("收到空帧")
|
||
return None, None
|
||
|
||
# 如果有ROI，处理ROI区域
|
||
if roi and roi != (0, 0, 0, 0):
|
||
x, y, w, h = roi
|
||
if x >= 0 and y >= 0 and w > 0 and h > 0 and \
|
||
x + w <= frame.shape[1] and y + h <= frame.shape[0]:
|
||
frame_roi = frame[y:y + h, x:x + w]
|
||
else:
|
||
# print("ROI 超出图像范围")
|
||
frame_roi = frame
|
||
else:
|
||
frame_roi = frame
|
||
|
||
# 运行YOLO检测
|
||
results = self.model(
|
||
source=frame_roi,
|
||
conf=self.predict_config['conf_thres'],
|
||
iou=self.predict_config['iou_thres'],
|
||
imgsz=self.predict_config['imgsz'],
|
||
device=self.device,
|
||
verbose=False
|
||
)
|
||
|
||
# 在图像上绘制检测结果
|
||
annotated_frame = frame.copy()
|
||
|
||
if len(results) > 0 and results[0].boxes is not None and len(results[0].boxes) > 0:
|
||
for box, conf, cls in zip(results[0].boxes.xyxy,
|
||
results[0].boxes.conf,
|
||
results[0].boxes.cls):
|
||
class_name = results[0].names[int(cls)]
|
||
x1, y1, x2, y2 = map(int, box)
|
||
|
||
# 如果使用ROI，调整坐标
|
||
# if roi and roi != (0, 0, 0, 0):
|
||
if roi and len(roi) == 4:
|
||
# x1, y1 = x1 + roi[0], y1 + roi[1]
|
||
# x2, y2 = x2 + roi[0], y2 + roi[1]
|
||
x1 += roi[0]
|
||
y1 += roi[1]
|
||
x2 += roi[0]
|
||
y2 += roi[1]
|
||
|
||
# 绘制边界框
|
||
cv2.rectangle(annotated_frame,
|
||
(x1, y1),
|
||
(x2, y2),
|
||
(0, 255, 0),
|
||
self.predict_config['line_width'])
|
||
|
||
# 添加中文标签
|
||
# pil_img = Image.fromarray(cv2.cvtColor(annotated_frame, cv2.COLOR_BGR2RGB))
|
||
# draw = ImageDraw.Draw(pil_img)
|
||
|
||
# 根据类别显示不同的中文标签
|
||
class_name = results[0].names[int(cls)].lower()
|
||
label_map = {
|
||
"helmet": "安全帽",
|
||
"person": "人员",
|
||
"safevest": "工服",
|
||
"smoke": "吸烟",
|
||
"animal": "异物入侵",
|
||
"cellphone": "玩手机",
|
||
"fire": "起火"
|
||
}
|
||
|
||
label = label_map.get(class_name, class_name)
|
||
|
||
# draw.text((x1, y1 - 30), f"{label} {conf:.2f}",
|
||
# font=self.font, fill=(0, 255, 0))
|
||
# annotated_frame = cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR)
|
||
|
||
# 绘制ROI区域
|
||
if roi and roi != (0, 0, 0, 0):
|
||
x, y, w, h = roi
|
||
cv2.rectangle(annotated_frame, (x, y), (x + w, y + h), (0, 0, 255), 2)
|
||
|
||
return annotated_frame, results[0]
|
||
|
||
except Exception as e:
|
||
print(f"处理帧时出错: {str(e)}")
|
||
return frame, None
|
||
|
||
|