使用 PyTorch Hub 加载 YOLOv5

最新推荐文章于 2024-09-18 20:48:51 发布

水榭花非梦

最新推荐文章于 2024-09-18 20:48:51 发布

阅读量3.2k

点赞数 33

文章标签： pytorch YOLO 人工智能

本文链接：https://blog.youkuaiyun.com/m0_68220374/article/details/138425500

版权

本文详细介绍了如何在Python中使用PyTorchHub加载预训练的YOLOv5模型进行对象检测，包括设置推理参数、设备管理、批处理推理、屏幕截图推理、多GPU并行和不同模型格式转换。还涉及训练和裁剪结果的处理。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

简单示例

此示例从 PyTorch Hub 中加载一个预训练的 YOLOv5s 模型作为模型，并传入一张图片进行推理。

import torch

model = torch.hub.load('ultralytics/yolov5', 'yolov5s')

image = 'https://ultralytics.com/images/zidane.jpg'

results = model(image)

results.pandas().xyxy[0]
#      xmin    ymin    xmax   ymax  confidence  class    name
# 0  749.50   43.50  1148.0  704.5    0.874023      0  person
# 1  433.50  433.50   517.5  714.5    0.687988     27     tie
# 2  114.75  195.75  1095.0  708.0    0.624512      0  person
# 3  986.00  304.00  1028.0  420.0    0.286865     27     tie

详细示例

该示例展示了使用 PIL 和 OpenCV 图像源的批处理推理。结果可以打印到控制台、保存到 runs/hub、在支持的环境中显示到屏幕，并以张量或 pandas 数据帧的形式返回。

import cv2
import torch
from PIL import Image

model = torch.hub.load('ultralytics/yolov5', 'yolov5s')

for f in 'zidane.jpg', 'bus.jpg':
    torch.hub.download_url_to_file('https://ultralytics.com/images/' + f, f)  # 下载两张图片
img1 = Image.open('zidane.jpg')  # PIL 图片
img2 = cv2.imread('bus.jpg')[..., ::-1]  # OpenCV 图片 (BGR to RGB)

results = model([img1, img2], size=640)  # 推理结果

results.print()
results.save()  # 将模型的推理结果保存为一个 CSV 文件或者一个 JSON 文件

results.xyxy[0]  # img1的推理结果，是一个一行六列的张量，分别为边界框的坐标（xmin,ymin,xmax,ymax）,置信度（confidence），类别（class_id）
results.pandas().xyxy[0]  # 将这些结果转换为一个 Pandas DataFrame，以表格形式存储数据。如下所示
#      xmin    ymin    xmax   ymax  confidence  class    name
# 0  749.50   43.50  1148.0  704.5    0.874023      0  person
# 1  433.50  433.50   517.5  714.5    0.687988     27     tie
# 2  114.75  195.75  1095.0  708.0    0.624512      0  person
# 3  986.00  304.00  1028.0  420.0    0.286865     27     tie

推理设置

YOLOv5 模型包含多种推理属性，如置信度阈值、IoU阈值等，可通过以下方式设置：

model.conf = 0.25  # 置信度阈值，只有当模型预测的置信度高于这个值时，预测结果才会被被用于后续处理。
iou = 0.45  # 非最大值抑制（NMS）过程中使用的交并比（IoU）阈值。NMS 用于合并重叠的检测框，移除多余的预测。当两个边界框的IoU大于45% 时，才会被认为是重叠的。
agnostic = False  # 设置为 True，所有类别的边界框都会一起进行 NMS 处理；设置为 False，每个类别的边界框将独立进行 NMS 处理。
multi_label = False  # 这个参数指定是否允许每个边界框有多于一个的类别标签。如果设置为 True，则允许一个边界框被分配给多个类别。
classes = None  # 可选参数，用于指定一个类列表，模型将只输出这些类别的检测结果。如果设置为 None，则不过滤类别，所有检测结果都会被输出。
max_det = 1000  # 每张图片最大检测数量的上限。
amp = False  # 这是自动混合精度（AMP）的一个标志，用于在模型推理时使用半精度浮点数（FP16）以加速计算并减少内存使用。设置为 False，则不使用 AMP

results = model(img, size=320)  # 使用设置好的参数对图片 img 进行了推理，size=320 指定模型输入图片的尺寸为 320x320 像素

设备

模型创建后可传输到任何设备：

