def _gen_line_pred_chars(self, line_prob, img_width, max_img_width): """ Get the predicted characters. :param line_prob: with shape of [seq_length, num_classes] :param img_width: :param max_img_width: :return: """ class_ids = np.argmax(line_prob, axis=-1) class_ids *= np.max(line_prob, axis=-1) > 0.5 # Delete low confidence result if img_width < max_img_width: comp_ratio = self._hp.seq_len_cmpr_ratio end_idx = img_width // comp_ratio if end_idx < len(class_ids): class_ids[end_idx:] = 0 prediction, start_end_idx = CtcMetrics.ctc_label(class_ids.tolist()) alphabet = self._alphabet res = [ alphabet[p] if alphabet[p] != '<space>' else ' ' for p in prediction ] return res
def ocr_for_single_line(self, img_fp): """ Recognize characters from an image with characters with only one line :param img_fp: image file path; or gray image mx.nd.NDArray; or gray image np.ndarray :return: charector list, such as ['你', '好'] """ hp = deepcopy(self._hp) if isinstance(img_fp, str) and os.path.isfile(img_fp): img = read_ocr_img(img_fp) elif isinstance(img_fp, mx.nd.NDArray) or isinstance( img_fp, np.ndarray): img = img_fp else: raise TypeError('Inappropriate argument type.') img = rescale_img(img, hp) init_state_names, init_state_arrays = lstm_init_states(batch_size=1, hp=hp) sample = SimpleBatch(data_names=['data'] + init_state_names, data=[mx.nd.array([img])] + init_state_arrays) mod = self._get_module(hp, sample) mod.forward(sample) prob = mod.get_outputs()[0].asnumpy() prediction, start_end_idx = CtcMetrics.ctc_label( np.argmax(prob, axis=-1).tolist()) # print(start_end_idx) alphabet = self._alphabet res = [alphabet[p] for p in prediction] return res
def _gen_line_pred_chars(self, line_prob, img_width, max_img_width): """ Get the predicted characters. :param line_prob: with shape of [seq_length, num_classes] :param img_width: :param max_img_width: :return: """ drop = [1 if l.max() > 0.8 else 0 for l in line_prob] class_ids_ = np.argmax(line_prob, axis=-1) class_ids = [] for c, d in zip(class_ids_, drop): class_ids += [c] if d == 1 else [6425] if img_width < max_img_width: comp_ratio = self._hp.seq_len_cmpr_ratio end_idx = img_width // comp_ratio if end_idx < len(class_ids): class_ids[end_idx:] = 0 #prediction, start_end_idx = CtcMetrics.ctc_label(class_ids.tolist()) prediction, start_end_idx = CtcMetrics.ctc_label(class_ids) alphabet = self._alphabet res = [ alphabet[p] if alphabet[p] != '<space>' else ' ' for p in prediction ] return res
def _gen_line_pred_chars(self, line_prob, img_width, max_img_width): """ Get the predicted characters. :param line_prob: with shape of [seq_length, num_classes] :param img_width: :param max_img_width: :return: """ # DCMMC: Greedy decoder for CTC class_ids = np.argmax(line_prob, axis=-1) if img_width < max_img_width: comp_ratio = self._hp.seq_len_cmpr_ratio end_idx = img_width // comp_ratio # DCMMC: 原来照片是 right padding 的... # 而我的数据集是 left and right padding if end_idx < len(class_ids): class_ids[end_idx:] = 0 prediction, start_end_idx = CtcMetrics.ctc_label(class_ids.tolist()) alphabet = self._alphabet res = [ alphabet[p] if alphabet[p] != '<space>' else ' ' for p in prediction ] return res
def main(): parser = argparse.ArgumentParser() 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("--file", help="Path to the CAPTCHA image file") parser.add_argument("--prefix", help="Checkpoint prefix [Default 'ocr']", default='./models/model') parser.add_argument("--epoch", help="Checkpoint epoch [Default 100]", type=int, default=20) parser.add_argument('--charset_file', type=str, help='存储了每个字对应哪个id的关系.') args = parser.parse_args() if args.dataset == 'cn_ocr': hp = Hyperparams() img = read_ocr_img(args.file, hp) else: hp = Hyperparams2() img = read_captcha_img(args.file, hp) # init_state_names, init_state_arrays = lstm_init_states(batch_size=1, hp=hp) # import pdb; pdb.set_trace() sample = SimpleBatch(data_names=['data'], data=[mx.nd.array([img])]) network = crnn_lstm(hp) mod = load_module(args.prefix, args.epoch, sample.data_names, sample.provide_data, network=network) mod.forward(sample) prob = mod.get_outputs()[0].asnumpy() prediction, start_end_idx = CtcMetrics.ctc_label( np.argmax(prob, axis=-1).tolist()) if args.charset_file: alphabet, _ = read_charset(args.charset_file) res = [alphabet[p] for p in prediction] print("Predicted Chars:", res) else: # Predictions are 1 to 10 for digits 0 to 9 respectively (prediction 0 means no-digit) prediction = [p - 1 for p in prediction] print("Digits:", prediction) return
def _gen_line_pred_chars(self, line_prob, img_width, max_img_width): """ Get the predicted characters. :param line_prob: with shape of [seq_length, num_classes] :param img_width: :param max_img_width: :return: """ class_ids = np.argmax(line_prob, axis=-1) # idxs = list(zip(range(len(class_ids)), class_ids)) # probs = [line_prob[e[0], e[1]] for e in idxs] if img_width < max_img_width: comp_ratio = self._hp.seq_len_cmpr_ratio end_idx = img_width // comp_ratio if end_idx < len(class_ids): class_ids[end_idx:] = 0 prediction, start_end_idx = CtcMetrics.ctc_label(class_ids.tolist()) # print(start_end_idx) alphabet = self._alphabet res = [alphabet[p] for p in prediction] # res = self._insert_space_char(res, start_end_idx) return res
def test_ctc_metrics(input, expected): input = list(map(int, list(input))) expected = list(map(int, list(expected))) p, _ = CtcMetrics.ctc_label(input) assert expected == p