### 使用SSD、FPN、R-CNN或YOLO预训练模型进行目标检测的完整Python代码示例 以下提供了基于四种不同预训练目标检测网络（SSD、FPN、R-CNN和YOLO）的完整Python代码示例，用于对图像进行目标探测。 --- #### 1. SSD（Single Shot MultiBox Detector） SSD 是一种单阶段目标检测算法，能够同时完成边界框预测和类别分类[^3]。 ```python import torch from torchvision import models, transforms from PIL import Image import matplotlib.pyplot as plt import numpy as np # 加载预训练的 SSD 模型 model = models.detection.ssd300_vgg16(pretrained=True) model.eval() # 图像预处理 transform = transforms.Compose([ transforms.Resize((300, 300)), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) # 加载并转换图像 image_path = "example.jpg" # 替换为你的图像路径 image = Image.open(image_path).convert("RGB") image_tensor = transform(image).unsqueeze(0) # 模型推理 with torch.no_grad(): predictions = model(image_tensor) # 可视化结果 def plot_image_with_boxes(image, boxes): plt.figure() img = np.array(image) plt.imshow(img) ax = plt.gca() for box in boxes: x_min, y_min, x_max, y_max = box rect = plt.Rectangle((x_min, y_min), x_max - x_min, y_max - y_min, fill=False, edgecolor='red', linewidth=2) ax.add_patch(rect) plt.show() boxes = predictions[0]['boxes'].cpu().numpy() plot_image_with_boxes(image, boxes) ``` --- #### 2. FPN（Feature Pyramid Network）结合 Faster R-CNN FPN 是一种用于增强特征提取能力的架构，通常与 Faster R-CNN 结合使用[^2]。 ```python import torch from torchvision import models, transforms from PIL import Image import matplotlib.pyplot as plt import numpy as np # 加载预训练的 Faster R-CNN 模型（带 FPN） model = models.detection.fasterrcnn_resnet50_fpn(pretrained=True) model.eval() # 图像预处理 transform = transforms.Compose([ transforms.ToTensor() ]) # 加载并转换图像 image_path = "example.jpg" # 替换为你的图像路径 image = Image.open(image_path).convert("RGB") image_tensor = transform(image).unsqueeze(0) # 模型推理 with torch.no_grad(): predictions = model(image_tensor) # 可视化结果 def plot_image_with_boxes(image, boxes): plt.figure() img = np.array(image) plt.imshow(img) ax = plt.gca() for box in boxes: x_min, y_min, x_max, y_max = box rect = plt.Rectangle((x_min, y_min), x_max - x_min, y_max - y_min, fill=False, edgecolor='red', linewidth=2) ax.add_patch(rect) plt.show() boxes = predictions[0]['boxes'].cpu().numpy() plot_image_with_boxes(image, boxes) ``` --- #### 3. R-CNN 系列（Faster R-CNN 示例） R-CNN 系列包括多种变体，如 Faster R-CNN 和 Mask R-CNN。以下是一个使用 Faster R-CNN 的示例[^1]。 ```python import torch from torchvision import models, transforms from PIL import Image import matplotlib.pyplot as plt import numpy as np # 加载预训练的 Faster R-CNN 模型 model = models.detection.fasterrcnn_resnet50_fpn(pretrained=True) model.eval() # 图像预处理 transform = transforms.Compose([ transforms.ToTensor() ]) # 加载并转换图像 image_path = "example.jpg" # 替换为你的图像路径 image = Image.open(image_path).convert("RGB") image_tensor = transform(image).unsqueeze(0) # 模型推理 with torch.no_grad(): predictions = model(image_tensor) # 可视化结果 def plot_image_with_boxes(image, boxes): plt.figure() img = np.array(image) plt.imshow(img) ax = plt.gca() for box in boxes: x_min, y_min, x_max, y_max = box rect = plt.Rectangle((x_min, y_min), x_max - x_min, y_max - y_min, fill=False, edgecolor='red', linewidth=2) ax.add_patch(rect) plt.show() boxes = predictions[0]['boxes'].cpu().numpy() plot_image_with_boxes(image, boxes) ``` --- #### 4. YOLO（You Only Look Once） YOLO 是一种快速且高效的单阶段目标检测算法[^3]。 ```python import torch from yolov5 import detect # 加载预训练的 YOLOv5 模型 model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True) # 加载并推理图像 image_path = "example.jpg" # 替换为你的图像路径 results = model(image_path) # 显示结果 results.show() ``` --- ### 注意事项 - **模型选择**：根据任务需求选择合适的模型。例如，如果需要高精度，可以选择 Mask R-CNN；如果需要实时性，可以选择 YOLO 或 SSD。 - **数据预处理**：确保输入图像符合模型要求的尺寸和格式。 - **后处理**：对模型输出进行非极大值抑制（NMS）等操作，以过滤冗余检测结果。 ---