glenn-jocher commited on
Commit
1935266
1 Parent(s): b098600

Scope imports for torch.hub.list() improvement (#3144)

Browse files
Files changed (1) hide show
  1. 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 create(name, pretrained=True, channels=3, classes=80, autoshape=True, verbose=True):
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 create(path, autoshape=autoshape, verbose=verbose)
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 create('yolov5s', pretrained, channels, classes, autoshape, verbose)
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 create('yolov5m', pretrained, channels, classes, autoshape, verbose)
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 create('yolov5l', pretrained, channels, classes, autoshape, verbose)
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 create('yolov5x', pretrained, channels, classes, autoshape, verbose)
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 create('yolov5s6', pretrained, channels, classes, autoshape, verbose)
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 create('yolov5m6', pretrained, channels, classes, autoshape, verbose)
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 create('yolov5l6', pretrained, channels, classes, autoshape, verbose)
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 create('yolov5x6', pretrained, channels, classes, autoshape, verbose)
106
 
107
 
108
  if __name__ == '__main__':
109
- model = create(name='yolov5s', pretrained=True, channels=3, classes=80, autoshape=True, verbose=True) # pretrained
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