pip install wired_table_rec lineless_table_rec table_cls
pip install rapidocr 

from pathlib import Path

from wired_table_rec.utils.utils import VisTable
from table_cls import TableCls
from wired_table_rec.main import WiredTableInput, WiredTableRecognition
from lineless_table_rec.main import LinelessTableInput, LinelessTableRecognition
from rapidocr import RapidOCR


if __name__ == "__main__":
    # Init
    wired_input = WiredTableInput()
    lineless_input = LinelessTableInput()
    wired_engine = WiredTableRecognition(wired_input)
    lineless_engine = LinelessTableRecognition(lineless_input)
    viser = VisTable()
    # 默认小yolo模型(0.1s)，可切换为精度更高yolox(0.25s),更快的qanything(0.07s)模型或paddle模型(0.03s)
    table_cls = TableCls()
    img_path = f"tests/test_files/table.jpg"

    cls, elasp = table_cls(img_path)
    if cls == "wired":
        table_engine = wired_engine
    else:
        table_engine = lineless_engine

    # 使用RapidOCR输入
    ocr_engine = RapidOCR()
    rapid_ocr_output = ocr_engine(img_path, return_word_box=True)
    ocr_result = list(
        zip(rapid_ocr_output.boxes, rapid_ocr_output.txts, rapid_ocr_output.scores)
    )
    table_results = table_engine(
        img_path, ocr_result=ocr_result
    )

    # 使用单字识别
    # word_results = rapid_ocr_output.word_results
    # ocr_result = [
    #     [word_result[2], word_result[0], word_result[1]] for word_result in word_results
    # ]
    # table_results = table_engine(
    #     img_path, ocr_result=ocr_result, enhance_box_line=False
    # )

    # Save
    # save_dir = Path("outputs")
    # save_dir.mkdir(parents=True, exist_ok=True)
    # 
    # save_html_path = f"outputs/{Path(img_path).stem}.html"
    # save_drawed_path = f"outputs/{Path(img_path).stem}_table_vis{Path(img_path).suffix}"
    # save_logic_path = (
    #     f"outputs/{Path(img_path).stem}_table_vis_logic{Path(img_path).suffix}"
    # )

    # Visualize table rec result
    # vis_imged = viser(
    #     img_path, table_results, save_html_path, save_drawed_path, save_logic_path
    # )

# 将单字box转换为行识别同样的结构)
from rapidocr import RapidOCR
img_path = "tests/test_files/wired/table4.jpg"
ocr_engine = RapidOCR()
rapid_ocr_output = ocr_engine(img_path, return_word_box=True)
word_results = rapid_ocr_output.word_results
ocr_result = [
    [word_result[2], word_result[0], word_result[1]] for word_result in word_results
]

import cv2

img_path = f'tests/test_files/wired/squeeze_error.jpeg'
from wired_table_rec.utils import ImageOrientationCorrector

img_orientation_corrector = ImageOrientationCorrector()
img = cv2.imread(img_path)
img = img_orientation_corrector(img)
cv2.imwrite(f'img_rotated.jpg', img)

pip install rapid-table-det

import os
import cv2
from rapid_table_det.utils import img_loader, visuallize, extract_table_img
from rapid_table_det.inference import TableDetector
table_det = TableDetector()
img_path = f"tests/test_files/chip.jpg"
result, elapse = table_det(img_path)
img = img_loader(img_path)
extract_img = img.copy()
#可能有多表格
for i, res in enumerate(result):
    box = res["box"]
    lt, rt, rb, lb = res["lt"], res["rt"], res["rb"], res["lb"]
    # 带识别框和左上角方向位置
    img = visuallize(img, box, lt, rt, rb, lb)
    # 透视变换提取表格图片
    wrapped_img = extract_table_img(extract_img.copy(), lt, rt, rb, lb)
#     cv2.imwrite(f"{out_dir}/{file_name}-extract-{i}.jpg", wrapped_img)
# cv2.imwrite(f"{out_dir}/{file_name}-visualize.jpg", img)

# 输入(WiredTableInput/LinelessTableInput)
@dataclass
class WiredTableInput:
    model_type: Optional[str] = "unet" #unet/cycle_center_net
    model_path: Union[str, Path, None, Dict[str, str]] = None
    use_cuda: bool = False
    device: str = "cpu"
    
@dataclass
class LinelessTableInput:
    model_type: Optional[str] = "lore" #lore
    model_path: Union[str, Path, None, Dict[str, str]] = None
    use_cuda: bool = False
    device: str = "cpu"
    
# 输出(WiredTableOutput/LinelessTableOutput)
@dataclass
class WiredTableOutput:
    pred_html: Optional[str] = None
    cell_bboxes: Optional[np.ndarray] = None
    logic_points: Optional[np.ndarray] = None
    elapse: Optional[float] = None
    
@dataclass
class LinelessTableOutput:
    pred_html: Optional[str] = None
    cell_bboxes: Optional[np.ndarray] = None
    logic_points: Optional[np.ndarray] = None
    elapse: Optional[float] = None

wired_table_rec = WiredTableRecognition(WiredTableInput())
table_results = wired_table_rec(
    img, # 图片 Union[str, np.ndarray, bytes, Path, PIL.Image.Image]
    ocr_result, # 输入rapidOCR识别结果，不传默认使用内部rapidocr模型
    enhance_box_line=True, # 识别框切割增强(关闭避免多余切割，开启减少漏切割)，默认为True
    col_threshold=15, # 识别框左边界x坐标差值小于col_threshold的默认同列
    row_threshold=10, # 识别框上边界y坐标差值小于row_threshold的默认同行
    rotated_fix=True, # wiredV2支持，轻度旋转(-45°~45°)矫正，默认为True
    need_ocr=True, # 是否进行OCR识别, 默认为True
)
lineless_table_rec = LinelessTableRecognition(LinelessTableInput())
table_results = lineless_table_rec(
    img, # 图片 Union[str, np.ndarray, bytes, Path, PIL.Image.Image]
    ocr_result, # 输入rapidOCR识别结果，不传默认使用内部rapidocr模型
    need_ocr=True, # 是否进行OCR识别, 默认为True
)