539 lines
20 KiB
Python
539 lines
20 KiB
Python
import cv2
|
||
import torch
|
||
import numpy as np
|
||
import concurrent.futures
|
||
from typing import List, Dict, Tuple, Any, Optional, Union
|
||
from dataclasses import dataclass
|
||
from collections import defaultdict
|
||
|
||
from PIL import Image, ImageDraw, ImageFont
|
||
|
||
# 尝试使用PIL绘制中文
|
||
font = None
|
||
FONT_PATHS = [
|
||
"simhei.ttf", # 黑体
|
||
"/usr/share/fonts/truetype/wqy/wqy-microhei.ttc", # 文泉驿微米黑
|
||
"C:/Windows/Fonts/simhei.ttf", # Windows黑体
|
||
"/System/Library/Fonts/PingFang.ttc" # macOS苹方
|
||
]
|
||
|
||
for path in FONT_PATHS:
|
||
try:
|
||
font = ImageFont.truetype(path, 20)
|
||
break
|
||
except:
|
||
continue
|
||
|
||
|
||
@dataclass
|
||
class ModelConfig:
|
||
"""模型配置数据类"""
|
||
model: Any
|
||
config: Any
|
||
tags: Optional[Dict[str, Dict]] = None
|
||
name: str = "未命名"
|
||
id: Any = None
|
||
device: str = "cuda:0"
|
||
conf_thres: float = 0.25
|
||
iou_thres: float = 0.45
|
||
half: bool = False
|
||
key_valid: Any = None
|
||
model_hash: Any = None
|
||
imgsz: Any = 640
|
||
|
||
def __post_init__(self):
|
||
"""初始化后处理"""
|
||
if self.tags is None:
|
||
self.tags = {}
|
||
|
||
|
||
@dataclass
|
||
class Detection:
|
||
"""检测结果数据类"""
|
||
model_idx: int
|
||
model_name: str
|
||
boxes: np.ndarray
|
||
confidences: np.ndarray
|
||
class_ids: np.ndarray
|
||
class_names: List[str]
|
||
tags: Dict[str, Dict]
|
||
raw_result: Any = None # 新增：保存原始结果对象
|
||
|
||
|
||
class DetectionVisualizer:
|
||
"""检测结果可视化器"""
|
||
|
||
# 预定义模型颜色
|
||
MODEL_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), # 橄榄色
|
||
]
|
||
|
||
def __init__(self, use_pil: bool = True):
|
||
"""
|
||
初始化可视化器
|
||
|
||
Args:
|
||
use_pil: 是否使用PIL绘制（支持中文）
|
||
"""
|
||
self.use_pil = use_pil
|
||
self.font = font
|
||
|
||
def should_draw_detection(self, class_id: int, confidence: float,
|
||
tags: Dict[str, Dict]) -> Tuple[bool, Optional[Tuple[int, int, int]]]:
|
||
"""
|
||
判断是否应该绘制检测框
|
||
|
||
Args:
|
||
class_id: 类别ID
|
||
confidence: 置信度
|
||
tags: 标签配置字典
|
||
|
||
Returns:
|
||
(是否绘制, 颜色)
|
||
"""
|
||
class_id_str = str(class_id)
|
||
|
||
# 如果标签配置为空，默认不绘制任何标签（因为需求是只绘制配置的标签）
|
||
if not tags:
|
||
return False, None
|
||
|
||
# 如果标签不在配置中，不绘制
|
||
tag_config = tags.get(class_id_str)
|
||
if not tag_config:
|
||
return False, None
|
||
|
||
# 检查select标记
|
||
if not tag_config.get('select', True):
|
||
return False, None
|
||
|
||
# 检查置信度阈值
|
||
reliability = tag_config.get('reliability', 0)
|
||
if confidence < reliability:
|
||
return False, None
|
||
|
||
# 获取自定义颜色
|
||
color = tag_config.get('color')
|
||
if color and isinstance(color, (list, tuple)) and len(color) >= 3:
|
||
return True, tuple(color[:3])
|
||
|
||
return True, None
|
||
|
||
def draw_with_pil(self, frame: np.ndarray, detections: List[Detection],
|
||
confidence_threshold: float) -> np.ndarray:
|
||
"""使用PIL绘制检测结果（支持中文）"""
|
||
# 转换到PIL格式
|
||
frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
|
||
pil_image = Image.fromarray(frame_rgb)
|
||
draw = ImageDraw.Draw(pil_image)
|
||
|
||
for det in detections:
|
||
model_idx = det.model_idx
|
||
boxes = det.boxes
|
||
confidences = det.confidences
|
||
class_names = det.class_names
|
||
class_ids = det.class_ids
|
||
tags = det.tags
|
||
|
||
for i, (box, conf, cls_id, cls_name) in enumerate(
|
||
zip(boxes, confidences, class_ids, class_names)):
|
||
|
||
if conf < confidence_threshold:
|
||
continue
|
||
|
||
should_draw, custom_color = self.should_draw_detection(
|
||
cls_id, conf, tags)
|
||
if not should_draw:
|
||
continue
|
||
|
||
# 使用自定义颜色或模型颜色
|
||
color = custom_color or self.MODEL_COLORS[model_idx % len(self.MODEL_COLORS)]
|
||
|
||
x1, y1, x2, y2 = map(int, box[:4])
|
||
|
||
# 标签文本
|
||
label = f"{cls_name} {conf:.2f}"
|
||
|
||
# 绘制矩形框
|
||
draw.rectangle([x1, y1, x2, y2], outline=color, width=2)
|
||
|
||
# 绘制标签背景
|
||
if self.font:
|
||
try:
|
||
text_bbox = draw.textbbox((x1, y1 - 25), label, font=self.font)
|
||
draw.rectangle(text_bbox, fill=color)
|
||
draw.text((x1, y1 - 25), label, fill=(255, 255, 255), font=self.font)
|
||
except:
|
||
# 字体失败，回退到OpenCV
|
||
cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)
|
||
cv2.putText(frame, label, (x1, y1 - 10),
|
||
cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
|
||
else:
|
||
cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)
|
||
cv2.putText(frame, label, (x1, y1 - 10),
|
||
cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
|
||
|
||
# 转换回OpenCV格式
|
||
return cv2.cvtColor(np.array(pil_image), cv2.COLOR_RGB2BGR)
|
||
|
||
def draw_with_opencv(self, frame: np.ndarray, detections: List[Detection],
|
||
confidence_threshold: float) -> np.ndarray:
|
||
"""使用OpenCV绘制检测结果"""
|
||
frame_drawn = frame.copy()
|
||
|
||
for det in detections:
|
||
model_idx = det.model_idx
|
||
model_name = det.model_name
|
||
boxes = det.boxes
|
||
confidences = det.confidences
|
||
class_names = det.class_names
|
||
class_ids = det.class_ids
|
||
tags = det.tags
|
||
|
||
for i, (box, conf, cls_id, cls_name) in enumerate(
|
||
zip(boxes, confidences, class_ids, class_names)):
|
||
|
||
if conf < confidence_threshold:
|
||
continue
|
||
|
||
should_draw, custom_color = self.should_draw_detection(
|
||
cls_id, conf, tags)
|
||
if not should_draw:
|
||
continue
|
||
|
||
# 使用自定义颜色或模型颜色
|
||
color = custom_color or self.MODEL_COLORS[model_idx % len(self.MODEL_COLORS)]
|
||
|
||
x1, y1, x2, y2 = map(int, box[:4])
|
||
|
||
# 绘制矩形框
|
||
cv2.rectangle(frame_drawn, (x1, y1), (x2, y2), color, 2)
|
||
|
||
# 绘制标签
|
||
label = f"{model_name}: {cls_name} {conf:.2f}"
|
||
|
||
# 计算文本大小
|
||
(text_width, text_height), baseline = cv2.getTextSize(
|
||
label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 2
|
||
)
|
||
|
||
# 绘制标签背景
|
||
cv2.rectangle(frame_drawn,
|
||
(x1, y1 - text_height - 10),
|
||
(x1 + text_width, y1),
|
||
color, -1)
|
||
|
||
# 绘制文本
|
||
cv2.putText(frame_drawn, label, (x1, y1 - 5),
|
||
cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 2)
|
||
|
||
return frame_drawn
|
||
|
||
def draw(self, frame: np.ndarray, detections: List[Detection],
|
||
confidence_threshold: float) -> np.ndarray:
|
||
"""绘制检测结果"""
|
||
if self.use_pil:
|
||
try:
|
||
return self.draw_with_pil(frame, detections, confidence_threshold)
|
||
except Exception as e:
|
||
print(f"PIL绘制失败，使用OpenCV: {e}")
|
||
return self.draw_with_opencv(frame, detections, confidence_threshold)
|
||
else:
|
||
return self.draw_with_opencv(frame, detections, confidence_threshold)
|
||
|
||
|
||
class YOLOModelWrapper:
|
||
"""YOLO模型包装器"""
|
||
|
||
@staticmethod
|
||
def infer_yolov8(model: Any, frame: np.ndarray, conf_thres: float,
|
||
iou_thres: float, imgsz: int, half: bool, device: str):
|
||
"""YOLOv8推理"""
|
||
with torch.no_grad():
|
||
results = model(frame, conf=conf_thres, verbose=False,
|
||
iou=iou_thres, imgsz=imgsz, half=half, device=device)
|
||
|
||
if not results or len(results) == 0:
|
||
return np.array([]), np.array([]), np.array([]), [], None
|
||
|
||
result = results[0]
|
||
if not hasattr(result, 'boxes'):
|
||
return np.array([]), np.array([]), np.array([]), [], None
|
||
|
||
boxes = result.boxes.xyxy.cpu().numpy()
|
||
confidences = result.boxes.conf.cpu().numpy()
|
||
class_ids = result.boxes.cls.cpu().numpy().astype(int)
|
||
class_names = [result.names.get(cid, str(cid)) for cid in class_ids]
|
||
|
||
return boxes, confidences, class_ids, class_names, result
|
||
|
||
@staticmethod
|
||
def infer_yolov5(model: Any, frame: np.ndarray):
|
||
"""YOLOv5推理"""
|
||
results = model(frame)
|
||
|
||
if not hasattr(results, 'xyxy'):
|
||
return np.array([]), np.array([]), np.array([]), [], None
|
||
|
||
detections = results.xyxy[0].cpu().numpy()
|
||
if len(detections) == 0:
|
||
return np.array([]), np.array([]), np.array([]), [], None
|
||
|
||
boxes = detections[:, :4]
|
||
confidences = detections[:, 4]
|
||
class_ids = detections[:, 5].astype(int)
|
||
|
||
if hasattr(results, 'names'):
|
||
class_names = [results.names.get(cid, str(cid)) for cid in class_ids]
|
||
else:
|
||
class_names = [str(cid) for cid in class_ids]
|
||
|
||
return boxes, confidences, class_ids, class_names, results
|
||
|
||
|
||
def multi_model_inference(
|
||
models_config: List[Union[Dict, ModelConfig]],
|
||
frame: np.ndarray,
|
||
confidence_threshold: float = 0.25,
|
||
parallel: bool = True,
|
||
use_pil: bool = True,
|
||
use_plot_for_single: bool = True # 新增参数：是否对单个模型使用plot绘制
|
||
) -> Tuple[np.ndarray, List[Detection]]:
|
||
"""
|
||
多模型并行推理
|
||
|
||
Args:
|
||
models_config: 模型配置列表
|
||
frame: 视频帧 (BGR格式)
|
||
confidence_threshold: 全局置信度阈值
|
||
parallel: 是否并行推理
|
||
use_pil: 是否使用PIL绘制
|
||
use_plot_for_single: 当只有一个模型时，是否使用result.plot()绘制
|
||
|
||
Returns:
|
||
(绘制完成的帧, 检测结果列表)
|
||
"""
|
||
# 转换为ModelConfig对象
|
||
configs = []
|
||
for cfg in models_config:
|
||
if isinstance(cfg, dict):
|
||
configs.append(ModelConfig(**cfg))
|
||
else:
|
||
configs.append(cfg)
|
||
|
||
def single_model_inference(model_cfg: ModelConfig, model_idx: int) -> Detection:
|
||
"""单个模型的推理函数"""
|
||
try:
|
||
model = model_cfg.model
|
||
|
||
# 根据模型类型进行推理
|
||
if hasattr(model, 'predict'): # YOLOv8
|
||
boxes, confidences, class_ids, class_names, raw_result = YOLOModelWrapper.infer_yolov8(
|
||
model, frame, model_cfg.conf_thres, model_cfg.iou_thres,
|
||
model_cfg.imgsz, model_cfg.half, model_cfg.device
|
||
)
|
||
elif hasattr(model, '__call__'): # YOLOv5
|
||
boxes, confidences, class_ids, class_names, raw_result = YOLOModelWrapper.infer_yolov5(model, frame)
|
||
else:
|
||
raise ValueError(f"不支持的模型类型: {type(model)}")
|
||
|
||
return Detection(
|
||
model_idx=model_idx,
|
||
model_name=model_cfg.name,
|
||
boxes=boxes,
|
||
confidences=confidences,
|
||
class_ids=class_ids,
|
||
class_names=class_names,
|
||
tags=model_cfg.tags,
|
||
raw_result=raw_result
|
||
)
|
||
|
||
except Exception as e:
|
||
print(f"模型 {model_idx} ({model_cfg.name}) 推理失败: {e}")
|
||
return Detection(
|
||
model_idx=model_idx,
|
||
model_name=model_cfg.name,
|
||
boxes=np.array([]),
|
||
confidences=np.array([]),
|
||
class_ids=np.array([]),
|
||
class_names=[],
|
||
tags=model_cfg.tags,
|
||
raw_result=None
|
||
)
|
||
|
||
# 并行推理
|
||
if parallel and len(configs) > 1:
|
||
detections = []
|
||
with concurrent.futures.ThreadPoolExecutor(max_workers=len(configs)) as executor:
|
||
# 提交任务
|
||
future_to_idx = {
|
||
executor.submit(single_model_inference, cfg, idx): idx
|
||
for idx, cfg in enumerate(configs)
|
||
}
|
||
|
||
# 收集结果
|
||
for future in concurrent.futures.as_completed(future_to_idx):
|
||
try:
|
||
detection = future.result(timeout=2.0)
|
||
detections.append(detection)
|
||
except concurrent.futures.TimeoutError:
|
||
idx = future_to_idx[future]
|
||
print(f"模型 {idx} 推理超时")
|
||
# 创建空的检测结果
|
||
detections.append(Detection(
|
||
model_idx=idx,
|
||
model_name=configs[idx].name,
|
||
boxes=np.array([]),
|
||
confidences=np.array([]),
|
||
class_ids=np.array([]),
|
||
class_names=[],
|
||
tags=configs[idx].tags,
|
||
raw_result=None
|
||
))
|
||
except Exception as e:
|
||
print(f"模型推理异常: {e}")
|
||
else:
|
||
# 顺序推理
|
||
detections = [single_model_inference(cfg, idx)
|
||
for idx, cfg in enumerate(configs)]
|
||
|
||
# 按照模型索引排序
|
||
detections.sort(key=lambda x: x.model_idx)
|
||
|
||
# 新增：如果只有一个模型且配置使用plot绘制，则使用result.plot()
|
||
if len(configs) == 1 and use_plot_for_single:
|
||
single_detection = detections[0]
|
||
if single_detection.raw_result is not None:
|
||
try:
|
||
# 检查是否有plot方法
|
||
if hasattr(single_detection.raw_result, 'plot'):
|
||
# 使用plot方法绘制结果
|
||
frame_drawn = single_detection.raw_result.plot()
|
||
# 确保返回的是numpy数组
|
||
if not isinstance(frame_drawn, np.ndarray):
|
||
frame_drawn = np.array(frame_drawn)
|
||
# plot方法通常返回RGB图像，转换为BGR
|
||
if len(frame_drawn.shape) == 3 and frame_drawn.shape[2] == 3:
|
||
frame_drawn = cv2.cvtColor(frame_drawn, cv2.COLOR_RGB2BGR)
|
||
print(f"使用 {single_detection.model_name} 的 plot() 方法绘制结果")
|
||
return frame_drawn, detections
|
||
else:
|
||
print(f"模型 {single_detection.model_name} 的结果对象没有 plot() 方法，使用自定义绘制")
|
||
except Exception as e:
|
||
print(f"使用 plot() 方法绘制失败: {e}，回退到自定义绘制")
|
||
|
||
# 绘制结果
|
||
visualizer = DetectionVisualizer(use_pil=use_pil)
|
||
frame_drawn = visualizer.draw(frame.copy(), detections, confidence_threshold)
|
||
|
||
return frame_drawn, detections
|
||
|
||
|
||
# 兼容旧接口
|
||
def multi_model_inference_legacy(_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]]:
|
||
"""旧接口兼容函数"""
|
||
frame_drawn, detections = multi_model_inference(
|
||
_models, frame, confidence_threshold, parallel, use_plot_for_single=use_plot_for_single
|
||
)
|
||
|
||
# 转换为旧格式
|
||
old_detections = []
|
||
for det in detections:
|
||
old_detections.append({
|
||
'model_idx': det.model_idx,
|
||
'model_name': det.model_name,
|
||
'boxes': det.boxes,
|
||
'confidences': det.confidences,
|
||
'class_ids': det.class_ids,
|
||
'class_names': det.class_names,
|
||
'tags': det.tags
|
||
})
|
||
|
||
return frame_drawn, old_detections
|
||
|
||
|
||
if __name__ == '__main__':
|
||
from ultralytics import YOLO
|
||
|
||
print('加载模型中')
|
||
model_paths = [
|
||
r"F:\PyModelScope\Yolov\models\yolov8m.pt",
|
||
r"F:\PyModelScope\Yolov\models\car.pt"
|
||
]
|
||
|
||
print("预热模型...")
|
||
model_list = []
|
||
for i, path in enumerate(model_paths):
|
||
model = YOLO(path)
|
||
model.to('cuda:0' if torch.cuda.is_available() else 'cpu')
|
||
model.eval()
|
||
|
||
# 创建模型配置
|
||
model_config = ModelConfig(
|
||
model=model,
|
||
name=f'model_{i}',
|
||
device='cuda:0',
|
||
conf_thres=0.45,
|
||
iou_thres=0.45,
|
||
half=False,
|
||
imgsz=1920,
|
||
tags={
|
||
"0": {"name": "汽车", "reliability": 0.4, "select": True, "color": [0, 255, 0]},
|
||
"1": {"name": "行人", "reliability": 0.3, "select": True, "color": [255, 0, 0]},
|
||
"2": {"name": "自行车", "reliability": 0.5, "select": False, "color": [0, 0, 255]}
|
||
},
|
||
config=None
|
||
)
|
||
model_list.append(model_config)
|
||
|
||
print("模型预热完成")
|
||
image_path = r"F:\PyModelScope\Yolov\images\444.png"
|
||
frame = cv2.imread(image_path)
|
||
|
||
# 测试单个模型的情况 - 使用plot绘制
|
||
print("\n=== 测试单个模型 (使用plot绘制) ===")
|
||
single_model_list = [model_list[0]] # 只使用第一个模型
|
||
frame_drawn_single, detections_single = multi_model_inference(
|
||
single_model_list, frame, use_plot_for_single=True
|
||
)
|
||
print(f"单个模型检测结果: {len(detections_single[0].boxes)} 个目标")
|
||
cv2.imwrite("uploads/result_single_plot.jpg", frame_drawn_single)
|
||
print("结果已保存到 uploads/result_single_plot.jpg")
|
||
|
||
# 测试单个模型的情况 - 强制使用自定义绘制
|
||
print("\n=== 测试单个模型 (强制使用自定义绘制) ===")
|
||
frame_drawn_single_custom, detections_single_custom = multi_model_inference(
|
||
single_model_list, frame, use_plot_for_single=False
|
||
)
|
||
print(f"单个模型自定义绘制结果: {len(detections_single_custom[0].boxes)} 个目标")
|
||
cv2.imwrite("uploads/result_single_custom.jpg", frame_drawn_single_custom)
|
||
print("结果已保存到 uploads/result_single_custom.jpg")
|
||
|
||
# 测试多个模型的情况
|
||
print("\n=== 测试多个模型 (使用自定义绘制) ===")
|
||
frame_drawn_multi, detections_multi = multi_model_inference(
|
||
model_list, frame, use_plot_for_single=True # 即使设为True，多个模型也会使用自定义绘制
|
||
)
|
||
print(f"多个模型检测结果:")
|
||
for det in detections_multi:
|
||
print(f" 模型 {det.model_name}: 检测到 {len(det.boxes)} 个目标")
|
||
for box, conf, cls_id, cls_name in zip(det.boxes, det.confidences,
|
||
det.class_ids, det.class_names):
|
||
print(f"box:{box},conf:{conf},cls_id:{cls_id},cls_name:{cls_name}")
|
||
if conf >= 0.25: # 全局阈值
|
||
should_draw, color = DetectionVisualizer().should_draw_detection(
|
||
cls_id, conf, det.tags)
|
||
status = "绘制" if should_draw else "不绘制"
|
||
print(f" {cls_name} (置信度: {conf:.2f}): {status}")
|
||
|
||
cv2.imwrite("uploads/result_multi.jpg", frame_drawn_multi)
|
||
print("结果已保存到 uploads/result_multi.jpg") |