(cfg, type)
| 61 | |
| 62 | |
| 63 | def main(cfg, type): |
| 64 | _cfg = cfg.cfg |
| 65 | logger = get_logger() |
| 66 | global_config = _cfg['Global'] |
| 67 | export_config = _cfg['Export'] |
| 68 | # build post process |
| 69 | post_process_class = build_post_process(_cfg['PostProcess']) |
| 70 | char_num = len(getattr(post_process_class, 'character')) |
| 71 | cfg['Architecture']['Decoder']['out_channels'] = char_num |
| 72 | model = build_model(_cfg['Architecture']) |
| 73 | |
| 74 | load_ckpt(model, _cfg) |
| 75 | model.eval() |
| 76 | |
| 77 | export_dir = export_config.get('export_dir', '') |
| 78 | if not export_dir: |
| 79 | export_dir = os.path.join(global_config.get('output_dir', 'output'), |
| 80 | 'export') |
| 81 | |
| 82 | if _cfg['Architecture']['algorithm'] in ['Distillation' |
| 83 | ]: # distillation model |
| 84 | _cfg['PostProcess'][ |
| 85 | 'name'] = post_process_class.__class__.__base__.__name__ |
| 86 | for model_name in model.model_list: |
| 87 | sub_model_save_path = os.path.join(export_dir, model_name) |
| 88 | export_single_model( |
| 89 | model.model_list[model_name], |
| 90 | _cfg, |
| 91 | sub_model_save_path, |
| 92 | export_config, |
| 93 | logger, |
| 94 | type, |
| 95 | ) |
| 96 | else: |
| 97 | export_single_model(model, _cfg, export_dir, export_config, logger, |
| 98 | type) |
| 99 | |
| 100 | |
| 101 | def parse_args(): |
no test coverage detected