jnhg_ai_code/utils/loggers/comet/comet_utils.py

import logging
import os
from urllib.parse import urlparse

try:
    import comet_ml
except (ModuleNotFoundError, ImportError):
    comet_ml = None

import yaml

logger = logging.getLogger(__name__)

COMET_PREFIX = 'comet://'
COMET_MODEL_NAME = os.getenv('COMET_MODEL_NAME', 'yolov5')
COMET_DEFAULT_CHECKPOINT_FILENAME = os.getenv('COMET_DEFAULT_CHECKPOINT_FILENAME', 'last.pt')


def download_model_checkpoint(opt, experiment):
    model_dir = f'{opt.project}/{experiment.name}'
    os.makedirs(model_dir, exist_ok=True)

    model_name = COMET_MODEL_NAME
    model_asset_list = experiment.get_model_asset_list(model_name)

    if len(model_asset_list) == 0:
        logger.error(f'COMET ERROR: No checkpoints found for model name : {model_name}')
        return

    model_asset_list = sorted(
        model_asset_list,
        key=lambda x: x['step'],
        reverse=True,
    )
    logged_checkpoint_map = {asset['fileName']: asset['assetId'] for asset in model_asset_list}

    resource_url = urlparse(opt.weights)
    checkpoint_filename = resource_url.query

    if checkpoint_filename:
        asset_id = logged_checkpoint_map.get(checkpoint_filename)
    else:
        asset_id = logged_checkpoint_map.get(COMET_DEFAULT_CHECKPOINT_FILENAME)
        checkpoint_filename = COMET_DEFAULT_CHECKPOINT_FILENAME

    if asset_id is None:
        logger.error(f'COMET ERROR: Checkpoint {checkpoint_filename} not found in the given Experiment')
        return

    try:
        logger.info(f'COMET INFO: Downloading checkpoint {checkpoint_filename}')
        asset_filename = checkpoint_filename

        model_binary = experiment.get_asset(asset_id, return_type='binary', stream=False)
        model_download_path = f'{model_dir}/{asset_filename}'
        with open(model_download_path, 'wb') as f:
            f.write(model_binary)

        opt.weights = model_download_path

    except Exception as e:
        logger.warning('COMET WARNING: Unable to download checkpoint from Comet')
        logger.exception(e)


def set_opt_parameters(opt, experiment):
    """Update the opts Namespace with parameters
    from Comet's ExistingExperiment when resuming a run

    Args:
        opt (argparse.Namespace): Namespace of command line options
        experiment (comet_ml.APIExperiment): Comet API Experiment object
    """
    asset_list = experiment.get_asset_list()
    resume_string = opt.resume

    for asset in asset_list:
        if asset['fileName'] == 'opt.yaml':
            asset_id = asset['assetId']
            asset_binary = experiment.get_asset(asset_id, return_type='binary', stream=False)
            opt_dict = yaml.safe_load(asset_binary)
            for key, value in opt_dict.items():
                setattr(opt, key, value)
            opt.resume = resume_string

    # Save hyperparameters to YAML file
    # Necessary to pass checks in training script
    save_dir = f'{opt.project}/{experiment.name}'
    os.makedirs(save_dir, exist_ok=True)

    hyp_yaml_path = f'{save_dir}/hyp.yaml'
    with open(hyp_yaml_path, 'w') as f:
        yaml.dump(opt.hyp, f)
    opt.hyp = hyp_yaml_path


def check_comet_weights(opt):
    """Downloads model weights from Comet and updates the
    weights path to point to saved weights location

    Args:
        opt (argparse.Namespace): Command Line arguments passed
            to YOLOv5 training script

    Returns:
        None/bool: Return True if weights are successfully downloaded
            else return None
    """
    if comet_ml is None:
        return

    if isinstance(opt.weights, str):
        if opt.weights.startswith(COMET_PREFIX):
            api = comet_ml.API()
            resource = urlparse(opt.weights)
            experiment_path = f'{resource.netloc}{resource.path}'
            experiment = api.get(experiment_path)
            download_model_checkpoint(opt, experiment)
            return True

    return None


def check_comet_resume(opt):
    """Restores run parameters to its original state based on the model checkpoint
    and logged Experiment parameters.

    Args:
        opt (argparse.Namespace): Command Line arguments passed
            to YOLOv5 training script

    Returns:
        None/bool: Return True if the run is restored successfully
            else return None
    """
    if comet_ml is None:
        return

    if isinstance(opt.resume, str):
        if opt.resume.startswith(COMET_PREFIX):
            api = comet_ml.API()
            resource = urlparse(opt.resume)
            experiment_path = f'{resource.netloc}{resource.path}'
            experiment = api.get(experiment_path)
            set_opt_parameters(opt, experiment)
            download_model_checkpoint(opt, experiment)

            return True

    return None