This commit is contained in:
Glenn Jocher 2018-12-04 19:17:03 +01:00
parent fd6619d773
commit be8603b2dd
3 changed files with 9 additions and 6 deletions

View File

@ -35,11 +35,11 @@ def main(opt):
model = Darknet(opt.cfg, opt.img_size) model = Darknet(opt.cfg, opt.img_size)
weights_path = f_path + 'weights/yolov3.pt' weights_path = f_path + 'weights/yolov3.pt'
if weights_path.endswith('.weights'): # saved in darknet format if weights_path.endswith('.pt'): # pytorch format
load_weights(model, weights_path)
else: # endswith('.pt'), saved in pytorch format
if weights_path.endswith('weights/yolov3.pt') and not os.path.isfile(weights_path): if weights_path.endswith('weights/yolov3.pt') and not os.path.isfile(weights_path):
os.system('wget https://storage.googleapis.com/ultralytics/yolov3.pt -O ' + weights_path) os.system('wget https://storage.googleapis.com/ultralytics/yolov3.pt -O ' + weights_path)
else: # darknet format
load_weights(model, weights_path)
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'])

View File

@ -35,12 +35,12 @@ def main(opt):
model = Darknet(opt.cfg, opt.img_size) model = Darknet(opt.cfg, opt.img_size)
# Load weights # Load weights
if opt.weights_path.endswith('.weights'): # darknet format if opt.weights_path.endswith('.pt'): # pytorch format
load_weights(model, opt.weights_path)
elif opt.weights_path.endswith('.pt'): # pytorch format
checkpoint = torch.load(opt.weights_path, map_location='cpu') checkpoint = torch.load(opt.weights_path, map_location='cpu')
model.load_state_dict(checkpoint['model']) model.load_state_dict(checkpoint['model'])
del checkpoint del checkpoint
else: # darknet format
load_weights(model, opt.weights_path)
model.to(device).eval() model.to(device).eval()

View File

@ -13,6 +13,9 @@ python3 detect.py
# Test # Test
python3 test.py -img_size 416 -weights_path weights/latest.pt python3 test.py -img_size 416 -weights_path weights/latest.pt
# Test Darknet
python3 test.py -img_size 416 -weights_path ../darknet/backup/yolov3.backup
# Download and Test # Download and Test
sudo rm -rf yolov3 && git clone https://github.com/ultralytics/yolov3 && cd yolov3 sudo rm -rf yolov3 && git clone https://github.com/ultralytics/yolov3 && cd yolov3
wget https://pjreddie.com/media/files/yolov3.weights -P weights wget https://pjreddie.com/media/files/yolov3.weights -P weights