faisalhr1997/codeformer

1import numpy as np2import os3import random4import time5import torch6from os import path as osp7 8from .dist_util import master_only9from .logger import get_root_logger10 11 12def set_random_seed(seed):13    """Set random seeds."""14    random.seed(seed)15    np.random.seed(seed)16    torch.manual_seed(seed)17    torch.cuda.manual_seed(seed)18    torch.cuda.manual_seed_all(seed)19 20 21def get_time_str():22    return time.strftime('%Y%m%d_%H%M%S', time.localtime())23 24 25def mkdir_and_rename(path):26    """mkdirs. If path exists, rename it with timestamp and create a new one.27 28    Args:29        path (str): Folder path.30    """31    if osp.exists(path):32        new_name = path + '_archived_' + get_time_str()33        print(f'Path already exists. Rename it to {new_name}', flush=True)34        os.rename(path, new_name)35    os.makedirs(path, exist_ok=True)36 37 38@master_only39def make_exp_dirs(opt):40    """Make dirs for experiments."""41    path_opt = opt['path'].copy()42    if opt['is_train']:43        mkdir_and_rename(path_opt.pop('experiments_root'))44    else:45        mkdir_and_rename(path_opt.pop('results_root'))46    for key, path in path_opt.items():47        if ('strict_load' not in key) and ('pretrain_network' not in key) and ('resume' not in key):48            os.makedirs(path, exist_ok=True)49 50 51def scandir(dir_path, suffix=None, recursive=False, full_path=False):52    """Scan a directory to find the interested files.53 54    Args:55        dir_path (str): Path of the directory.56        suffix (str | tuple(str), optional): File suffix that we are57            interested in. Default: None.58        recursive (bool, optional): If set to True, recursively scan the59            directory. Default: False.60        full_path (bool, optional): If set to True, include the dir_path.61            Default: False.62 63    Returns:64        A generator for all the interested files with relative pathes.65    """66 67    if (suffix is not None) and not isinstance(suffix, (str, tuple)):68        raise TypeError('"suffix" must be a string or tuple of strings')69 70    root = dir_path71 72    def _scandir(dir_path, suffix, recursive):73        for entry in os.scandir(dir_path):74            if not entry.name.startswith('.') and entry.is_file():75                if full_path:76                    return_path = entry.path77                else:78                    return_path = osp.relpath(entry.path, root)79 80                if suffix is None:81                    yield return_path82                elif return_path.endswith(suffix):83                    yield return_path84            else:85                if recursive:86                    yield from _scandir(entry.path, suffix=suffix, recursive=recursive)87                else:88                    continue89 90    return _scandir(dir_path, suffix=suffix, recursive=recursive)91 92 93def check_resume(opt, resume_iter):94    """Check resume states and pretrain_network paths.95 96    Args:97        opt (dict): Options.98        resume_iter (int): Resume iteration.99    """100    logger = get_root_logger()101    if opt['path']['resume_state']:102        # get all the networks103        networks = [key for key in opt.keys() if key.startswith('network_')]104        flag_pretrain = False105        for network in networks:106            if opt['path'].get(f'pretrain_{network}') is not None:107                flag_pretrain = True108        if flag_pretrain:109            logger.warning('pretrain_network path will be ignored during resuming.')110        # set pretrained model paths111        for network in networks:112            name = f'pretrain_{network}'113            basename = network.replace('network_', '')114            if opt['path'].get('ignore_resume_networks') is None or (basename115                                                                     not in opt['path']['ignore_resume_networks']):116                opt['path'][name] = osp.join(opt['path']['models'], f'net_{basename}_{resume_iter}.pth')117                logger.info(f"Set {name} to {opt['path'][name]}")118 119 120def sizeof_fmt(size, suffix='B'):121    """Get human readable file size.122 123    Args:124        size (int): File size.125        suffix (str): Suffix. Default: 'B'.126 127    Return:128        str: Formated file siz.129    """130    for unit in ['', 'K', 'M', 'G', 'T', 'P', 'E', 'Z']:131        if abs(size) < 1024.0:132            return f'{size:3.1f} {unit}{suffix}'133        size /= 1024.0134    return f'{size:3.1f} Y{suffix}'135