PyTorch 目标检测教程

PyTorch 目标检测教程

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

1. 目标检测概述

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

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

2. 官方文档链接


3. Faster R-CNN 模型

3.1 加载 Faster R-CNN 模型并进行推理
python 复制代码
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。下面是如何定义自定义数据集并在上面进行训练。

python 复制代码
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
python 复制代码
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 模型进行推理
python 复制代码
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 仓库并安装依赖项。

bash 复制代码
git clone https://github.com/ultralytics/yolov5
cd yolov5
pip install -r requirements.txt
5.2 使用 YOLOv5 进行推理

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

python 复制代码
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 格式的数据集,然后使用以下命令进行训练。

bash 复制代码
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 是理想的选择。

相关推荐
ZOMI酱16 分钟前
【AI系统】模型转换基本介绍
人工智能
CodeDevMaster28 分钟前
LangChain之检索增强生成RAG
人工智能·python·llm
数学人学c语言1 小时前
yolov11剪枝
pytorch·python·深度学习
今天又是学习1 小时前
深度学习5
人工智能·深度学习
新加坡内哥谈技术2 小时前
RAG架构类型
大数据·人工智能·语言模型·chatgpt
Topstip2 小时前
iOS 19 重大更新泄露,将带来更“聪明”的 Siri 挑战 ChatGPT
人工智能·ios·ai·chatgpt
Nerinic2 小时前
深度学习基础1
人工智能·深度学习
数字扫地僧2 小时前
深度学习与知识图谱嵌入的结合:从理论到实践
人工智能·深度学习·知识图谱
真理Eternal2 小时前
手搓人工智能—聚类分析(下)谱系聚类与K-mean聚类
人工智能·机器学习
ZOMI酱2 小时前
【AI系统】昇腾 AI 架构介绍
人工智能·架构