快速调用Yolov5模型检检测图片

前提:未修改模型结构

1.快速调用官方的Yolov5预模型

import torch

# 使用torch.hub加载yolov5的预训练模型训练
model = torch.hub.load('ultralytics/yolov5', 'yolov5s')  # or yolov5m, yolov5x, custom

# 进行模型调用测试
img_path = './6800.jpg'  # or file, PIL, OpenCV, numpy, multiple
results = model(img_path) # 得到预测结果

print(results.xyxy) # 输出预测出的bbox_list
results.show() # 预测结果展示

2.快速调用自己训练好的的Yolov5预模型(有pt文件即可)

import torch

# 使用torch.hub加载yolov5的预训练模型训练
model = torch.hub.load('ultralytics/yolov5', 'yolov5s')  # or yolov5m, yolov5x, custom

# 加载自己训练好的模型及相关参数
cpkt = torch.load("./best.pt",map_location=torch.device("cuda:0"))

# 将预训练的模型的骨干替换成自己训练好的
yolov5_load = model
yolov5_load.model = cpkt["model"]

# 进行模型调用测试
img_path = './6800.jpg'  # or file, PIL, OpenCV, numpy, multiple
results = yolov5_load(img_path) # 得到预测结果

print(results.xyxy) # 输出预测出的bbox_list
results.show() # 预测结果展示

参考资料

1. https://github.com/ultralytics/yolov5