Kevin Hu
commited on
Commit
·
1955ace
1
Parent(s):
54cd6b8
Fix t_recognizer.py after model updating. (#4330)
Browse files### What problem does this PR solve?
#4230
### Type of change
- [x] Bug Fix (non-breaking change which fixes an issue)
- deepdoc/vision/t_recognizer.py +11 -18
deepdoc/vision/t_recognizer.py
CHANGED
|
@@ -23,8 +23,7 @@ sys.path.insert(
|
|
| 23 |
'../../')))
|
| 24 |
|
| 25 |
from deepdoc.vision.seeit import draw_box
|
| 26 |
-
from deepdoc.vision import
|
| 27 |
-
from api.utils.file_utils import get_project_base_directory
|
| 28 |
import argparse
|
| 29 |
import re
|
| 30 |
import numpy as np
|
|
@@ -33,13 +32,7 @@ import numpy as np
|
|
| 33 |
def main(args):
|
| 34 |
images, outputs = init_in_out(args)
|
| 35 |
if args.mode.lower() == "layout":
|
| 36 |
-
|
| 37 |
-
detr = Recognizer(
|
| 38 |
-
labels,
|
| 39 |
-
"layout",
|
| 40 |
-
os.path.join(
|
| 41 |
-
get_project_base_directory(),
|
| 42 |
-
"rag/res/deepdoc/"))
|
| 43 |
if args.mode.lower() == "tsr":
|
| 44 |
labels = TableStructureRecognizer.labels
|
| 45 |
detr = TableStructureRecognizer()
|
|
@@ -64,7 +57,7 @@ def main(args):
|
|
| 64 |
|
| 65 |
def get_table_html(img, tb_cpns, ocr):
|
| 66 |
boxes = ocr(np.array(img))
|
| 67 |
-
boxes =
|
| 68 |
[{"x0": b[0][0], "x1": b[1][0],
|
| 69 |
"top": b[0][1], "text": t[0],
|
| 70 |
"bottom": b[-1][1],
|
|
@@ -75,26 +68,26 @@ def get_table_html(img, tb_cpns, ocr):
|
|
| 75 |
|
| 76 |
def gather(kwd, fzy=10, ption=0.6):
|
| 77 |
nonlocal boxes
|
| 78 |
-
eles =
|
| 79 |
[r for r in tb_cpns if re.match(kwd, r["label"])], fzy)
|
| 80 |
-
eles =
|
| 81 |
-
return
|
| 82 |
|
| 83 |
headers = gather(r".*header$")
|
| 84 |
rows = gather(r".* (row|header)")
|
| 85 |
spans = gather(r".*spanning")
|
| 86 |
clmns = sorted([r for r in tb_cpns if re.match(
|
| 87 |
r"table column$", r["label"])], key=lambda x: x["x0"])
|
| 88 |
-
clmns =
|
| 89 |
|
| 90 |
for b in boxes:
|
| 91 |
-
ii =
|
| 92 |
if ii is not None:
|
| 93 |
b["R"] = ii
|
| 94 |
b["R_top"] = rows[ii]["top"]
|
| 95 |
b["R_bott"] = rows[ii]["bottom"]
|
| 96 |
|
| 97 |
-
ii =
|
| 98 |
if ii is not None:
|
| 99 |
b["H_top"] = headers[ii]["top"]
|
| 100 |
b["H_bott"] = headers[ii]["bottom"]
|
|
@@ -102,13 +95,13 @@ def get_table_html(img, tb_cpns, ocr):
|
|
| 102 |
b["H_right"] = headers[ii]["x1"]
|
| 103 |
b["H"] = ii
|
| 104 |
|
| 105 |
-
ii =
|
| 106 |
if ii is not None:
|
| 107 |
b["C"] = ii
|
| 108 |
b["C_left"] = clmns[ii]["x0"]
|
| 109 |
b["C_right"] = clmns[ii]["x1"]
|
| 110 |
|
| 111 |
-
ii =
|
| 112 |
if ii is not None:
|
| 113 |
b["H_top"] = spans[ii]["top"]
|
| 114 |
b["H_bott"] = spans[ii]["bottom"]
|
|
|
|
| 23 |
'../../')))
|
| 24 |
|
| 25 |
from deepdoc.vision.seeit import draw_box
|
| 26 |
+
from deepdoc.vision import LayoutRecognizer, TableStructureRecognizer, OCR, init_in_out
|
|
|
|
| 27 |
import argparse
|
| 28 |
import re
|
| 29 |
import numpy as np
|
|
|
|
| 32 |
def main(args):
|
| 33 |
images, outputs = init_in_out(args)
|
| 34 |
if args.mode.lower() == "layout":
|
| 35 |
+
detr = LayoutRecognizer("layout")
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 36 |
if args.mode.lower() == "tsr":
|
| 37 |
labels = TableStructureRecognizer.labels
|
| 38 |
detr = TableStructureRecognizer()
|
|
|
|
| 57 |
|
| 58 |
def get_table_html(img, tb_cpns, ocr):
|
| 59 |
boxes = ocr(np.array(img))
|
| 60 |
+
boxes = LayoutRecognizer.sort_Y_firstly(
|
| 61 |
[{"x0": b[0][0], "x1": b[1][0],
|
| 62 |
"top": b[0][1], "text": t[0],
|
| 63 |
"bottom": b[-1][1],
|
|
|
|
| 68 |
|
| 69 |
def gather(kwd, fzy=10, ption=0.6):
|
| 70 |
nonlocal boxes
|
| 71 |
+
eles = LayoutRecognizer.sort_Y_firstly(
|
| 72 |
[r for r in tb_cpns if re.match(kwd, r["label"])], fzy)
|
| 73 |
+
eles = LayoutRecognizer.layouts_cleanup(boxes, eles, 5, ption)
|
| 74 |
+
return LayoutRecognizer.sort_Y_firstly(eles, 0)
|
| 75 |
|
| 76 |
headers = gather(r".*header$")
|
| 77 |
rows = gather(r".* (row|header)")
|
| 78 |
spans = gather(r".*spanning")
|
| 79 |
clmns = sorted([r for r in tb_cpns if re.match(
|
| 80 |
r"table column$", r["label"])], key=lambda x: x["x0"])
|
| 81 |
+
clmns = LayoutRecognizer.layouts_cleanup(boxes, clmns, 5, 0.5)
|
| 82 |
|
| 83 |
for b in boxes:
|
| 84 |
+
ii = LayoutRecognizer.find_overlapped_with_threashold(b, rows, thr=0.3)
|
| 85 |
if ii is not None:
|
| 86 |
b["R"] = ii
|
| 87 |
b["R_top"] = rows[ii]["top"]
|
| 88 |
b["R_bott"] = rows[ii]["bottom"]
|
| 89 |
|
| 90 |
+
ii = LayoutRecognizer.find_overlapped_with_threashold(b, headers, thr=0.3)
|
| 91 |
if ii is not None:
|
| 92 |
b["H_top"] = headers[ii]["top"]
|
| 93 |
b["H_bott"] = headers[ii]["bottom"]
|
|
|
|
| 95 |
b["H_right"] = headers[ii]["x1"]
|
| 96 |
b["H"] = ii
|
| 97 |
|
| 98 |
+
ii = LayoutRecognizer.find_horizontally_tightest_fit(b, clmns)
|
| 99 |
if ii is not None:
|
| 100 |
b["C"] = ii
|
| 101 |
b["C_left"] = clmns[ii]["x0"]
|
| 102 |
b["C_right"] = clmns[ii]["x1"]
|
| 103 |
|
| 104 |
+
ii = LayoutRecognizer.find_overlapped_with_threashold(b, spans, thr=0.3)
|
| 105 |
if ii is not None:
|
| 106 |
b["H_top"] = spans[ii]["top"]
|
| 107 |
b["H_bott"] = spans[ii]["bottom"]
|