def ocr(self, img_fp):
"""
:param img_fp: image file path; or color image mx.nd.NDArray or np.ndarray,
with shape (height, width, 3), and the channels should be RGB formatted.
:return: List(List(Char)), such as:
[['第', '一', '行'], ['第', '二', '行'], ['第', '三', '行']]
"""
if isinstance(img_fp, str):
if not os.path.isfile(img_fp):
raise FileNotFoundError(img_fp)
img = mx.image.imread(img_fp, 1).asnumpy()
elif isinstance(img_fp, mx.nd.NDArray):
img = img_fp.asnumpy()
elif isinstance(img_fp, np.ndarray):
img = img_fp
else:
raise TypeError('Inappropriate argument type.')
if min(img.shape[0], img.shape[1]) < 2:
return ''
if img.mean() < 145: # 把黑底白字的图片对调为白底黑字
img = 255 - img
line_imgs = line_split(img, blank=True)
line_img_list = [line_img for line_img, _ in line_imgs]
line_chars_list = self.ocr_for_single_lines(line_img_list)
return line_chars_list
def test_ocr_for_single_lines(img_fp, expected):
ocr = CNOCR
root_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
img_fp = os.path.join(root_dir, 'examples', img_fp)
img = mx.image.imread(img_fp, 1).asnumpy()
line_imgs = line_split(img, blank=True)
line_img_list = [line_img for line_img, _ in line_imgs]
pred = ocr.ocr_for_single_lines(line_img_list)
print('\n')
print("Predicted Chars:", pred)
assert expected == pred
line_img_list = [nd.array(line_img) for line_img in line_img_list]
pred = ocr.ocr_for_single_lines(line_img_list)
print("Predicted Chars:", pred)
assert expected == pred
def test_ocr_for_single_lines(img_fp, expected):
ocr = CNOCR
root_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
img_fp = os.path.join(root_dir, 'examples', img_fp)
img = mx.image.imread(img_fp, 1).asnumpy()
if img.mean() < 145: # 把黑底白字的图片对调为白底黑字
img = 255 - img
line_imgs = line_split(img, blank=True)
line_img_list = [line_img for line_img, _ in line_imgs]
pred = ocr.ocr_for_single_lines(line_img_list)
print('\n')
print_preds(pred)
assert cal_score(pred, expected) >= 0.9
line_img_list = [nd.array(line_img) for line_img in line_img_list]
pred = ocr.ocr_for_single_lines(line_img_list)
print_preds(pred)
assert cal_score(pred, expected) >= 0.9
def ocr(self, img_fp):
"""
:param img_fp: image file path; or color image mx.nd.NDArray or np.ndarray,
with shape (height, width, 3), and the channels should be RGB formatted.
:return: List(List(Letter)), such as:
[['第', '一', '行'], ['第', '二', '行'], ['第', '三', '行']]
"""
if isinstance(img_fp, str) and os.path.isfile(img_fp):
img = mx.image.imread(img_fp, 1).asnumpy()
elif isinstance(img_fp, mx.nd.NDArray) or isinstance(img_fp, np.ndarray):
img = img_fp
else:
raise TypeError('Inappropriate argument type.')
if min(img.shape[0], img.shape[1]) < 2:
return ''
line_imgs = line_split(img, blank=True)
line_chars_list = []
for line_idx, (line_img, _) in enumerate(line_imgs):
line_img = np.array(Image.fromarray(line_img).convert('L'))
line_chars = self.ocr_for_single_line(line_img)
line_chars_list.append(line_chars)
return line_chars_list
