PyTorch 目标检测教程

PyTorch 目标检测教程

本教程将介绍如何在 PyTorch 中使用几种常见的目标检测模型，包括 Faster R-CNN、SSD 以及 YOLO (You Only Look Once)。我们将涵盖预训练模型的使用、推理、微调，以及自定义数据集上的训练。

1. 目标检测概述

目标检测任务不仅要识别图像中的物体类别，还要精确定位物体的边界框。在此任务中，每个模型输出一个物体类别标签和一个边界框。

常见目标检测模型：

• Faster R-CNN：基于区域提议方法，具有较高的检测精度。
• SSD (Single Shot MultiBox Detector)：单阶段检测器，速度较快，适合实时应用。
• YOLO (You Only Look Once)：单阶段检测器，具有极快的检测速度，适合大规模实时检测。

2. 官方文档链接

• PyTorch 官方文档
• Torchvision 模型
• YOLO 官方实现 （YOLOv5 及 PyTorch 版 YOLO）

3. Faster R-CNN 模型

3.1 加载 Faster R-CNN 模型并进行推理

import torch
from torchvision.models.detection import fasterrcnn_resnet50_fpn
from torchvision.transforms import functional as F
from PIL import Image
import matplotlib.pyplot as plt
import matplotlib.patches as patches

# 加载预训练的 Faster R-CNN 模型
model = fasterrcnn_resnet50_fpn(pretrained=True)
model.eval()  # 切换为评估模式

# 加载和预处理图像
image_path = "test_image.jpg"
image = Image.open(image_path)
image_tensor = F.to_tensor(image).unsqueeze(0)  # 转换为张量并添加批次维度

# 将模型和输入图像移动到 GPU（如果可用）
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)
image_tensor = image_tensor.to(device)

# 进行推理
with torch.no_grad():
    predictions = model(image_tensor)

# 获取预测的边界框和类别标签
boxes = predictions[0]['boxes'].cpu().numpy()  # 预测的边界框
labels = predictions[0]['labels'].cpu().numpy()  # 预测的类别标签
scores = predictions[0]['scores'].cpu().numpy()  # 预测的分数

# 可视化预测结果
fig, ax = plt.subplots(1)
ax.imshow(image)

# 设置阈值，只显示高置信度的检测结果
threshold = 0.5
for box, label, score in zip(boxes, labels, scores):
    if score > threshold:
        rect = patches.Rectangle((box[0], box[1]), box[2] - box[0], box[3] - box[1], linewidth=2, edgecolor='r', facecolor='none')
        ax.add_patch(rect)
        ax.text(box[0], box[1], f'{label}: {score:.2f}', color='white', fontsize=12, bbox=dict(facecolor='red', alpha=0.5))

plt.show()

3.2 微调 Faster R-CNN 模型

你可以使用自定义的数据集微调 Faster R-CNN。下面是如何定义自定义数据集并在上面进行训练。

import torch
from torch.utils.data import Dataset
from PIL import Image

# 自定义数据集类
class CustomDataset(Dataset):
    def __init__(self, image_paths, annotations):
        self.image_paths = image_paths
        self.annotations = annotations

    def __len__(self):
        return len(self.image_paths)

    def __getitem__(self, idx):
        image = Image.open(self.image_paths[idx]).convert("RGB")
        boxes, labels = self.annotations[idx]
        boxes = torch.as_tensor(boxes, dtype=torch.float32)
        labels = torch.as_tensor(labels, dtype=torch.int64)

        target = {}
        target["boxes"] = boxes
        target["labels"] = labels

        image = F.to_tensor(image)
        return image, target

from torch.utils.data import DataLoader
from torchvision.models.detection import fasterrcnn_resnet50_fpn

# 加载自定义数据集
dataset = CustomDataset(image_paths=["img1.jpg", "img2.jpg"], annotations=[([[10, 20, 200, 300]], [1]), ([[30, 40, 180, 220]], [2])])
dataloader = DataLoader(dataset, batch_size=2, shuffle=True, collate_fn=lambda x: tuple(zip(*x)))

