glenn-jocher commited on
Commit
2435bfe
·
unverified ·
1 Parent(s): c6b5bfc

Add URL download to check_file() (#3330)

Browse files

* Add URL file download to check_file()

* cleanup

* pathlib bug fix

Files changed (1) hide show
  1. utils/general.py +11 -4
utils/general.py CHANGED
@@ -173,12 +173,19 @@ def check_imshow():
173
 
174
 
175
  def check_file(file):
176
- # Search for file if not found
177
- if Path(file).is_file() or file == '':
 
178
  return file
179
- else:
 
 
 
 
 
 
180
  files = glob.glob('./**/' + file, recursive=True) # find file
181
- assert len(files), f'File Not Found: {file}' # assert file was found
182
  assert len(files) == 1, f"Multiple files match '{file}', specify exact path: {files}" # assert unique
183
  return files[0] # return file
184
 
 
173
 
174
 
175
  def check_file(file):
176
+ # Search/download file (if necessary) and return path
177
+ file = str(file) # convert to str()
178
+ if Path(file).is_file() or file == '': # exists
179
  return file
180
+ elif file.startswith(('http://', 'https://')): # download
181
+ url, file = file, Path(file).name
182
+ print(f'Downloading {url} to {file}...')
183
+ torch.hub.download_url_to_file(url, file)
184
+ assert Path(file).exists() and Path(file).stat().st_size > 0, f'File download failed: {url}' # check
185
+ return file
186
+ else: # search
187
  files = glob.glob('./**/' + file, recursive=True) # find file
188
+ assert len(files), f'File not found: {file}' # assert file was found
189
  assert len(files) == 1, f"Multiple files match '{file}', specify exact path: {files}" # assert unique
190
  return files[0] # return file
191