hub

1
import torch
2

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

1
import torch
2

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

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
                       )

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")

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

1
import torch
2

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