first commit

app/yolov5/utils/loggers/clearml/clearml_utils.py

@@ -0,0 +1,156 @@
+"""Main Logger class for ClearML experiment tracking."""
+import glob
+import re
+from pathlib import Path
+
+import numpy as np
+import yaml
+
+from app.yolov5.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='YOLOv5',
+                task_name='training',
+                tags=['YOLOv5'],
+                output_uri=True,
+                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')
+
+            # 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)