model.cpu()  # 将模型的所有参数和缓冲区迁移到 CPU 上。通常用于在模型不需要大量计算或当 GPU 资源不足时进行推理。
model.cuda()  # 将模型迁移到第一个可用的 GPU 上。当有多个 GPU时，想使用特定的 GPU，需要指定 GPU 的索引。
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")  #使用 torch.device 可以更明确地指定设备，并且可以轻松地更改运行设备而不需要修改代码的其他部分。
model.to(device)  # 将模型迁移到任意指定的设备上。device 参数可以是 torch.device 对象，它允许你指定模型运行的计算设备，可以是 CPU 或者 GPU。
# 在迁移模型到不同的设备时，重要的是要注意输入数据也必须发送到相同的设备，否则会出现设备不匹配的错误

或者直接在创建模型时指定设备：

model = torch.hub.load('ultralytics/yolov5', 'yolov5s', device='cpu')

静默输出

model = torch.hub.load('ultralytics/yolov5', 'yolov5s', _verbose=False)  
# _verbose用于控制加载模型时的输出信息量，设置为 False 时，模型将“静默”加载，即不会打印出额外的日志信息。如果希望在加载过程中看到更多的信息，可以将此参数设置为 True。

输入通道

使用 4 个输入通道（而不是默认的 3 个）加载预训练的 YOLOv5s 模型：

model = torch.hub.load('ultralytics/yolov5', 'yolov5s', channels=4)

在这种情况下，模型将由预训练的权重组成，但第一个输入层除外，因为它不再与预训练的输入层形状相同。输入层仍将由随机权重初始化。

类别数

加载预训练的 YOLOv5s 模型，输出类数为 10 个，而不是默认的 80 个：

model = torch.hub.load('ultralytics/yolov5', 'yolov5s', classes=10)

在这种情况下，模型将由预训练的权重组成，但输出层除外，因为输出层的形状不再与预训练的输出层相同。输出层仍由随机权重初始化。

强制重新加载

如果在上述步骤中遇到问题，设置 force_reload=True 可能会有所帮助，它可以丢弃现有缓存，并强制从 PyTorch Hub 重新下载最新的 YOLOv5 版本。

model = torch.hub.load('ultralytics/yolov5', 'yolov5s', force_reload=True)

屏幕截图推理

在桌面屏幕上运行推理：

import torch
from PIL import ImageGrab

model = torch.hub.load('ultralytics/yolov5', 'yolov5s')

img = ImageGrab.grab()  # 用于截取当前屏幕的截图。这个函数会捕获整个屏幕的内容，并返回一个 PIL 图像对象。

results = model(img)

多 GPU 推理

YOLOv5 模型可以通过线程推理并行加载到多个 GPU 上：

import torch
import threading

def run(model, img):
    results = model(img)
    results.save()

model0 = torch.hub.load('ultralytics/yolov5', 'yolov5s', device=0)
model1 = torch.hub.load('ultralytics/yolov5', 'yolov5s', device=1)

threading.Thread(target=run, args=[model0, 'https://ultralytics.com/images/zidane.jpg'], daemon=True).start()
threading.Thread(target=run, args=[model1, 'https://ultralytics.com/images/bus.jpg'], daemon=True).start()
# 使用 threading.Thread() 创建两个线程，并将 run 函数和模型推理所需的参数传递给它们。
# daemon=True 参数使得线程成为守护线程，这意味着当主程序退出时，这些线程也会自动退出。
# 调用 .start() 方法启动线程，线程将异步执行

示例展示了如何使用 PyTorch 的多线程功能来并行执行模型推理。这里创建了两个线程，每个线程可以使用不同的 YOLOv5 模型对不同的图像进行对象检测。

训练

加载 YOLOv5 模型进行训练而非推理。

import torch

model = torch.hub.load('ultralytics/yolov5', 'yolov5s', autoshape=False)
# autoshape=False 参数指示加载模型时不自动调整输入的形状，这通常用于当您需要对输入数据的形状进行手动控制时。
model = torch.hub.load('ultralytics/yolov5', 'yolov5s', autoshape=False, pretrained=False)  
# pretrained=False 参数确保模型不会自动下载并加载预训练权重，这样就可以从头开始训练模型。

当加载未训练的模型时，需要自己提供数据，并使用适当的训练循环来训练模型。

注意：autoshape=False 参数在 YOLOv5 的某些版本中可能是不合法的，因为模型通常期望以特定的输入尺寸来接收数据。

base64 编码

results = model(img)  # 模型推理

