一、准备
PyTorch安装请点这里
二、简单示例
这里使用轻量级yolov5s模型。
import torch
model = torch.hub.load('ultralytics/yolov5', 'yolov5s')
img = 'https://ultralytics.com/images/zidane.jpg'
results = model(img)
三、详细示例
代码中同时使用PIL和OpenCV，识别结果保存在runs/hub目录下。
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')
img2 = cv2.imread('bus.jpg')[:, :, ::-1]
imgs = [img1, img2]
results = model(imgs, size=640)
results.print()
results.save()
results.xyxy[0]
results.pandas().xyxy[0]
四、参数设置
这里参数主要是指置信度阈值，NMS loU阈值，分类筛选器等模型属性参数。
model.conf = 0.25
model.iou = 0.45
model.classes = None
results = model(imgs, size=320)
五、输入通道设置
加载YOLOv5s模型输入通道数默认值为3，可以通过以下方式修改。
model = torch.hub.load('ultralytics/yolov5', 'yolov5s', channels=4)
六、分类设置
YOLOv5模型默认分类数为80，可以通过以下方式修改。
model = torch.hub.load('ultralytics/yolov5', 'yolov5s', classes=10)
七、强制重新加载
使用force_reload=True可以帮助清理缓存并且强制更新下载最新YOLOv5版本。
model = torch.hub.load('ultralytics/yolov5', 'yolov5s', force_reload=True)
八、训练
加载YOLOv5模型如果是为了训练，可以设置autoshape=False。
加载模型并随机初始化权值可以设置pretrained=False。
model = torch.hub.load('ultralytics/yolov5', 'yolov5s', autoshape=False)
model = torch.hub.load('ultralytics/yolov5', 'yolov5s', autoshape=False, pretrained=False)
九、Base64结果
示例如下：
results = model(imgs)
results.imgs
results.render()
for img in results.imgs:
buffered = BytesIO()
img_base64 = Image.fromarray(img)
img_base64.save(buffered, format="JPEG")
print(base64.b64encode(buffered.getvalue()).decode('utf-8'))
十、JSON结果
示例如下：
results = model(imgs)
results.pandas().xyxy[0].to_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”}, {“xmin”:115.25,”ymin”:195.75,”xmax”:1096.0,”ymax”:708.0,”confidence”:0.6254882812,”class”:0,”name”:”person”}, {“xmin”:986.0,”ymin”:304.0,”xmax”:1028.0,”ymax”:420.0,”confidence”:0.2873535156,”class”:27,”name”:”tie”} ]
Original: https://blog.csdn.net/u011922698/article/details/123268070
Author: icodekang
Title: 通过PyTorch Hub加载YOLOv5
原创文章受到原创版权保护。转载请注明出处：https://www.johngo689.com/709875/
转载文章受原作者版权保护。转载请注明原作者出处！