This commit is contained in:
Glenn Jocher 2019-12-06 13:44:13 -08:00
parent 6067b22605
commit ef133382c5
2 changed files with 16 additions and 17 deletions

View File

@ -436,10 +436,11 @@ def convert(cfg='cfg/yolov3-spp.cfg', weights='weights/yolov3-spp.weights'):
def attempt_download(weights): def attempt_download(weights):
# Attempt to download pretrained weights if not found locally # Attempt to download pretrained weights if not found locally
msg = weights + ' missing, download from https://drive.google.com/open?id=1LezFG5g3BCW6iYaV89B2i64cqEUZD7e0'
if weights and not os.path.isfile(weights):
file = Path(weights).name file = Path(weights).name
msg = weights + ' missing, try downloading from https://drive.google.com/open?id=1LezFG5g3BCW6iYaV89B2i64cqEUZD7e0'
r = 1 # error value
if weights and not os.path.isfile(weights):
d = {'yolov3-spp.weights': '16lYS4bcIdM2HdmyJBVDOvt3Trx6N3W2R', d = {'yolov3-spp.weights': '16lYS4bcIdM2HdmyJBVDOvt3Trx6N3W2R',
'yolov3.weights': '1uTlyDWlnaqXcsKOktP5aH_zRDbfcDp-y', 'yolov3.weights': '1uTlyDWlnaqXcsKOktP5aH_zRDbfcDp-y',
'yolov3-tiny.weights': '1CCF-iNIIkYesIDzaPvdwlcf7H9zSsKZQ', 'yolov3-tiny.weights': '1CCF-iNIIkYesIDzaPvdwlcf7H9zSsKZQ',
@ -451,17 +452,14 @@ def attempt_download(weights):
'ultralytics49.pt': '158g62Vs14E3aj7oPVPuEnNZMKFNgGyNq', 'ultralytics49.pt': '158g62Vs14E3aj7oPVPuEnNZMKFNgGyNq',
'ultralytics68.pt': '1Jm8kqnMdMGUUxGo8zMFZMJ0eaPwLkxSG'} 'ultralytics68.pt': '1Jm8kqnMdMGUUxGo8zMFZMJ0eaPwLkxSG'}
if weights in d: if file in d:
gdrive_download(id=d[weights], name=weights) r = gdrive_download(id=d[file], name=weights)
else: # download from pjreddie.com else: # download from pjreddie.com
try:
url = 'https://pjreddie.com/media/files/' + file url = 'https://pjreddie.com/media/files/' + file
print('Downloading ' + url) print('Downloading ' + url)
os.system('curl -f ' + url + ' -o ' + weights) r = os.system('curl -f ' + url + ' -o ' + weights)
except IOError:
print(msg)
os.system('rm ' + weights) # remove partial downloads
if os.path.getsize(weights) < 5E6: # weights < 5MB (too small), download failed # Error check
os.remove(weights) # delete corrupted weightsfile if not (r == 0 and os.path.exists(weights) and os.path.getsize(weights) > 1E6): # weights exist and > 1MB
assert os.path.exists(weights), msg # download missing weights from Google Drive os.system('rm ' + weights) # remove partial downloads
raise Exception(msg)

View File

@ -30,11 +30,11 @@ def gdrive_download(id='1HaXkef9z6y5l4vUnCYgdmEAj61c6bfWO', name='coco.zip'):
s = 'curl -f -L -o %s https://drive.google.com/uc?export=download&id=%s' % (name, id) s = 'curl -f -L -o %s https://drive.google.com/uc?export=download&id=%s' % (name, id)
r = os.system(s) r = os.system(s)
# Check for errors # Error check
if r != 0: if r != 0:
os.system('rm ' + name) # remove partial downloads os.system('rm ' + name) # remove partial downloads
print('ERROR: Download failure. ') print('ERROR: Download failure ')
return return r
# Unzip if archive # Unzip if archive
if name.endswith('.zip'): if name.endswith('.zip'):
@ -43,6 +43,7 @@ def gdrive_download(id='1HaXkef9z6y5l4vUnCYgdmEAj61c6bfWO', name='coco.zip'):
os.remove(name) # remove zip to free space os.remove(name) # remove zip to free space
print('Done (%.1fs)' % (time.time() - t)) print('Done (%.1fs)' % (time.time() - t))
return r
def upload_blob(bucket_name, source_file_name, destination_blob_name): def upload_blob(bucket_name, source_file_name, destination_blob_name):