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 模型并进行推理
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. 总结
- Faster R-CNN:基于区域提议方法,具有高精度,但速度相对较慢。适合需要高精度的场景。
- SSD:单阶段检测器,速度较快,适合实时检测任务。
- YOLOv5:速度极快,适合大规模实时检测应用。
PyTorch 和 torchvision
提供了丰富的目标检测模型,可以用于快速的推理或在自定义数据集上进行微调。对于实时需求较高的应用,YOLOv5 是不错的选择,而对于精度要求更高的场景,Faster R-CNN 是理想的选择。