目标检测数据集

一、为什么需要专门的目标检测数据集

目标检测和图像分类不一样，图像分类只需要知道图片里的物体是什么类别，而目标检测不仅要知道类别，还要知道物体的位置，所以目标检测的数据集里不仅有图片，还有每个物体的边界框信息。

但是目标检测领域没有像 MNIST、Fashion-MNIST 那样的小型标准数据集，为了快速测试目标检测模型，我们可以用一个简单的香蕉检测数据集来学习。

二、香蕉检测数据集介绍

这个数据集是人工制作的，步骤是：

1. 拍摄了很多香蕉的照片，生成了 1000 张不同角度、不同大小的香蕉图片。
2. 把这些香蕉图片放到不同的背景图片里，每个背景图片里只放一个香蕉。
3. 给每个香蕉标注了边界框，也就是香蕉在图片里的位置。

这个数据集分为训练集和验证集：训练集有 1000 张图片，验证集有 100 张图片。

三、下载数据集

我们可以直接从网上下载这个数据集，用 d2l 的download_extract函数就可以下载并解压：

import os
import pandas as pd
import torch
import torchvision
from d2l import torch as d2l

# 下载并解压数据集
data_dir = d2l.download_extract('banana-detection')

四、读取数据集

1. 读取 CSV 标签文件

数据集里的标签信息存在 CSV 文件里，每个图片的标签包含：

• 图片名称
• 物体的类别（这里只有香蕉，类别是 0）
• 边界框的左上角坐标（x1, y1）
• 边界框的右下角坐标（x2, y2）

我们可以用read_data_bananas函数来读取图片和标签：

def read_data_bananas(is_train=True):
    """读取香蕉检测数据集的图片和标签"""
    # 确定数据集的路径
    data_dir = d2l.download_extract('banana-detection')
    csv_fname = os.path.join(data_dir, 'bananas_train' if is_train else 'bananas_val', 'label.csv')
    # 读取CSV文件
    csv_data = pd.read_csv(csv_fname)
    csv_data = csv_data.set_index('img_name')
    
    images, targets = [], []
    for img_name, target in csv_data.iterrows():
        # 读取图片
        img_path = os.path.join(data_dir, 'bananas_train' if is_train else 'bananas_val', 'images', f'{img_name}')
        images.append(torchvision.io.read_image(img_path))
        # 标签里包含（类别, 左上角x, 左上角y, 右下角x, 右下角y）
        targets.append(list(target))
    
    # 把标签转换成张量，并且把坐标除以256，缩放到0-1之间
    return images, torch.tensor(targets).unsqueeze(1) / 256

2. 创建自定义数据集类

为了方便使用 PyTorch 的数据加载器，我们可以创建一个自定义的数据集类BananasDataset：

class BananasDataset(torch.utils.data.Dataset):
    """自定义的香蕉检测数据集类"""
    def __init__(self, is_train):
        # 读取图片和标签
        self.features, self.labels = read_data_bananas(is_train)
        print(f'读取了{len(self.features)}张{"训练" if is_train else "验证"}图片')
    
    def __getitem__(self, idx):
        # 返回第idx张图片和标签
        return (self.features[idx].float(), self.labels[idx])
    
    def __len__(self):
        # 返回数据集的大小
        return len(self.features)

3. 创建数据加载器

最后，我们可以用load_data_bananas函数来创建训练集和验证集的数据加载器：

def load_data_bananas(batch_size):
    """加载香蕉检测数据集，返回数据加载器"""
    train_iter = torch.utils.data.DataLoader(BananasDataset(is_train=True),
                                           batch_size, shuffle=True)
    val_iter = torch.utils.data.DataLoader(BananasDataset(is_train=False),
                                         batch_size)
    return train_iter, val_iter

五、目标检测数据集的特点

1. 数据形状

我们可以查看一下数据加载器返回的小批量数据的形状：

batch_size = 32
train_iter, _ = load_data_bananas(batch_size)
batch = next(iter(train_iter))
print('图片的形状:', batch[0].shape)
print('标签的形状:', batch[1].shape)

运行结果是：

图片的形状: torch.Size([32, 3, 256, 256])
标签的形状: torch.Size([32, 1, 5])

• 图片的形状是(批量大小, 通道数, 高度, 宽度)，和图像分类的图片形状一样。

• 标签的形状是(批量大小, m, 5)，其中 m 是每张图片里最多的边界框数量，这里每张图片只有一个香蕉，所以 m=1。每个标签是一个长度为 5 的数组：

  • 第一个元素是物体的类别，这里是 0（香蕉），如果是 - 1 的话，代表这个边界框是用来填充的无效框。

  • 后面四个元素是边界框的坐标，已经缩放到 0-1 之间了。

2. 展示图片和边界框

我们可以展示一些图片和它们的边界框，看看数据集的样子：

edge_size = 256
imgs = (batch[0][0:10].permute(0, 2, 3, 1)) / 255
axes = d2l.show_images(imgs, 2, 5, scale=2)
for ax, label in zip(axes, batch[1][0:10]):
    # 把坐标转换回图片的尺寸
    bbox = label[0][1:5] * edge_size
    d2l.show_bboxes(ax, [bbox], colors=['w'])

这样我们就能看到 10 张图片，每张图片里的香蕉都被白色的边界框框住了，这些香蕉的角度、大小、位置都不一样。

六、小结

1. 目标检测数据集里不仅有图片，还有每个物体的边界框信息，这和图像分类数据集不一样。

2. 香蕉检测数据集是一个简单的人工数据集，适合用来快速测试目标检测模型。

3. 读取目标检测数据集的时候，需要同时读取图片和标签，并且把标签转换成适合模型输入的形状。

4. 目标检测的标签形状是(批量大小, m, 5)，m 是每张图片里最多的边界框数量，每个标签包含类别和边界框坐标。

（注：文档部分内容可能由 AI 生成）