基于Matplotlib实现YOLOv5目标检测结果的实时棋盘式绘图技术求助

我明白你现在的需求：用自定义YOLOv5模型检测视频帧里的目标，然后把这些目标实时绘制在类似棋盘的区域里，而且想要x轴对应图像高度、y轴对应图像宽度。结合你的代码，我来一步步帮你实现这个功能：

1. 先理清坐标转换逻辑 #

YOLOv5输出的xyxy格式坐标（xmin, ymin, xmax, ymax）默认是基于图像的**宽度（x轴：左→右，0到图像宽度）和高度（y轴：上→下，0到图像高度）**的。如果你想让x轴对应图像高度、y轴对应图像宽度，需要对坐标做转换：

• 假设原图像尺寸是(img_height, img_width)（注意：PIL图像的size是(width, height)，所以要反过来）
• 转换后的x坐标 = 原y坐标（ymin/ymax）
• 转换后的y坐标 = img_width - 原x坐标（xmin/xmax）（可根据棋盘布局调整方向）

2. 初始化实时绘图环境 #

Matplotlib的交互式模式（plt.ion()）是实现实时绘图的关键，我们需要先初始化画布、设置棋盘背景和坐标轴范围：

import matplotlib.pyplot as plt
import matplotlib.patches as patches

# 初始化绘图
plt.ion()  # 开启交互式模式
fig, ax = plt.subplots()

# 设置棋盘背景（黑白格子示例，可自定义）
def create_board_background(img_height, img_width, grid_size=50):
    ax.clear()
    # 设置坐标轴范围：x对应图像高度，y对应图像宽度
    ax.set_xlim(0, img_height)
    ax.set_ylim(0, img_width)
    # 绘制网格线
    for x in range(0, img_height, grid_size):
        ax.axvline(x=x, color='gray', linestyle='-', linewidth=0.5)
    for y in range(0, img_width, grid_size):
        ax.axhline(y=y, color='gray', linestyle='-', linewidth=0.5)
    # 填充黑白格子
    for i in range(0, img_height // grid_size):
        for j in range(0, img_width // grid_size):
            rect_color = 'white' if (i + j) % 2 == 0 else 'lightgray'
            rect = patches.Rectangle((i*grid_size, j*grid_size), grid_size, grid_size, facecolor=rect_color)
            ax.add_patch(rect)
    ax.set_aspect('equal')  # 保持坐标轴比例一致

# 先获取一次截图确定图像尺寸
initial_screenshot = pyautogui.screenshot(region=(window.left, window.top, window.width, window.height))
img_width, img_height = initial_screenshot.size  # PIL图像size是(width, height)
create_board_background(img_height, img_width)
plt.show()

3. 修改推理循环，加入实时绘图逻辑 #

把绘图代码嵌入到你的循环里，每次检测到目标后更新画布上的标记：

import pyautogui
import time

# 假设model已提前加载完成
while True:
    current_screen_shot = pyautogui.screenshot(region=(window.left, window.top, window.width, window.height))
    img_width, img_height = current_screen_shot.size  # 实时获取窗口尺寸，防止大小变化
    
    # 模型推理
    results = model(current_screen_shot, size=640)
    df = results.pandas().xyxy[0]
    
    # 重置棋盘背景，清除上一帧的标记
    create_board_background(img_height, img_width)
    
    if not df.empty:
        print("Got results!")
        # 遍历每个检测到的目标
        for idx, row in df.iterrows():
            # 获取原YOLO检测坐标和信息
            xmin, ymin, xmax, ymax = row['xmin'], row['ymin'], row['xmax'], row['ymax']
            conf = row['confidence']
            cls_name = row['name']
            
            # 坐标转换：适配x轴对应图像高度、y轴对应图像宽度的需求
            new_x1 = ymin
            new_y1 = img_width - xmax
            new_x2 = ymax
            new_y2 = img_width - xmin
            
            # 绘制目标矩形框
            rect = patches.Rectangle((new_x1, new_y1), new_x2 - new_x1, new_y2 - new_y1, 
                                    linewidth=2, edgecolor='red', facecolor='none')
            ax.add_patch(rect)
            # 添加类别和置信度标签
            ax.text(new_x1, new_y1, f"{cls_name}: {conf:.2f}", 
                    bbox=dict(facecolor='yellow', alpha=0.5), fontsize=8)
    else:
        print("No results!")
    
    # 更新画布实现实时显示
    plt.draw()
    plt.pause(0.01)  # 短暂暂停让Matplotlib处理绘图事件
    time.sleep(2)  # 你的原有延迟，可根据需求调整

# 循环结束后关闭交互式模式
plt.ioff()
plt.close()

关键注意点 #

• 坐标轴方向调整：如果你的棋盘y轴不需要反转，直接把new_y1和new_y2换成xmin和xmax即可，完全匹配你的棋盘布局需求。
• 实时效率优化：plt.pause(0.01)是实时更新的核心，它会触发画布刷新；如果检测速度较慢，可以适当增大这个值。
• 残留标记清理：每次循环都重置棋盘背景，避免上一帧的检测标记残留。

内容的提问来源于stack exchange，提问作者eddy