Update google_utils.py (ultralytics#2900)

fhkiel-mlaip · Apr 22, 2021 · 48fd9fe · 48fd9fe
1 parent f5f65ad
commit 48fd9fe
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/utils/google_utils.py b/utils/google_utils.py
@@ -26,8 +26,12 @@ def attempt_download(file, repo='ultralytics/yolov5'):
             assets = [x['name'] for x in response['assets']]  # release assets, i.e. ['yolov5s.pt', 'yolov5m.pt', ...]
             tag = response['tag_name']  # i.e. 'v1.0'
         except:  # fallback plan
-            assets = ['yolov5s.pt', 'yolov5m.pt', 'yolov5l.pt', 'yolov5x.pt']
-            tag = subprocess.check_output('git tag', shell=True).decode().split()[-1]
+            assets = ['yolov5s.pt', 'yolov5m.pt', 'yolov5l.pt', 'yolov5x.pt',
+                      'yolov5s6.pt', 'yolov5m6.pt', 'yolov5l6.pt', 'yolov5x6.pt']
+            try:
+                tag = subprocess.check_output('git tag', shell=True, stderr=subprocess.STDOUT).decode().split()[-1]
+            except:
+                tag = 'v5.0'  # current release
 
         name = file.name
         if name in assets: