Просмотр исходного кода

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

### What problem does this PR solve?

#4230

### Type of change

- [x] Bug Fix (non-breaking change which fixes an issue)
tags/v0.16.0
Kevin Hu 10 месяцев назад
Родитель
Сommit
59a78408be
Аккаунт пользователя с таким Email не найден
1 измененных файлов: 11 добавлений и 18 удалений
  1. 11
    18
      deepdoc/vision/t_recognizer.py

+ 11
- 18
deepdoc/vision/t_recognizer.py Просмотреть файл

'../../'))) '../../')))


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 api.utils.file_utils import get_project_base_directory
from deepdoc.vision import LayoutRecognizer, TableStructureRecognizer, OCR, init_in_out
import argparse import argparse
import re import re
import numpy as np 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 = Recognizer(
labels,
"layout",
os.path.join(
get_project_base_directory(),
"rag/res/deepdoc/"))
detr = LayoutRecognizer("layout")
if args.mode.lower() == "tsr": if args.mode.lower() == "tsr":
labels = TableStructureRecognizer.labels labels = TableStructureRecognizer.labels
detr = TableStructureRecognizer() detr = TableStructureRecognizer()


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],


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)
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$") 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"]
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"]

Загрузка…
Отмена
Сохранить