276 lines
11 KiB
Python
276 lines
11 KiB
Python
|
# -*- encoding: utf-8 -*-
|
|||
|
|
# @Author: SWHL
|
|||
|
|
# @Contact: liekkaskono@163.com
|
|||
|
|
import argparse
|
|||
|
|
import importlib
|
|||
|
|
import logging
|
|||
|
|
import time
|
|||
|
|
import traceback
|
|||
|
|
from pathlib import Path
|
|||
|
|
from typing import List, Optional, Tuple, Union, Dict, Any
|
|||
|
|
import numpy as np
|
|||
|
|
import cv2
|
|||
|
|
|
|||
|
|
from wired_table_rec.table_line_rec import TableLineRecognition
|
|||
|
|
from wired_table_rec.table_line_rec_plus import TableLineRecognitionPlus
|
|||
|
|
from .table_recover import TableRecover
|
|||
|
|
from .utils import InputType, LoadImage
|
|||
|
|
from .utils_table_recover import (
|
|||
|
|
match_ocr_cell,
|
|||
|
|
plot_html_table,
|
|||
|
|
box_4_2_poly_to_box_4_1,
|
|||
|
|
get_rotate_crop_image,
|
|||
|
|
sorted_ocr_boxes,
|
|||
|
|
gather_ocr_list_by_row,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
cur_dir = Path(__file__).resolve().parent
|
|||
|
|
default_model_path = cur_dir / "models" / "cycle_center_net_v1.onnx"
|
|||
|
|
default_model_path_v2 = cur_dir / "models" / "cycle_center_net_v2.onnx"
|
|||
|
|
|
|||
|
|
def convert_float32_to_float(data):
|
|||
|
|
if isinstance(data, list):
|
|||
|
|
return [convert_float32_to_float(item) for item in data]
|
|||
|
|
elif isinstance(data, np.ndarray):
|
|||
|
|
return data.tolist() # 将 NumPy 数组转换为 Python 列表
|
|||
|
|
elif isinstance(data, np.float32):
|
|||
|
|
return float(data) # 将 float32 转换为 float
|
|||
|
|
else:
|
|||
|
|
return data # 其他类型保持不变
|
|||
|
|
|
|||
|
|
|
|||
|
|
class WiredTableRecognition:
|
|||
|
|
def __init__(self, table_model_path: Union[str, Path] = None, version="v2"):
|
|||
|
|
self.load_img = LoadImage()
|
|||
|
|
if version == "v2":
|
|||
|
|
model_path = table_model_path if table_model_path else default_model_path_v2
|
|||
|
|
self.table_line_rec = TableLineRecognitionPlus(str(model_path))
|
|||
|
|
else:
|
|||
|
|
model_path = table_model_path if table_model_path else default_model_path
|
|||
|
|
self.table_line_rec = TableLineRecognition(str(model_path))
|
|||
|
|
|
|||
|
|
self.table_recover = TableRecover()
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
self.ocr = importlib.import_module("rapidocr_onnxruntime").RapidOCR()
|
|||
|
|
except ModuleNotFoundError:
|
|||
|
|
self.ocr = None
|
|||
|
|
|
|||
|
|
def __call__(
|
|||
|
|
self,
|
|||
|
|
img: InputType,
|
|||
|
|
ocr_result: Optional[List[Union[List[List[float]], str, str]]] = None,
|
|||
|
|
**kwargs,
|
|||
|
|
) -> Tuple[str, float, Any, Any, Any, Any, Any]:
|
|||
|
|
if self.ocr is None and ocr_result is None:
|
|||
|
|
raise ValueError(
|
|||
|
|
"One of two conditions must be met: ocr_result is not empty, or rapidocr_onnxruntime is installed."
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
s = time.perf_counter()
|
|||
|
|
rec_again = True
|
|||
|
|
need_ocr = True
|
|||
|
|
col_threshold = 15
|
|||
|
|
row_threshold = 10
|
|||
|
|
if kwargs:
|
|||
|
|
rec_again = kwargs.get("rec_again", True)
|
|||
|
|
need_ocr = kwargs.get("need_ocr", True)
|
|||
|
|
col_threshold = kwargs.get("col_threshold", 15)
|
|||
|
|
row_threshold = kwargs.get("row_threshold", 10)
|
|||
|
|
img = self.load_img(img)
|
|||
|
|
polygons, rotated_polygons = self.table_line_rec(img, **kwargs)
|
|||
|
|
if polygons is None:
|
|||
|
|
logging.warning("polygons is None.")
|
|||
|
|
return "", 0.0, None, None, None, None, None
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
table_res, logi_points = self.table_recover(
|
|||
|
|
rotated_polygons, row_threshold, col_threshold
|
|||
|
|
)
|
|||
|
|
# 将坐标由逆时针转为顺时针方向，后续处理与无线表格对齐
|
|||
|
|
polygons[:, 1, :], polygons[:, 3, :] = (
|
|||
|
|
polygons[:, 3, :].copy(),
|
|||
|
|
polygons[:, 1, :].copy(),
|
|||
|
|
)
|
|||
|
|
if not need_ocr:
|
|||
|
|
sorted_polygons, idx_list = sorted_ocr_boxes(
|
|||
|
|
[box_4_2_poly_to_box_4_1(box) for box in polygons]
|
|||
|
|
)
|
|||
|
|
return (
|
|||
|
|
"",
|
|||
|
|
time.perf_counter() - s,
|
|||
|
|
sorted_polygons,
|
|||
|
|
logi_points[idx_list],
|
|||
|
|
[],
|
|||
|
|
{},
|
|||
|
|
[],
|
|||
|
|
)
|
|||
|
|
if ocr_result is None and need_ocr:
|
|||
|
|
ocr_result, _ = self.ocr(img)
|
|||
|
|
cell_box_det_map, not_match_orc_boxes = match_ocr_cell(ocr_result, polygons)
|
|||
|
|
# 如果有识别框没有ocr结果，直接进行rec补充
|
|||
|
|
cell_box_det_map = self.re_rec(img, polygons, cell_box_det_map, rec_again)
|
|||
|
|
# 转换为中间格式，修正识别框坐标,将物理识别框，逻辑识别框，ocr识别框整合为dict，方便后续处理
|
|||
|
|
t_rec_ocr_list = self.transform_res(cell_box_det_map, polygons, logi_points)
|
|||
|
|
# 将每个单元格中的ocr识别结果排序和同行合并，输出的html能完整保留文字的换行格式
|
|||
|
|
t_rec_ocr_list = self.sort_and_gather_ocr_res(t_rec_ocr_list)
|
|||
|
|
# cell_box_map =
|
|||
|
|
logi_points = [t_box_ocr["t_logic_box"] for t_box_ocr in t_rec_ocr_list]
|
|||
|
|
cell_box_det_map = {
|
|||
|
|
i: [ocr_box_and_text[1] for ocr_box_and_text in t_box_ocr["t_ocr_res"]]
|
|||
|
|
for i, t_box_ocr in enumerate(t_rec_ocr_list)
|
|||
|
|
}
|
|||
|
|
# cell_box_det_map_with_cor = {
|
|||
|
|
# i: [[ocr_box_and_text[0],ocr_box_and_text[1]] for ocr_box_and_text in t_box_ocr["t_ocr_res"]]
|
|||
|
|
# for i, t_box_ocr in enumerate(t_rec_ocr_list)
|
|||
|
|
# }
|
|||
|
|
table_str = plot_html_table(logi_points, cell_box_det_map)
|
|||
|
|
ocr_boxes_res = [
|
|||
|
|
box_4_2_poly_to_box_4_1(ori_ocr[0]) for ori_ocr in ocr_result
|
|||
|
|
]
|
|||
|
|
sorted_ocr_boxes_res, _ = sorted_ocr_boxes(ocr_boxes_res)
|
|||
|
|
sorted_polygons = [box_4_2_poly_to_box_4_1(box) for box in polygons]
|
|||
|
|
sorted_logi_points = logi_points
|
|||
|
|
table_elapse = time.perf_counter() - s
|
|||
|
|
|
|||
|
|
# for i, box in enumerate(not_match_orc_boxes):
|
|||
|
|
# cell_box_det_map[len(cell_box_det_map) + 1] = box[1]
|
|||
|
|
cell_box_det_map_with_cor = {}
|
|||
|
|
# 遍历cell_box_det_map的每个key
|
|||
|
|
for i, box in enumerate(cell_box_det_map):
|
|||
|
|
cell_box_det_map_with_cor[i] = convert_float32_to_float(sorted_polygons[i]), cell_box_det_map[i]
|
|||
|
|
|
|||
|
|
# print("not_match_orc_boxes", not_match_orc_boxes)
|
|||
|
|
|
|||
|
|
except Exception:
|
|||
|
|
logging.warning(traceback.format_exc())
|
|||
|
|
return "", 0.0, None, None, None, None, None
|
|||
|
|
return (
|
|||
|
|
table_str,
|
|||
|
|
table_elapse,
|
|||
|
|
sorted_polygons,
|
|||
|
|
sorted_logi_points,
|
|||
|
|
sorted_ocr_boxes_res,
|
|||
|
|
cell_box_det_map_with_cor,
|
|||
|
|
not_match_orc_boxes
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
def transform_res(
|
|||
|
|
self,
|
|||
|
|
cell_box_det_map: Dict[int, List[any]],
|
|||
|
|
polygons: np.ndarray,
|
|||
|
|
logi_points: List[np.ndarray],
|
|||
|
|
) -> List[Dict[str, any]]:
|
|||
|
|
res = []
|
|||
|
|
for i in range(len(polygons)):
|
|||
|
|
ocr_res_list = cell_box_det_map.get(i)
|
|||
|
|
if not ocr_res_list:
|
|||
|
|
continue
|
|||
|
|
xmin = min([ocr_box[0][0][0] for ocr_box in ocr_res_list])
|
|||
|
|
ymin = min([ocr_box[0][0][1] for ocr_box in ocr_res_list])
|
|||
|
|
xmax = max([ocr_box[0][2][0] for ocr_box in ocr_res_list])
|
|||
|
|
ymax = max([ocr_box[0][2][1] for ocr_box in ocr_res_list])
|
|||
|
|
dict_res = {
|
|||
|
|
# xmin,xmax,ymin,ymax
|
|||
|
|
"t_box": [xmin, ymin, xmax, ymax],
|
|||
|
|
# row_start,row_end,col_start,col_end
|
|||
|
|
"t_logic_box": logi_points[i].tolist(),
|
|||
|
|
# [[xmin,xmax,ymin,ymax], text]
|
|||
|
|
"t_ocr_res": [
|
|||
|
|
[box_4_2_poly_to_box_4_1(ocr_det[0]), ocr_det[1]]
|
|||
|
|
for ocr_det in ocr_res_list
|
|||
|
|
],
|
|||
|
|
}
|
|||
|
|
res.append(dict_res)
|
|||
|
|
return res
|
|||
|
|
|
|||
|
|
def sort_and_gather_ocr_res(self, res):
|
|||
|
|
for i, dict_res in enumerate(res):
|
|||
|
|
_, sorted_idx = sorted_ocr_boxes(
|
|||
|
|
[ocr_det[0] for ocr_det in dict_res["t_ocr_res"]], threhold=0.3
|
|||
|
|
)
|
|||
|
|
dict_res["t_ocr_res"] = [dict_res["t_ocr_res"][i] for i in sorted_idx]
|
|||
|
|
dict_res["t_ocr_res"] = gather_ocr_list_by_row(
|
|||
|
|
dict_res["t_ocr_res"], threhold=0.3
|
|||
|
|
)
|
|||
|
|
return res
|
|||
|
|
|
|||
|
|
def re_rec(
|
|||
|
|
self,
|
|||
|
|
img: np.ndarray,
|
|||
|
|
sorted_polygons: np.ndarray,
|
|||
|
|
cell_box_map: Dict[int, List[str]],
|
|||
|
|
rec_again=True,
|
|||
|
|
) -> Dict[int, List[Any]]:
|
|||
|
|
"""找到poly对应为空的框，尝试将直接将poly框直接送到识别中"""
|
|||
|
|
for i in range(sorted_polygons.shape[0]):
|
|||
|
|
if cell_box_map.get(i):
|
|||
|
|
continue
|
|||
|
|
if not rec_again:
|
|||
|
|
box = sorted_polygons[i]
|
|||
|
|
cell_box_map[i] = [[box, "", 1]]
|
|||
|
|
continue
|
|||
|
|
crop_img = get_rotate_crop_image(img, sorted_polygons[i])
|
|||
|
|
pad_img = cv2.copyMakeBorder(
|
|||
|
|
crop_img, 5, 5, 100, 100, cv2.BORDER_CONSTANT, value=(255, 255, 255)
|
|||
|
|
)
|
|||
|
|
rec_res, _ = self.ocr(pad_img, use_det=False, use_cls=True, use_rec=True)
|
|||
|
|
box = sorted_polygons[i]
|
|||
|
|
text = [rec[0] for rec in rec_res]
|
|||
|
|
scores = [rec[1] for rec in rec_res]
|
|||
|
|
cell_box_map[i] = [[box, "".join(text), min(scores)]]
|
|||
|
|
return cell_box_map
|
|||
|
|
|
|||
|
|
def re_rec_high_precise(
|
|||
|
|
self,
|
|||
|
|
img: np.ndarray,
|
|||
|
|
sorted_polygons: np.ndarray,
|
|||
|
|
cell_box_map: Dict[int, List[str]],
|
|||
|
|
) -> Dict[int, List[any]]:
|
|||
|
|
"""找到poly对应为空的框，尝试将直接将poly框直接送到识别中"""
|
|||
|
|
#
|
|||
|
|
cell_box_map = {}
|
|||
|
|
for i in range(sorted_polygons.shape[0]):
|
|||
|
|
if cell_box_map.get(i):
|
|||
|
|
continue
|
|||
|
|
crop_img = get_rotate_crop_image(img, sorted_polygons[i])
|
|||
|
|
pad_img = cv2.copyMakeBorder(
|
|||
|
|
crop_img, 10, 10, 10, 10, cv2.BORDER_CONSTANT, value=(255, 255, 255)
|
|||
|
|
)
|
|||
|
|
rec_res, _ = self.ocr(pad_img, use_det=True, use_cls=True, use_rec=True)
|
|||
|
|
if not rec_res:
|
|||
|
|
det_boxes = [sorted_polygons[i]]
|
|||
|
|
text = [""]
|
|||
|
|
scores = [1.0]
|
|||
|
|
else:
|
|||
|
|
det_boxes = [rec[0] for rec in rec_res]
|
|||
|
|
text = [rec[1] for rec in rec_res]
|
|||
|
|
scores = [rec[2] for rec in rec_res]
|
|||
|
|
cell_box_map[i] = [
|
|||
|
|
[box, text, score] for box, text, score in zip(det_boxes, text, scores)
|
|||
|
|
]
|
|||
|
|
return cell_box_map
|
|||
|
|
|
|||
|
|
|
|||
|
|
def main():
|
|||
|
|
parser = argparse.ArgumentParser()
|
|||
|
|
parser.add_argument("-img", "--img_path", type=str, required=True)
|
|||
|
|
args = parser.parse_args()
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
ocr_engine = importlib.import_module("rapidocr_onnxruntime").RapidOCR()
|
|||
|
|
except ModuleNotFoundError as exc:
|
|||
|
|
raise ModuleNotFoundError(
|
|||
|
|
"Please install the rapidocr_onnxruntime by pip install rapidocr_onnxruntime."
|
|||
|
|
) from exc
|
|||
|
|
|
|||
|
|
table_rec = WiredTableRecognition()
|
|||
|
|
ocr_result, _ = ocr_engine(args.img_path)
|
|||
|
|
table_str, elapse = table_rec(args.img_path, ocr_result)
|
|||
|
|
print(table_str)
|
|||
|
|
print(f"cost: {elapse:.5f}")
|
|||
|
|
|
|||
|
|
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
main()
|