results.ims  # 一个包含原始图像的 NumPy 数组，这些图像是传递给模型进行推理的。
results.render()  # results.render() 方法将在图像上绘制检测到的边界框和标签
for im in results.ims:
    buffered = BytesIO()  # 创建一个 BytesIO 对象，它是一个用于读写文本和二进制数据的文件类。
    im_base64 = Image.fromarray(img)  # 使用 Image.fromarray(im) 将 NumPy 数组转换为一个 PIL Image 对象。
    im_base64.save(buffered, format="JPEG")  #调用 save() 方法将图像保存到 buffered 对象中，格式为 JPEG。
    print(base64.b64encode(buffered.getvalue()).decode('utf-8'))  # 使用 base64.b64encode() 函数将图像数据编码为 base64 字符串。使用 decode('utf-8') 将 base64 编码的字节字符串解码为 UTF-8 字符串。
    # print 语句输出的是包含检测结果的图像的 base64 编码字符串，这些字符串可以用于在不支持图像文件直接传输的环境中展示图像，例如在网页上或者在一些文本协议中。

裁剪结果

可将结果返回并保存为检测结果：

results = model(img)
crops = results.crop(save=True)
# results.crop() 是一个方法，根据检测到的边界框坐标裁剪原始图像，获取只包含检测对象的图像区域。这个方法返回一个字典，包含裁剪后的图像。
# save=True 参数指示方法在裁剪图像后将其保存到文件系统中。默认情况下，裁剪的图像会保存在 runs/detect/exp*/ 目录下，可以通过配置模型的方式来改变这个行为。
# 返回的 crops 字典通常包含以下内容：
# 每个检测到的对象都会被裁剪成一张单独的图像。
# 字典的键通常是检测到的对象的索引或者类别。
# 字典的值是保存裁剪图像的路径或者裁剪图像的数组。

Pandas 结果

结果可以作为 Pandas DataFrames 返回：

results = model(img)  
results.pandas().xyxy[0]

结果排序

结果可按列排序：

results = model(im)  # 
results.pandas().xyxy[0].sort_values('xmin')  # 对检测结果按照边界框的 xmin 坐标进行从左到右的排序

JSON 结果

使用 .to_json() 方法转换为 .pandas() 数据帧后，结果可以 JSON 格式返回。JSON 格式可以使用 orient 参数修改。

results = model(ims)
results.pandas().xyxy[0].to_json(orient="records")  
# .to_json(orient="records") 将这个 DataFrame 转换为 JSON 格式的字符串。参数 orient="records" 指定了 JSON 的结构，其中每个检测结果将作为一个单独的记录（对象）。
"""
[
  {
    "xmin": 749.5,
    "ymin": 43.5,
    "xmax": 1148.0,
    "ymax": 704.5,
    "confidence": 0.8740234375,
    "class": 0,
    "name": "person"
  },
  {
    "xmin": 433.5,
    "ymin": 433.5,
    "xmax": 517.5,
    "ymax": 714.5,
    "confidence": 0.6879882812,
    "class": 27,
    "name": "tie"
  }
]
"""

自定义模型

本例使用 PyTorch Hub 加载了一个本地仓库的 YOLOv5n 模型 "best.pt"。

model = torch.hub.load('C:\yolov5-6.0', 'custom', path = "C:/yolov5-6.0/runs/train/exp3/weights/best.pt", source='local')

TensorRT、ONNX 和 OpenVINO 模型

PyTorch Hub 支持大多数 YOLOv5 导出格式的推理，包括自定义训练的模型。

import torch

model = torch.hub.load('ultralytics/yolov5', 'custom', path='yolov5s.pt')  # PyTorch
model = torch.hub.load('ultralytics/yolov5', 'custom', path='yolov5s.torchscript')  # TorchScript
model = torch.hub.load('ultralytics/yolov5', 'custom', path='yolov5s.onnx')  # ONNX
model = torch.hub.load('ultralytics/yolov5', 'custom', path='yolov5s_openvino_model/')  # OpenVINO
model = torch.hub.load('ultralytics/yolov5', 'custom', path='yolov5s.engine')  # TensorRT
model = torch.hub.load('ultralytics/yolov5', 'custom', path='yolov5s.mlmodel')  # CoreML (macOS-only)
model = torch.hub.load('ultralytics/yolov5', 'custom', path='yolov5s.tflite')  # TFLite
model = torch.hub.load('ultralytics/yolov5', 'custom', path='yolov5s_paddle_model/')  # PaddlePaddle

提示：在 GPU 基准上，TensorRT 可能比 PyTorch 快 2-5 倍；在 CPU 基准上，ONNX 和 OpenVINO 可能比 PyTorch 快 2-3 倍

参考：PyTorch Hub - Ultralytics YOLOv8 Docs