当前位置: 首页 > 工具软件 > PyTorch Hub > 使用案例 >

【从PyTorch Hub加载YOLOv5运行推理】

孔鸿哲
2023-12-01

从PyTorch Hub加载YOLOv5运行推理
教程:https://github.com/ultralytics/yolov5/issues/36

import torch

# Model
# model = torch.hub.load('ultralytics/yolov5', 'yolov5s')  # or yolov5n - yolov5x6, custom 加载官网训练好的模型
model = torch.hub.load('ultralytics/yolov5', 'custom', path='runs/train/yolov5s-face_mask/weights/best.pt')  # 加载自己训练的模型

# Images
img = 'data/images/zidane.jpg'  # or file, Path, PIL, OpenCV, numpy, list, 
# imgs = ['data/images/zidane.jpg',  # filename
#         Path('data/images/zidane.jpg'),  # Path
#         # 'https://ultralytics.com/images/zidane.jpg',  # URI
#         cv2.imread('data/images/bus.jpg')[:, :, ::-1],  # OpenCV
#         Image.open('data/images/bus.jpg'),  # PIL
#         np.zeros((320, 640, 3))]  # numpy

# Inference
results = model(img)

# Results
results.print()  # or .show(), .save(), .crop(), .pandas(), etc.
# results.pandas().xyxy[0]  # Pandas DataFrame
results.show()
results.xyxy[0]  # im1 predictions (tensor)
# results.pandas().xyxy[0]  # im1 predictions (pandas)
results.save('./save_img/')
results.crop('./save_img/')
 类似资料: