glenn-jocher
commited on
Commit
•
1935266
1
Parent(s):
b098600
Scope imports for torch.hub.list() improvement (#3144)
Browse files- hubconf.py +15 -14
hubconf.py
CHANGED
@@ -9,16 +9,13 @@ from pathlib import Path
|
|
9 |
|
10 |
import torch
|
11 |
|
12 |
-
from models.yolo import Model, attempt_load
|
13 |
from utils.general import check_requirements, set_logging
|
14 |
-
from utils.google_utils import attempt_download
|
15 |
-
from utils.torch_utils import select_device
|
16 |
|
17 |
dependencies = ['torch', 'yaml']
|
18 |
check_requirements(Path(__file__).parent / 'requirements.txt', exclude=('tensorboard', 'pycocotools', 'thop'))
|
19 |
|
20 |
|
21 |
-
def
|
22 |
"""Creates a specified YOLOv5 model
|
23 |
|
24 |
Arguments:
|
@@ -32,6 +29,10 @@ def create(name, pretrained=True, channels=3, classes=80, autoshape=True, verbos
|
|
32 |
Returns:
|
33 |
YOLOv5 pytorch model
|
34 |
"""
|
|
|
|
|
|
|
|
|
35 |
set_logging(verbose=verbose)
|
36 |
fname = Path(name).with_suffix('.pt') # checkpoint filename
|
37 |
try:
|
@@ -62,51 +63,51 @@ def create(name, pretrained=True, channels=3, classes=80, autoshape=True, verbos
|
|
62 |
|
63 |
def custom(path='path/to/model.pt', autoshape=True, verbose=True):
|
64 |
# YOLOv5 custom or local model
|
65 |
-
return
|
66 |
|
67 |
|
68 |
def yolov5s(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True):
|
69 |
# YOLOv5-small model https://github.com/ultralytics/yolov5
|
70 |
-
return
|
71 |
|
72 |
|
73 |
def yolov5m(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True):
|
74 |
# YOLOv5-medium model https://github.com/ultralytics/yolov5
|
75 |
-
return
|
76 |
|
77 |
|
78 |
def yolov5l(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True):
|
79 |
# YOLOv5-large model https://github.com/ultralytics/yolov5
|
80 |
-
return
|
81 |
|
82 |
|
83 |
def yolov5x(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True):
|
84 |
# YOLOv5-xlarge model https://github.com/ultralytics/yolov5
|
85 |
-
return
|
86 |
|
87 |
|
88 |
def yolov5s6(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True):
|
89 |
# YOLOv5-small-P6 model https://github.com/ultralytics/yolov5
|
90 |
-
return
|
91 |
|
92 |
|
93 |
def yolov5m6(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True):
|
94 |
# YOLOv5-medium-P6 model https://github.com/ultralytics/yolov5
|
95 |
-
return
|
96 |
|
97 |
|
98 |
def yolov5l6(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True):
|
99 |
# YOLOv5-large-P6 model https://github.com/ultralytics/yolov5
|
100 |
-
return
|
101 |
|
102 |
|
103 |
def yolov5x6(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True):
|
104 |
# YOLOv5-xlarge-P6 model https://github.com/ultralytics/yolov5
|
105 |
-
return
|
106 |
|
107 |
|
108 |
if __name__ == '__main__':
|
109 |
-
model =
|
110 |
# model = custom(path='path/to/model.pt') # custom
|
111 |
|
112 |
# Verify inference
|
|
|
9 |
|
10 |
import torch
|
11 |
|
|
|
12 |
from utils.general import check_requirements, set_logging
|
|
|
|
|
13 |
|
14 |
dependencies = ['torch', 'yaml']
|
15 |
check_requirements(Path(__file__).parent / 'requirements.txt', exclude=('tensorboard', 'pycocotools', 'thop'))
|
16 |
|
17 |
|
18 |
+
def _create(name, pretrained=True, channels=3, classes=80, autoshape=True, verbose=True):
|
19 |
"""Creates a specified YOLOv5 model
|
20 |
|
21 |
Arguments:
|
|
|
29 |
Returns:
|
30 |
YOLOv5 pytorch model
|
31 |
"""
|
32 |
+
from models.yolo import Model, attempt_load
|
33 |
+
from utils.google_utils import attempt_download
|
34 |
+
from utils.torch_utils import select_device
|
35 |
+
|
36 |
set_logging(verbose=verbose)
|
37 |
fname = Path(name).with_suffix('.pt') # checkpoint filename
|
38 |
try:
|
|
|
63 |
|
64 |
def custom(path='path/to/model.pt', autoshape=True, verbose=True):
|
65 |
# YOLOv5 custom or local model
|
66 |
+
return _create(path, autoshape=autoshape, verbose=verbose)
|
67 |
|
68 |
|
69 |
def yolov5s(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True):
|
70 |
# YOLOv5-small model https://github.com/ultralytics/yolov5
|
71 |
+
return _create('yolov5s', pretrained, channels, classes, autoshape, verbose)
|
72 |
|
73 |
|
74 |
def yolov5m(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True):
|
75 |
# YOLOv5-medium model https://github.com/ultralytics/yolov5
|
76 |
+
return _create('yolov5m', pretrained, channels, classes, autoshape, verbose)
|
77 |
|
78 |
|
79 |
def yolov5l(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True):
|
80 |
# YOLOv5-large model https://github.com/ultralytics/yolov5
|
81 |
+
return _create('yolov5l', pretrained, channels, classes, autoshape, verbose)
|
82 |
|
83 |
|
84 |
def yolov5x(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True):
|
85 |
# YOLOv5-xlarge model https://github.com/ultralytics/yolov5
|
86 |
+
return _create('yolov5x', pretrained, channels, classes, autoshape, verbose)
|
87 |
|
88 |
|
89 |
def yolov5s6(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True):
|
90 |
# YOLOv5-small-P6 model https://github.com/ultralytics/yolov5
|
91 |
+
return _create('yolov5s6', pretrained, channels, classes, autoshape, verbose)
|
92 |
|
93 |
|
94 |
def yolov5m6(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True):
|
95 |
# YOLOv5-medium-P6 model https://github.com/ultralytics/yolov5
|
96 |
+
return _create('yolov5m6', pretrained, channels, classes, autoshape, verbose)
|
97 |
|
98 |
|
99 |
def yolov5l6(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True):
|
100 |
# YOLOv5-large-P6 model https://github.com/ultralytics/yolov5
|
101 |
+
return _create('yolov5l6', pretrained, channels, classes, autoshape, verbose)
|
102 |
|
103 |
|
104 |
def yolov5x6(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True):
|
105 |
# YOLOv5-xlarge-P6 model https://github.com/ultralytics/yolov5
|
106 |
+
return _create('yolov5x6', pretrained, channels, classes, autoshape, verbose)
|
107 |
|
108 |
|
109 |
if __name__ == '__main__':
|
110 |
+
model = _create(name='yolov5s', pretrained=True, channels=3, classes=80, autoshape=True, verbose=True) # pretrained
|
111 |
# model = custom(path='path/to/model.pt') # custom
|
112 |
|
113 |
# Verify inference
|