从代码学习深度学习 - 目标检测前置知识(二) PyTorch版

前言

大家好！欢迎来到“从代码学习深度学习-目标检测前置知识”的第二部分，我们将继续深入探讨目标检测的前置知识。在上一部分，我们已经了解了目标检测的基本概念和锚框的生成。今天，我们将重点关注多尺度目标检测的理念，学习如何生成不同尺度的锚框来适应不同大小的目标物体，并了解如何加载和处理自定义的目标检测数据集。本篇将结合 PyTorch 代码进行实践，帮助大家更直观地理解这些概念。

目标检测的一个核心挑战是如何有效地检测图像中大小各异的物体。简单的单尺度锚框生成策略可能难以覆盖所有情况。因此，多尺度检测应运而生，它允许我们在不同的特征图层级上生成不同尺寸和比例的锚框，从而更有效地捕捉从微小到巨大的各类目标。

此外，为了方便模型验证和快速迭代，我们还将介绍一个专门用于目标检测的小型数据集——香蕉检测数据集，并学习如何使用 PyTorch 的 Dataset 和 DataLoader 来加载和处理它。

让我们开始吧！

完整代码：下载链接

一、多尺度目标检测

在之前的讨论中，我们可能已经了解到，可以围绕图像中的每个像素生成锚框。然而，这种方法会产生海量的锚框，导致计算量巨大。为了解决这个问题，可以采取一些策略来减少锚框数量，例如：

1. 采样像素： 不再为每个像素生成锚框，而是在图像中均匀采样一部分像素作为锚框中心。
2. 多尺度生成： 在不同的特征图尺度上生成不同大小和数量的锚框。直观上，小目标在图像中可能出现的位置更密集、更多样，因此在检测小目标时，可以使用较小的锚框并采样更多的区域；而在检测大目标时，可以使用较大的锚框并采样较少的区域。

1.1 多尺度锚框

为了更具体地理解多尺度锚框，我们首先加载一张示例图片，并了解其基本信息。

%matplotlib inline
import torch
import matplotlib.pyplot as plt

# 读取图像文件 'img/03_catdog.jpg'
# 返回的 img 是一个形状为 (高, 宽, 通道数) 的 numpy 数组
img = plt.imread('img/03_catdog.jpg')

# 获取图像的高度和宽度
# h: 标量，表示图像的高度（像素数）
# w: 标量，表示图像的宽度（像素数）
# img.shape[:2] 返回的是一个包含两个元素的元组 (高度, 宽度)
h, w = img.shape[:2]

# 打印图像的高度和宽度
# h: 标量，表示图像的高度（像素数）
# w: 标量，表示图像的宽度（像素数）
print(f"图像尺寸: 高度 = {
     h} 像素, 宽度 = {
     w} 像素")

输出：

图像尺寸: 高度 = 561 像素, 宽度 = 728 像素

接下来，我们定义一个核心函数 multibox_prior，用于生成以特征图上每个像素为中心、具有不同形状（大小和宽高比）的锚框。

%matplotlib inline
import torch
torch.set_printoptions(2)  # 精简输出精度

def multibox_prior(data, sizes, ratios):
    """生成以每个像素为中心具有不同形状的锚框

    参数：
    data：输入图像张量，维度为(批量大小, 通道数, 高度, 宽度)
    sizes：锚框缩放比列表，元素个数为num_sizes，每个元素∈(0,1]
    ratios：锚框宽高比列表，元素个数为num_ratios，每个元素>0

    返回：
    输出张量，维度为(1, 像素总数*每像素锚框数, 4)，表示所有锚框的坐标
    """
    # 获取输入数据的高度和宽度
    # in_height, in_width: 标量
    in_height, in_width = data.shape[-2:]

    # 获取设备信息以及尺寸和比例的数量
    # device: 字符串; num_sizes, num_ratios: 标量
    device, num_sizes, num_ratios = data.device, len(sizes), len(ratios)

    # 计算每个像素点产生的锚框数量 = 尺寸数 + 宽高比数 - 1
    # boxes_per_pixel: 标量
    boxes_per_pixel = (num_sizes + num_ratios - 1)

    # 将尺寸和比例转换为张量
    # size_tensor: 维度为(num_sizes,)
    # ratio_tensor: 维度为(num_ratios,)
    size_tensor = torch.tensor(sizes, device=device)
    ratio_tensor = torch.tensor(ratios, device=device)

    # 为了将锚点移动到像素的中心，需要设置偏移量
    # 因为一个像素的高为1且宽为1，我们选择偏移中心0.5
    # offset_h, offset_w: 标量
    offset_h, offset_w = 0.5, 0.5

    # 计算高度和宽度方向上的步长（归一化）
    # steps_h, steps_w: 标量
    steps_h = 1.0 / in_height  # 在y轴上缩放步长
    steps_w = 1.0 / in_width   # 在x轴上缩放步长

    # 生成锚框的所有中心点
    # center_h: 维度为(in_height,)
    # center_w: 维度为(in_width,)
    center_h = (torch.arange(in_height, device=device) + offset_h) * steps_h
    center_w = (torch.arange(in_width, device=device) + offset_w) * steps_w

    # 使用meshgrid生成网格坐标
    # shift_y, shift_x: 维度均为(in_height, in_width)
    shift_y, shift_x = torch.meshgrid(center_h, center_w, indexing='ij')

    # 将坐标展平为一维
    # shift_y, shift_x: 展平后维度均为(in_height*in_width,)
    shift_y, shift_x = shift_y.reshape(-1), shift_x.reshape(-1)

    # 生成"boxes_per_pixel"个高和宽，
    # 之后用于创建锚框的四角坐标(xmin,ymin,xmax,ymax)

    # 计算锚框宽度：先计算尺寸与第一个比例的组合，再计算第一个尺寸与其余比例的组合
    # w: 维度为(num_sizes + num_ratios - 1,)
    w = torch.cat((size_tensor * torch.sqrt(ratio_tensor[0]),
                   sizes[0] * torch.sqrt(ratio_tensor[1:])))\
                   * in_height / in_width  # 处理矩形输入，调整宽度

    # 计算锚框高度：对应于宽度的计算方式
    # h: 维度为(num_sizes + num_ratios - 1,)
    h = torch.cat((size_tensor / torch.sqrt(ratio_tensor[0]),
                   sizes[0] / torch.sqrt(ratio_tensor[1:])))

    # 计算锚框的四个坐标相对于中心点的偏移量：(-w/2, -h/2, w/2, h/2)
    # anchor_manipulations: 维度为(in_height*in_width*boxes_per_pixel, 4)
    anchor_manipulations = torch.stack((-w, -h, w, h)).T.repeat(
                                        in_height * in_width, 1) / 2

    # 每个中心点都将有"boxes_per_pixel"个锚框，