Skip to content

Commit

Permalink
Fix t_recognizer.py after model updating. (#4330)
Browse files Browse the repository at this point in the history
### What problem does this PR solve?

#4230

### Type of change

- [x] Bug Fix (non-breaking change which fixes an issue)
  • Loading branch information
KevinHuSh authored Jan 2, 2025
1 parent df22ead commit 59a7840
Showing 1 changed file with 11 additions and 18 deletions.
29 changes: 11 additions & 18 deletions deepdoc/vision/t_recognizer.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,8 +23,7 @@
'../../')))

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
Expand All @@ -33,13 +32,7 @@
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()
Expand All @@ -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],
Expand All @@ -75,40 +68,40 @@ 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"]
b["H_left"] = headers[ii]["x0"]
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"]
Expand Down

0 comments on commit 59a7840

Please sign in to comment.