AI/ragflow
1
0
Fork 0
mirror of https://git.mirrors.martin98.com/https://github.com/infiniflow/ragflow.git synced 2025-06-04 11:24:00 +08:00
Code Issues Packages Projects Releases Wiki Activity

Fix t_recognizer.py after model updating. (#4330)

Browse Source 
### What problem does this PR solve?

#4230

### Type of change

- [x] Bug Fix (non-breaking change which fixes an issue)
This commit is contained in:
Kevin Hu 2025-01-02 17:00:11 +08:00 committed by GitHub
parent df22ead841
commit 59a78408be
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 11 additions and 18 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

29
deepdoc/vision/t_recognizer.py
View File

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