def parse_args(): # Parse command line arguments parser = argparse.ArgumentParser() default_model_prefix = os.path.join(data_dir(), 'models', 'model-v{}'.format(__version__)) parser.add_argument("--dataset", help="use which kind of dataset, captcha or cn_ocr", choices=['captcha', 'cn_ocr'], type=str, default='captcha') parser.add_argument("--data_root", help="Path to image files", type=str, default='/Users/king/Documents/WhatIHaveDone/Test/text_renderer/output/wechat_simulator') parser.add_argument("--train_file", help="Path to train txt file", type=str, default='/Users/king/Documents/WhatIHaveDone/Test/text_renderer/output/wechat_simulator/train.txt') parser.add_argument("--test_file", help="Path to test txt file", type=str, default='/Users/king/Documents/WhatIHaveDone/Test/text_renderer/output/wechat_simulator/test.txt') parser.add_argument("--cpu", help="Number of CPUs for training [Default 8]. Ignored if --gpu is specified.", type=int, default=2) parser.add_argument("--gpu", help="Number of GPUs for training [Default 0]", type=int) parser.add_argument('--load_epoch', type=int, help='load the model on an epoch using the model-load-prefix [Default: no trained model will be loaded]') parser.add_argument("--prefix", help="Checkpoint prefix [Default '{}']".format(default_model_prefix), default=default_model_prefix) parser.add_argument("--loss", help="'ctc' or 'warpctc' loss [Default 'ctc']", default='ctc') parser.add_argument("--num_proc", help="Number CAPTCHA generating processes [Default 4]", type=int, default=4) parser.add_argument("--font_path", help="Path to ttf font file or directory containing ttf files") return parser.parse_args()
def __init__(self, root=data_dir(), model_epoch=MODEL_EPOCE): self._model_dir = os.path.join(root, 'models') self._model_epoch = model_epoch self._assert_and_prepare_model_files(root) self._alphabet, _ = read_charset(os.path.join(self._model_dir, 'label_cn.txt')) self._hp = Hyperparams() self._mods = {}
def __init__(self, model_name='conv-lite-fc', model_epoch=None, cand_alphabet=None, root=data_dir(), gpus=0): """ :param model_name: 模型名称 :param model_epoch: 模型迭代次数 :param cand_alphabet: 待识别字符所在的候选集合。默认为 `None`,表示不限定识别字符范围 :param root: 模型文件所在的根目录。 Linux/Mac下默认值为 `~/.cnocr`,表示模型文件所处文件夹类似 `~/.cnocr/1.1.0/conv-lite-fc-0027`。 Windows下默认值为 ``。 """ check_model_name(model_name) self._model_name = model_name self._model_file_prefix = '{}-{}'.format(self.MODEL_FILE_PREFIX, model_name) self._model_epoch = model_epoch or AVAILABLE_MODELS[model_name][0] root = os.path.join(root, MODEL_VERSION) self._model_dir = os.path.join(root, self._model_name) self._assert_and_prepare_model_files() self._alphabet, inv_alph_dict = read_charset( os.path.join(self._model_dir, 'label_cn.txt')) self._cand_alph_idx = None if cand_alphabet is not None: self._cand_alph_idx = [0] + [ inv_alph_dict[word] for word in cand_alphabet ] self._cand_alph_idx.sort() self._hp = Hyperparams() self._hp._loss_type = None # infer mode # DCMMC: gpu context for mxnet if gpus > 0: self.context = [mx.context.gpu(i) for i in range(gpus)] else: self.context = [mx.context.cpu()] self._mod = self._get_module()
def __init__( self, model_name='densenet-lite-fc', model_epoch=None, cand_alphabet=None, root=data_dir(), context='cpu', name=None, ): """ :param model_name: 模型名称 :param model_epoch: 模型迭代次数 :param cand_alphabet: 待识别字符所在的候选集合。默认为 `None`,表示不限定识别字符范围 :param root: 模型文件所在的根目录。 Linux/Mac下默认值为 `~/.cnocr`,表示模型文件所处文件夹类似 `~/.cnocr/1.1.0/conv-lite-fc-0027`。 Windows下默认值为 ``。 :param context: 'cpu', or 'gpu'。表明预测时是使用CPU还是GPU。默认为CPU。 :param name: 正在初始化的这个实例名称。如果需要同时初始化多个实例,需要为不同的实例指定不同的名称。 """ check_model_name(model_name) self._model_name = model_name self._model_file_prefix = '{}-{}'.format(self.MODEL_FILE_PREFIX, model_name) self._model_epoch = model_epoch or AVAILABLE_MODELS[model_name][0] root = os.path.join(root, MODEL_VERSION) self._model_dir = os.path.join(root, self._model_name) self._assert_and_prepare_model_files() self._alphabet, self._inv_alph_dict = read_charset( os.path.join(self._model_dir, 'label_cn.txt')) self._cand_alph_idx = None self.set_cand_alphabet(cand_alphabet) self._hp = Hyperparams() self._hp._loss_type = None # infer mode # 传入''的话,也改成传入None self._net_prefix = None if name == '' else name self._mod = self._get_module(context)
def main(): charset_fp = os.path.join(data_dir(), 'models', 'label_cn.txt') alphabet, inv_alph_dict = read_charset(charset_fp) for idx in BAD_CHARS: print('idx: {}, char: {}'.format(idx, alphabet[idx]))
def parse_args(): # Parse command line arguments parser = argparse.ArgumentParser() parser.add_argument( "--emb_model_type", help="which embedding model to use", choices=EMB_MODEL_TYPES, type=str, default='conv-lite', ) parser.add_argument( "--seq_model_type", help='which sequence model to use', default='fc', type=str, choices=SEQ_MODEL_TYPES, ) parser.add_argument( "--train_file", help="Path to train txt file", type=str, default='data/sample-data-lst/train.txt', ) parser.add_argument( "--test_file", help="Path to test txt file", type=str, default='data/sample-data-lst/test.txt', ) parser.add_argument('--dataset', help='file path for dataset hdf5', type=str, required=True) parser.add_argument('--charset', help='file path for chat set of labels', type=str, required=True) parser.add_argument( '--debug', help='debug mode', action='store_true', ) parser.add_argument( "--use_train_image_aug", action='store_true', help="Whether to use image augmentation for training", ) parser.add_argument( "--gpu", help="Number of GPUs for training [Default 0, means using cpu]", type=int, default=0, ) parser.add_argument( "--optimizer", help="optimizer for training [Default: Adam]", type=str, default='Adam', ) parser.add_argument('--epoch', type=int, default=20, help='train epochs [Default: 20]') parser.add_argument( '--load_epoch', type=int, help= 'load the model on an epoch using the model-load-prefix [Default: no trained model will be loaded]', ) parser.add_argument('--lr', type=float, default=0.001, help='learning rate') parser.add_argument('--wd', type=float, default=0.0, help='weight decay factor [Default: 0.0]') parser.add_argument( '--clip_gradient', type=float, default=None, help= 'value for clip gradient [Default: None, means no gradient will be clip]', ) parser.add_argument( "--out_model_dir", help='output model directory', default=os.path.join(data_dir(), __version__), ) return parser.parse_args()
def __init__( self, model_name: str = 'densenet_lite_136-fc', *, cand_alphabet: Optional[Union[Collection, str]] = None, context: str = 'cpu', # ['cpu', 'gpu', 'cuda'] model_fp: Optional[str] = None, root: Union[str, Path] = data_dir(), **kwargs, ): """ 识别模型初始化函数。 Args: model_name (str): 模型名称。默认为 `densenet_lite_136-fc` cand_alphabet (Optional[Union[Collection, str]]): 待识别字符所在的候选集合。默认为 `None`,表示不限定识别字符范围 context (str): 'cpu', or 'gpu'。表明预测时是使用CPU还是GPU。默认为 `cpu` model_fp (Optional[str]): 如果不使用系统自带的模型,可以通过此参数直接指定所使用的模型文件('.ckpt' 文件) root (Union[str, Path]): 模型文件所在的根目录。 Linux/Mac下默认值为 `~/.cnocr`,表示模型文件所处文件夹类似 `~/.cnocr/2.1/densenet_lite_136-fc`。 Windows下默认值为 `C:/Users/<username>/AppData/Roaming/cnocr`。 **kwargs: 目前未被使用。 Examples: 使用默认参数: >>> ocr = CnOcr() 使用指定模型: >>> ocr = CnOcr(model_name='densenet_lite_136-fc') 识别时只考虑数字: >>> ocr = CnOcr(model_name='densenet_lite_136-fc', cand_alphabet='0123456789') """ if 'name' in kwargs: logger.warning( 'param `name` is useless and deprecated since version %s' % MODEL_VERSION) check_model_name(model_name) check_context(context) self._model_name = model_name if context == 'gpu': context = 'cuda' self.context = context self._model_file_prefix = '{}-{}'.format(self.MODEL_FILE_PREFIX, model_name) model_epoch = AVAILABLE_MODELS.get(model_name, [None])[0] if model_epoch is not None: self._model_file_prefix = '%s-epoch=%03d' % ( self._model_file_prefix, model_epoch, ) self._assert_and_prepare_model_files(model_fp, root) self._vocab, self._letter2id = read_charset(VOCAB_FP) self._candidates = None self.set_cand_alphabet(cand_alphabet) self._model = self._get_model(context)