Nelze vybrat více než 25 témat Téma musí začínat písmenem nebo číslem, může obsahovat pomlčky („-“) a může být dlouhé až 35 znaků.

154 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 = ['yolov5n.pt', 'yolov5s.pt', 'yolov5m.pt', 'yolov5l.pt', 'yolov5x.pt',
  57. 'yolov5n6.pt', 'yolov5s6.pt', 'yolov5m6.pt', 'yolov5l6.pt', 'yolov5x6.pt']
  58. try:
  59. tag = subprocess.check_output('git tag', shell=True, stderr=subprocess.STDOUT).decode().split()[-1]
  60. except Exception:
  61. tag = 'v6.0' # current release
  62. if name in assets:
  63. safe_download(file,
  64. url=f'https://github.com/{repo}/releases/download/{tag}/{name}',
  65. # url2=f'https://storage.googleapis.com/{repo}/ckpt/{name}', # backup url (optional)
  66. min_bytes=1E5,
  67. error_msg=f'{file} missing, try downloading from https://github.com/{repo}/releases/')
  68. return str(file)
  69. def gdrive_download(id='16TiPfZj7htmTyhntwcZyEEAejOUxuT6m', file='tmp.zip'):
  70. # Downloads a file from Google Drive. from yolov5.utils.downloads import *; gdrive_download()
  71. t = time.time()
  72. file = Path(file)
  73. cookie = Path('cookie') # gdrive cookie
  74. print(f'Downloading https://drive.google.com/uc?export=download&id={id} as {file}... ', end='')
  75. file.unlink(missing_ok=True) # remove existing file
  76. cookie.unlink(missing_ok=True) # remove existing cookie
  77. # Attempt file download
  78. out = "NUL" if platform.system() == "Windows" else "/dev/null"
  79. os.system(f'curl -c ./cookie -s -L "drive.google.com/uc?export=download&id={id}" > {out}')
  80. if os.path.exists('cookie'): # large file
  81. s = f'curl -Lb ./cookie "drive.google.com/uc?export=download&confirm={get_token()}&id={id}" -o {file}'
  82. else: # small file
  83. s = f'curl -s -L -o {file} "drive.google.com/uc?export=download&id={id}"'
  84. r = os.system(s) # execute, capture return
  85. cookie.unlink(missing_ok=True) # remove existing cookie
  86. # Error check
  87. if r != 0:
  88. file.unlink(missing_ok=True) # remove partial
  89. print('Download error ') # raise Exception('Download error')
  90. return r
  91. # Unzip if archive
  92. if file.suffix == '.zip':
  93. print('unzipping... ', end='')
  94. ZipFile(file).extractall(path=file.parent) # unzip
  95. file.unlink() # remove zip
  96. print(f'Done ({time.time() - t:.1f}s)')
  97. return r
  98. def get_token(cookie="./cookie"):
  99. with open(cookie) as f:
  100. for line in f:
  101. if "download" in line:
  102. return line.split()[-1]
  103. return ""
  104. # Google utils: https://cloud.google.com/storage/docs/reference/libraries ----------------------------------------------
  105. #
  106. #
  107. # def upload_blob(bucket_name, source_file_name, destination_blob_name):
  108. # # Uploads a file to a bucket
  109. # # https://cloud.google.com/storage/docs/uploading-objects#storage-upload-object-python
  110. #
  111. # storage_client = storage.Client()
  112. # bucket = storage_client.get_bucket(bucket_name)
  113. # blob = bucket.blob(destination_blob_name)
  114. #
  115. # blob.upload_from_filename(source_file_name)
  116. #
  117. # print('File {} uploaded to {}.'.format(
  118. # source_file_name,
  119. # destination_blob_name))
  120. #
  121. #
  122. # def download_blob(bucket_name, source_blob_name, destination_file_name):
  123. # # Uploads a blob from a bucket
  124. # storage_client = storage.Client()
  125. # bucket = storage_client.get_bucket(bucket_name)
  126. # blob = bucket.blob(source_blob_name)
  127. #
  128. # blob.download_to_filename(destination_file_name)
  129. #
  130. # print('Blob {} downloaded to {}.'.format(
  131. # source_blob_name,
  132. # destination_file_name))