diff --git a/utils/general.py b/utils/general.py index 842f28c60886..ac9b0cefd7c8 100755 --- a/utils/general.py +++ b/utils/general.py @@ -300,7 +300,7 @@ def git_describe(path=ROOT): # path must be a directory @TryExcept() @WorkingDirectory(ROOT) -def check_git_status(repo='ultralytics/yolov5'): +def check_git_status(repo='ultralytics/yolov5', branch='master'): # YOLOv5 status check, recommend 'git pull' if code is out of date url = f'https://github.com/{repo}' msg = f', for updates see {url}' @@ -316,10 +316,10 @@ def check_git_status(repo='ultralytics/yolov5'): remote = 'ultralytics' check_output(f'git remote add {remote} {url}', shell=True) check_output(f'git fetch {remote}', shell=True, timeout=5) # git fetch - branch = check_output('git rev-parse --abbrev-ref HEAD', shell=True).decode().strip() # checked out - n = int(check_output(f'git rev-list {branch}..{remote}/master --count', shell=True)) # commits behind + local_branch = check_output('git rev-parse --abbrev-ref HEAD', shell=True).decode().strip() # checked out + n = int(check_output(f'git rev-list {local_branch}..{remote}/{branch} --count', shell=True)) # commits behind if n > 0: - pull = 'git pull' if remote == 'origin' else f'git pull {remote} master' + pull = 'git pull' if remote == 'origin' else f'git pull {remote} {branch}' s += f"⚠️ YOLOv5 is out of date by {n} commit{'s' * (n > 1)}. Use `{pull}` or `git clone {url}` to update." else: s += f'up to date with {url} ✅'