This commit is contained in:
Glenn Jocher 2018-11-21 18:25:52 +01:00
parent 4e4b67b3c5
commit f1a94abafa
1 changed files with 3 additions and 0 deletions

View File

@ -36,6 +36,9 @@ def detect(opt):
if weights_path.endswith('.weights'): # saved in darknet format if weights_path.endswith('.weights'): # saved in darknet format
load_weights(model, weights_path) load_weights(model, weights_path)
else: # endswith('.pt'), saved in pytorch format else: # endswith('.pt'), saved in pytorch format
if weights_path == 'weights/yolov3.pt' and not os.path.isfile(weights_path):
os.system('wget https://storage.googleapis.com/ultralytics/yolov3.pt -P weights')
checkpoint = torch.load(weights_path, map_location='cpu') checkpoint = torch.load(weights_path, map_location='cpu')
model.load_state_dict(checkpoint['model']) model.load_state_dict(checkpoint['model'])
del checkpoint del checkpoint