yolov5 正样本可视化

效果展示

yolov5的学习重点有3个，分别是：

1. 网络模型
2. 正样本筛选
3. 损失函数

正样本的筛选过程在上一篇文章中已经详细解析了 yolov5 筛选正样本流程 代码多图详解 。本篇说明在这个过程中如何实现正样本可视化。首先来看看实现的效果：
原图：

网格可视化，不同尺度下匹配到的anchor不同。

anchor可视化

原图中的anchor网格可视化：

实现代码

原图添加标注框

原图中并没有标注框，为了更友好的查看，将标注框还原到原图上。
将标注框绘制在图片上，yolov5/utils/dataloaders.py 中load_mosaic函数加入下面绘制矩形框的功能

if labels.size:
    labels[:, 1:] = xywhn2xyxy(labels[:, 1:], w, h, padw, padh)  # normalized xywh to pixel xyxy format
    segments = [xyn2xy(x, w, h, padw, padh) for x in segments]
    for _, x_left, y_left, x_right, y_right in labels:
        cv2.rectangle(img4, (int(x_left), int(y_left)), (int(x_right), int(y_right)), (255, 0, 0), 2)

主要逻辑位置

在loss计算的 yolov5/utils/loss.py 函数 build_targets中，添加相关绘制函数

# feature map 上绘制正样本所在grid cell
anchor_grid_cell_in_feature_map(gij, shape[2])

# feature map 上绘制正样本anchor
anchor_in_feature_map(gij, anchors[a], shape[2])

# 绘制原图
origin_img(imgs)

# 原图上绘制正样本所在grid cell
grid_cell_in_image(gij, shape[2], imgs)

# 原图上绘制正样本anchor
anchor_in_image(gij, anchors[a], shape[2], imgs)

# Append
# indices 保存的内容是：image_id, anchor_id（0，1，2）, grid x刻度  grid y刻度。这里的刻度就是正样本
indices.append((b, a, gj.clamp_(0, shape[2] - 1), gi.clamp_(0, shape[3] - 1)))  # image, anchor, grid

# tbox保存的是gt中心相对于所在grid cell左上角偏移量。也会计算出gt中心相对扩展anchor的偏移量
tbox.append(torch.cat((gxy - gij, gwh), 1))  # box

保存原图

首先将原图保存起来做一个对比

def origin_img(imgs):
    import uuid
    import numpy as np
    from PIL import Image

    image_arr = imgs[0].cpu().numpy().transpose(1, 2, 0) * 255
    image = Image.fromarray(np.uint8(image_arr))
    image.save(f'positive_sample_img/{uuid.uuid4().hex}.png')

[网格图]正负样本所在网格

在划分为2020,4040,80*80的网格下，anchor所在网格的可视化

def anchor_grid_cell_in_feature_map(gij, scale):
    """feature map 网格图"""
    import uuid
    import matplotlib.pyplot as plt
    import numpy as np
    import matplotlib.patches as patches

    # 创建一个新的图形对象
    fig, ax = plt.subplots()

    # 设置网格的大小
    grid_size = scale

    # 设置x和y的范围
    x_range = np.arange(0, grid_size + 1, 1)
    y_range = np.arange(0, grid_size + 1, 1)

    # 绘制水平线
    for y in y_range:
        ax.plot([0, grid_size], [y, y], color='black', linewidth=0.5)

    # 绘制垂直线
    for x in x_range:
        ax.plot([x, x], [0, grid_size], color='black', linewidth=0.5)

    # 设置x和y轴的范围
    ax.set_xlim(0, grid_size)
    ax.set_ylim(grid_size, 0)

    # 设置相等的纵横比
    ax.set_aspect('equal')

    # 将x轴的标尺和标签放在顶部
    ax.xaxis.set_ticks_position('top')
    ax.xaxis.set_label_position('top')

    print(f"当前尺度：{scale}")
    for x, y in gij:
        print(f"x: {x.cpu()}, y: {y.cpu()}")
        rect = patches.Rectangle((x.cpu(), y.cpu()), 1, 1, linewidth=1, edgecolor='red', facecolor='red')
        ax.add_patch(rect)

    # 显示网格
    plt.grid(True)

    plt.savefig(f"positive_sample_img/feature_map_grid_cell{uuid.uuid4().hex[:5]}.png")

    # 显示图形
    plt.show()

