Commit
•
1f92422
1
Parent(s):
4949401
Github release assets model autodownload (#711)
Browse files* assets autodownload
* Update google_utils.py
* Update google_utils.py
* Update google_utils.py
* Update google_utils.py
* Update google_utils.py
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
- utils/google_utils.py +36 -26
utils/google_utils.py
CHANGED
@@ -6,40 +6,50 @@ import os
|
|
6 |
import platform
|
7 |
import time
|
8 |
from pathlib import Path
|
|
|
9 |
|
10 |
|
11 |
def attempt_download(weights):
|
12 |
# Attempt to download pretrained weights if not found locally
|
13 |
weights = weights.strip().replace("'", '')
|
14 |
-
|
15 |
-
|
16 |
-
|
17 |
-
|
18 |
-
|
19 |
-
|
20 |
-
|
21 |
-
|
22 |
-
|
23 |
-
|
24 |
-
|
25 |
-
|
26 |
-
if
|
27 |
-
|
28 |
-
|
29 |
-
|
30 |
-
|
31 |
-
|
32 |
-
|
33 |
-
|
34 |
-
|
35 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
36 |
os.remove(weights) if os.path.exists(weights) else None # remove partial downloads
|
37 |
-
|
|
|
|
|
38 |
|
39 |
|
40 |
def gdrive_download(id='1n_oKgR81BJtqk75b00eAjdv03qVCQn2f', name='coco128.zip'):
|
41 |
-
# Downloads a file from Google Drive
|
42 |
-
# from utils.google_utils import *; gdrive_download()
|
43 |
t = time.time()
|
44 |
|
45 |
print('Downloading https://drive.google.com/uc?export=download&id=%s as %s... ' % (id, name), end='')
|
@@ -53,7 +63,7 @@ def gdrive_download(id='1n_oKgR81BJtqk75b00eAjdv03qVCQn2f', name='coco128.zip'):
|
|
53 |
s = 'curl -Lb ./cookie "drive.google.com/uc?export=download&confirm=%s&id=%s" -o %s' % (get_token(), id, name)
|
54 |
else: # small file
|
55 |
s = 'curl -s -L -o %s "drive.google.com/uc?export=download&id=%s"' % (name, id)
|
56 |
-
r = os.system(s) # execute, capture return
|
57 |
os.remove('cookie') if os.path.exists('cookie') else None
|
58 |
|
59 |
# Error check
|
|
|
6 |
import platform
|
7 |
import time
|
8 |
from pathlib import Path
|
9 |
+
import torch
|
10 |
|
11 |
|
12 |
def attempt_download(weights):
|
13 |
# Attempt to download pretrained weights if not found locally
|
14 |
weights = weights.strip().replace("'", '')
|
15 |
+
file = Path(weights).name
|
16 |
+
|
17 |
+
msg = weights + ' missing, try downloading from https://github.com/ultralytics/yolov5/releases/'
|
18 |
+
models = ['yolov5s.pt', 'yolov5m.pt', 'yolov5l.pt', 'yolov5x.pt'] # available models
|
19 |
+
|
20 |
+
if file in models and not os.path.isfile(weights):
|
21 |
+
# Google Drive
|
22 |
+
# d = {'yolov5s.pt': '1R5T6rIyy3lLwgFXNms8whc-387H0tMQO',
|
23 |
+
# 'yolov5m.pt': '1vobuEExpWQVpXExsJ2w-Mbf3HJjWkQJr',
|
24 |
+
# 'yolov5l.pt': '1hrlqD1Wdei7UT4OgT785BEk1JwnSvNEV',
|
25 |
+
# 'yolov5x.pt': '1mM8aZJlWTxOg7BZJvNUMrTnA2AbeCVzS'}
|
26 |
+
# r = gdrive_download(id=d[file], name=weights) if file in d else 1
|
27 |
+
# if r == 0 and os.path.exists(weights) and os.path.getsize(weights) > 1E6: # check
|
28 |
+
# return
|
29 |
+
|
30 |
+
try: # GitHub
|
31 |
+
url = 'https://github.com/ultralytics/yolov5/releases/download/v2.0/' + file
|
32 |
+
print('Downloading %s to %s...' % (url, weights))
|
33 |
+
if platform.system() == 'Darwin': # avoid MacOS python requests certificate error
|
34 |
+
r = os.system('curl -L %s -o %s' % (url, weights))
|
35 |
+
else:
|
36 |
+
torch.hub.download_url_to_file(url, weights)
|
37 |
+
assert os.path.exists(weights) and os.path.getsize(weights) > 1E6 # check
|
38 |
+
except Exception as e: # GCP
|
39 |
+
print('Download error: %s' % e)
|
40 |
+
url = 'https://storage.googleapis.com/ultralytics/yolov5/ckpt/' + file
|
41 |
+
print('Downloading %s to %s...' % (url, weights))
|
42 |
+
r = os.system('curl -L %s -o %s' % (url, weights)) # torch.hub.download_url_to_file(url, weights)
|
43 |
+
finally:
|
44 |
+
if not (os.path.exists(weights) and os.path.getsize(weights) > 1E6): # check
|
45 |
os.remove(weights) if os.path.exists(weights) else None # remove partial downloads
|
46 |
+
print('ERROR: Download failure: %s' % msg)
|
47 |
+
print('')
|
48 |
+
return
|
49 |
|
50 |
|
51 |
def gdrive_download(id='1n_oKgR81BJtqk75b00eAjdv03qVCQn2f', name='coco128.zip'):
|
52 |
+
# Downloads a file from Google Drive. from utils.google_utils import *; gdrive_download()
|
|
|
53 |
t = time.time()
|
54 |
|
55 |
print('Downloading https://drive.google.com/uc?export=download&id=%s as %s... ' % (id, name), end='')
|
|
|
63 |
s = 'curl -Lb ./cookie "drive.google.com/uc?export=download&confirm=%s&id=%s" -o %s' % (get_token(), id, name)
|
64 |
else: # small file
|
65 |
s = 'curl -s -L -o %s "drive.google.com/uc?export=download&id=%s"' % (name, id)
|
66 |
+
r = os.system(s) # execute, capture return
|
67 |
os.remove('cookie') if os.path.exists('cookie') else None
|
68 |
|
69 |
# Error check
|