调用yolov5训好的本地模型pt

原创已于 2022-09-16 12:36:56 修改 · 7.1k 阅读

25 ·

CC 4.0 BY-SA版权

文章标签：

#pytorch #深度学习 #yolov5

于 2022-09-16 12:36:13 首次发布

检测识别专栏收录该内容

6 篇文章

订阅专栏

这篇博客介绍了如何利用Yolov5模型进行PDF文件中的对象检测。首先创建一个名为pdf2picture的文件夹，然后将预训练模型yolov5s-flash.pt复制到该文件夹。在yolov5项目中的hubconf.py定义入口点为'yolov5s_flash'。接着，在新文件pdf_yolov5_to_flash.py中加载模型并进行设备分配。确保模型、yaml配置文件和hubconf.py的入口点名称一致。

部署运行你感兴趣的模型镜像

1、新建一个文件夹pdf2picture

2、将yolov5项目中的模型yolov5s-flash.pt复制到该目录下

3、必须在yolov5项目中的hubconf.py中定义可调用(入口点)的名称：

4、新建的文件夹下建议一个pdf_yolov5_to_flash.py：

import torch

def yolov5_model():
    # GPU
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    # Model
    model = torch.hub.load('../yolov5',   
                          'yolov5s_flash',
                           pretrained=True,
                           source='local')  # or yolov5m, yolov5l, yolov5x, custom
    model = model.to(device)
    return model

model = yolov5_model()

注意：

yolov5s-flash.pt也必须在yolov5项目的根目录下

model目录下的yaml名称要与pt文件名、hubconf.py入口点名称一致，这里统一是“yolov5s_flash”