# 加载预训练的 Faster R-CNN 模型
model = fasterrcnn_resnet50_fpn(pretrained=True)
num_classes = 3
in_features = model.roi_heads.box_predictor.cls_score.in_features
model.roi_heads.box_predictor = torch.nn.Linear(in_features, num_classes)

# 定义优化器并进行训练
optimizer = torch.optim.SGD(model.parameters(), lr=0.005, momentum=0.9, weight_decay=0.0005)

# 训练模型
model.train()
for epoch in range(5):
    for images, targets in dataloader:
        images = list(image.to(device) for image 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()

    print(f'Epoch [{epoch+1}/5], Loss: {losses.item():.4f}')

4. SSD 模型

SSD（Single Shot MultiBox Detector）是一种速度较快的单阶段检测器。torchvision 也提供了预训练的 SSD 模型。

4.1 使用预训练 SSD 模型进行推理

import torch
from torchvision.models.detection import ssd300_vgg16
from torchvision.transforms import functional as F
from PIL import Image
import matplotlib.pyplot as plt
import matplotlib.patches as patches

# 加载预训练的 SSD 模型
model = ssd300_vgg16(pretrained=True)
model.eval()

# 加载和预处理图像
image_path = "test_image.jpg"
image = Image.open(image_path).convert("RGB")
image_tensor = F.to_tensor(image).unsqueeze(0)

# 进行推理
with torch.no_grad():
    predictions = model(image_tensor)

# 获取预测的边界框和类别标签
boxes = predictions[0]['boxes'].cpu().numpy()
labels = predictions[0]['labels'].cpu().numpy()
scores = predictions[0]['scores'].cpu().numpy()

# 可视化预测结果
fig, ax = plt.subplots(1)
ax.imshow(image)

threshold = 0.5
for box, label, score in zip(boxes, labels, scores):
    if score > threshold:
        rect = patches.Rectangle((box[0], box[1]), box[2] - box[0], box[3] - box[1], linewidth=2, edgecolor='r', facecolor='none')
        ax.add_patch(rect)
        ax.text(box[0], box[1], f'{label}: {score:.2f}', color='white', fontsize=12, bbox=dict(facecolor='red', alpha=0.5))

plt.show()

5. YOLO 模型

YOLO（You Only Look Once）是单阶段目标检测器，具有非常快的检测速度。YOLOv5 是目前广泛使用的 YOLO 变体，官方提供了 PyTorch 版 YOLOv5。

5.1 安装 YOLOv5

首先，克隆 YOLOv5 的 GitHub 仓库并安装依赖项。

git clone https://github.com/ultralytics/yolov5
cd yolov5
pip install -r requirements.txt

5.2 使用 YOLOv5 进行推理

YOLOv5 提供了一个简单的 API，用于进行推理和训练。

import torch

# 加载预训练的 YOLOv5 模型
model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)

# 加载图像并进行推理
image_path = 'test_image.jpg'
results = model(image_path)

# 打印预测结果并显示
results.print()  # 打印检测到的结果
results.show()   # 显示带有检测框的图像

5.3 微调 YOLOv5 模型

你可以使用 YOLOv5 提供的工具在自定义数据集上训练模型。首先，准备符合 YOLO 格式的数据集，然后使用以下命令进行训练。

python train.py --img 640 --batch 16 --epochs 100 --data custom_dataset.yaml --weights yolov5s.pt

custom_dataset.yaml 文件

需要指定训练和验证集路径，以及类别信息。

6. 总结

1. Faster R-CNN：基于区域提议方法，具有高精度，但速度相对较慢。适合需要高精度的场景。
2. SSD：单阶段检测器，速度较快，适合实时检测任务。
3. YOLOv5：速度极快，适合大规模实时检测应用。

PyTorch 和 torchvision 提供了丰富的目标检测模型，可以用于快速的推理或在自定义数据集上进行微调。对于实时需求较高的应用，YOLOv5 是不错的选择，而对于精度要求更高的场景，Faster R-CNN 是理想的选择。