【AI 进阶笔记】cv 小白初上手 Kaggle Person-Collecting-Waste COCO Dataset
原创
【AI 进阶笔记】cv 小白初上手 Kaggle Person-Collecting-Waste COCO Dataset
原创
1. 引言
在计算机视觉的领域，目标检测（Object Detection）是一个非常重要的任务。与传统的分类任务不同，目标检测不仅要求模型判断图片中有哪些物体，还需要精确地确定物体在图像中的位置。Faster R-CNN 是一种强大的目标检测模型，广泛应用于图像识别、自动驾驶、视频分析等场景。
最近刚刚学习了 Faster R-CNN，为了深入理解其原理和实现，决定亲自动手实践，于是从 Kaggle 找了个数据集（Person-Collecting-Waste COCO Dataset）练练手。这篇博客将详细讲解 Faster R-CNN 的工作原理，并分享从数据准备、模型训练、到推理与可视化的完整过程。
2. Faster R-CNN 原理解析
2.1 Faster R-CNN 的整体架构
Faster R-CNN 是基于卷积神经网络（CNN）的一种目标检测框架，它的核心思想是通过一系列层来生成候选区域、进行目标分类和位置回归。Faster R-CNN 由以下几部分组成：
- 特征提取（Backbone）：使用卷积神经网络（如 ResNet-50 或 ResNet-101）作为特征提取网络，提取图像的高层次特征。
- 区域生成网络（Region Proposal Network, RPN）：RPN 网络负责在图像的特征图上生成候选框，并为每个候选框打分，评估其是否包含目标物体。
- ROI Pooling：由于候选框大小不一致，需要使用 ROI Pooling 层将每个候选框裁剪并调整为相同大小，以便输入到全连接层进行分类和回归。
- 目标分类与边界框回归：最后，通过分类器预测每个候选框的类别，并通过回归网络预测目标的精确边界框。
这些模块协同工作，最终将图像中的目标检测出来，并给出每个目标的类别及其边界框位置。
下图展示了 Faster R-CNN 的整体流程：
图中（这个图不是很生动），输入图像首先通过 backbone 提取特征，然后通过 RPN 生成候选框，并使用 ROI Pooling 进行统一尺寸处理，最后通过分类和回归网络输出最终的检测结果。
2.2 RPN 详解
RPN（Region Proposal Network）是 Faster R-CNN 中一个关键的部分，它的作用是生成候选区域。RPN 会在特征图的每个像素点生成多个不同尺度的候选框（Anchor），并根据目标框的真实标签（Ground Truth）进行分类和回归。
RPN 通过一个小型的卷积网络来计算每个 Anchor 的得分和偏移量：
- 前景（Positive）：与真实目标框的 IoU（Intersection over Union）大于 0.7 的 Anchor。
- 背景（Negative）：与真实目标框的 IoU 小于 0.3 的 Anchor。
通过这种方式，RPN 能够从大量的候选框中筛选出一些最有可能包含物体的区域，从而提高后续检测的效率和准确性。
2.3 ROI Pooling 和分类回归
经过 RPN 之后，得到的候选框大小不一，为了让这些框能够输入到后续的全连接层进行分类和回归，Faster R-CNN 使用了 ROI Pooling（或 ROI Align）方法。ROI Pooling 的作用是将不同尺寸的候选框转换成相同尺寸的特征块。
具体来说，ROI Pooling 会将每个候选框映射到一个固定大小的特征图（例如 7x7），并将其中的每个区域进行池化操作，以保留最重要的特征信息。经过 ROI Pooling 后，特征图会传递给两个全连接层：
- 分类分支：预测每个候选框的类别。
- 回归分支：预测每个候选框的精确边界框坐标。
这样，模型就能同时完成目标分类和位置回归的任务，最终输出多个带有类别和位置的预测框。
3. 数据准备与预处理
3.1 COCO 数据格式解析
COCO 数据集是计算机视觉领域广泛使用的数据集，通常用于目标检测、图像分割等任务。COCO 数据集的标注格式包含了每个目标的边界框（bbox）和类别信息：
{
"annotations": [
{
"bbox": [x_min, y_min, width, height],
"category_id": 1
}
]
}
其中，bbox 是目标的边界框，它由左上角的 (x_min, y_min) 坐标和目标的宽度和高度 [width, height] 组成。我们需要将这个格式转换成 Faster R-CNN 需要的 [x_min, y_min, x_max, y_max] 格式，其中 x_max = x_min + width 和 y_max = y_min + height。
3.2 自定义数据集类
为了加载和处理 COCO 数据，我们可以使用 torchvision.datasets.CocoDetection 类来读取数据，并且封装成一个自定义的数据集类 CustomDataset：
import os
import torch
import torchvision.transforms as transforms
import torchvision.datasets as datasets
class CustomDataset(torch.utils.data.Dataset):
def __init__(self, root, split, transforms):
self.root = os.path.join(root, split) # 数据集路径
self.transforms = transforms # 数据预处理操作
self.dataset = datasets.CocoDetection(self.root, f"{self.root}/_annotations.coco.json", transform=self.transforms)
def __getitem__(self, idx):
img, target = self.dataset[idx] # 获取图片和标注
boxes = []
labels = []
for obj in target:
x_min, y_min, width, height = obj["bbox"]
x_max = x_min + width
y_max = y_min + height
# 过滤掉无效的 bbox
if width > 0 and height > 0:
boxes.append([x_min, y_min, x_max, y_max])
labels.append(obj["category_id"])
target = {
"boxes": torch.tensor(boxes, dtype=torch.float32),
"labels": torch.tensor(labels, dtype=torch.int64)
}
return img, target
def __len__(self):
return len(self.dataset)这个数据集类对每张图片的标注进行了转换，并且将每张图片与其对应的目标框一一对应。这样，在训练时我们就可以直接使用这个自定义的数据集类来读取并处理数据。
4. 训练 Faster R-CNN
4.1 加载数据集
在加载数据集时，我们将训练集、验证集和测试集分别加载，并进行批量化处理：
train_dataset = CustomDataset("data", "train", transforms.ToTensor())
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=2, shuffle=True, collate_fn=lambda x: tuple(zip(*x)))4.2 定义 Faster R-CNN 模型
这里使用 torchvision 提供的预训练 Faster R-CNN 模型，并根据数据集的类别数量进行调整：
import torchvision
from torchvision.models.detection.faster_rcnn import FasterRCNN_ResNet50_FPN_Weights
model = torchvision.models.detection.fasterrcnn_resnet50_fpn(weights=FasterRCNN_ResNet50_FPN_Weights.DEFAULT)
num_classes = 2 # 背景 + 1 个类别
in_features = model.roi_heads.box_predictor.cls_score.in_features
model.roi_heads.box_predictor = torchvision.models.detection.faster_rcnn.FastRCNNPredictor(in_features, num_classes)通过修改 roi_heads.box_predictor 来适配我们自己数据集的类别数量。这样模型就能够进行目标分类与回归。
4.3 训练模型
在训练模型时，我们使用 Adam 优化器，并进行损失计算与反向传播：
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
num_epochs = 10
for epoch in range(num_epochs):
model.train()
total_loss = 0
for images, targets in train_loader:
images = [img.to(device) for img in images]
targets = [{k: v.to(device) for k, v in t.items()} for t in targets]
loss_dict = model(images, targets)
losses = sum(loss for loss in loss_dict.values())
optimizer.zero_grad()
losses.backward()
optimizer.step()
total_loss += losses.item()
print(f"Epoch {epoch+1}, Train Loss: {total_loss / len(train_loader)}")数据量很小，没一会就训练完了。
如果和我在本地训练慢的，可以试试腾讯云 Cloud Studio 的高性能工作空间，把模型加载到 GPU，这样就会训练很快啦！
5. 训练分析
5.1 绘制损失曲线
通过绘制训练过程中的损失曲线，我们可以直观地看到模型的训练进展情况：
import matplotlib.pyplot as plt
plt.plot(train_losses, label="Train Loss")
plt.plot(valid_losses, label="Valid Loss")
plt.xlabel("Epoch")
plt.ylabel("Loss")
plt.legend()
plt.show()可以看到 Loss 正常下降
在完成训练后，我们可以通过以下代码对新图片进行预测，并展示检测结果：
def predict(model, image_path):
model.eval()
img = Image.open(image_path).convert("RGB")
img_tensor = transforms.ToTensor()(img).unsqueeze(0).to(device)
with torch.no_grad():
predictions = model(img_tensor)
return predictions
也可以保存模型，方便下次使用。
# 保存模型
torch.save(model.state_dict(), "faster_rcnn_trained.pth")
print("模型已保存")希望这篇文章对你有所帮助！下次见！🚀
原创声明：本文系作者授权腾讯云开发者社区发表，未经许可，不得转载。
如有侵权，请联系 cloudcommunity@tencent.com 删除。
原创声明：本文系作者授权腾讯云开发者社区发表，未经许可，不得转载。
如有侵权，请联系 cloudcommunity@tencent.com 删除。
腾讯云开发者
