Upload New File

yolov5/models/utils/loggers/comet/__init__.py

+import glob
+import json
+import logging
+import os
+import sys
+from pathlib import Path
+
+logger = logging.getLogger(__name__)
+
+FILE = Path(__file__).resolve()
+ROOT = FILE.parents[3]  # YOLOv5 root directory
+if str(ROOT) not in sys.path:
+    sys.path.append(str(ROOT))  # add ROOT to PATH
+
+try:
+    import comet_ml
+
+    # Project Configuration
+    config = comet_ml.config.get_config()
+    COMET_PROJECT_NAME = config.get_string(os.getenv('COMET_PROJECT_NAME'), 'comet.project_name', default='yolov5')
+except (ModuleNotFoundError, ImportError):
+    comet_ml = None
+    COMET_PROJECT_NAME = None
+
+import PIL
+import torch
+import torchvision.transforms as T
+import yaml
+
+from utils.dataloaders import img2label_paths
+from utils.general import check_dataset, scale_boxes, xywh2xyxy
+from utils.metrics import box_iou
+
+COMET_PREFIX = 'comet://'
+
+COMET_MODE = os.getenv('COMET_MODE', 'online')
+
+# Model Saving Settings
+COMET_MODEL_NAME = os.getenv('COMET_MODEL_NAME', 'yolov5')
+
+# Dataset Artifact Settings
+COMET_UPLOAD_DATASET = os.getenv('COMET_UPLOAD_DATASET', 'false').lower() == 'true'
+
+# Evaluation Settings
+COMET_LOG_CONFUSION_MATRIX = os.getenv('COMET_LOG_CONFUSION_MATRIX', 'true').lower() == 'true'
+COMET_LOG_PREDICTIONS = os.getenv('COMET_LOG_PREDICTIONS', 'true').lower() == 'true'
+COMET_MAX_IMAGE_UPLOADS = int(os.getenv('COMET_MAX_IMAGE_UPLOADS', 100))
+
+# Confusion Matrix Settings
+CONF_THRES = float(os.getenv('CONF_THRES', 0.001))
+IOU_THRES = float(os.getenv('IOU_THRES', 0.6))
+
+# Batch Logging Settings
+COMET_LOG_BATCH_METRICS = os.getenv('COMET_LOG_BATCH_METRICS', 'false').lower() == 'true'
+COMET_BATCH_LOGGING_INTERVAL = os.getenv('COMET_BATCH_LOGGING_INTERVAL', 1)
+COMET_PREDICTION_LOGGING_INTERVAL = os.getenv('COMET_PREDICTION_LOGGING_INTERVAL', 1)
+COMET_LOG_PER_CLASS_METRICS = os.getenv('COMET_LOG_PER_CLASS_METRICS', 'false').lower() == 'true'
+
+RANK = int(os.getenv('RANK', -1))
+
+to_pil = T.ToPILImage()
+
+
+class CometLogger:
+    """Log metrics, parameters, source code, models and much more
+    with Comet
+    """
+
+    def __init__(self, opt, hyp, run_id=None, job_type='Training', **experiment_kwargs) -> None:
+        self.job_type = job_type
+        self.opt = opt
+        self.hyp = hyp
+
+        # Comet Flags
+        self.comet_mode = COMET_MODE
+
+        self.save_model = opt.save_period > -1
+        self.model_name = COMET_MODEL_NAME
+
+        # Batch Logging Settings
+        self.log_batch_metrics = COMET_LOG_BATCH_METRICS
+        self.comet_log_batch_interval = COMET_BATCH_LOGGING_INTERVAL
+
+        # Dataset Artifact Settings
+        self.upload_dataset = self.opt.upload_dataset if self.opt.upload_dataset else COMET_UPLOAD_DATASET
+        self.resume = self.opt.resume
+
+        # Default parameters to pass to Experiment objects
+        self.default_experiment_kwargs = {
+            'log_code': False,
+            'log_env_gpu': True,
+            'log_env_cpu': True,
+            'project_name': COMET_PROJECT_NAME,}
+        self.default_experiment_kwargs.update(experiment_kwargs)
+        self.experiment = self._get_experiment(self.comet_mode, run_id)
+
+        self.data_dict = self.check_dataset(self.opt.data)
+        self.class_names = self.data_dict['names']
+        self.num_classes = self.data_dict['nc']
+
+        self.logged_images_count = 0
+        self.max_images = COMET_MAX_IMAGE_UPLOADS
+
+        if run_id is None:
+            self.experiment.log_other('Created from', 'YOLOv5')
+            if not isinstance(self.experiment, comet_ml.OfflineExperiment):
+                workspace, project_name, experiment_id = self.experiment.url.split('/')[-3:]
+                self.experiment.log_other(
+                    'Run Path',
+                    f'{workspace}/{project_name}/{experiment_id}',
+                )
+            self.log_parameters(vars(opt))
+            self.log_parameters(self.opt.hyp)
+            self.log_asset_data(
+                self.opt.hyp,
+                name='hyperparameters.json',
+                metadata={'type': 'hyp-config-file'},
+            )
+            self.log_asset(
+                f'{self.opt.save_dir}/opt.yaml',
+                metadata={'type': 'opt-config-file'},
+            )
+
+        self.comet_log_confusion_matrix = COMET_LOG_CONFUSION_MATRIX
+
+        if hasattr(self.opt, 'conf_thres'):
+            self.conf_thres = self.opt.conf_thres
+        else:
+            self.conf_thres = CONF_THRES
+        if hasattr(self.opt, 'iou_thres'):
+            self.iou_thres = self.opt.iou_thres
+        else:
+            self.iou_thres = IOU_THRES
+
+        self.log_parameters({'val_iou_threshold': self.iou_thres, 'val_conf_threshold': self.conf_thres})
+
+        self.comet_log_predictions = COMET_LOG_PREDICTIONS
+        if self.opt.bbox_interval == -1:
+            self.comet_log_prediction_interval = 1 if self.opt.epochs < 10 else self.opt.epochs // 10
+        else:
+            self.comet_log_prediction_interval = self.opt.bbox_interval
+
+        if self.comet_log_predictions:
+            self.metadata_dict = {}
+            self.logged_image_names = []
+
+        self.comet_log_per_class_metrics = COMET_LOG_PER_CLASS_METRICS
+
+        self.experiment.log_others({
+            'comet_mode': COMET_MODE,
+            'comet_max_image_uploads': COMET_MAX_IMAGE_UPLOADS,
+            'comet_log_per_class_metrics': COMET_LOG_PER_CLASS_METRICS,
+            'comet_log_batch_metrics': COMET_LOG_BATCH_METRICS,
+            'comet_log_confusion_matrix': COMET_LOG_CONFUSION_MATRIX,
+            'comet_model_name': COMET_MODEL_NAME,})
+
+        # Check if running the Experiment with the Comet Optimizer
+        if hasattr(self.opt, 'comet_optimizer_id'):
+            self.experiment.log_other('optimizer_id', self.opt.comet_optimizer_id)
+            self.experiment.log_other('optimizer_objective', self.opt.comet_optimizer_objective)
+            self.experiment.log_other('optimizer_metric', self.opt.comet_optimizer_metric)
+            self.experiment.log_other('optimizer_parameters', json.dumps(self.hyp))
+
+    def _get_experiment(self, mode, experiment_id=None):
+        if mode == 'offline':
+            if experiment_id is not None:
+                return comet_ml.ExistingOfflineExperiment(
+                    previous_experiment=experiment_id,
+                    **self.default_experiment_kwargs,
+                )
+
+            return comet_ml.OfflineExperiment(**self.default_experiment_kwargs,)
+
+        else:
+            try:
+                if experiment_id is not None:
+                    return comet_ml.ExistingExperiment(
+                        previous_experiment=experiment_id,
+                        **self.default_experiment_kwargs,
+                    )
+
+                return comet_ml.Experiment(**self.default_experiment_kwargs)
+
+            except ValueError:
+                logger.warning('COMET WARNING: '
+                               'Comet credentials have not been set. '
+                               'Comet will default to offline logging. '
+                               'Please set your credentials to enable online logging.')
+                return self._get_experiment('offline', experiment_id)
+
+        return
+
+    def log_metrics(self, log_dict, **kwargs):
+        self.experiment.log_metrics(log_dict, **kwargs)
+
+    def log_parameters(self, log_dict, **kwargs):
+        self.experiment.log_parameters(log_dict, **kwargs)
+
+    def log_asset(self, asset_path, **kwargs):
+        self.experiment.log_asset(asset_path, **kwargs)
+
+    def log_asset_data(self, asset, **kwargs):
+        self.experiment.log_asset_data(asset, **kwargs)
+
+    def log_image(self, img, **kwargs):
+        self.experiment.log_image(img, **kwargs)
+
+    def log_model(self, path, opt, epoch, fitness_score, best_model=False):
+        if not self.save_model:
+            return
+
+        model_metadata = {
+            'fitness_score': fitness_score[-1],
+            'epochs_trained': epoch + 1,
+            'save_period': opt.save_period,
+            'total_epochs': opt.epochs,}
+
+        model_files = glob.glob(f'{path}/*.pt')
+        for model_path in model_files:
+            name = Path(model_path).name
+
+            self.experiment.log_model(
+                self.model_name,
+                file_or_folder=model_path,
+                file_name=name,
+                metadata=model_metadata,
+                overwrite=True,
+            )
+
+    def check_dataset(self, data_file):
+        with open(data_file) as f:
+            data_config = yaml.safe_load(f)
+
+        if data_config['path'].startswith(COMET_PREFIX):
+            path = data_config['path'].replace(COMET_PREFIX, '')
+            data_dict = self.download_dataset_artifact(path)
+
+            return data_dict
+
+        self.log_asset(self.opt.data, metadata={'type': 'data-config-file'})
+
+        return check_dataset(data_file)
+
+    def log_predictions(self, image, labelsn, path, shape, predn):
+        if self.logged_images_count >= self.max_images:
+            return
+        detections = predn[predn[:, 4] > self.conf_thres]
+        iou = box_iou(labelsn[:, 1:], detections[:, :4])
+        mask, _ = torch.where(iou > self.iou_thres)
+        if len(mask) == 0:
+            return
+
+        filtered_detections = detections[mask]
+        filtered_labels = labelsn[mask]
+
+        image_id = path.split('/')[-1].split('.')[0]
+        image_name = f'{image_id}_curr_epoch_{self.experiment.curr_epoch}'
+        if image_name not in self.logged_image_names:
+            native_scale_image = PIL.Image.open(path)
+            self.log_image(native_scale_image, name=image_name)
+            self.logged_image_names.append(image_name)
+
+        metadata = []
+        for cls, *xyxy in filtered_labels.tolist():
+            metadata.append({
+                'label': f'{self.class_names[int(cls)]}-gt',
+                'score': 100,
+                'box': {
+                    'x': xyxy[0],
+                    'y': xyxy[1],
+                    'x2': xyxy[2],
+                    'y2': xyxy[3]},})
+        for *xyxy, conf, cls in filtered_detections.tolist():
+            metadata.append({
+                'label': f'{self.class_names[int(cls)]}',
+                'score': conf * 100,
+                'box': {
+                    'x': xyxy[0],
+                    'y': xyxy[1],
+                    'x2': xyxy[2],
+                    'y2': xyxy[3]},})
+
+        self.metadata_dict[image_name] = metadata
+        self.logged_images_count += 1
+
+        return
+
+    def preprocess_prediction(self, image, labels, shape, pred):
+        nl, _ = labels.shape[0], pred.shape[0]
+
+        # Predictions
+        if self.opt.single_cls:
+            pred[:, 5] = 0
+
+        predn = pred.clone()
+        scale_boxes(image.shape[1:], predn[:, :4], shape[0], shape[1])
+
+        labelsn = None
+        if nl:
+            tbox = xywh2xyxy(labels[:, 1:5])  # target boxes
+            scale_boxes(image.shape[1:], tbox, shape[0], shape[1])  # native-space labels
+            labelsn = torch.cat((labels[:, 0:1], tbox), 1)  # native-space labels
+            scale_boxes(image.shape[1:], predn[:, :4], shape[0], shape[1])  # native-space pred
+
+        return predn, labelsn
+
+    def add_assets_to_artifact(self, artifact, path, asset_path, split):
+        img_paths = sorted(glob.glob(f'{asset_path}/*'))
+        label_paths = img2label_paths(img_paths)
+
+        for image_file, label_file in zip(img_paths, label_paths):
+            image_logical_path, label_logical_path = map(lambda x: os.path.relpath(x, path), [image_file, label_file])
+
+            try:
+                artifact.add(image_file, logical_path=image_logical_path, metadata={'split': split})
+                artifact.add(label_file, logical_path=label_logical_path, metadata={'split': split})
+            except ValueError as e:
+                logger.error('COMET ERROR: Error adding file to Artifact. Skipping file.')
+                logger.error(f'COMET ERROR: {e}')
+                continue
+
+        return artifact
+
+    def upload_dataset_artifact(self):
+        dataset_name = self.data_dict.get('dataset_name', 'yolov5-dataset')
+        path = str((ROOT / Path(self.data_dict['path'])).resolve())
+
+        metadata = self.data_dict.copy()
+        for key in ['train', 'val', 'test']:
+            split_path = metadata.get(key)
+            if split_path is not None:
+                metadata[key] = split_path.replace(path, '')
+
+        artifact = comet_ml.Artifact(name=dataset_name, artifact_type='dataset', metadata=metadata)
+        for key in metadata.keys():
+            if key in ['train', 'val', 'test']:
+                if isinstance(self.upload_dataset, str) and (key != self.upload_dataset):
+                    continue
+
+                asset_path = self.data_dict.get(key)
+                if asset_path is not None:
+                    artifact = self.add_assets_to_artifact(artifact, path, asset_path, key)
+
+        self.experiment.log_artifact(artifact)
+
+        return
+
+    def download_dataset_artifact(self, artifact_path):
+        logged_artifact = self.experiment.get_artifact(artifact_path)
+        artifact_save_dir = str(Path(self.opt.save_dir) / logged_artifact.name)
+        logged_artifact.download(artifact_save_dir)
+
+        metadata = logged_artifact.metadata
+        data_dict = metadata.copy()
+        data_dict['path'] = artifact_save_dir
+
+        metadata_names = metadata.get('names')
+        if type(metadata_names) == dict:
+            data_dict['names'] = {int(k): v for k, v in metadata.get('names').items()}
+        elif type(metadata_names) == list:
+            data_dict['names'] = {int(k): v for k, v in zip(range(len(metadata_names)), metadata_names)}
+        else:
+            raise "Invalid 'names' field in dataset yaml file. Please use a list or dictionary"
+
+        data_dict = self.update_data_paths(data_dict)
+        return data_dict
+
+    def update_data_paths(self, data_dict):
+        path = data_dict.get('path', '')
+
+        for split in ['train', 'val', 'test']:
+            if data_dict.get(split):
+                split_path = data_dict.get(split)
+                data_dict[split] = (f'{path}/{split_path}' if isinstance(split, str) else [
+                    f'{path}/{x}' for x in split_path])
+
+        return data_dict
+
+    def on_pretrain_routine_end(self, paths):
+        if self.opt.resume:
+            return
+
+        for path in paths:
+            self.log_asset(str(path))
+
+        if self.upload_dataset:
+            if not self.resume:
+                self.upload_dataset_artifact()
+
+        return
+
+    def on_train_start(self):
+        self.log_parameters(self.hyp)
+
+    def on_train_epoch_start(self):
+        return
+
+    def on_train_epoch_end(self, epoch):
+        self.experiment.curr_epoch = epoch
+
+        return
+
+    def on_train_batch_start(self):
+        return
+
+    def on_train_batch_end(self, log_dict, step):
+        self.experiment.curr_step = step
+        if self.log_batch_metrics and (step % self.comet_log_batch_interval == 0):
+            self.log_metrics(log_dict, step=step)
+
+        return
+
+    def on_train_end(self, files, save_dir, last, best, epoch, results):
+        if self.comet_log_predictions:
+            curr_epoch = self.experiment.curr_epoch
+            self.experiment.log_asset_data(self.metadata_dict, 'image-metadata.json', epoch=curr_epoch)
+
+        for f in files:
+            self.log_asset(f, metadata={'epoch': epoch})
+        self.log_asset(f'{save_dir}/results.csv', metadata={'epoch': epoch})
+
+        if not self.opt.evolve:
+            model_path = str(best if best.exists() else last)
+            name = Path(model_path).name
+            if self.save_model:
+                self.experiment.log_model(
+                    self.model_name,
+                    file_or_folder=model_path,
+                    file_name=name,
+                    overwrite=True,
+                )
+
+        # Check if running Experiment with Comet Optimizer
+        if hasattr(self.opt, 'comet_optimizer_id'):
+            metric = results.get(self.opt.comet_optimizer_metric)
+            self.experiment.log_other('optimizer_metric_value', metric)
+
+        self.finish_run()
+
+    def on_val_start(self):
+        return
+
+    def on_val_batch_start(self):
+        return
+
+    def on_val_batch_end(self, batch_i, images, targets, paths, shapes, outputs):
+        if not (self.comet_log_predictions and ((batch_i + 1) % self.comet_log_prediction_interval == 0)):
+            return
+
+        for si, pred in enumerate(outputs):
+            if len(pred) == 0:
+                continue
+
+            image = images[si]
+            labels = targets[targets[:, 0] == si, 1:]
+            shape = shapes[si]
+            path = paths[si]
+            predn, labelsn = self.preprocess_prediction(image, labels, shape, pred)
+            if labelsn is not None:
+                self.log_predictions(image, labelsn, path, shape, predn)
+
+        return
+
+    def on_val_end(self, nt, tp, fp, p, r, f1, ap, ap50, ap_class, confusion_matrix):
+        if self.comet_log_per_class_metrics:
+            if self.num_classes > 1:
+                for i, c in enumerate(ap_class):
+                    class_name = self.class_names[c]
+                    self.experiment.log_metrics(
+                        {
+                            'mAP@.5': ap50[i],
+                            'mAP@.5:.95': ap[i],
+                            'precision': p[i],
+                            'recall': r[i],
+                            'f1': f1[i],
+                            'true_positives': tp[i],
+                            'false_positives': fp[i],
+                            'support': nt[c]},
+                        prefix=class_name)
+
+        if self.comet_log_confusion_matrix:
+            epoch = self.experiment.curr_epoch
+            class_names = list(self.class_names.values())
+            class_names.append('background')
+            num_classes = len(class_names)
+
+            self.experiment.log_confusion_matrix(
+                matrix=confusion_matrix.matrix,
+                max_categories=num_classes,
+                labels=class_names,
+                epoch=epoch,
+                column_label='Actual Category',
+                row_label='Predicted Category',
+                file_name=f'confusion-matrix-epoch-{epoch}.json',
+            )
+
+    def on_fit_epoch_end(self, result, epoch):
+        self.log_metrics(result, epoch=epoch)
+
+    def on_model_save(self, last, epoch, final_epoch, best_fitness, fi):
+        if ((epoch + 1) % self.opt.save_period == 0 and not final_epoch) and self.opt.save_period != -1:
+            self.log_model(last.parent, self.opt, epoch, fi, best_model=best_fitness == fi)
+
+    def on_params_update(self, params):
+        self.log_parameters(params)
+
+    def finish_run(self):
+        self.experiment.end()