hub

2023-04-23

list

查看可用模型

1
import torch
2
3
torch.hub.list("ultralytics/yolov5")
4
# 返回: ['custom', 'yolov5l', 'yolov5l6', 'yolov5m', 'yolov5m6', 'yolov5n', 'yolov5n6', 'yolov5s', 'yolov5s6', 'yolov5x', 'yolov5x6']

help

查看模型帮助

1
import torch
2
3
torch.hub.help("pytorch/vision", "resnet18")

load

加载模型

1
import torch
2
3
# 加载
4
# github 仓库和模型
5
model = torch.hub.load("ultralytics/yolov5", "yolov5s")
6
7
# 本地模型
8
model = torch.hub.load(
9
"ultralytics/yolov5",
10
"custom",
11
path="path/best.pt"
12
)
13
14
# 本地版本和模型
15
model = torch.hub.load(
16
"path/yolov5",
17
"custom",
18
path="path/best.pt",
19
source="local"
20
)

参数

  • device: GPU
  • _verbose: 静默加载

设置

1
# 设置
2
model.conf = 0.25 # 置信度
3
model.classes = [0] # 分类

使用

1
res = model("file.jpg")
2
3
res.print()
4
res.save()
5
res.show() # 预览图片
6
res.names: 所有 classess
7
res.xyxy[0]
8
res.files: 文件名 # 列表形式
9
10
# pandas 输出
11
res.pandas().xyxy[0]
12
res.pandas().xyxy[0].sort_values("xmin") # 排序从左到右
13
# json 输出
14
res.pandas().xyxy[0].to_json(orient="records")

pandas 包含:

  • names: 所有 classess #字典
  • files: 文件名 #列表
  • xyxy
  • xyxyn: 包含归一化
  • xywhn: 归一化
  • xywhn
  • n
  • t
  • s

加载到设备

1
model.cpu() # CPU
2
model.cuda() # GPU
3
model.to(device) # i.e. device=torch.device(0)

download_url_to_file

1
import torch
2
3
# 下载文件到本地
4
torch.hub.download_url_to_file('http://url/img.jpg', '/save_path/file_name')

参考