diff --git a/deepdoc/vision/t_recognizer.py b/deepdoc/vision/t_recognizer.py index a5efc5dc..d0c0cb49 100644 --- a/deepdoc/vision/t_recognizer.py +++ b/deepdoc/vision/t_recognizer.py @@ -23,8 +23,7 @@ sys.path.insert( '../../'))) from deepdoc.vision.seeit import draw_box -from deepdoc.vision import Recognizer, LayoutRecognizer, TableStructureRecognizer, OCR, init_in_out -from api.utils.file_utils import get_project_base_directory +from deepdoc.vision import LayoutRecognizer, TableStructureRecognizer, OCR, init_in_out import argparse import re import numpy as np @@ -33,13 +32,7 @@ import numpy as np def main(args): images, outputs = init_in_out(args) if args.mode.lower() == "layout": - labels = LayoutRecognizer.labels - detr = Recognizer( - labels, - "layout", - os.path.join( - get_project_base_directory(), - "rag/res/deepdoc/")) + detr = LayoutRecognizer("layout") if args.mode.lower() == "tsr": labels = TableStructureRecognizer.labels detr = TableStructureRecognizer() @@ -64,7 +57,7 @@ def main(args): def get_table_html(img, tb_cpns, ocr): boxes = ocr(np.array(img)) - boxes = Recognizer.sort_Y_firstly( + boxes = LayoutRecognizer.sort_Y_firstly( [{"x0": b[0][0], "x1": b[1][0], "top": b[0][1], "text": t[0], "bottom": b[-1][1], @@ -75,26 +68,26 @@ def get_table_html(img, tb_cpns, ocr): def gather(kwd, fzy=10, ption=0.6): nonlocal boxes - eles = Recognizer.sort_Y_firstly( + eles = LayoutRecognizer.sort_Y_firstly( [r for r in tb_cpns if re.match(kwd, r["label"])], fzy) - eles = Recognizer.layouts_cleanup(boxes, eles, 5, ption) - return Recognizer.sort_Y_firstly(eles, 0) + eles = LayoutRecognizer.layouts_cleanup(boxes, eles, 5, ption) + return LayoutRecognizer.sort_Y_firstly(eles, 0) headers = gather(r".*header$") rows = gather(r".* (row|header)") spans = gather(r".*spanning") clmns = sorted([r for r in tb_cpns if re.match( r"table column$", r["label"])], key=lambda x: x["x0"]) - clmns = Recognizer.layouts_cleanup(boxes, clmns, 5, 0.5) + clmns = LayoutRecognizer.layouts_cleanup(boxes, clmns, 5, 0.5) for b in boxes: - ii = Recognizer.find_overlapped_with_threashold(b, rows, thr=0.3) + ii = LayoutRecognizer.find_overlapped_with_threashold(b, rows, thr=0.3) if ii is not None: b["R"] = ii b["R_top"] = rows[ii]["top"] b["R_bott"] = rows[ii]["bottom"] - ii = Recognizer.find_overlapped_with_threashold(b, headers, thr=0.3) + ii = LayoutRecognizer.find_overlapped_with_threashold(b, headers, thr=0.3) if ii is not None: b["H_top"] = headers[ii]["top"] b["H_bott"] = headers[ii]["bottom"] @@ -102,13 +95,13 @@ def get_table_html(img, tb_cpns, ocr): b["H_right"] = headers[ii]["x1"] b["H"] = ii - ii = Recognizer.find_horizontally_tightest_fit(b, clmns) + ii = LayoutRecognizer.find_horizontally_tightest_fit(b, clmns) if ii is not None: b["C"] = ii b["C_left"] = clmns[ii]["x0"] b["C_right"] = clmns[ii]["x1"] - ii = Recognizer.find_overlapped_with_threashold(b, spans, thr=0.3) + ii = LayoutRecognizer.find_overlapped_with_threashold(b, spans, thr=0.3) if ii is not None: b["H_top"] = spans[ii]["top"] b["H_bott"] = spans[ii]["bottom"]