[网格图] 正负样本anchor框

在划分为2020,4040,80*80的网格图上可视化anchor框

def anchor_in_feature_map(gij, anchors, scale):
    """feature map 绘制anchor ，三种尺度"""
    import uuid
    import matplotlib.pyplot as plt
    import numpy as np
    from PIL import Image
    import matplotlib.patches as patches

    # 创建一个新的图形对象
    fig, ax = plt.subplots()

    # 设置网格的大小
    grid_size = scale

    # 设置x和y的范围
    x_range = np.arange(0, grid_size + 1, 1)
    y_range = np.arange(0, grid_size + 1, 1)

    # 绘制水平线
    for y in y_range:
        ax.plot([0, grid_size], [y, y], color='black', linewidth=0.5)

    # 绘制垂直线
    for x in x_range:
        ax.plot([x, x], [0, grid_size], color='black', linewidth=0.5)

    # 设置x和y轴的范围
    ax.set_xlim(0, grid_size)
    ax.set_ylim(grid_size, 0)

    # 设置相等的纵横比
    ax.set_aspect('equal')

    # 将x轴的标尺和标签放在顶部
    ax.xaxis.set_ticks_position('top')
    ax.xaxis.set_label_position('top')

    print(f"当前尺度：{scale}")
    for index in range(len(gij.cpu())):
        x, y = gij.cpu()[index]
        width, height = anchors.cpu()[index]
        new_x = x - width / 2
        new_y = y - height / 2
        print(f"x: {new_x}, y: {new_y}")

        rect = patches.Rectangle((new_x, new_y), width, height, linewidth=1, edgecolor='red', facecolor='none')
        ax.add_patch(rect)

    plt.savefig(f"positive_sample_img/feature_map_anchor_{uuid.uuid4().hex[:5]}.png")

    # 显示网格
    plt.grid(True)

[原图] 正负样本所在网格

在原图rezise到640*640之后，可视化出anchor所在的网格。

def grid_cell_in_image(gij, scale, imgs):
    """原图 正负样本网格"""
    import uuid
    import matplotlib.pyplot as plt
    from PIL import Image
    import numpy as np
    import matplotlib.patches as patches

    # 读取图片
    image_arr = imgs[0].cpu().numpy().transpose(1, 2, 0) * 255
    image = Image.fromarray(np.uint8(image_arr))
    image_width, image_height = image.size

    # 创建一个Matplotlib图形对象
    fig, ax = plt.subplots()

    # 显示图片
    ax.imshow(image)
    # 设置网格的大小
    grid_size = image_width / scale  # 你可以根据需要调整网格的大小

    # 设置x和y的范围
    x_range = np.arange(0, image_width + 1, grid_size)
    y_range = np.arange(0, image_height + 1, grid_size)

    # 绘制水平线
    for y in y_range:
        ax.plot([0, image_width], [y, y], color='black', linewidth=0.2)

    # 绘制垂直线
    for x in x_range:
        ax.plot([x, x], [0, image_height], color='black', linewidth=0.2)

    # 设置x和y轴的范围以匹配图像尺寸
    ax.set_xlim(0, image_width)
    ax.set_ylim(image_width, 0)  # 注意：y轴方向与图像坐标系相反

    # 将x轴的标尺和标签放在顶部
    ax.xaxis.set_ticks_position('top')
    ax.xaxis.set_label_position('top')

    print(f"当前尺度：{scale}")
    for index in range(len(gij.cpu())):
        x, y = gij.cpu()[index]
        new_x = x * grid_size
        new_y = y * grid_size
        print(f"x: {new_x}, y: {new_y}")
        rect = patches.Rectangle((new_x, new_y), 1 * grid_size, 1 * grid_size, linewidth=1, edgecolor='red', facecolor='red')
        ax.add_patch(rect)

    plt.savefig(f"positive_sample_img/image_gird_cell_{uuid.uuid4().hex[:5]}.png")

    # 显示图形
    plt.show()

