YOLOv5-手写数字识别：四、在本地部署训练模型

本文最后更新于452 天前，其中的信息可能已经过时，如有错误请联系站长。

上一篇讲了如何在OpenBayes训练YOLOv5模型，而我们往往希望能在实际应用中使用训练好的模型。本篇我们讲述将训练好的YOLOv5模型部署至本地的方法。

一、在OpenBayes下载训练好的模型

上一篇我们得知，训练好的模型放到了YOLOv5本地的/runs/train/exp/weights路径中。发现有best.pt和last.pt两个模型。两个模型的区别我们已经讲过，直接下载这两个模型即可。

如果不太能独立训练模型的话，可以直接使用博主训练的模型。点击此处下载。

二、在本地配置YOLOv5环境

我们的目标是离线环境下运行模型，而YOLOv5的离线模型依赖YOLOv5项目环境。在本系列的第零篇文章时，我们在本地创建了两个环境，一个是用于标注图像的Labelimg，一个就是预留给YOLOv5的环境。我们前往YOLOv5的Github项目地址，将YOLOv5项目下载至本地并解压。

YOLOv5下载地址

我们解压YOLOv5项目，然后激活我们之前创建的yolov5环境，转移到YOLOv5文件夹下，安装其所需依赖库。防止遗忘我再把指令放在下面（这里假设你已经到达了YOLOv5的文件夹下，不要再问为什么找不到requirements.txt诸如此类的问题了）：

conda activate YOLOv5
pip install -r requirements.txt

等待一段时间安装依赖完成，我们还需要一些额外的依赖库。YOLOv5的模型支持使用PyTorch来部署，所以我们也用这种方法来部署模型。首先我们安装PyTorch的库：

pip install pytorch

文件比较大，所以要等一会。安装完成后，YOLOv5的本地运行环境至此搭建完成。

三、在本地使用模型进行推理

我们在本地的Yolov5文件夹下创建一个main.py来运行模型。首先确保我们在OpenBayes训练好的模型转移到你本地的Yolov5文件夹下，这一步不做演示。

接下来我们在main.py中写入以下内容：

import cv2
import torch

# 以下部分十分关键！！！
# 由于我们是在OpenBayes的工作环境中训练的模型
# 而所用的生产环境是Linux系统
# 所以我们要重定向PosixPath到WindowsPath
import pathlib
temp = pathlib.PosixPath
pathlib.PosixPath = pathlib.WindowsPath

# 加载模型
model = torch.hub.load('.', 'custom', path='./best.pt',source='local')  
# 第一个参数'.'意味着你的本地yolov5路径。'.'代表其就在main.py所在目录下。
# 第二个参数'custom'代表使用自制模型。
# 第三个参数 path 代表你的模型的位置。这里认为best.pt就在同main.py所在的文件夹里。
# 第四个参数 source 代表你使用本地的yolov5库。

变量model便是你训练好的模型，我们可以用model进行很多工作。下面是其中几个的示例，直接在main.py中接着写就行：

3.1 使用本地图片进行推理

# 使用 OpenCV 加载你要识别的图像
img_path = './xxx.jpg'  # 这里填你的图像路径
img = cv2.imread(img_path)  # 使用 OpenCV 读取图像

# 将 BGR 图像转换为 RGB
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

# 执行识别
results = model(img)
results.print()  # 打印结果
results.show()  # 在默认的图片查看器中显示结果图像

# 保存图像（如果需要的话）
for img in results.render():
    cv2.imwrite('./result.jpg', cv2.cvtColor(img, cv2.COLOR_RGB2BGR))

3.2 使用本地视频进行推理

# 打开视频文件
video_path = './xxx.mp4'  # 使用你的视频路径替换这里
cap = cv2.VideoCapture(video_path)

# 获取视频的基本参数
fps = int(cap.get(cv2.CAP_PROP_FPS))  # 获取视频的帧率
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))  # 获取视频宽度
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))  # 获取视频高度

# 定义视频编码器和创建 VideoWriter 对象
fourcc = cv2.VideoWriter_fourcc(*'XVID')  # 定义编码器
out = cv2.VideoWriter('./results/output.avi', fourcc, fps, (width, height))  # 输出文件名和参数

while cap.isOpened():
    ret, frame = cap.read()  # 读取一帧图像
    if not ret:
        break

    # 将 BGR 图像转换为 RGB，因为模型需要 RGB 图像
    frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

    # 执行推理
    results = model(frame_rgb)
    # 获取推理后的图像帧
    result_frame = results.render()[0]
    # 将 RGB 图像转换回 BGR，以便正确保存
    result_frame_bgr = cv2.cvtColor(result_frame, cv2.COLOR_RGB2BGR)
    
    # 写入帧到视频文件
    out.write(result_frame_bgr)

    # 可选：显示结果图像
    cv2.imshow('Inference', result_frame_bgr)

    if cv2.waitKey(1) & 0xFF == ord('q'):  # 按 q 退出
        break

# 释放资源
cap.release()
out.release()  # 关闭视频写入对象
cv2.destroyAllWindows()

3.3 使用本地摄像头进行推理

# 使用 OpenCV 捕获视频流
cap = cv2.VideoCapture(0)  # 0 通常是默认摄像头的ID

while cap.isOpened():
    ret, frame = cap.read()  # 读取一帧图像
    if not ret:
        break

    # 将 BGR 图像转换为 RGB，因为模型需要 RGB 图像
    frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

    # 执行推理
    results = model(frame_rgb)
    
    # 可选：显示结果图像
    for img in results.render():
        cv2.imshow('Inference', cv2.cvtColor(img, cv2.COLOR_RGB2BGR))

    if cv2.waitKey(1) & 0xFF == ord('q'):  # 按 q 退出
        break

cap.release()
cv2.destroyAllWindows()

四、结语

至此，YOLOv5-手写数字识别系列完结。虽然说是手写数字识别，但几乎没有多少涉及手写数字的内容……总之感谢各位的支持，本次系列的完成也让博主认识到自己的语言和知识面的各种薄弱之处。

希望这个教程对各位读者有所帮助，也请各位多多灌注喵立知，我们下个项目再见！