Spaces:
Build error
Build error
| """Main Logger class for ClearML experiment tracking.""" | |
| import glob | |
| import re | |
| from pathlib import Path | |
| import numpy as np | |
| import yaml | |
| from utils.plots import Annotator, colors | |
| try: | |
| import clearml | |
| from clearml import Dataset, Task | |
| assert hasattr(clearml, '__version__') # verify package import not local dir | |
| except (ImportError, AssertionError): | |
| clearml = None | |
| def construct_dataset(clearml_info_string): | |
| """Load in a clearml dataset and fill the internal data_dict with its contents. | |
| """ | |
| dataset_id = clearml_info_string.replace('clearml://', '') | |
| dataset = Dataset.get(dataset_id=dataset_id) | |
| dataset_root_path = Path(dataset.get_local_copy()) | |
| # We'll search for the yaml file definition in the dataset | |
| yaml_filenames = list(glob.glob(str(dataset_root_path / '*.yaml')) + glob.glob(str(dataset_root_path / '*.yml'))) | |
| if len(yaml_filenames) > 1: | |
| raise ValueError('More than one yaml file was found in the dataset root, cannot determine which one contains ' | |
| 'the dataset definition this way.') | |
| elif len(yaml_filenames) == 0: | |
| raise ValueError('No yaml definition found in dataset root path, check that there is a correct yaml file ' | |
| 'inside the dataset root path.') | |
| with open(yaml_filenames[0]) as f: | |
| dataset_definition = yaml.safe_load(f) | |
| assert set(dataset_definition.keys()).issuperset( | |
| {'train', 'test', 'val', 'nc', 'names'} | |
| ), "The right keys were not found in the yaml file, make sure it at least has the following keys: ('train', 'test', 'val', 'nc', 'names')" | |
| data_dict = dict() | |
| data_dict['train'] = str( | |
| (dataset_root_path / dataset_definition['train']).resolve()) if dataset_definition['train'] else None | |
| data_dict['test'] = str( | |
| (dataset_root_path / dataset_definition['test']).resolve()) if dataset_definition['test'] else None | |
| data_dict['val'] = str( | |
| (dataset_root_path / dataset_definition['val']).resolve()) if dataset_definition['val'] else None | |
| data_dict['nc'] = dataset_definition['nc'] | |
| data_dict['names'] = dataset_definition['names'] | |
| return data_dict | |
| class ClearmlLogger: | |
| """Log training runs, datasets, models, and predictions to ClearML. | |
| This logger sends information to ClearML at app.clear.ml or to your own hosted server. By default, | |
| this information includes hyperparameters, system configuration and metrics, model metrics, code information and | |
| basic data metrics and analyses. | |
| By providing additional command line arguments to train.py, datasets, | |
| models and predictions can also be logged. | |
| """ | |
| def __init__(self, opt, hyp): | |
| """ | |
| - Initialize ClearML Task, this object will capture the experiment | |
| - Upload dataset version to ClearML Data if opt.upload_dataset is True | |
| arguments: | |
| opt (namespace) -- Commandline arguments for this run | |
| hyp (dict) -- Hyperparameters for this run | |
| """ | |
| self.current_epoch = 0 | |
| # Keep tracked of amount of logged images to enforce a limit | |
| self.current_epoch_logged_images = set() | |
| # Maximum number of images to log to clearML per epoch | |
| self.max_imgs_to_log_per_epoch = 16 | |
| # Get the interval of epochs when bounding box images should be logged | |
| self.bbox_interval = opt.bbox_interval | |
| self.clearml = clearml | |
| self.task = None | |
| self.data_dict = None | |
| if self.clearml: | |
| self.task = Task.init( | |
| project_name=opt.project if opt.project != 'runs/train' else 'YOLOv5', | |
| task_name=opt.name if opt.name != 'exp' else 'Training', | |
| tags=['YOLOv5'], | |
| output_uri=True, | |
| reuse_last_task_id=opt.exist_ok, | |
| auto_connect_frameworks={'pytorch': False} | |
| # We disconnect pytorch auto-detection, because we added manual model save points in the code | |
| ) | |
| # ClearML's hooks will already grab all general parameters | |
| # Only the hyperparameters coming from the yaml config file | |
| # will have to be added manually! | |
| self.task.connect(hyp, name='Hyperparameters') | |
| self.task.connect(opt, name='Args') | |
| # Make sure the code is easily remotely runnable by setting the docker image to use by the remote agent | |
| self.task.set_base_docker('ultralytics/yolov5:latest', | |
| docker_arguments='--ipc=host -e="CLEARML_AGENT_SKIP_PYTHON_ENV_INSTALL=1"', | |
| docker_setup_bash_script='pip install clearml') | |
| # Get ClearML Dataset Version if requested | |
| if opt.data.startswith('clearml://'): | |
| # data_dict should have the following keys: | |
| # names, nc (number of classes), test, train, val (all three relative paths to ../datasets) | |
| self.data_dict = construct_dataset(opt.data) | |
| # Set data to data_dict because wandb will crash without this information and opt is the best way | |
| # to give it to them | |
| opt.data = self.data_dict | |
| def log_debug_samples(self, files, title='Debug Samples'): | |
| """ | |
| Log files (images) as debug samples in the ClearML task. | |
| arguments: | |
| files (List(PosixPath)) a list of file paths in PosixPath format | |
| title (str) A title that groups together images with the same values | |
| """ | |
| for f in files: | |
| if f.exists(): | |
| it = re.search(r'_batch(\d+)', f.name) | |
| iteration = int(it.groups()[0]) if it else 0 | |
| self.task.get_logger().report_image(title=title, | |
| series=f.name.replace(it.group(), ''), | |
| local_path=str(f), | |
| iteration=iteration) | |
| def log_image_with_boxes(self, image_path, boxes, class_names, image, conf_threshold=0.25): | |
| """ | |
| Draw the bounding boxes on a single image and report the result as a ClearML debug sample. | |
| arguments: | |
| image_path (PosixPath) the path the original image file | |
| boxes (list): list of scaled predictions in the format - [xmin, ymin, xmax, ymax, confidence, class] | |
| class_names (dict): dict containing mapping of class int to class name | |
| image (Tensor): A torch tensor containing the actual image data | |
| """ | |
| if len(self.current_epoch_logged_images) < self.max_imgs_to_log_per_epoch and self.current_epoch >= 0: | |
| # Log every bbox_interval times and deduplicate for any intermittend extra eval runs | |
| if self.current_epoch % self.bbox_interval == 0 and image_path not in self.current_epoch_logged_images: | |
| im = np.ascontiguousarray(np.moveaxis(image.mul(255).clamp(0, 255).byte().cpu().numpy(), 0, 2)) | |
| annotator = Annotator(im=im, pil=True) | |
| for i, (conf, class_nr, box) in enumerate(zip(boxes[:, 4], boxes[:, 5], boxes[:, :4])): | |
| color = colors(i) | |
| class_name = class_names[int(class_nr)] | |
| confidence_percentage = round(float(conf) * 100, 2) | |
| label = f'{class_name}: {confidence_percentage}%' | |
| if conf > conf_threshold: | |
| annotator.rectangle(box.cpu().numpy(), outline=color) | |
| annotator.box_label(box.cpu().numpy(), label=label, color=color) | |
| annotated_image = annotator.result() | |
| self.task.get_logger().report_image(title='Bounding Boxes', | |
| series=image_path.name, | |
| iteration=self.current_epoch, | |
| image=annotated_image) | |
| self.current_epoch_logged_images.add(image_path) | |