[原图] 正负样本anchor框

在原图中画出anchor框的位置

def anchor_in_image(gij, anchors, scale, imgs):
    """原图上绘制anchor"""
    import uuid
    import matplotlib.pyplot as plt
    from PIL import Image
    import numpy as np
    import matplotlib.patches as patches

    # 读取图片
    image_arr = imgs[0].cpu().numpy().transpose(1, 2, 0) * 255
    image = Image.fromarray(np.uint8(image_arr))
    image_width, image_height = image.size

    # 创建一个Matplotlib图形对象
    fig, ax = plt.subplots()

    # 显示图片
    ax.imshow(image)
    # 设置网格的大小
    grid_size = image_width / scale  # 你可以根据需要调整网格的大小

    # 设置x和y的范围
    x_range = np.arange(0, image_width + 1, grid_size)
    y_range = np.arange(0, image_height + 1, grid_size)

    # 绘制水平线
    for y in y_range:
        ax.plot([0, image_width], [y, y], color='black', linewidth=0.2)

    # 绘制垂直线
    for x in x_range:
        ax.plot([x, x], [0, image_height], color='black', linewidth=0.2)

    # 设置x和y轴的范围以匹配图像尺寸
    ax.set_xlim(0, image_width)
    ax.set_ylim(image_width, 0)  # 注意：y轴方向与图像坐标系相反

    # 将x轴的标尺和标签放在顶部
    ax.xaxis.set_ticks_position('top')
    ax.xaxis.set_label_position('top')

    print(f"当前尺度：{scale}")
    for index in range(len(gij.cpu())):
        x, y = gij.cpu()[index]
        width, height = anchors.cpu()[index]
        new_x = x * grid_size - width * grid_size / 2
        new_y = y * grid_size - height * grid_size / 2
        print(f"x: {new_x}, y: {new_y}")

        rect = patches.Rectangle((new_x, new_y), width * grid_size, height * grid_size, linewidth=1, edgecolor='red', facecolor='none')
        ax.add_patch(rect)

    plt.savefig(f"positive_sample_img/image_anchor_{uuid.uuid4().hex[:5]}.png")

    # 显示图形
    plt.show()

正样本筛选规则

yolov5中正样本筛选的规则如下，分析可视化中如何体现这些规则。

1. 跨anchor预测

一个网格有3个anchor，只要符合anchor匹配的宽高比规则，3个anchor都有可以匹配成为正样本。
可以从比较密集的anchor中看出，一个中心点可以匹配多个anchor。

2. 跨grid预测

假设一个GT框落在了某个预测分支的某个网格内，则该网格有左、上、右、下4个邻域网格，根据GT框的中心位置，将最近的2个邻域网格也作为预测网格，也即一个GT框可以由3个网格来预测。
从所有的anchor网格图中都能看出，网格可视化都是3个网格在一起的。当一个anchor网格选中之后，肯定有其他两个网格也被选中。

3. 跨分支预测

假设一个GT框可以和2个甚至3个预测分支上的anchor匹配，则这2个或3个预测分支都可以预测该GT框。即一个GT框可以在3个预测分支上匹配正样本，在每一个分支上重复anchor匹配和grid匹配的步骤，最终可以得到某个GT 匹配到的所有正样本。

可以看到在3种不同的尺度下，黄色框中的对象都匹配到了。

总结

通过可视化anchor，可以加深对yolov5中正样本筛选的理解。而实现可视化本身也是一件有趣的事情。