PyTorch 目標(biāo)檢測教程

本教程將介紹如何在 PyTorch 中使用幾種常見的目標(biāo)檢測模型，包括 Faster R-CNN、SSD 以及 YOLO (You Only Look Once)。我們將涵蓋預(yù)訓(xùn)練模型的使用、推理、微調(diào)，以及自定義數(shù)據(jù)集上的訓(xùn)練。

1. 目標(biāo)檢測概述

目標(biāo)檢測任務(wù)不僅要識別圖像中的物體類別，還要精確定位物體的邊界框。在此任務(wù)中，每個(gè)模型輸出一個(gè)物體類別標(biāo)簽和一個(gè)邊界框。

常見目標(biāo)檢測模型：

• Faster R-CNN：基于區(qū)域提議方法，具有較高的檢測精度。
• SSD (Single Shot MultiBox Detector)：單階段檢測器，速度較快，適合實(shí)時(shí)應(yīng)用。
• YOLO (You Only Look Once)：單階段檢測器，具有極快的檢測速度，適合大規(guī)模實(shí)時(shí)檢測。

2. 官方文檔鏈接

• PyTorch 官方文檔
• Torchvision 模型
• YOLO 官方實(shí)現(xiàn) （YOLOv5 及 PyTorch 版 YOLO）

---

3. Faster R-CNN 模型

3.1 加載 Faster R-CNN 模型并進(jìn)行推理

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# 加載預(yù)訓(xùn)練的 Faster R-CNN 模型
model = fasterrcnn_resnet50_fpn(pretrained=True)
model.eval()  # 切換為評估模式# 加載和預(yù)處理圖像
image_path = "test_image.jpg"
image = Image.open(image_path)
image_tensor = F.to_tensor(image).unsqueeze(0)  # 轉(zhuǎn)換為張量并添加批次維度# 將模型和輸入圖像移動(dòng)到 GPU（如果可用）
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)
image_tensor = image_tensor.to(device)# 進(jìn)行推理
with torch.no_grad():predictions = model(image_tensor)# 獲取預(yù)測的邊界框和類別標(biāo)簽
boxes = predictions[0]['boxes'].cpu().numpy()  # 預(yù)測的邊界框
labels = predictions[0]['labels'].cpu().numpy()  # 預(yù)測的類別標(biāo)簽
scores = predictions[0]['scores'].cpu().numpy()  # 預(yù)測的分?jǐn)?shù)# 可視化預(yù)測結(jié)果
fig, ax = plt.subplots(1)
ax.imshow(image)# 設(shè)置閾值，只顯示高置信度的檢測結(jié)果
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 微調(diào) Faster R-CNN 模型

你可以使用自定義的數(shù)據(jù)集微調(diào) Faster R-CNN。下面是如何定義自定義數(shù)據(jù)集并在上面進(jìn)行訓(xùn)練。

import torch
from torch.utils.data import Dataset
from PIL import Image# 自定義數(shù)據(jù)集類
class CustomDataset(Dataset):def __init__(self, image_paths, annotations):self.image_paths = image_pathsself.annotations = annotationsdef __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"] = boxestarget["labels"] = labelsimage = F.to_tensor(image)return image, target

from torch.utils.data import DataLoader
from torchvision.models.detection import fasterrcnn_resnet50_fpn# 加載自定義數(shù)據(jù)集
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)))# 加載預(yù)訓(xùn)練的 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)# 定義優(yōu)化器并進(jìn)行訓(xùn)練
optimizer = torch.optim.SGD(model.parameters(), lr=0.005, momentum=0.9, weight_decay=0.0005)# 訓(xùn)練模型
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 也提供了預(yù)訓(xùn)練的 SSD 模型。

4.1 使用預(yù)訓(xùn)練 SSD 模型進(jìn)行推理

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# 加載預(yù)訓(xùn)練的 SSD 模型
model = ssd300_vgg16(pretrained=True)
model.eval()# 加載和預(yù)處理圖像
image_path = "test_image.jpg"
image = Image.open(image_path).convert("RGB")
image_tensor = F.to_tensor(image).unsqueeze(0)# 進(jìn)行推理
with torch.no_grad():predictions = model(image_tensor)# 獲取預(yù)測的邊界框和類別標(biāo)簽
boxes = predictions[0]['boxes'].cpu().numpy()
labels = predictions[0]['labels'].cpu().numpy()
scores = predictions[0]['scores'].cpu().numpy()# 可視化預(yù)測結(jié)果
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）是單階段目標(biāo)檢測器，具有非?？斓臋z測速度。YOLOv5 是目前廣泛使用的 YOLO 變體，官方提供了 PyTorch 版 YOLOv5。

5.1 安裝 YOLOv5

首先，克隆 YOLOv5 的 GitHub 倉庫并安裝依賴項(xiàng)。

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

5.2 使用 YOLOv5 進(jìn)行推理

YOLOv5 提供了一個(gè)簡單的 API，用于進(jìn)行推理和訓(xùn)練。

import torch# 加載預(yù)訓(xùn)練的 YOLOv5 模型
model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)# 加載圖像并進(jìn)行推理
image_path = 'test_image.jpg'
results = model(image_path)# 打印預(yù)測結(jié)果并顯示
results.print()  # 打印檢測到的結(jié)果
results.show()   # 顯示帶有檢測框的圖像

5.3 微調(diào) YOLOv5 模型

你可以使用 YOLOv5 提供的工具在自定義數(shù)據(jù)集上訓(xùn)練模型。首先，準(zhǔn)備符合 YOLO 格式的數(shù)據(jù)集，然后使用以下命令進(jìn)行訓(xùn)練。

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

custom_dataset.yaml 文件

需要指定訓(xùn)練和驗(yàn)證集路徑，以及類別信息。

---

6. 總結(jié)

1. Faster R-CNN：基于區(qū)域提議方法，具有高精度，但速度相對較慢。適合需要高精度的場景。
2. SSD：單階段檢測器，速度較快，適合實(shí)時(shí)檢測任務(wù)。
3. YOLOv5：速度極快，適合大規(guī)模實(shí)時(shí)檢測應(yīng)用。

PyTorch 和 torchvision 提供了豐富的目標(biāo)檢測模型，可以用于快速的推理或在自定義數(shù)據(jù)集上進(jìn)行微調(diào)。對于實(shí)時(shí)需求較高的應(yīng)用，YOLOv5 是不錯(cuò)的選擇，而對于精度要求更高的場景，Faster R-CNN 是理想的選擇。