Update hubconf.py (#1210)
Browse files- 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', '
|
30 |
try:
|
31 |
model = Model(config, channels, classes)
|
32 |
if pretrained:
|
33 |
-
|
34 |
-
attempt_download(
|
35 |
-
|
|
|
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 |
-
|
|
|
|
|
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
|