Nevar pievienot vairāk kā 25 tēmas Tēmai ir jāsākas ar burtu vai ciparu, tā var saturēt domu zīmes ('-') un var būt līdz 35 simboliem gara.

144 rindas
5.8KB

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