radames/Text2Human-API

1import os2import os.path as osp3from collections import OrderedDict4 5import yaml6 7 8def ordered_yaml():9    """Support OrderedDict for yaml.10 11    Returns:12        yaml Loader and Dumper.13    """14    try:15        from yaml import CDumper as Dumper16        from yaml import CLoader as Loader17    except ImportError:18        from yaml import Dumper, Loader19 20    _mapping_tag = yaml.resolver.BaseResolver.DEFAULT_MAPPING_TAG21 22    def dict_representer(dumper, data):23        return dumper.represent_dict(data.items())24 25    def dict_constructor(loader, node):26        return OrderedDict(loader.construct_pairs(node))27 28    Dumper.add_representer(OrderedDict, dict_representer)29    Loader.add_constructor(_mapping_tag, dict_constructor)30    return Loader, Dumper31 32 33def parse(opt_path, is_train=True):34    """Parse option file.35 36    Args:37        opt_path (str): Option file path.38        is_train (str): Indicate whether in training or not. Default: True.39 40    Returns:41        (dict): Options.42    """43    with open(opt_path, mode='r') as f:44        Loader, _ = ordered_yaml()45        opt = yaml.load(f, Loader=Loader)46 47    gpu_list = ','.join(str(x) for x in opt['gpu_ids'])48    if opt.get('set_CUDA_VISIBLE_DEVICES', None):49        os.environ['CUDA_VISIBLE_DEVICES'] = gpu_list50        print('export CUDA_VISIBLE_DEVICES=' + gpu_list, flush=True)51    else:52        print('gpu_list: ', gpu_list, flush=True)53 54    opt['is_train'] = is_train55 56    # paths57    opt['path'] = {}58    opt['path']['root'] = osp.abspath(59        osp.join(__file__, osp.pardir, osp.pardir))60    if is_train:61        experiments_root = osp.join(opt['path']['root'], 'experiments',62                                    opt['name'])63        opt['path']['experiments_root'] = experiments_root64        opt['path']['models'] = osp.join(experiments_root, 'models')65        opt['path']['log'] = experiments_root66        opt['path']['visualization'] = osp.join(experiments_root,67                                                'visualization')68 69        # change some options for debug mode70        if 'debug' in opt['name']:71            opt['debug'] = True72            opt['val_freq'] = 173            opt['print_freq'] = 174            opt['save_checkpoint_freq'] = 175    else:  # test76        results_root = osp.join(opt['path']['root'], 'results', opt['name'])77        opt['path']['results_root'] = results_root78        opt['path']['log'] = results_root79        opt['path']['visualization'] = osp.join(results_root, 'visualization')80 81    return opt82 83 84def dict2str(opt, indent_level=1):85    """dict to string for printing options.86 87    Args:88        opt (dict): Option dict.89        indent_level (int): Indent level. Default: 1.90 91    Return:92        (str): Option string for printing.93    """94    msg = ''95    for k, v in opt.items():96        if isinstance(v, dict):97            msg += ' ' * (indent_level * 2) + k + ':[\n'98            msg += dict2str(v, indent_level + 1)99            msg += ' ' * (indent_level * 2) + ']\n'100        else:101            msg += ' ' * (indent_level * 2) + k + ': ' + str(v) + '\n'102    return msg103 104 105class NoneDict(dict):106    """None dict. It will return none if key is not in the dict."""107 108    def __missing__(self, key):109        return None110 111 112def dict_to_nonedict(opt):113    """Convert to NoneDict, which returns None for missing keys.114 115    Args:116        opt (dict): Option dict.117 118    Returns:119        (dict): NoneDict for options.120    """121    if isinstance(opt, dict):122        new_opt = dict()123        for key, sub_opt in opt.items():124            new_opt[key] = dict_to_nonedict(sub_opt)125        return NoneDict(**new_opt)126    elif isinstance(opt, list):127        return [dict_to_nonedict(sub_opt) for sub_opt in opt]128    else:129        return opt130