glenn-jocher commited on
Commit
7f16406
·
unverified ·
1 Parent(s): d3dad42

Update hubconf.py (#1210)

Browse files
Files changed (1) hide show
  1. hubconf.py +15 -5
hubconf.py CHANGED
@@ -11,8 +11,11 @@ import os
11
  import torch
12
 
13
  from models.yolo import Model
 
14
  from utils.google_utils import attempt_download
15
 
 
 
16
 
17
  def create(name, pretrained, channels, classes):
18
  """Creates a specified YOLOv5 model
@@ -26,16 +29,19 @@ def create(name, pretrained, channels, classes):
26
  Returns:
27
  pytorch model
28
  """
29
- config = os.path.join(os.path.dirname(__file__), 'models', '%s.yaml' % name) # model.yaml path
30
  try:
31
  model = Model(config, channels, classes)
32
  if pretrained:
33
- ckpt = '%s.pt' % name # checkpoint filename
34
- attempt_download(ckpt) # download if not found locally
35
- state_dict = torch.load(ckpt, map_location=torch.device('cpu'))['model'].float().state_dict() # to FP32
 
36
  state_dict = {k: v for k, v in state_dict.items() if model.state_dict()[k].shape == v.shape} # filter
37
  model.load_state_dict(state_dict, strict=False) # load
38
- # model = model.autoshape() # cv2/PIL/np/torch inference: predictions = model(Image.open('image.jpg'))
 
 
39
  return model
40
 
41
  except Exception as e:
@@ -98,3 +104,7 @@ def yolov5x(pretrained=False, channels=3, classes=80):
98
  pytorch model
99
  """
100
  return create('yolov5x', pretrained, channels, classes)
 
 
 
 
 
11
  import torch
12
 
13
  from models.yolo import Model
14
+ from utils.general import set_logging
15
  from utils.google_utils import attempt_download
16
 
17
+ set_logging()
18
+
19
 
20
  def create(name, pretrained, channels, classes):
21
  """Creates a specified YOLOv5 model
 
29
  Returns:
30
  pytorch model
31
  """
32
+ config = os.path.join(os.path.dirname(__file__), 'models', f'{name}.yaml') # model.yaml path
33
  try:
34
  model = Model(config, channels, classes)
35
  if pretrained:
36
+ fname = f'{name}.pt' # checkpoint filename
37
+ attempt_download(fname) # download if not found locally
38
+ ckpt = torch.load(fname, map_location=torch.device('cpu')) # load
39
+ state_dict = ckpt['model'].float().state_dict() # to FP32
40
  state_dict = {k: v for k, v in state_dict.items() if model.state_dict()[k].shape == v.shape} # filter
41
  model.load_state_dict(state_dict, strict=False) # load
42
+ if len(ckpt['model'].names) == classes:
43
+ model.names = ckpt['model'].names # set class names attribute
44
+ # model = model.autoshape() # for autoshaping of PIL/cv2/np inputs and NMS
45
  return model
46
 
47
  except Exception as e:
 
104
  pytorch model
105
  """
106
  return create('yolov5x', pretrained, channels, classes)
107
+
108
+
109
+ if __name__ == '__main__':
110
+ model = create(name='yolov5s', pretrained=True, channels=3, classes=80) # example