117 lines
3.4 KiB
Python
117 lines
3.4 KiB
Python
import cv2
|
||
import numpy as np
|
||
from datetime import datetime
|
||
import os
|
||
|
||
class VisualizationUtils:
|
||
@staticmethod
|
||
def draw_detection_results(
|
||
image: np.ndarray,
|
||
persons: list,
|
||
distances: list = None,
|
||
save_path: str = "output",
|
||
show_confidence: bool = True
|
||
) -> np.ndarray:
|
||
"""
|
||
在图像上绘制检测结果和距离信息
|
||
|
||
Args:
|
||
image: 原始图像
|
||
persons: 检测到的人物列表，每个元素包含bbox和confidence
|
||
distances: 对应的距离列表（可选）
|
||
save_path: 保存路径
|
||
show_confidence: 是否显示置信度
|
||
|
||
Returns:
|
||
绘制了检测结果的图像
|
||
"""
|
||
# 创建图像副本
|
||
vis_image = image.copy()
|
||
|
||
# 定义颜色和字体
|
||
BOX_COLOR = (0, 255, 0) # 绿色边框
|
||
TEXT_COLOR = (255, 255, 255) # 白色文字
|
||
FONT = cv2.FONT_HERSHEY_SIMPLEX
|
||
FONT_SCALE = 0.6
|
||
THICKNESS = 2
|
||
|
||
# 遍历每个检测结果
|
||
for idx, person in enumerate(persons):
|
||
bbox = person['bbox']
|
||
conf = person['confidence']
|
||
|
||
# 绘制边界框
|
||
cv2.rectangle(
|
||
vis_image,
|
||
(bbox[0], bbox[1]),
|
||
(bbox[2], bbox[3]),
|
||
BOX_COLOR,
|
||
THICKNESS
|
||
)
|
||
|
||
# 准备显示文本
|
||
text_items = []
|
||
if show_confidence:
|
||
text_items.append(f"Conf: {conf:.2f}")
|
||
|
||
if distances and idx < len(distances):
|
||
dist_meters = distances[idx] / 1000 # 转换为米
|
||
text_items.append(f"Dist: {dist_meters:.2f}m")
|
||
|
||
# 绘制文本
|
||
text = " | ".join(text_items)
|
||
if text:
|
||
# 获取文本大小
|
||
(text_width, text_height), _ = cv2.getTextSize(
|
||
text, FONT, FONT_SCALE, THICKNESS
|
||
)
|
||
|
||
# 绘制文本背景
|
||
cv2.rectangle(
|
||
vis_image,
|
||
(bbox[0], bbox[1] - text_height - 10),
|
||
(bbox[0] + text_width + 10, bbox[1]),
|
||
BOX_COLOR,
|
||
-1 # 填充矩形
|
||
)
|
||
|
||
# 绘制文本
|
||
cv2.putText(
|
||
vis_image,
|
||
text,
|
||
(bbox[0] + 5, bbox[1] - 5),
|
||
FONT,
|
||
FONT_SCALE,
|
||
TEXT_COLOR,
|
||
THICKNESS
|
||
)
|
||
|
||
return vis_image
|
||
|
||
@staticmethod
|
||
def save_detection_image(
|
||
image: np.ndarray,
|
||
save_path: str = "output"
|
||
) -> str:
|
||
"""
|
||
保存检测结果图像
|
||
|
||
Args:
|
||
image: 要保存的图像
|
||
save_path: 保存路径
|
||
|
||
Returns:
|
||
保存的文件路径
|
||
"""
|
||
# 确保输出目录存在
|
||
os.makedirs(save_path, exist_ok=True)
|
||
|
||
# 生成文件名（使用时间戳）
|
||
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
|
||
filename = f"detection_{timestamp}.jpg"
|
||
filepath = os.path.join(save_path, filename)
|
||
|
||
# 保存图像
|
||
cv2.imwrite(filepath, image)
|
||
|
||
return filepath |