mirror of
https://github.com/RapidAI/RapidOCR.git
synced 2026-11-10 13:03:58 +08:00
feat(rapidocr_openvino): support to char rec and bump to v1.4.0
This commit is contained in:
1 parent
4e9b7ccfb9
commit
62bc4871a0
9 files changed
+423
-89
No files matched your search
@@ -87,7 +87,7 @@ jobs:
|
||||
|
||||
cd ..
|
||||
python -m pip install --upgrade pip
|
||||
python setup_openvino.py bdist_wheel "{{ github.event.head_commit.message }}"
|
||||
python setup_openvino.py bdist_wheel "${{ github.event.head_commit.message }}"
|
||||
mv dist ../
|
||||
|
||||
- name: Publish distribution 📦 to PyPI
|
||||
|
||||
@@ -8,7 +8,6 @@ repos:
|
||||
"--recursive",
|
||||
"--in-place",
|
||||
"--remove-all-unused-imports",
|
||||
"--remove-unused-variable",
|
||||
"--ignore-init-module-imports",
|
||||
]
|
||||
files: \.py$
|
||||
|
||||
@@ -0,0 +1,4 @@
|
||||
# -*- encoding: utf-8 -*-
|
||||
# @Author: SWHL
|
||||
# @Contact: liekkaskono@163.com
|
||||
from .main import CalRecBoxes
|
||||
@@ -0,0 +1,260 @@
|
||||
# -*- encoding: utf-8 -*-
|
||||
# @Author: SWHL / Joker1212
|
||||
# @Contact: liekkaskono@163.com
|
||||
import copy
|
||||
import math
|
||||
from typing import Any, List, Optional, Tuple
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
|
||||
class CalRecBoxes:
|
||||
"""计算识别文字的汉字单字和英文单词的坐标框。代码借鉴自PaddlePaddle/PaddleOCR和fanqie03/char-detection"""
|
||||
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
imgs: Optional[List[np.ndarray]],
|
||||
dt_boxes: Optional[List[np.ndarray]],
|
||||
rec_res: Optional[List[Any]],
|
||||
):
|
||||
res = []
|
||||
for img, box, rec_res in zip(imgs, dt_boxes, rec_res):
|
||||
direction = self.get_box_direction(box)
|
||||
|
||||
rec_txt, rec_conf, rec_word_info = rec_res[0], rec_res[1], rec_res[2]
|
||||
h, w = img.shape[:2]
|
||||
img_box = np.array([[0, 0], [w, 0], [w, h], [0, h]])
|
||||
word_box_content_list, word_box_list = self.cal_ocr_word_box(
|
||||
rec_txt, img_box, rec_word_info
|
||||
)
|
||||
word_box_list = self.adjust_box_overlap(copy.deepcopy(word_box_list))
|
||||
word_box_list = self.reverse_rotate_crop_image(
|
||||
copy.deepcopy(box), word_box_list, direction
|
||||
)
|
||||
res.append([rec_txt, rec_conf, word_box_list, word_box_content_list])
|
||||
return res
|
||||
|
||||
@staticmethod
|
||||
def get_box_direction(box: np.ndarray) -> str:
|
||||
direction = "w"
|
||||
img_crop_width = int(
|
||||
max(
|
||||
np.linalg.norm(box[0] - box[1]),
|
||||
np.linalg.norm(box[2] - box[3]),
|
||||
)
|
||||
)
|
||||
img_crop_height = int(
|
||||
max(
|
||||
np.linalg.norm(box[0] - box[3]),
|
||||
np.linalg.norm(box[1] - box[2]),
|
||||
)
|
||||
)
|
||||
if img_crop_height * 1.0 / img_crop_width >= 1.5:
|
||||
direction = "h"
|
||||
return direction
|
||||
|
||||
@staticmethod
|
||||
def cal_ocr_word_box(
|
||||
rec_txt: str, box: np.ndarray, rec_word_info: List[Tuple[str, List[int]]]
|
||||
) -> Tuple[List[str], List[List[int]]]:
|
||||
"""Calculate the detection frame for each word based on the results of recognition and detection of ocr
|
||||
汉字坐标是单字的
|
||||
英语坐标是单词级别的
|
||||
"""
|
||||
|
||||
col_num, word_list, word_col_list, state_list = rec_word_info
|
||||
box = box.tolist()
|
||||
bbox_x_start = box[0][0]
|
||||
bbox_x_end = box[1][0]
|
||||
bbox_y_start = box[0][1]
|
||||
bbox_y_end = box[2][1]
|
||||
|
||||
cell_width = (bbox_x_end - bbox_x_start) / col_num
|
||||
word_box_list = []
|
||||
word_box_content_list = []
|
||||
cn_width_list = []
|
||||
cn_col_list = []
|
||||
for word, word_col, state in zip(word_list, word_col_list, state_list):
|
||||
if state == "cn":
|
||||
if len(word_col) != 1:
|
||||
char_seq_length = (word_col[-1] - word_col[0] + 1) * cell_width
|
||||
char_width = char_seq_length / (len(word_col) - 1)
|
||||
cn_width_list.append(char_width)
|
||||
cn_col_list += word_col
|
||||
word_box_content_list += word
|
||||
else:
|
||||
cell_x_start = bbox_x_start + int(word_col[0] * cell_width)
|
||||
cell_x_end = bbox_x_start + int((word_col[-1] + 1) * cell_width)
|
||||
cell = [
|
||||
[cell_x_start, bbox_y_start],
|
||||
[cell_x_end, bbox_y_start],
|
||||
[cell_x_end, bbox_y_end],
|
||||
[cell_x_start, bbox_y_end],
|
||||
]
|
||||
word_box_list.append(cell)
|
||||
word_box_content_list.append("".join(word))
|
||||
|
||||
if len(cn_col_list) != 0:
|
||||
if len(cn_width_list) != 0:
|
||||
avg_char_width = np.mean(cn_width_list)
|
||||
else:
|
||||
avg_char_width = (bbox_x_end - bbox_x_start) / len(rec_txt)
|
||||
|
||||
for center_idx in cn_col_list:
|
||||
center_x = (center_idx + 0.5) * cell_width
|
||||
cell_x_start = max(int(center_x - avg_char_width / 2), 0) + bbox_x_start
|
||||
cell_x_end = (
|
||||
min(int(center_x + avg_char_width / 2), bbox_x_end - bbox_x_start)
|
||||
+ bbox_x_start
|
||||
)
|
||||
cell = [
|
||||
[cell_x_start, bbox_y_start],
|
||||
[cell_x_end, bbox_y_start],
|
||||
[cell_x_end, bbox_y_end],
|
||||
[cell_x_start, bbox_y_end],
|
||||
]
|
||||
word_box_list.append(cell)
|
||||
sorted_word_box_list = sorted(word_box_list, key=lambda box: box[0][0])
|
||||
return word_box_content_list, sorted_word_box_list
|
||||
|
||||
@staticmethod
|
||||
def adjust_box_overlap(
|
||||
word_box_list: List[List[List[int]]],
|
||||
) -> List[List[List[int]]]:
|
||||
# 调整bbox有重叠的地方
|
||||
for i in range(len(word_box_list) - 1):
|
||||
cur, nxt = word_box_list[i], word_box_list[i + 1]
|
||||
if cur[1][0] > nxt[0][0]: # 有交集
|
||||
distance = abs(cur[1][0] - nxt[0][0])
|
||||
cur[1][0] -= distance / 2
|
||||
cur[2][0] -= distance / 2
|
||||
nxt[0][0] += distance / 2
|
||||
nxt[3][0] += distance / 2
|
||||
return word_box_list
|
||||
|
||||
def reverse_rotate_crop_image(
|
||||
self,
|
||||
bbox_points: np.ndarray,
|
||||
word_points_list: List[List[List[int]]],
|
||||
direction: str = "w",
|
||||
) -> List[List[List[int]]]:
|
||||
"""
|
||||
get_rotate_crop_image的逆操作
|
||||
img为原图
|
||||
part_img为crop后的图
|
||||
bbox_points为part_img中对应在原图的bbox, 四个点,左上,右上,右下,左下
|
||||
part_points为在part_img中的点[(x, y), (x, y)]
|
||||
"""
|
||||
bbox_points = np.float32(bbox_points)
|
||||
|
||||
left = int(np.min(bbox_points[:, 0]))
|
||||
top = int(np.min(bbox_points[:, 1]))
|
||||
bbox_points[:, 0] = bbox_points[:, 0] - left
|
||||
bbox_points[:, 1] = bbox_points[:, 1] - top
|
||||
|
||||
img_crop_width = int(np.linalg.norm(bbox_points[0] - bbox_points[1]))
|
||||
img_crop_height = int(np.linalg.norm(bbox_points[0] - bbox_points[3]))
|
||||
|
||||
pts_std = np.array(
|
||||
[
|
||||
[0, 0],
|
||||
[img_crop_width, 0],
|
||||
[img_crop_width, img_crop_height],
|
||||
[0, img_crop_height],
|
||||
]
|
||||
).astype(np.float32)
|
||||
M = cv2.getPerspectiveTransform(bbox_points, pts_std)
|
||||
_, IM = cv2.invert(M)
|
||||
|
||||
new_word_points_list = []
|
||||
for word_points in word_points_list:
|
||||
new_word_points = []
|
||||
for point in word_points:
|
||||
new_point = point
|
||||
if direction == "h":
|
||||
new_point = self.s_rotate(
|
||||
math.radians(-90), new_point[0], new_point[1], 0, 0
|
||||
)
|
||||
new_point[0] = new_point[0] + img_crop_width
|
||||
|
||||
p = np.float32(new_point + [1])
|
||||
x, y, z = np.dot(IM, p)
|
||||
new_point = [x / z, y / z]
|
||||
|
||||
new_point = [int(new_point[0] + left), int(new_point[1] + top)]
|
||||
new_word_points.append(new_point)
|
||||
new_word_points = self.order_points(new_word_points)
|
||||
new_word_points_list.append(new_word_points)
|
||||
return new_word_points_list
|
||||
|
||||
@staticmethod
|
||||
def s_rotate(angle, valuex, valuey, pointx, pointy):
|
||||
"""绕pointx,pointy顺时针旋转
|
||||
https://blog.csdn.net/qq_38826019/article/details/84233397
|
||||
"""
|
||||
valuex = np.array(valuex)
|
||||
valuey = np.array(valuey)
|
||||
sRotatex = (
|
||||
(valuex - pointx) * math.cos(angle)
|
||||
+ (valuey - pointy) * math.sin(angle)
|
||||
+ pointx
|
||||
)
|
||||
sRotatey = (
|
||||
(valuey - pointy) * math.cos(angle)
|
||||
- (valuex - pointx) * math.sin(angle)
|
||||
+ pointy
|
||||
)
|
||||
return [sRotatex, sRotatey]
|
||||
|
||||
@staticmethod
|
||||
def order_points(box: List[List[int]]) -> List[List[int]]:
|
||||
"""矩形框顺序排列"""
|
||||
box = np.array(box).reshape((-1, 2))
|
||||
center_x, center_y = np.mean(box[:, 0]), np.mean(box[:, 1])
|
||||
if np.any(box[:, 0] == center_x) and np.any(
|
||||
box[:, 1] == center_y
|
||||
): # 有两点横坐标相等,有两点纵坐标相等,菱形
|
||||
p1 = box[np.where(box[:, 0] == np.min(box[:, 0]))]
|
||||
p2 = box[np.where(box[:, 1] == np.min(box[:, 1]))]
|
||||
p3 = box[np.where(box[:, 0] == np.max(box[:, 0]))]
|
||||
p4 = box[np.where(box[:, 1] == np.max(box[:, 1]))]
|
||||
elif np.all(box[:, 0] == center_x): # 四个点的横坐标都相同
|
||||
y_sort = np.argsort(box[:, 1])
|
||||
p1 = box[y_sort[0]]
|
||||
p2 = box[y_sort[1]]
|
||||
p3 = box[y_sort[2]]
|
||||
p4 = box[y_sort[3]]
|
||||
elif np.any(box[:, 0] == center_x) and np.all(
|
||||
box[:, 1] != center_y
|
||||
): # 只有两点横坐标相等,先上下再左右
|
||||
p12, p34 = (
|
||||
box[np.where(box[:, 1] < center_y)],
|
||||
box[np.where(box[:, 1] > center_y)],
|
||||
)
|
||||
p1, p2 = (
|
||||
p12[np.where(p12[:, 0] == np.min(p12[:, 0]))],
|
||||
p12[np.where(p12[:, 0] == np.max(p12[:, 0]))],
|
||||
)
|
||||
p3, p4 = (
|
||||
p34[np.where(p34[:, 0] == np.max(p34[:, 0]))],
|
||||
p34[np.where(p34[:, 0] == np.min(p34[:, 0]))],
|
||||
)
|
||||
else: # 只有两点纵坐标相等,或者是没有相等的,先左右再上下
|
||||
p14, p23 = (
|
||||
box[np.where(box[:, 0] < center_x)],
|
||||
box[np.where(box[:, 0] > center_x)],
|
||||
)
|
||||
p1, p4 = (
|
||||
p14[np.where(p14[:, 1] == np.min(p14[:, 1]))],
|
||||
p14[np.where(p14[:, 1] == np.max(p14[:, 1]))],
|
||||
)
|
||||
p2, p3 = (
|
||||
p23[np.where(p23[:, 1] == np.min(p23[:, 1]))],
|
||||
p23[np.where(p23[:, 1] == np.max(p23[:, 1]))],
|
||||
)
|
||||
|
||||
return np.array([p1, p2, p3, p4]).reshape((-1, 2)).tolist()
|
||||
@@ -36,10 +36,12 @@ class TextRecognizer:
|
||||
self.character_dict_path = dict_path if dict_path else DEFAULT_DICT_PATH
|
||||
self.postprocess_op = CTCLabelDecode(character_path=self.character_dict_path)
|
||||
|
||||
self.infer = OpenVINOInferSession(config)
|
||||
self.session = OpenVINOInferSession(config)
|
||||
|
||||
def __call__(
|
||||
self, img_list: Union[np.ndarray, List[np.ndarray]]
|
||||
self,
|
||||
img_list: Union[np.ndarray, List[np.ndarray]],
|
||||
return_word_box: bool = False,
|
||||
) -> Tuple[List[Tuple[str, float]], float]:
|
||||
if isinstance(img_list, np.ndarray):
|
||||
img_list = [img_list]
|
||||
@@ -57,11 +59,16 @@ class TextRecognizer:
|
||||
elapse = 0
|
||||
for beg_img_no in range(0, img_num, batch_num):
|
||||
end_img_no = min(img_num, beg_img_no + batch_num)
|
||||
max_wh_ratio = 0
|
||||
|
||||
# Parameter Alignment for PaddleOCR
|
||||
imgC, imgH, imgW = self.rec_image_shape[:3]
|
||||
max_wh_ratio = imgW / imgH
|
||||
wh_ratio_list = []
|
||||
for ino in range(beg_img_no, end_img_no):
|
||||
h, w = img_list[indices[ino]].shape[0:2]
|
||||
wh_ratio = w * 1.0 / h
|
||||
max_wh_ratio = max(max_wh_ratio, wh_ratio)
|
||||
wh_ratio_list.append(wh_ratio)
|
||||
|
||||
norm_img_batch = []
|
||||
for ino in range(beg_img_no, end_img_no):
|
||||
@@ -70,8 +77,13 @@ class TextRecognizer:
|
||||
norm_img_batch = np.concatenate(norm_img_batch).astype(np.float32)
|
||||
|
||||
starttime = time.time()
|
||||
preds = self.infer(norm_img_batch)
|
||||
rec_result = self.postprocess_op(preds)
|
||||
preds = self.session(norm_img_batch)
|
||||
rec_result = self.postprocess_op(
|
||||
preds,
|
||||
return_word_box,
|
||||
wh_ratio_list=wh_ratio_list,
|
||||
max_wh_ratio=max_wh_ratio,
|
||||
)
|
||||
|
||||
for rno, one_res in enumerate(rec_result):
|
||||
rec_res[indices[beg_img_no + rno]] = one_res
|
||||
|
||||
@@ -16,10 +16,19 @@ class CTCLabelDecode:
|
||||
self.character = self.get_character(character, character_path)
|
||||
self.dict = {char: i for i, char in enumerate(self.character)}
|
||||
|
||||
def __call__(self, preds: np.ndarray) -> List[Tuple[str, float]]:
|
||||
def __call__(
|
||||
self, preds: np.ndarray, return_word_box: bool = False, **kwargs
|
||||
) -> List[Tuple[str, float]]:
|
||||
preds_idx = preds.argmax(axis=2)
|
||||
preds_prob = preds.max(axis=2)
|
||||
text = self.decode(preds_idx, preds_prob, is_remove_duplicate=True)
|
||||
text = self.decode(
|
||||
preds_idx, preds_prob, return_word_box, is_remove_duplicate=True
|
||||
)
|
||||
if return_word_box:
|
||||
for rec_idx, rec in enumerate(text):
|
||||
wh_ratio = kwargs["wh_ratio_list"][rec_idx]
|
||||
max_wh_ratio = kwargs["max_wh_ratio"]
|
||||
rec[2][0] = rec[2][0] * (wh_ratio / max_wh_ratio)
|
||||
return text
|
||||
|
||||
def get_character(
|
||||
@@ -67,6 +76,7 @@ class CTCLabelDecode:
|
||||
self,
|
||||
text_index: np.ndarray,
|
||||
text_prob: Optional[np.ndarray] = None,
|
||||
return_word_box: bool = False,
|
||||
is_remove_duplicate: bool = False,
|
||||
) -> List[Tuple[str, float]]:
|
||||
"""convert text-index into text-label."""
|
||||
@@ -74,28 +84,99 @@ class CTCLabelDecode:
|
||||
ignored_tokens = self.get_ignored_tokens()
|
||||
batch_size = len(text_index)
|
||||
for batch_idx in range(batch_size):
|
||||
char_list, conf_list = [], []
|
||||
cur_pred_ids = text_index[batch_idx]
|
||||
for idx, cur_idx in enumerate(cur_pred_ids):
|
||||
if cur_idx in ignored_tokens:
|
||||
continue
|
||||
selection = np.ones(len(text_index[batch_idx]), dtype=bool)
|
||||
if is_remove_duplicate:
|
||||
selection[1:] = text_index[batch_idx][1:] != text_index[batch_idx][:-1]
|
||||
|
||||
if is_remove_duplicate:
|
||||
# only for predict
|
||||
if idx > 0 and cur_pred_ids[idx - 1] == cur_idx:
|
||||
continue
|
||||
for ignored_token in ignored_tokens:
|
||||
selection &= text_index[batch_idx] != ignored_token
|
||||
|
||||
char_list.append(self.character[int(cur_idx)])
|
||||
if text_prob is not None:
|
||||
conf_list = text_prob[batch_idx][selection]
|
||||
else:
|
||||
conf_list = [1] * len(selection)
|
||||
|
||||
if text_prob is None:
|
||||
conf_list.append(1)
|
||||
else:
|
||||
conf_list.append(text_prob[batch_idx][idx])
|
||||
if len(conf_list) == 0:
|
||||
conf_list = [0]
|
||||
|
||||
char_list = [
|
||||
self.character[text_id] for text_id in text_index[batch_idx][selection]
|
||||
]
|
||||
text = "".join(char_list)
|
||||
result_list.append((text, np.mean(conf_list if any(conf_list) else [0])))
|
||||
if return_word_box:
|
||||
word_list, word_col_list, state_list = self.get_word_info(
|
||||
text, selection
|
||||
)
|
||||
result_list.append(
|
||||
(
|
||||
text,
|
||||
np.mean(conf_list).tolist(),
|
||||
[
|
||||
len(text_index[batch_idx]),
|
||||
word_list,
|
||||
word_col_list,
|
||||
state_list,
|
||||
],
|
||||
)
|
||||
)
|
||||
else:
|
||||
result_list.append((text, np.mean(conf_list).tolist()))
|
||||
return result_list
|
||||
|
||||
@staticmethod
|
||||
def get_word_info(
|
||||
text: str, selection: np.ndarray
|
||||
) -> Tuple[List[List[str]], List[List[int]], List[str]]:
|
||||
"""
|
||||
Group the decoded characters and record the corresponding decoded positions.
|
||||
from https://github.com/PaddlePaddle/PaddleOCR/blob/fbba2178d7093f1dffca65a5b963ec277f1a6125/ppocr/postprocess/rec_postprocess.py#L70
|
||||
|
||||
Args:
|
||||
text: the decoded text
|
||||
selection: the bool array that identifies which columns of features are decoded as non-separated characters
|
||||
Returns:
|
||||
word_list: list of the grouped words
|
||||
word_col_list: list of decoding positions corresponding to each character in the grouped word
|
||||
state_list: list of marker to identify the type of grouping words, including two types of grouping words:
|
||||
- 'cn': continous chinese characters (e.g., 你好啊)
|
||||
- 'en&num': continous english characters (e.g., hello), number (e.g., 123, 1.123), or mixed of them connected by '-' (e.g., VGG-16)
|
||||
"""
|
||||
state = None
|
||||
word_content = []
|
||||
word_col_content = []
|
||||
word_list = []
|
||||
word_col_list = []
|
||||
state_list = []
|
||||
valid_col = np.where(selection == True)[0]
|
||||
|
||||
for c_i, char in enumerate(text):
|
||||
if "\u4e00" <= char <= "\u9fff":
|
||||
c_state = "cn"
|
||||
else:
|
||||
c_state = "en&num"
|
||||
|
||||
if state == None:
|
||||
state = c_state
|
||||
|
||||
if state != c_state:
|
||||
if len(word_content) != 0:
|
||||
word_list.append(word_content)
|
||||
word_col_list.append(word_col_content)
|
||||
state_list.append(state)
|
||||
word_content = []
|
||||
word_col_content = []
|
||||
state = c_state
|
||||
|
||||
word_content.append(char)
|
||||
word_col_content.append(int(valid_col[c_i]))
|
||||
|
||||
if len(word_content) != 0:
|
||||
word_list.append(word_content)
|
||||
word_col_list.append(word_col_content)
|
||||
state_list.append(state)
|
||||
|
||||
return word_list, word_col_list, state_list
|
||||
|
||||
@staticmethod
|
||||
def get_ignored_tokens() -> List[int]:
|
||||
return [0] # for ctc blank
|
||||
@@ -8,6 +8,7 @@ Global:
|
||||
width_height_ratio: 8
|
||||
max_side_len: 2000
|
||||
min_side_len: 30
|
||||
return_word_box: false
|
||||
|
||||
inference_num_threads: &infer_num_threads -1
|
||||
|
||||
|
||||
@@ -8,6 +8,7 @@ from typing import Any, Dict, List, Optional, Tuple, Union
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from .cal_rec_boxes import CalRecBoxes
|
||||
from .ch_ppocr_cls import TextClassifier
|
||||
from .ch_ppocr_det import TextDetector
|
||||
from .ch_ppocr_rec import TextRecognizer
|
||||
@@ -60,6 +61,8 @@ class RapidOCR:
|
||||
self.max_side_len = global_config["max_side_len"]
|
||||
self.min_side_len = global_config["min_side_len"]
|
||||
|
||||
self.cal_rec_boxes = CalRecBoxes()
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
img_content: Union[str, np.ndarray, bytes, Path],
|
||||
@@ -71,12 +74,12 @@ class RapidOCR:
|
||||
use_det = self.use_det if use_det is None else use_det
|
||||
use_cls = self.use_cls if use_cls is None else use_cls
|
||||
use_rec = self.use_rec if use_rec is None else use_rec
|
||||
|
||||
return_word_box = False
|
||||
if kwargs:
|
||||
box_thresh = kwargs.get("box_thresh", 0.5)
|
||||
unclip_ratio = kwargs.get("unclip_ratio", 1.6)
|
||||
text_score = kwargs.get("text_score", 0.5)
|
||||
|
||||
return_word_box = kwargs.get("return_word_box", False)
|
||||
self.text_det.postprocess_op.box_thresh = box_thresh
|
||||
self.text_det.postprocess_op.unclip_ratio = unclip_ratio
|
||||
self.text_score = text_score
|
||||
@@ -103,7 +106,17 @@ class RapidOCR:
|
||||
img, cls_res, cls_elapse = self.text_cls(img)
|
||||
|
||||
if use_rec:
|
||||
rec_res, rec_elapse = self.text_rec(img)
|
||||
rec_res, rec_elapse = self.text_rec(img, return_word_box)
|
||||
|
||||
if dt_boxes is not None and rec_res is not None and return_word_box:
|
||||
rec_res = self.cal_rec_boxes(img, dt_boxes, rec_res)
|
||||
for rec_res_i in rec_res:
|
||||
if rec_res_i[2]:
|
||||
rec_res_i[2] = (
|
||||
self._get_origin_points(rec_res_i[2], op_record, raw_h, raw_w)
|
||||
.astype(np.int32)
|
||||
.tolist()
|
||||
)
|
||||
|
||||
if dt_boxes is not None and rec_res is not None:
|
||||
dt_boxes = self._get_origin_points(dt_boxes, op_record, raw_h, raw_w)
|
||||
@@ -237,7 +250,7 @@ class RapidOCR:
|
||||
raw_h: int,
|
||||
raw_w: int,
|
||||
) -> np.ndarray:
|
||||
dt_boxes_array = np.array(dt_boxes)
|
||||
dt_boxes_array = np.array(dt_boxes).astype(np.float32)
|
||||
for op in reversed(list(op_record.keys())):
|
||||
v = op_record[op]
|
||||
if "padding" in op:
|
||||
@@ -284,9 +297,11 @@ class RapidOCR:
|
||||
if not dt_boxes or not rec_res or len(dt_boxes) <= 0:
|
||||
return None, None
|
||||
|
||||
ocr_res = [
|
||||
[box.tolist(), res[0], res[1]] for box, res in zip(dt_boxes, rec_res)
|
||||
], [det_elapse, cls_elapse, rec_elapse]
|
||||
ocr_res = [[box.tolist(), *res] for box, res in zip(dt_boxes, rec_res)], [
|
||||
det_elapse,
|
||||
cls_elapse,
|
||||
rec_elapse,
|
||||
]
|
||||
return ocr_res
|
||||
|
||||
def filter_result(
|
||||
@@ -299,7 +314,7 @@ class RapidOCR:
|
||||
|
||||
filter_boxes, filter_rec_res = [], []
|
||||
for box, rec_reuslt in zip(dt_boxes, rec_res):
|
||||
text, score = rec_reuslt
|
||||
text, score = rec_reuslt[0], rec_reuslt[1]
|
||||
if float(score) >= self.text_score:
|
||||
filter_boxes.append(box)
|
||||
filter_rec_res.append(rec_reuslt)
|
||||
|
||||
+19
-57
@@ -12,7 +12,7 @@ root_dir = Path(__file__).resolve().parent.parent
|
||||
sys.path.append(str(root_dir))
|
||||
|
||||
from rapidocr_openvino import LoadImageError, RapidOCR
|
||||
from tests.base_module import BaseModule, download_file
|
||||
from tests.base_module import download_file
|
||||
|
||||
engine = RapidOCR()
|
||||
tests_dir = root_dir / "tests" / "test_files"
|
||||
@@ -199,61 +199,23 @@ def test_input_three_ndim_one_channel():
|
||||
|
||||
result, _ = engine(img)
|
||||
|
||||
assert result[0][1] == "正品促销"
|
||||
assert len(result) == 17
|
||||
assert len(result) >= 17
|
||||
|
||||
|
||||
def test_det():
|
||||
module_name = "ch_ppocr_det"
|
||||
class_name = "TextDetector"
|
||||
|
||||
base = BaseModule(package_name)
|
||||
TextDetector = base.init_module(module_name, class_name)
|
||||
|
||||
yaml_path = base.package_dir / module_name / "config.yaml"
|
||||
config = base.read_yaml(str(yaml_path))
|
||||
config["model_path"] = str(base.package_dir / config["model_path"])
|
||||
|
||||
text_det = TextDetector(config)
|
||||
img_path = base.tests_dir / "test_files" / "text_det.jpg"
|
||||
img = cv2.imread(str(img_path))
|
||||
dt_boxes, elapse = text_det(img)
|
||||
assert dt_boxes.shape == (18, 4, 2)
|
||||
|
||||
|
||||
def test_cls():
|
||||
module_name = "ch_ppocr_cls"
|
||||
class_name = "TextClassifier"
|
||||
|
||||
base = BaseModule(package_name=package_name)
|
||||
TextClassifier = base.init_module(module_name, class_name)
|
||||
|
||||
yaml_path = base.package_dir / module_name / "config.yaml"
|
||||
config = base.read_yaml(str(yaml_path))
|
||||
config["model_path"] = str(base.package_dir / config["model_path"])
|
||||
|
||||
text_cls = TextClassifier(config)
|
||||
|
||||
img_path = base.tests_dir / "test_files" / "text_cls.jpg"
|
||||
img = cv2.imread(str(img_path))
|
||||
result = text_cls([img])
|
||||
assert result[1][0][0] == "180"
|
||||
|
||||
|
||||
def test_rec():
|
||||
module_name = "ch_ppocr_rec"
|
||||
class_name = "TextRecognizer"
|
||||
|
||||
base = BaseModule(package_name)
|
||||
TextRecognizer = base.init_module(module_name, class_name)
|
||||
|
||||
yaml_path = base.package_dir / module_name / "config.yaml"
|
||||
config = base.read_yaml(str(yaml_path))
|
||||
config["model_path"] = str(base.package_dir / config["model_path"])
|
||||
|
||||
text_rec = TextRecognizer(config)
|
||||
|
||||
img_path = base.tests_dir / "test_files" / "text_rec.jpg"
|
||||
img = cv2.imread(str(img_path))
|
||||
rec_res, elapse = text_rec(img)
|
||||
assert rec_res[0][0] == "韩国小馆"
|
||||
# @pytest.mark.parametrize(
|
||||
# "img_name,words",
|
||||
# [
|
||||
# (
|
||||
# "black_font_color_transparent.png",
|
||||
# ["我", "是", "中", "国", "人"],
|
||||
# ),
|
||||
# (
|
||||
# "text_vertical_words.png",
|
||||
# ["已", "取", "之", "時", "不", "參", "一", "人", "見", "而"],
|
||||
# ),
|
||||
# ],
|
||||
# )
|
||||
# def test_word_ocr(img_name: str, words: List[str]):
|
||||
# img_path = tests_dir / img_name
|
||||
# result, _ = engine(img_path, return_word_box=True)
|
||||
# assert result[0][4] == words
|
||||
Reference in new issue
Block a user