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.

179 line
6.9KB

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