sczhou/CodeFormer

1import os2import re3import random4import time5import torch6import numpy as np7from os import path as osp8 9from .dist_util import master_only10from .logger import get_root_logger11 12IS_HIGH_VERSION = [int(m) for m in list(re.findall(r"^([0-9]+)\.([0-9]+)\.([0-9]+)([^0-9][a-zA-Z0-9]*)?(\+git.*)?$",\13    torch.__version__)[0][:3])] >= [1, 12, 0]14 15def gpu_is_available():16    if IS_HIGH_VERSION:17        if torch.backends.mps.is_available():18            return True19    return True if torch.cuda.is_available() and torch.backends.cudnn.is_available() else False20 21def get_device(gpu_id=None):22    if gpu_id is None:23        gpu_str = ''24    elif isinstance(gpu_id, int):25        gpu_str = f':{gpu_id}'26    else:27        raise TypeError('Input should be int value.')28 29    if IS_HIGH_VERSION:30        if torch.backends.mps.is_available():31            return torch.device('mps'+gpu_str)32    return torch.device('cuda'+gpu_str if torch.cuda.is_available() and torch.backends.cudnn.is_available() else 'cpu')33 34 35def set_random_seed(seed):36    """Set random seeds."""37    random.seed(seed)38    np.random.seed(seed)39    torch.manual_seed(seed)40    torch.cuda.manual_seed(seed)41    torch.cuda.manual_seed_all(seed)42 43 44def get_time_str():45    return time.strftime('%Y%m%d_%H%M%S', time.localtime())46 47 48def mkdir_and_rename(path):49    """mkdirs. If path exists, rename it with timestamp and create a new one.50 51    Args:52        path (str): Folder path.53    """54    if osp.exists(path):55        new_name = path + '_archived_' + get_time_str()56        print(f'Path already exists. Rename it to {new_name}', flush=True)57        os.rename(path, new_name)58    os.makedirs(path, exist_ok=True)59 60 61@master_only62def make_exp_dirs(opt):63    """Make dirs for experiments."""64    path_opt = opt['path'].copy()65    if opt['is_train']:66        mkdir_and_rename(path_opt.pop('experiments_root'))67    else:68        mkdir_and_rename(path_opt.pop('results_root'))69    for key, path in path_opt.items():70        if ('strict_load' not in key) and ('pretrain_network' not in key) and ('resume' not in key):71            os.makedirs(path, exist_ok=True)72 73 74def scandir(dir_path, suffix=None, recursive=False, full_path=False):75    """Scan a directory to find the interested files.76 77    Args:78        dir_path (str): Path of the directory.79        suffix (str | tuple(str), optional): File suffix that we are80            interested in. Default: None.81        recursive (bool, optional): If set to True, recursively scan the82            directory. Default: False.83        full_path (bool, optional): If set to True, include the dir_path.84            Default: False.85 86    Returns:87        A generator for all the interested files with relative pathes.88    """89 90    if (suffix is not None) and not isinstance(suffix, (str, tuple)):91        raise TypeError('"suffix" must be a string or tuple of strings')92 93    root = dir_path94 95    def _scandir(dir_path, suffix, recursive):96        for entry in os.scandir(dir_path):97            if not entry.name.startswith('.') and entry.is_file():98                if full_path:99                    return_path = entry.path100                else:101                    return_path = osp.relpath(entry.path, root)102 103                if suffix is None:104                    yield return_path105                elif return_path.endswith(suffix):106                    yield return_path107            else:108                if recursive:109                    yield from _scandir(entry.path, suffix=suffix, recursive=recursive)110                else:111                    continue112 113    return _scandir(dir_path, suffix=suffix, recursive=recursive)114 115 116def check_resume(opt, resume_iter):117    """Check resume states and pretrain_network paths.118 119    Args:120        opt (dict): Options.121        resume_iter (int): Resume iteration.122    """123    logger = get_root_logger()124    if opt['path']['resume_state']:125        # get all the networks126        networks = [key for key in opt.keys() if key.startswith('network_')]127        flag_pretrain = False128        for network in networks:129            if opt['path'].get(f'pretrain_{network}') is not None:130                flag_pretrain = True131        if flag_pretrain:132            logger.warning('pretrain_network path will be ignored during resuming.')133        # set pretrained model paths134        for network in networks:135            name = f'pretrain_{network}'136            basename = network.replace('network_', '')137            if opt['path'].get('ignore_resume_networks') is None or (basename138                                                                     not in opt['path']['ignore_resume_networks']):139                opt['path'][name] = osp.join(opt['path']['models'], f'net_{basename}_{resume_iter}.pth')140                logger.info(f"Set {name} to {opt['path'][name]}")141 142 143def sizeof_fmt(size, suffix='B'):144    """Get human readable file size.145 146    Args:147        size (int): File size.148        suffix (str): Suffix. Default: 'B'.149 150    Return:151        str: Formated file siz.152    """153    for unit in ['', 'K', 'M', 'G', 'T', 'P', 'E', 'Z']:154        if abs(size) < 1024.0:155            return f'{size:3.1f} {unit}{suffix}'156        size /= 1024.0157    return f'{size:3.1f} Y{suffix}'