diff --git a/requirements.txt b/requirements.txt index 70dd7ce53ba3..85eb839df8a0 100644 --- a/requirements.txt +++ b/requirements.txt @@ -2,6 +2,7 @@ # Usage: pip install -r requirements.txt # Base ------------------------------------------------------------------------ +gitpython ipython # interactive notebook matplotlib>=3.2.2 numpy>=1.18.5 diff --git a/train.py b/train.py index bbbd6d07db00..6fa33f47d100 100644 --- a/train.py +++ b/train.py @@ -47,7 +47,7 @@ from utils.callbacks import Callbacks from utils.dataloaders import create_dataloader from utils.downloads import attempt_download, is_url -from utils.general import (LOGGER, TQDM_BAR_FORMAT, check_amp, check_dataset, check_file, check_git_status, +from utils.general import (GIT, LOGGER, TQDM_BAR_FORMAT, check_amp, check_dataset, check_file, check_git_status, check_img_size, check_requirements, check_suffix, check_yaml, colorstr, get_latest_run, increment_path, init_seeds, intersect_dicts, labels_to_class_weights, labels_to_image_weights, methods, one_cycle, print_args, print_mutation, strip_optimizer, @@ -376,6 +376,7 @@ def train(hyp, opt, device, callbacks): # hyp is path/to/hyp.yaml or hyp dictio 'updates': ema.updates, 'optimizer': optimizer.state_dict(), 'opt': vars(opt), + 'git': GIT, 'date': datetime.now().isoformat()} # Save last, best and delete diff --git a/utils/general.py b/utils/general.py index 58181f00568d..57b6e4e78166 100644 --- a/utils/general.py +++ b/utils/general.py @@ -29,6 +29,7 @@ from zipfile import ZipFile, is_zipfile import cv2 +import git import IPython import numpy as np import pandas as pd @@ -344,6 +345,22 @@ def check_git_status(repo='ultralytics/yolov5', branch='master'): LOGGER.info(s) +@WorkingDirectory(ROOT) +def check_git(path='.'): + # YOLOv5 git check, return git {remote, branch, commit} + try: + repo = git.Repo(path) + remote = repo.remotes.origin.url.replace('.git', '') # i.e. 'https://github.com/ultralytics/yolov5' + commit = repo.head.commit.hexsha # i.e. '3134699c73af83aac2a481435550b968d5792c0d' + try: + branch = repo.active_branch.name # i.e. 'main' + except TypeError: # not on any branch + branch = None # i.e. 'detached HEAD' state + return {'remote': remote, 'branch': branch, 'commit': commit} + except git.exc.InvalidGitRepositoryError: # path is not a git dir + return {'remote': None, 'branch': None, 'commit': None} + + def check_python(minimum='3.7.0'): # Check current python version vs. required python version check_version(platform.python_version(), minimum, name='Python ', hard=True) @@ -1121,4 +1138,4 @@ def imshow(path, im): cv2.imread, cv2.imwrite, cv2.imshow = imread, imwrite, imshow # redefine # Variables ------------------------------------------------------------------------------------------------------------ -NCOLS = 0 if is_docker() else shutil.get_terminal_size().columns # terminal window size for tqdm +GIT = check_git() # repo, branch, commit