From a7b193d77bd5bd9302c6ad5fbf92d67fe734c515 Mon Sep 17 00:00:00 2001 From: zcxGGmu <72263081+zcxGGmu@users.noreply.github.com> Date: Sat, 18 Jul 2026 18:22:33 +0800 Subject: [PATCH] fix: align pdf table structure coordinates (#17016) ### What problem does this PR solve? Table structure recognition rows, columns, headers, and spans are produced in cropped table image coordinates, while OCR boxes are matched later in page-cumulative coordinates. Comparing those boxes without normalization can skip or misassign table row and column metadata. Closes #16992. ### What is changed? - Map TSR components from cropped or rotated table-image coordinates back into page-cumulative coordinates before matching OCR boxes. - Reuse one inverse rotation transform for rotated OCR boxes and TSR components. - Keep TSR layout ids in the same `table-N` form used by table OCR boxes. - Sort columns by mapped page x-coordinate after coordinate normalization. - Add focused unit coverage for page offsets, zoom scaling, and 90/180/270 degree rotated tables. ### Type of change - [x] Bug fix - [x] Test coverage ### How has this been tested? - `uv run --group test pytest test/unit_test/deepdoc/parser/test_pdf_parser_table_coordinates.py -q` - `uv run --no-sync --group test pytest --confcutdir=test/unit_test/deepdoc/parser test/unit_test/deepdoc/parser/test_pdf_parser_table_coordinates.py -q` - `uv run ruff check deepdoc/parser/pdf_parser.py test/unit_test/deepdoc/parser/test_pdf_parser_table_coordinates.py` - `uv run --no-sync python -m py_compile deepdoc/parser/pdf_parser.py test/unit_test/deepdoc/parser/test_pdf_parser_table_coordinates.py` - `git diff --check` A later dependency-sync attempt was blocked while resolving the `en-core-web-sm` wheel from GitHub, and the repository-level unit-test conftest can try to download missing NLTK `wordnet` data when it is not already present locally. The focused parser test above does not require that data fixture. --------- Co-authored-by: zq --- build.sh | 5 +- deepdoc/parser/pdf_parser.py | 87 +++-- internal/deepdoc/parser/pdf/table_extract.go | 4 +- .../deepdoc/parser/pdf/table_extract_test.go | 97 ++++++ internal/deepdoc/parser/pdf/util/crop.go | 23 ++ internal/deepdoc/parser/pdf/util/crop_test.go | 27 ++ .../test_pdf_parser_table_coordinates.py | 306 ++++++++++++++++++ 7 files changed, 511 insertions(+), 38 deletions(-) create mode 100644 internal/deepdoc/parser/pdf/table_extract_test.go create mode 100644 test/unit_test/deepdoc/parser/test_pdf_parser_table_coordinates.py diff --git a/build.sh b/build.sh index 04c683cd51..24b4bffd9f 100755 --- a/build.sh +++ b/build.sh @@ -376,8 +376,9 @@ setup_cgo_env() { ;; Darwin) export CGO_LDFLAGS="$CGO_LDFLAGS \ - -framework CoreFoundation -framework Security \ - -framework SystemConfiguration -liconv -lresolv" + -framework CoreGraphics -framework CoreFoundation \ + -framework Security -framework SystemConfiguration \ + -liconv -lresolv -lc++" ;; esac diff --git a/deepdoc/parser/pdf_parser.py b/deepdoc/parser/pdf_parser.py index 5953829291..802286690e 100644 --- a/deepdoc/parser/pdf_parser.py +++ b/deepdoc/parser/pdf_parser.py @@ -455,6 +455,18 @@ class RAGFlowPdfParser: return best_angle, best_img, results + @staticmethod + def _map_clockwise_rotated_point_to_original(x, y, angle, width, height): + if angle == 0: + return x, y + if angle == 90: + return y, height - x + if angle == 180: + return width - x, height - y + if angle == 270: + return width - y, x + return x, y + def _table_transformer_job(self, ZM, auto_rotate=True): """ Process table structure recognition. @@ -480,7 +492,7 @@ class RAGFlowPdfParser: assert len(self.page_layout) == len(self.page_images) # Collect layout info for all tables - table_layouts = [] # [(page, table_layout, left, top, right, bott), ...] + table_layouts = [] table_index = 0 for p, tbls in enumerate(self.page_layout): # for page @@ -488,16 +500,17 @@ class RAGFlowPdfParser: tbcnt.append(len(tbls)) if not tbls: continue - for tb in tbls: # for table + for page_table_index, tb in enumerate(tbls): # for table left, top, right, bott = tb["x0"] - MARGIN, tb["top"] - MARGIN, tb["x1"] + MARGIN, tb["bottom"] + MARGIN left *= ZM top *= ZM right *= ZM bott *= ZM - pos.append((left, top, p, table_index)) # Add page and table_index + layoutno = f"table-{page_table_index}" + pos.append((left, top, p, table_index, layoutno)) # Record table layout info - table_layouts.append({"page": p, "table_index": table_index, "layout": tb, "coords": (left, top, right, bott)}) + table_layouts.append({"page": p, "table_index": table_index, "layoutno": layoutno, "layout": tb, "coords": (left, top, right, bott)}) # Crop table image table_img = self.page_images[p].crop((left, top, right, bott)) @@ -538,7 +551,28 @@ class RAGFlowPdfParser: if auto_rotate: self._ocr_rotated_tables(ZM, table_layouts, recos, tbcnt) - # Process TSR results (keep original logic but handle rotated coordinates) + def _map_tsr_component_to_page_space(component, table_pos): + crop_left, crop_top, page, table_index, _ = table_pos + rotation_info = self.table_rotations.get(table_index, {}) + angle = rotation_info.get("best_angle", 0) + original_pos = rotation_info.get("original_pos", (crop_left, crop_top, crop_left, crop_top)) + width = original_pos[2] - original_pos[0] + height = original_pos[3] - original_pos[1] + points = [ + (component["x0_rotated"], component["top_rotated"]), + (component["x1_rotated"], component["top_rotated"]), + (component["x0_rotated"], component["bottom_rotated"]), + (component["x1_rotated"], component["bottom_rotated"]), + ] + mapped = [self._map_clockwise_rotated_point_to_original(x, y, angle, width, height) for x, y in points] + xs = [p[0] for p in mapped] + ys = [p[1] for p in mapped] + component["x0"] = min(xs) / ZM + crop_left / ZM + component["x1"] = max(xs) / ZM + crop_left / ZM + component["top"] = min(ys) / ZM + crop_top / ZM + self.page_cum_height[page] + component["bottom"] = max(ys) / ZM + crop_top / ZM + self.page_cum_height[page] + + # Process TSR results and align structure boxes with page-cumulative OCR boxes. tbcnt = np.cumsum(tbcnt) for i in range(len(tbcnt) - 1): # for page pg = [] @@ -551,11 +585,10 @@ class RAGFlowPdfParser: it["top_rotated"] = it["top"] it["bottom_rotated"] = it["bottom"] - # For rotated tables, coordinate transformation to page space requires rotation - # Since we already re-OCR'd on rotated image, keep simple processing here it["pn"] = poss[j][2] # page number - it["layoutno"] = j + it["layoutno"] = poss[j][4] it["table_index"] = poss[j][3] # table index + _map_tsr_component_to_page_space(it, poss[j]) pg.append(it) self.tb_cpns.extend(pg) @@ -568,7 +601,7 @@ class RAGFlowPdfParser: headers = gather(r".*header$") rows = gather(r".* (row|header)") spans = gather(r".*spanning") - clmns = sorted([r for r in self.tb_cpns if re.match(r"table column$", r["label"])], key=lambda x: (x["pn"], x["layoutno"], x["x0_rotated"] if "x0_rotated" in x else x["x0"])) + clmns = sorted([r for r in self.tb_cpns if re.match(r"table column$", r["label"])], key=lambda x: (x["pn"], x["layoutno"], x["x0"])) clmns = Recognizer.layouts_cleanup(self.boxes, clmns, 5, 0.5) for b in self.boxes: @@ -648,28 +681,12 @@ class RAGFlowPdfParser: insert_at += 1 return insert_at - def _map_rotated_point(x, y, angle, width, height): - # Map a point from rotated image coords back to original image coords. - if angle == 0: - return x, y - if angle == 90: - # clockwise 90: original->rotated (x', y') = (y, width - x) - # inverse: - return width - y, x - if angle == 180: - return width - x, height - y - if angle == 270: - # clockwise 270: original->rotated (x', y') = (height - y, x) - # inverse: - return y, height - x - return x, y - - def _insert_ocr_boxes(ocr_results, page_index, table_x0, table_top, insert_at, table_index, best_angle, table_w_px, table_h_px): + def _insert_ocr_boxes(ocr_results, page_index, crop_left, crop_top, insert_at, table_index, layoutno, best_angle, table_w_px, table_h_px): added = 0 for bbox, (text, conf) in ocr_results: if conf < 0.5: continue - mapped = [_map_rotated_point(p[0], p[1], best_angle, table_w_px, table_h_px) for p in bbox] + mapped = [self._map_clockwise_rotated_point_to_original(p[0], p[1], best_angle, table_w_px, table_h_px) for p in bbox] x_coords = [p[0] for p in mapped] y_coords = [p[1] for p in mapped] box_x0 = min(x_coords) / ZM @@ -678,13 +695,13 @@ class RAGFlowPdfParser: box_bottom = max(y_coords) / ZM new_box = { "text": text, - "x0": box_x0 + table_x0, - "x1": box_x1 + table_x0, - "top": box_top + table_top + self.page_cum_height[page_index], - "bottom": box_bottom + table_top + self.page_cum_height[page_index], + "x0": box_x0 + crop_left / ZM, + "x1": box_x1 + crop_left / ZM, + "top": box_top + crop_top / ZM + self.page_cum_height[page_index], + "bottom": box_bottom + crop_top / ZM + self.page_cum_height[page_index], "page_number": page_index + self.page_from, "layout_type": "table", - "layoutno": f"table-{table_index}", + "layoutno": layoutno, "_rotated": True, "_rotation_angle": best_angle, "_table_index": table_index, @@ -702,6 +719,7 @@ class RAGFlowPdfParser: table_index = tbl_info["table_index"] page = tbl_info["page"] layout = tbl_info["layout"] + layoutno = tbl_info["layoutno"] left, top, right, bott = tbl_info["coords"] rotation_info = self.table_rotations.get(table_index, {}) @@ -738,10 +756,11 @@ class RAGFlowPdfParser: added = _insert_ocr_boxes( ocr_results, page, - table_x0, - table_top, + left, + top, insert_at, table_index, + layoutno, best_angle, table_w_px, table_h_px, diff --git a/internal/deepdoc/parser/pdf/table_extract.go b/internal/deepdoc/parser/pdf/table_extract.go index acb4b050c1..3f9030914c 100644 --- a/internal/deepdoc/parser/pdf/table_extract.go +++ b/internal/deepdoc/parser/pdf/table_extract.go @@ -108,8 +108,8 @@ func (p *Parser) processOneTable(ctx context.Context, pageImg image.Image, boxes p.ocrTableCells(ctx, cells, tsrImg, docAnalyzer) } for i := range cells { - cells[i].X0, cells[i].Y0 = util.MapRotatedPointToOriginal(cells[i].X0, cells[i].Y0, bestAngle, origW, origH) - cells[i].X1, cells[i].Y1 = util.MapRotatedPointToOriginal(cells[i].X1, cells[i].Y1, bestAngle, origW, origH) + cells[i].X0, cells[i].Y0, cells[i].X1, cells[i].Y1 = util.MapRotatedRectToOriginal( + cells[i].X0, cells[i].Y0, cells[i].X1, cells[i].Y1, bestAngle, origW, origH) } } firstCellTop := 1e9 diff --git a/internal/deepdoc/parser/pdf/table_extract_test.go b/internal/deepdoc/parser/pdf/table_extract_test.go new file mode 100644 index 0000000000..5c8faf46c7 --- /dev/null +++ b/internal/deepdoc/parser/pdf/table_extract_test.go @@ -0,0 +1,97 @@ +package pdf + +import ( + "context" + "image" + "testing" + + tbl "ragflow/internal/deepdoc/parser/pdf/table" + pdf "ragflow/internal/deepdoc/parser/pdf/type" +) + +type orientationScoringDoc struct{} + +func (d *orientationScoringDoc) DLA(_ context.Context, _ image.Image) ([]pdf.DLARegion, error) { + return nil, nil +} + +func (d *orientationScoringDoc) TSR(_ context.Context, _ image.Image) ([]pdf.TSRCell, error) { + return nil, nil +} + +func (d *orientationScoringDoc) OCRDetect(_ context.Context, img image.Image) ([]pdf.OCRBox, error) { + regions := 1 + if img.Bounds().Dy() > img.Bounds().Dx() { + regions = 5 + } + boxes := make([]pdf.OCRBox, regions) + for i := range boxes { + x0 := float64((i + 1) * 10) + boxes[i] = pdf.OCRBox{ + X0: x0, Y0: 10, + X1: x0 + 5, Y1: 10, + X2: x0 + 5, Y2: 30, + X3: x0, Y3: 30, + } + } + return boxes, nil +} + +func (d *orientationScoringDoc) OCRRecognize(_ context.Context, _ image.Image) ([]pdf.OCRText, error) { + return nil, nil +} + +func (d *orientationScoringDoc) Health() bool { return true } + +type staticTableBuilder struct { + cells []pdf.TSRCell +} + +func (b *staticTableBuilder) Name() string { return "static" } + +func (b *staticTableBuilder) DetectCells(_ context.Context, _ image.Image) ([]pdf.TSRCell, error) { + return append([]pdf.TSRCell(nil), b.cells...), nil +} + +func (b *staticTableBuilder) GroupCells(cells []pdf.TSRCell) [][]pdf.TSRCell { + if len(cells) == 0 { + return nil + } + return [][]pdf.TSRCell{{cells[0]}} +} + +func TestProcessOneTable_AutoRotateNormalizesCellBounds(t *testing.T) { + autoRotate := true + cfg := pdf.DefaultParserConfig() + cfg.AutoRotateTables = &autoRotate + cfg.SkipOCR = true + p := NewParser(cfg) + + pageImg := image.NewRGBA(image.Rect(0, 0, 320, 220)) + boxes := []pdf.TextBox{ + {X0: 10, X1: 60, Top: 10, Bottom: 30, Text: "cell", LayoutType: pdf.LayoutTypeTable}, + } + match := tbl.TableMatch{ + Region: pdf.DLARegion{X0: 10, Y0: 10, X1: 210, Y1: 110, Label: pdf.LayoutTypeTable}, + BoxIdx: []int{0}, + } + builder := &staticTableBuilder{ + cells: []pdf.TSRCell{ + {X0: 10, Y0: 20, X1: 60, Y1: 80, Label: "table row"}, + }, + } + + item := p.processOneTable(context.Background(), pageImg, boxes, 0, &orientationScoringDoc{}, builder, match, pdf.DlaScale) + if len(item.Cells) != 1 { + t.Fatalf("cells = %d, want 1", len(item.Cells)) + } + + got := item.Cells[0] + if got.X0 != 20 || got.Y0 != 45 || got.X1 != 80 || got.Y1 != 95 { + t.Errorf("cell bounds = (%.0f,%.0f,%.0f,%.0f), want (20,45,80,95)", + got.X0, got.Y0, got.X1, got.Y1) + } + if got.X0 > got.X1 || got.Y0 > got.Y1 { + t.Fatalf("cell bounds are inverted: (%.0f,%.0f,%.0f,%.0f)", got.X0, got.Y0, got.X1, got.Y1) + } +} diff --git a/internal/deepdoc/parser/pdf/util/crop.go b/internal/deepdoc/parser/pdf/util/crop.go index e5ff629af6..f770fe1bb3 100644 --- a/internal/deepdoc/parser/pdf/util/crop.go +++ b/internal/deepdoc/parser/pdf/util/crop.go @@ -535,6 +535,29 @@ func MapRotatedPointToOriginal(x, y float64, angle int, origW, origH int) (float } } +// MapRotatedRectToOriginal maps a rotated-image rectangle back into original +// image coordinates and normalizes the resulting bounds. For 90°/270° rotation, +// mapping only two diagonal corners can invert X/Y bounds; mapping all four +// corners preserves the enclosing rectangle. +func MapRotatedRectToOriginal(x0, y0, x1, y1 float64, angle int, origW, origH int) (float64, float64, float64, float64) { + points := [][2]float64{ + {x0, y0}, + {x1, y0}, + {x0, y1}, + {x1, y1}, + } + minX, minY := math.Inf(1), math.Inf(1) + maxX, maxY := math.Inf(-1), math.Inf(-1) + for _, p := range points { + x, y := MapRotatedPointToOriginal(p[0], p[1], angle, origW, origH) + minX = math.Min(minX, x) + minY = math.Min(minY, y) + maxX = math.Max(maxX, x) + maxY = math.Max(maxY, y) + } + return minX, minY, maxX, maxY +} + // CropImageRegion crops a pdf.DLARegion from an image with a 3% margin // (matching Python's _table_transformer_job: w*0.03, h*0.03). func CropImageRegion(img image.Image, r pdf.DLARegion) (image.Image, error) { diff --git a/internal/deepdoc/parser/pdf/util/crop_test.go b/internal/deepdoc/parser/pdf/util/crop_test.go index 6e1b0d3e24..f9a8b1b061 100644 --- a/internal/deepdoc/parser/pdf/util/crop_test.go +++ b/internal/deepdoc/parser/pdf/util/crop_test.go @@ -245,6 +245,33 @@ func TestMapRotatedPointToOriginal(t *testing.T) { } } +func TestMapRotatedRectToOriginal_NormalizesBounds(t *testing.T) { + tests := []struct { + name string + angle int + wantX0, wantY0 float64 + wantX1, wantY1 float64 + }{ + {name: "zero", angle: 0, wantX0: 10, wantY0: 20, wantX1: 60, wantY1: 80}, + {name: "ninety", angle: 90, wantX0: 20, wantY0: 39, wantX1: 80, wantY1: 89}, + {name: "one-eighty", angle: 180, wantX0: 139, wantY0: 19, wantX1: 189, wantY1: 79}, + {name: "two-seventy", angle: 270, wantX0: 119, wantY0: 10, wantX1: 179, wantY1: 60}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + gotX0, gotY0, gotX1, gotY1 := MapRotatedRectToOriginal(10, 20, 60, 80, tt.angle, 200, 100) + if gotX0 != tt.wantX0 || gotY0 != tt.wantY0 || gotX1 != tt.wantX1 || gotY1 != tt.wantY1 { + t.Errorf("got (%.0f,%.0f,%.0f,%.0f), want (%.0f,%.0f,%.0f,%.0f)", + gotX0, gotY0, gotX1, gotY1, tt.wantX0, tt.wantY0, tt.wantX1, tt.wantY1) + } + if gotX0 > gotX1 || gotY0 > gotY1 { + t.Fatalf("mapped rectangle is inverted: (%.0f,%.0f,%.0f,%.0f)", gotX0, gotY0, gotX1, gotY1) + } + }) + } +} + func colorEqual(a, b color.Color) bool { ar, ag, ab, aa := a.RGBA() br, bg, bb, ba := b.RGBA() diff --git a/test/unit_test/deepdoc/parser/test_pdf_parser_table_coordinates.py b/test/unit_test/deepdoc/parser/test_pdf_parser_table_coordinates.py new file mode 100644 index 0000000000..e0a3ab2d3a --- /dev/null +++ b/test/unit_test/deepdoc/parser/test_pdf_parser_table_coordinates.py @@ -0,0 +1,306 @@ +import importlib.util +import sys +from pathlib import Path +from types import ModuleType, SimpleNamespace + +import pytest + + +def _load_pdf_parser(monkeypatch): + repo_root = Path(__file__).resolve().parents[4] + + _stub_module(monkeypatch, "pdfplumber") + _stub_module(monkeypatch, "pypdf", PdfReader=object) + _stub_module(monkeypatch, "huggingface_hub", snapshot_download=lambda **_kwargs: "") + _stub_module(monkeypatch, "xgboost", Booster=object) + _stub_module(monkeypatch, "sklearn") + _stub_module(monkeypatch, "sklearn.cluster", KMeans=object) + _stub_module(monkeypatch, "sklearn.metrics", silhouette_score=lambda *_args, **_kwargs: 0) + + common_mod = _stub_module(monkeypatch, "common") + common_mod.__path__ = [str(repo_root / "common")] + _stub_module(monkeypatch, "common.constants", MAXIMUM_PAGE_NUMBER=1024) + _stub_module(monkeypatch, "common.file_utils", get_project_base_directory=lambda: str(repo_root)) + _stub_module(monkeypatch, "common.settings", PARALLEL_DEVICES=1) + _stub_module(monkeypatch, "common.misc_utils", thread_pool_exec=lambda fn, *args, **kwargs: fn(*args, **kwargs)) + + deepdoc_mod = _stub_module(monkeypatch, "deepdoc") + deepdoc_mod.__path__ = [str(repo_root / "deepdoc")] + parser_mod = _stub_module(monkeypatch, "deepdoc.parser") + parser_mod.__path__ = [str(repo_root / "deepdoc" / "parser")] + _stub_module(monkeypatch, "deepdoc.parser.utils", extract_pdf_outlines=lambda *_args, **_kwargs: []) + _stub_module( + monkeypatch, + "deepdoc.vision", + OCR=object, + AscendLayoutRecognizer=object, + LayoutRecognizer=object, + Recognizer=_FakeRecognizer, + TableStructureRecognizer=object, + ) + + rag_mod = _stub_module(monkeypatch, "rag") + rag_mod.__path__ = [str(repo_root / "rag")] + _stub_module(monkeypatch, "rag.nlp", rag_tokenizer=SimpleNamespace(tokenize=lambda text: text)) + prompts_mod = _stub_module(monkeypatch, "rag.prompts") + prompts_mod.__path__ = [str(repo_root / "rag" / "prompts")] + _stub_module(monkeypatch, "rag.prompts.generator", vision_llm_describe_prompt="") + + module_name = "test_pdf_parser_unit_module" + module_path = repo_root / "deepdoc" / "parser" / "pdf_parser.py" + spec = importlib.util.spec_from_file_location(module_name, module_path) + module = importlib.util.module_from_spec(spec) + monkeypatch.setitem(sys.modules, module_name, module) + spec.loader.exec_module(module) + return module + + +def _stub_module(monkeypatch, name, **attrs): + module = ModuleType(name) + for key, value in attrs.items(): + setattr(module, key, value) + monkeypatch.setitem(sys.modules, name, module) + return module + + +class _FakeRecognizer: + @staticmethod + def sort_Y_firstly(arr, _threshold): + return sorted(arr, key=lambda item: (item["top"], item["x0"])) + + @staticmethod + def layouts_cleanup(_boxes, layouts, _far=2, _thr=0.7): + return layouts + + @staticmethod + def overlapped_area(a, b, ratio=True): + x0 = max(a["x0"], b["x0"]) + x1 = min(a["x1"], b["x1"]) + top = max(a["top"], b["top"]) + bottom = min(a["bottom"], b["bottom"]) + if x1 <= x0 or bottom <= top: + return 0 + area = (x1 - x0) * (bottom - top) + if ratio: + area /= (a["x1"] - a["x0"]) * (a["bottom"] - a["top"]) + return area + + @staticmethod + def find_overlapped_with_threshold(box, boxes, thr=0.3): + best_i = None + best = thr + best_reverse = 0 + for i, candidate in enumerate(boxes): + overlap = _FakeRecognizer.overlapped_area(box, candidate) + reverse = _FakeRecognizer.overlapped_area(candidate, box) + if (overlap, reverse) < (best, best_reverse): + continue + best_i = i + best = overlap + best_reverse = reverse + return best_i + + @staticmethod + def find_horizontally_tightest_fit(box, boxes): + min_distance = 1000000 + min_i = None + for i, candidate in enumerate(boxes): + if box.get("layoutno", "0") != candidate.get("layoutno", "0"): + continue + distance = min( + abs(box["x0"] - candidate["x0"]), + abs(box["x1"] - candidate["x1"]), + abs(box["x0"] + box["x1"] - candidate["x1"] - candidate["x0"]) / 2, + ) + if distance < min_distance: + min_distance = distance + min_i = i + return min_i + + +class _FakeImage: + def __init__(self, width=300, height=400): + self.size = (width, height) + + def crop(self, box): + left, top, right, bottom = box + return _FakeImage(right - left, bottom - top) + + def __array__(self, dtype=None): + import numpy as np + + return np.zeros((int(self.size[1]), int(self.size[0]), 3), dtype=dtype or np.uint8) + + +class _FakeTableDetector: + def __init__(self, zoom, angle=0, crop_width=140, crop_height=90): + self.zoom = zoom + self.angle = angle + self.crop_width = crop_width * zoom + self.crop_height = crop_height * zoom + + def __call__(self, _imgs): + z = self.zoom + rows = [ + _scale_bbox((15, 20, 125, 35), z), + _scale_bbox((15, 50, 125, 65), z), + ] + columns = [ + _scale_bbox((15, 20, 55, 65), z), + _scale_bbox((80, 20, 125, 65), z), + ] + rows = [_rotate_bbox_clockwise(row, self.angle, self.crop_width, self.crop_height) for row in rows] + columns = [_rotate_bbox_clockwise(column, self.angle, self.crop_width, self.crop_height) for column in columns] + return [ + [ + _component("table row", rows[0]), + _component("table row", rows[1]), + _component("table column", columns[0]), + _component("table column", columns[1]), + ] + ] + + +class _FakeOcr: + def __init__(self, angle, crop_width=140, crop_height=90): + self.angle = angle + self.crop_width = crop_width + self.crop_height = crop_height + + def __call__(self, _img_array): + boxes = [ + ("A1", (15, 20, 55, 35)), + ("B2", (80, 50, 125, 65)), + ] + return [ + ( + _bbox_points(_rotate_bbox_clockwise(bbox, self.angle, self.crop_width, self.crop_height)), + (text, 0.99), + ) + for text, bbox in boxes + ] + + +def _component(label, bbox): + x0, top, x1, bottom = bbox + return {"label": label, "x0": x0, "x1": x1, "top": top, "bottom": bottom} + + +def _scale_bbox(bbox, zoom): + x0, top, x1, bottom = bbox + return x0 * zoom, top * zoom, x1 * zoom, bottom * zoom + + +def _bbox_points(bbox): + x0, top, x1, bottom = bbox + return [(x0, top), (x1, top), (x1, bottom), (x0, bottom)] + + +def _rotate_bbox_clockwise(bbox, angle, width, height): + points = [_rotate_point_clockwise(x, y, angle, width, height) for x, y in _bbox_points(bbox)] + xs = [p[0] for p in points] + ys = [p[1] for p in points] + return min(xs), min(ys), max(xs), max(ys) + + +def _rotate_point_clockwise(x, y, angle, width, height): + if angle == 0: + return x, y + if angle == 90: + return height - y, x + if angle == 180: + return width - x, height - y + if angle == 270: + return y, width - x + raise ValueError(f"unsupported angle: {angle}") + + +@pytest.mark.p1 +@pytest.mark.parametrize(("page_index", "page_offset", "zoom"), [(0, 0, 1), (1, 500, 2)]) +def test_table_transformer_maps_tsr_crop_coordinates_to_page_coordinates(monkeypatch, page_index, page_offset, zoom): + module = _load_pdf_parser(monkeypatch) + parser = module.RAGFlowPdfParser.__new__(module.RAGFlowPdfParser) + parser.page_from = 0 + parser.page_cum_height = [0] if page_index == 0 else [0, page_offset] + parser.page_images = [_FakeImage() for _ in range(page_index + 1)] + parser.page_layout = [[] for _ in range(page_index + 1)] + parser.page_layout[page_index] = [{"type": "table", "x0": 100, "top": 200, "x1": 220, "bottom": 270}] + parser.tbl_det = _FakeTableDetector(zoom) + parser.boxes = [ + { + "text": "A1", + "layout_type": "table", + "layoutno": "table-0", + "page_number": page_index, + "x0": 105, + "x1": 145, + "top": page_offset + 210, + "bottom": page_offset + 225, + }, + { + "text": "B2", + "layout_type": "table", + "layoutno": "table-0", + "page_number": page_index, + "x0": 170, + "x1": 215, + "top": page_offset + 240, + "bottom": page_offset + 255, + }, + ] + + parser._table_transformer_job(ZM=zoom, auto_rotate=False) + + assert [box["R"] for box in parser.boxes] == [0, 1] + assert [box["R_top"] for box in parser.boxes] == [page_offset + 210, page_offset + 240] + assert [box["C"] for box in parser.boxes] == [0, 1] + assert [box["C_left"] for box in parser.boxes] == [105, 170] + + +@pytest.mark.p1 +@pytest.mark.parametrize("angle", [90, 180, 270]) +def test_table_transformer_keeps_rotated_ocr_and_tsr_coordinates_aligned(monkeypatch, angle): + module = _load_pdf_parser(monkeypatch) + parser = module.RAGFlowPdfParser.__new__(module.RAGFlowPdfParser) + parser.page_from = 0 + parser.page_cum_height = [0] + parser.page_images = [_FakeImage()] + parser.page_layout = [[{"type": "table", "x0": 100, "top": 200, "x1": 220, "bottom": 270}]] + parser.tbl_det = _FakeTableDetector(zoom=1, angle=angle) + parser.ocr = _FakeOcr(angle) + parser._evaluate_table_orientation = lambda table_img: ( + angle, + _FakeImage(table_img.size[1], table_img.size[0]) if angle in (90, 270) else _FakeImage(*table_img.size), + {}, + ) + parser.boxes = [ + { + "text": "old A1", + "layout_type": "table", + "layoutno": "table-0", + "page_number": 0, + "x0": 105, + "x1": 145, + "top": 210, + "bottom": 225, + }, + { + "text": "old B2", + "layout_type": "table", + "layoutno": "table-0", + "page_number": 0, + "x0": 170, + "x1": 215, + "top": 240, + "bottom": 255, + }, + ] + + parser._table_transformer_job(ZM=1, auto_rotate=True) + + assert [box["text"] for box in parser.boxes] == ["A1", "B2"] + assert [box["layoutno"] for box in parser.boxes] == ["table-0", "table-0"] + assert [box["R"] for box in parser.boxes] == [0, 1] + assert [box["R_top"] for box in parser.boxes] == [210, 240] + assert [box["C"] for box in parser.boxes] == [0, 1] + assert [box["C_left"] for box in parser.boxes] == [105, 170]