Add retry=3 to download() (#7313)

* Add `retry=3` to `download()` * Update general.py * Update general.py * Update general.py * Update VOC.yaml * Update VisDrone.yaml
ultralytics · Apr 6, 2022 · 32661f7 · 32661f7
1 parent f735458
commit 32661f7
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 8 deletions.
diff --git a/data/VOC.yaml b/data/VOC.yaml
@@ -62,7 +62,7 @@ download: |
   urls = [url + 'VOCtrainval_06-Nov-2007.zip',  # 446MB, 5012 images
           url + 'VOCtest_06-Nov-2007.zip',  # 438MB, 4953 images
           url + 'VOCtrainval_11-May-2012.zip']  # 1.95GB, 17126 images
-  download(urls, dir=dir / 'images', delete=False, threads=3)
+  download(urls, dir=dir / 'images', delete=False, curl=True, threads=3)
 
   # Convert
   path = dir / f'images/VOCdevkit'

diff --git a/data/VisDrone.yaml b/data/VisDrone.yaml
@@ -54,7 +54,7 @@ download: |
           'https://github.com/ultralytics/yolov5/releases/download/v1.0/VisDrone2019-DET-val.zip',
           'https://github.com/ultralytics/yolov5/releases/download/v1.0/VisDrone2019-DET-test-dev.zip',
           'https://github.com/ultralytics/yolov5/releases/download/v1.0/VisDrone2019-DET-test-challenge.zip']
-  download(urls, dir=dir, threads=4)
+  download(urls, dir=dir, curl=True, threads=4)
 
   # Convert
   for d in 'VisDrone2019-DET-train', 'VisDrone2019-DET-val', 'VisDrone2019-DET-test-dev':

diff --git a/utils/general.py b/utils/general.py
@@ -497,20 +497,32 @@ def url2file(url):
     return file
 
 
-def download(url, dir='.', unzip=True, delete=True, curl=False, threads=1):
+def download(url, dir='.', unzip=True, delete=True, curl=False, threads=1, retry=3):
     # Multi-threaded file download and unzip function, used in data.yaml for autodownload
     def download_one(url, dir):
         # Download 1 file
+        success = True
         f = dir / Path(url).name  # filename
         if Path(url).is_file():  # exists in current path
             Path(url).rename(f)  # move to dir
         elif not f.exists():
             LOGGER.info(f'Downloading {url} to {f}...')
-            if curl:
-                os.system(f"curl -L '{url}' -o '{f}' --retry 9 -C -")  # curl download, retry and resume on fail
-            else:
-                torch.hub.download_url_to_file(url, f, progress=threads == 1)  # torch download
-        if unzip and f.suffix in ('.zip', '.gz'):
+            for i in range(retry + 1):
+                if curl:
+                    s = 'sS' if threads > 1 else ''  # silent
+                    r = os.system(f"curl -{s}L '{url}' -o '{f}' --retry 9 -C -")  # curl download
+                    success = r == 0
+                else:
+                    torch.hub.download_url_to_file(url, f, progress=threads == 1)  # torch download
+                    success = f.is_file()
+                if success:
+                    break
+                elif i < retry:
+                    LOGGER.warning(f'Download failure, retrying {i + 1}/{retry} {url}...')
+                else:
+                    LOGGER.warning(f'Failed to download {url}...')
+
+        if unzip and success and f.suffix in ('.zip', '.gz'):
             LOGGER.info(f'Unzipping {f}...')
             if f.suffix == '.zip':
                 ZipFile(f).extractall(path=dir)  # unzip