From ef133382c5ac2c1b5704b113c25c5cbf4786a7ce Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Fri, 6 Dec 2019 13:44:13 -0800
Subject: [PATCH] updates

---
 models.py             | 26 ++++++++++++--------------
 utils/google_utils.py |  7 ++++---
 2 files changed, 16 insertions(+), 17 deletions(-)

diff --git a/models.py b/models.py
index b61e398b..6dd1aa6f 100755
--- a/models.py
+++ b/models.py
@@ -436,10 +436,11 @@ def convert(cfg='cfg/yolov3-spp.cfg', weights='weights/yolov3-spp.weights'):
 
 def attempt_download(weights):
     # Attempt to download pretrained weights if not found locally
+    file = Path(weights).name
+    msg = weights + ' missing, try downloading from https://drive.google.com/open?id=1LezFG5g3BCW6iYaV89B2i64cqEUZD7e0'
+    r = 1  # error value
 
-    msg = weights + ' missing, download from https://drive.google.com/open?id=1LezFG5g3BCW6iYaV89B2i64cqEUZD7e0'
     if weights and not os.path.isfile(weights):
-        file = Path(weights).name
         d = {'yolov3-spp.weights': '16lYS4bcIdM2HdmyJBVDOvt3Trx6N3W2R',
              'yolov3.weights': '1uTlyDWlnaqXcsKOktP5aH_zRDbfcDp-y',
              'yolov3-tiny.weights': '1CCF-iNIIkYesIDzaPvdwlcf7H9zSsKZQ',
@@ -451,17 +452,14 @@ def attempt_download(weights):
              'ultralytics49.pt': '158g62Vs14E3aj7oPVPuEnNZMKFNgGyNq',
              'ultralytics68.pt': '1Jm8kqnMdMGUUxGo8zMFZMJ0eaPwLkxSG'}
 
-        if weights in d:
-            gdrive_download(id=d[weights], name=weights)
+        if file in d:
+            r = gdrive_download(id=d[file], name=weights)
         else:  # download from pjreddie.com
-            try:
-                url = 'https://pjreddie.com/media/files/' + file
-                print('Downloading ' + url)
-                os.system('curl -f ' + url + ' -o ' + weights)
-            except IOError:
-                print(msg)
-                os.system('rm ' + weights)  # remove partial downloads
+            url = 'https://pjreddie.com/media/files/' + file
+            print('Downloading ' + url)
+            r = os.system('curl -f ' + url + ' -o ' + weights)
 
-        if os.path.getsize(weights) < 5E6:  # weights < 5MB (too small), download failed
-            os.remove(weights)  # delete corrupted weightsfile
-        assert os.path.exists(weights), msg  # download missing weights from Google Drive
+    # Error check
+    if not (r == 0 and os.path.exists(weights) and os.path.getsize(weights) > 1E6):  # weights exist and > 1MB
+        os.system('rm ' + weights)  # remove partial downloads
+        raise Exception(msg)
diff --git a/utils/google_utils.py b/utils/google_utils.py
index fc2bc34f..92fde590 100644
--- a/utils/google_utils.py
+++ b/utils/google_utils.py
@@ -30,11 +30,11 @@ def gdrive_download(id='1HaXkef9z6y5l4vUnCYgdmEAj61c6bfWO', name='coco.zip'):
         s = 'curl -f -L -o %s https://drive.google.com/uc?export=download&id=%s' % (name, id)
         r = os.system(s)
 
-    # Check for errors
+    # Error check
     if r != 0:
         os.system('rm ' + name)  # remove partial downloads
-        print('ERROR: Download failure. ')
-        return
+        print('ERROR: Download failure ')
+        return r
 
     # Unzip if archive
     if name.endswith('.zip'):
@@ -43,6 +43,7 @@ def gdrive_download(id='1HaXkef9z6y5l4vUnCYgdmEAj61c6bfWO', name='coco.zip'):
         os.remove(name)  # remove zip to free space
 
     print('Done (%.1fs)' % (time.time() - t))
+    return r
 
 
 def upload_blob(bucket_name, source_file_name, destination_blob_name):