Skip to content

Commit

Permalink
Enable direct --weights URL definition (ultralytics#3373)
Browse files Browse the repository at this point in the history
* Enable direct `--weights URL` definition

@kalenmike this PR will enable direct --weights URL definition. Example use case:
```
python train.py --weights https://storage.googleapis.com/bucket/dir/model.pt
```

* cleanup

* bug fixes

* weights = attempt_download(weights)

* Update experimental.py

* Update hubconf.py

* return bug fix

* comment mirror

* min_bytes

(cherry picked from commit ba6f3f9)
  • Loading branch information
glenn-jocher authored and Joshua Friedrich committed Jun 14, 2021
1 parent 09dd0da commit c316d44
Showing 1 changed file with 7 additions and 10 deletions.
17 changes: 7 additions & 10 deletions utils/google_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,6 @@
import platform
import subprocess
import time
import urllib
from pathlib import Path

import requests
Expand All @@ -20,32 +19,30 @@ def gsutil_getsize(url=''):
def safe_download(file, url, url2=None, min_bytes=1E0, error_msg=''):
# Attempts to download file from url or url2, checks and removes incomplete downloads < min_bytes
file = Path(file)
assert_msg = f"Downloaded file '{file}' does not exist or size is < min_bytes={min_bytes}"
try: # url1
try: # GitHub
print(f'Downloading {url} to {file}...')
torch.hub.download_url_to_file(url, str(file))
assert file.exists() and file.stat().st_size > min_bytes, assert_msg # check
except Exception as e: # url2
assert file.exists() and file.stat().st_size > min_bytes # check
except Exception as e: # GCP
file.unlink(missing_ok=True) # remove partial downloads
print(f'ERROR: {e}\nRe-attempting {url2 or url} to {file}...')
print(f'Download error: {e}\nRe-attempting {url2 or url} to {file}...')
os.system(f"curl -L '{url2 or url}' -o '{file}' --retry 3 -C -") # curl download, retry and resume on fail
finally:
if not file.exists() or file.stat().st_size < min_bytes: # check
file.unlink(missing_ok=True) # remove partial downloads
print(f"ERROR: {assert_msg}\n{error_msg}")
print(f'ERROR: Download failure: {error_msg or url}')
print('')


def attempt_download(file, repo='ultralytics/yolov5'): # from utils.google_utils import *; attempt_download()
def attempt_download(file, repo='ultralytics/yolov5'):
# Attempt file download if does not exist
file = Path(str(file).strip().replace("'", ''))

if not file.exists():
# URL specified
name = Path(urllib.parse.unquote(str(file))).name # decode '%2F' to '/' etc.
name = file.name
if str(file).startswith(('http:/', 'https:/')): # download
url = str(file).replace(':/', '://') # Pathlib turns :// -> :/
name = name.split('?')[0] # parse authentication https://url.com/file.txt?auth...
safe_download(file=name, url=url, min_bytes=1E5)
return name

Expand Down

0 comments on commit c316d44

Please sign in to comment.