You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

157 lines
6.1KB

  1. # YOLOv5 🚀 by Ultralytics, GPL-3.0 license
  2. """
  3. Download utils
  4. """
  5. import os
  6. import platform
  7. import subprocess
  8. import time
  9. import urllib
  10. from pathlib import Path
  11. from zipfile import ZipFile
  12. import requests
  13. import torch
  14. def gsutil_getsize(url=''):
  15. # gs://bucket/file size https://cloud.google.com/storage/docs/gsutil/commands/du
  16. s = subprocess.check_output(f'gsutil du {url}', shell=True).decode('utf-8')
  17. return eval(s.split(' ')[0]) if len(s) else 0 # bytes
  18. def safe_download(file, url, url2=None, min_bytes=1E0, error_msg=''):
  19. # Attempts to download file from url or url2, checks and removes incomplete downloads < min_bytes
  20. file = Path(file)
  21. assert_msg = f"Downloaded file '{file}' does not exist or size is < min_bytes={min_bytes}"
  22. try: # url1
  23. print(f'Downloading {url} to {file}...')
  24. torch.hub.download_url_to_file(url, str(file))
  25. assert file.exists() and file.stat().st_size > min_bytes, assert_msg # check
  26. except Exception as e: # url2
  27. file.unlink(missing_ok=True) # remove partial downloads
  28. print(f'ERROR: {e}\nRe-attempting {url2 or url} to {file}...')
  29. os.system(f"curl -L '{url2 or url}' -o '{file}' --retry 3 -C -") # curl download, retry and resume on fail
  30. finally:
  31. if not file.exists() or file.stat().st_size < min_bytes: # check
  32. file.unlink(missing_ok=True) # remove partial downloads
  33. print(f"ERROR: {assert_msg}\n{error_msg}")
  34. print('')
  35. def attempt_download(file, repo='ultralytics/yolov5'): # from utils.downloads import *; attempt_download()
  36. # Attempt file download if does not exist
  37. file = Path(str(file).strip().replace("'", ''))
  38. if not file.exists():
  39. # URL specified
  40. name = Path(urllib.parse.unquote(str(file))).name # decode '%2F' to '/' etc.
  41. if str(file).startswith(('http:/', 'https:/')): # download
  42. url = str(file).replace(':/', '://') # Pathlib turns :// -> :/
  43. file = name.split('?')[0] # parse authentication https://url.com/file.txt?auth...
  44. if Path(file).is_file():
  45. print(f'Found {url} locally at {file}') # file already exists
  46. else:
  47. safe_download(file=file, url=url, min_bytes=1E5)
  48. return file
  49. # GitHub assets
  50. file.parent.mkdir(parents=True, exist_ok=True) # make parent dir (if required)
  51. try:
  52. response = requests.get(f'https://api.github.com/repos/{repo}/releases/latest').json() # github api
  53. assets = [x['name'] for x in response['assets']] # release assets, i.e. ['yolov5s.pt', 'yolov5m.pt', ...]
  54. tag = response['tag_name'] # i.e. 'v1.0'
  55. except Exception: # fallback plan
  56. assets = [
  57. 'yolov5n.pt', 'yolov5s.pt', 'yolov5m.pt', 'yolov5l.pt', 'yolov5x.pt', 'yolov5n6.pt', 'yolov5s6.pt',
  58. 'yolov5m6.pt', 'yolov5l6.pt', 'yolov5x6.pt']
  59. try:
  60. tag = subprocess.check_output('git tag', shell=True, stderr=subprocess.STDOUT).decode().split()[-1]
  61. except Exception:
  62. tag = 'v6.0' # current release
  63. if name in assets:
  64. safe_download(
  65. file,
  66. url=f'https://github.com/{repo}/releases/download/{tag}/{name}',
  67. # url2=f'https://storage.googleapis.com/{repo}/ckpt/{name}', # backup url (optional)
  68. min_bytes=1E5,
  69. error_msg=f'{file} missing, try downloading from https://github.com/{repo}/releases/')
  70. return str(file)
  71. def gdrive_download(id='16TiPfZj7htmTyhntwcZyEEAejOUxuT6m', file='tmp.zip'):
  72. # Downloads a file from Google Drive. from yolov5.utils.downloads import *; gdrive_download()
  73. t = time.time()
  74. file = Path(file)
  75. cookie = Path('cookie') # gdrive cookie
  76. print(f'Downloading https://drive.google.com/uc?export=download&id={id} as {file}... ', end='')
  77. file.unlink(missing_ok=True) # remove existing file
  78. cookie.unlink(missing_ok=True) # remove existing cookie
  79. # Attempt file download
  80. out = "NUL" if platform.system() == "Windows" else "/dev/null"
  81. os.system(f'curl -c ./cookie -s -L "drive.google.com/uc?export=download&id={id}" > {out}')
  82. if os.path.exists('cookie'): # large file
  83. s = f'curl -Lb ./cookie "drive.google.com/uc?export=download&confirm={get_token()}&id={id}" -o {file}'
  84. else: # small file
  85. s = f'curl -s -L -o {file} "drive.google.com/uc?export=download&id={id}"'
  86. r = os.system(s) # execute, capture return
  87. cookie.unlink(missing_ok=True) # remove existing cookie
  88. # Error check
  89. if r != 0:
  90. file.unlink(missing_ok=True) # remove partial
  91. print('Download error ') # raise Exception('Download error')
  92. return r
  93. # Unzip if archive
  94. if file.suffix == '.zip':
  95. print('unzipping... ', end='')
  96. ZipFile(file).extractall(path=file.parent) # unzip
  97. file.unlink() # remove zip
  98. print(f'Done ({time.time() - t:.1f}s)')
  99. return r
  100. def get_token(cookie="./cookie"):
  101. with open(cookie) as f:
  102. for line in f:
  103. if "download" in line:
  104. return line.split()[-1]
  105. return ""
  106. # Google utils: https://cloud.google.com/storage/docs/reference/libraries ----------------------------------------------
  107. #
  108. #
  109. # def upload_blob(bucket_name, source_file_name, destination_blob_name):
  110. # # Uploads a file to a bucket
  111. # # https://cloud.google.com/storage/docs/uploading-objects#storage-upload-object-python
  112. #
  113. # storage_client = storage.Client()
  114. # bucket = storage_client.get_bucket(bucket_name)
  115. # blob = bucket.blob(destination_blob_name)
  116. #
  117. # blob.upload_from_filename(source_file_name)
  118. #
  119. # print('File {} uploaded to {}.'.format(
  120. # source_file_name,
  121. # destination_blob_name))
  122. #
  123. #
  124. # def download_blob(bucket_name, source_blob_name, destination_file_name):
  125. # # Uploads a blob from a bucket
  126. # storage_client = storage.Client()
  127. # bucket = storage_client.get_bucket(bucket_name)
  128. # blob = bucket.blob(source_blob_name)
  129. #
  130. # blob.download_to_filename(destination_file_name)
  131. #
  132. # print('Blob {} downloaded to {}.'.format(
  133. # source_blob_name,
  134. # destination_file_name))