zjc
/
yolov5
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

Update hubconf.py (#1210)

This commit is contained in:
Glenn Jocher 2020-10-26 01:06:54 +01:00 committed by GitHub
parent d3dad42256
commit 7f1640695b
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 15 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

20
hubconf.py
View File

@ -11,8 +11,11 @@ import os
import torch
from models.yolo import Model
from utils.general import set_logging
from utils.google_utils import attempt_download
set_logging()
def create(name, pretrained, channels, classes):
"""Creates a specified YOLOv5 model
@ -26,16 +29,19 @@ def create(name, pretrained, channels, classes):
Returns:
pytorch model
"""
config = os.path.join(os.path.dirname(__file__), 'models', '%s.yaml' % name) # model.yaml path
config = os.path.join(os.path.dirname(__file__), 'models', f'{name}.yaml') # model.yaml path
try:
model = Model(config, channels, classes)
if pretrained:
ckpt = '%s.pt' % name # checkpoint filename
attempt_download(ckpt) # download if not found locally
state_dict = torch.load(ckpt, map_location=torch.device('cpu'))['model'].float().state_dict() # to FP32
fname = f'{name}.pt' # checkpoint filename
attempt_download(fname) # download if not found locally
ckpt = torch.load(fname, map_location=torch.device('cpu')) # load
state_dict = ckpt['model'].float().state_dict() # to FP32
state_dict = {k: v for k, v in state_dict.items() if model.state_dict()[k].shape == v.shape} # filter
model.load_state_dict(state_dict, strict=False) # load
# model = model.autoshape() # cv2/PIL/np/torch inference: predictions = model(Image.open('image.jpg'))
if len(ckpt['model'].names) == classes:
model.names = ckpt['model'].names # set class names attribute
# model = model.autoshape() # for autoshaping of PIL/cv2/np inputs and NMS
return model
except Exception as e:
@ -98,3 +104,7 @@ def yolov5x(pretrained=False, channels=3, classes=80):
pytorch model
"""
return create('yolov5x', pretrained, channels, classes)
if __name__ == '__main__':
model = create(name='yolov5s', pretrained=True, channels=3, classes=80) # example