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 是理想的选择。

相关推荐
杰说新技术4 分钟前
Meta AI最新推出的长视频语言理解多模态模型LongVU分享
人工智能·aigc
说私域7 分钟前
基于开源 AI 智能名片、S2B2C 商城小程序的用户获取成本优化分析
人工智能·小程序
东胜物联27 分钟前
探寻5G工业网关市场,5G工业网关品牌解析
人工智能·嵌入式硬件·5g
皓74138 分钟前
服饰电商行业知识管理的创新实践与知识中台的重要性
大数据·人工智能·科技·数据分析·零售
985小水博一枚呀1 小时前
【深度学习滑坡制图|论文解读3】基于融合CNN-Transformer网络和深度迁移学习的遥感影像滑坡制图方法
人工智能·深度学习·神经网络·cnn·transformer
AltmanChan1 小时前
大语言模型安全威胁
人工智能·安全·语言模型
985小水博一枚呀1 小时前
【深度学习滑坡制图|论文解读2】基于融合CNN-Transformer网络和深度迁移学习的遥感影像滑坡制图方法
人工智能·深度学习·神经网络·cnn·transformer·迁移学习
数据与后端架构提升之路1 小时前
从神经元到神经网络:深度学习的进化之旅
人工智能·神经网络·学习
爱技术的小伙子2 小时前
【ChatGPT】如何通过逐步提示提高ChatGPT的细节描写
人工智能·chatgpt
咔叽布吉2 小时前
【论文阅读笔记】CamoFormer: Masked Separable Attention for Camouflaged Object Detection
论文阅读·笔记·目标检测