Fix `torch.hub.list('ultralytics/yolov5')` pathlib bug (#3921)
Browse files- hubconf.py +5 -5
hubconf.py
CHANGED
@@ -4,12 +4,9 @@ Usage:
|
|
4 |
import torch
|
5 |
model = torch.hub.load('ultralytics/yolov5', 'yolov5s')
|
6 |
"""
|
7 |
-
from pathlib import Path
|
8 |
|
9 |
import torch
|
10 |
|
11 |
-
FILE = Path(__file__).absolute()
|
12 |
-
|
13 |
|
14 |
def _create(name, pretrained=True, channels=3, classes=80, autoshape=True, verbose=True, device=None):
|
15 |
"""Creates a specified YOLOv5 model
|
@@ -26,15 +23,18 @@ def _create(name, pretrained=True, channels=3, classes=80, autoshape=True, verbo
|
|
26 |
Returns:
|
27 |
YOLOv5 pytorch model
|
28 |
"""
|
|
|
|
|
29 |
from models.yolo import Model, attempt_load
|
30 |
from utils.general import check_requirements, set_logging
|
31 |
from utils.google_utils import attempt_download
|
32 |
from utils.torch_utils import select_device
|
33 |
|
34 |
-
|
|
|
35 |
set_logging(verbose=verbose)
|
36 |
|
37 |
-
save_dir = Path('') if str(name).endswith('.pt') else
|
38 |
path = (save_dir / name).with_suffix('.pt') # checkpoint path
|
39 |
try:
|
40 |
device = select_device(('0' if torch.cuda.is_available() else 'cpu') if device is None else device)
|
|
|
4 |
import torch
|
5 |
model = torch.hub.load('ultralytics/yolov5', 'yolov5s')
|
6 |
"""
|
|
|
7 |
|
8 |
import torch
|
9 |
|
|
|
|
|
10 |
|
11 |
def _create(name, pretrained=True, channels=3, classes=80, autoshape=True, verbose=True, device=None):
|
12 |
"""Creates a specified YOLOv5 model
|
|
|
23 |
Returns:
|
24 |
YOLOv5 pytorch model
|
25 |
"""
|
26 |
+
from pathlib import Path
|
27 |
+
|
28 |
from models.yolo import Model, attempt_load
|
29 |
from utils.general import check_requirements, set_logging
|
30 |
from utils.google_utils import attempt_download
|
31 |
from utils.torch_utils import select_device
|
32 |
|
33 |
+
file = Path(__file__).absolute()
|
34 |
+
check_requirements(requirements=file.parent / 'requirements.txt', exclude=('tensorboard', 'thop', 'opencv-python'))
|
35 |
set_logging(verbose=verbose)
|
36 |
|
37 |
+
save_dir = Path('') if str(name).endswith('.pt') else file.parent
|
38 |
path = (save_dir / name).with_suffix('.pt') # checkpoint path
|
39 |
try:
|
40 |
device = select_device(('0' if torch.cuda.is_available() else 'cpu') if device is None else device)
|