langparse 0.1.0__py3-none-any.whl

langparse/engines/pdf/deepdoc/recognizer.py

@@ -0,0 +1,418 @@
+#
+#  Copyright 2025 The InfiniFlow Authors. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+#
+import gc
+import logging
+import math
+import numpy as np
+import cv2
+from functools import cmp_to_key
+
+
+from .model_loader import default_model_dir
+from .operators import *  # noqa: F403
+from .operators import preprocess
+from . import operators
+from .ocr import load_model
+
+
+class Recognizer:
+    def __init__(self, label_list, task_name, model_dir=None):
+        """
+        If you have trouble downloading HuggingFace models, -_^ this might help!!
+
+        For Linux:
+        export HF_ENDPOINT=https://hf-mirror.com
+
+        For Windows:
+        Good luck
+        ^_-
+
+        """
+        if not model_dir:
+            model_dir = str(default_model_dir())
+        self.ort_sess, self.run_options = load_model(model_dir, task_name)
+        self.input_names = [node.name for node in self.ort_sess.get_inputs()]
+        self.output_names = [node.name for node in self.ort_sess.get_outputs()]
+        self.input_shape = self.ort_sess.get_inputs()[0].shape[2:4]
+        self.label_list = label_list
+
+    @staticmethod
+    def sort_Y_firstly(arr, threshold):
+        def cmp(c1, c2):
+            diff = c1["top"] - c2["top"]
+            if abs(diff) < threshold:
+                diff = c1["x0"] - c2["x0"]
+            return diff
+
+        arr = sorted(arr, key=cmp_to_key(cmp))
+        return arr
+
+    @staticmethod
+    def sort_X_firstly(arr, threshold):
+        def cmp(c1, c2):
+            diff = c1["x0"] - c2["x0"]
+            if abs(diff) < threshold:
+                diff = c1["top"] - c2["top"]
+            return diff
+
+        arr = sorted(arr, key=cmp_to_key(cmp))
+        return arr
+
+    @staticmethod
+    def sort_C_firstly(arr, thr=0):
+        # sort using y1 first and then x1
+        # sorted(arr, key=lambda r: (r["x0"], r["top"]))
+        arr = Recognizer.sort_X_firstly(arr, thr)
+        for i in range(len(arr) - 1):
+            for j in range(i, -1, -1):
+                # restore the order using th
+                if "C" not in arr[j] or "C" not in arr[j + 1]:
+                    continue
+                if arr[j + 1]["C"] < arr[j]["C"] or (arr[j + 1]["C"] == arr[j]["C"] and arr[j + 1]["top"] < arr[j]["top"]):
+                    tmp = arr[j]
+                    arr[j] = arr[j + 1]
+                    arr[j + 1] = tmp
+        return arr
+
+    @staticmethod
+    def sort_R_firstly(arr, thr=0):
+        # sort using y1 first and then x1
+        # sorted(arr, key=lambda r: (r["top"], r["x0"]))
+        arr = Recognizer.sort_Y_firstly(arr, thr)
+        for i in range(len(arr) - 1):
+            for j in range(i, -1, -1):
+                if "R" not in arr[j] or "R" not in arr[j + 1]:
+                    continue
+                if arr[j + 1]["R"] < arr[j]["R"] or (arr[j + 1]["R"] == arr[j]["R"] and arr[j + 1]["x0"] < arr[j]["x0"]):
+                    tmp = arr[j]
+                    arr[j] = arr[j + 1]
+                    arr[j + 1] = tmp
+        return arr
+
+    @staticmethod
+    def overlapped_area(a, b, ratio=True):
+        tp, btm, x0, x1 = a["top"], a["bottom"], a["x0"], a["x1"]
+        if b["x0"] > x1 or b["x1"] < x0:
+            return 0
+        if b["bottom"] < tp or b["top"] > btm:
+            return 0
+        x0_ = max(b["x0"], x0)
+        x1_ = min(b["x1"], x1)
+        assert x0_ <= x1_, "Bbox mismatch! T:{},B:{},X0:{},X1:{} ==> {}".format(tp, btm, x0, x1, b)
+        tp_ = max(b["top"], tp)
+        btm_ = min(b["bottom"], btm)
+        assert tp_ <= btm_, "Bbox mismatch! T:{},B:{},X0:{},X1:{} => {}".format(tp, btm, x0, x1, b)
+        ov = (btm_ - tp_) * (x1_ - x0_) if x1 - x0 != 0 and btm - tp != 0 else 0
+        if ov > 0 and ratio:
+            ov /= (x1 - x0) * (btm - tp)
+        return ov
+
+    @staticmethod
+    def layouts_cleanup(boxes, layouts, far=2, thr=0.7):
+        def not_overlapped(a, b):
+            return any([a["x1"] < b["x0"], a["x0"] > b["x1"], a["bottom"] < b["top"], a["top"] > b["bottom"]])
+
+        i = 0
+        while i + 1 < len(layouts):
+            j = i + 1
+            while j < min(i + far, len(layouts)) and (layouts[i].get("type", "") != layouts[j].get("type", "") or not_overlapped(layouts[i], layouts[j])):
+                j += 1
+            if j >= min(i + far, len(layouts)):
+                i += 1
+                continue
+            if Recognizer.overlapped_area(layouts[i], layouts[j]) < thr and Recognizer.overlapped_area(layouts[j], layouts[i]) < thr:
+                i += 1
+                continue
+
+            if layouts[i].get("score") and layouts[j].get("score"):
+                if layouts[i]["score"] > layouts[j]["score"]:
+                    layouts.pop(j)
+                else:
+                    layouts.pop(i)
+                continue
+
+            area_i, area_i_1 = 0, 0
+            for b in boxes:
+                if not not_overlapped(b, layouts[i]):
+                    area_i += Recognizer.overlapped_area(b, layouts[i], False)
+                if not not_overlapped(b, layouts[j]):
+                    area_i_1 += Recognizer.overlapped_area(b, layouts[j], False)
+
+            if area_i > area_i_1:
+                layouts.pop(j)
+            else:
+                layouts.pop(i)
+
+        return layouts
+
+    def create_inputs(self, imgs, im_info):
+        """generate input for different model type
+        Args:
+            imgs (list(numpy)): list of images (np.ndarray)
+            im_info (list(dict)): list of image info
+        Returns:
+            inputs (dict): input of model
+        """
+        inputs = {}
+
+        im_shape = []
+        scale_factor = []
+        if len(imgs) == 1:
+            inputs["image"] = np.array((imgs[0],)).astype("float32")
+            inputs["im_shape"] = np.array((im_info[0]["im_shape"],)).astype("float32")
+            inputs["scale_factor"] = np.array((im_info[0]["scale_factor"],)).astype("float32")
+            return inputs
+
+        im_shape = np.array([info["im_shape"] for info in im_info], dtype="float32")
+        scale_factor = np.array([info["scale_factor"] for info in im_info], dtype="float32")
+
+        inputs["im_shape"] = np.concatenate(im_shape, axis=0)
+        inputs["scale_factor"] = np.concatenate(scale_factor, axis=0)
+
+        imgs_shape = [[e.shape[1], e.shape[2]] for e in imgs]
+        max_shape_h = max([e[0] for e in imgs_shape])
+        max_shape_w = max([e[1] for e in imgs_shape])
+        padding_imgs = []
+        for img in imgs:
+            im_c, im_h, im_w = img.shape[:]
+            padding_im = np.zeros((im_c, max_shape_h, max_shape_w), dtype=np.float32)
+            padding_im[:, :im_h, :im_w] = img
+            padding_imgs.append(padding_im)
+        inputs["image"] = np.stack(padding_imgs, axis=0)
+        return inputs
+
+    @staticmethod
+    def find_overlapped(box, boxes_sorted_by_y, naive=False):
+        if not boxes_sorted_by_y:
+            return
+        bxs = boxes_sorted_by_y
+        s, e, ii = 0, len(bxs), 0
+        while s < e and not naive:
+            ii = (e + s) // 2
+            pv = bxs[ii]
+            if box["bottom"] < pv["top"]:
+                e = ii
+                continue
+            if box["top"] > pv["bottom"]:
+                s = ii + 1
+                continue
+            break
+        while s < ii:
+            if box["top"] > bxs[s]["bottom"]:
+                s += 1
+            break
+        while e - 1 > ii:
+            if box["bottom"] < bxs[e - 1]["top"]:
+                e -= 1
+            break
+
+        max_overlapped_i, max_overlapped = None, 0
+        for i in range(s, e):
+            ov = Recognizer.overlapped_area(bxs[i], box)
+            if ov <= max_overlapped:
+                continue
+            max_overlapped_i = i
+            max_overlapped = ov
+
+        return max_overlapped_i
+
+    @staticmethod
+    def find_horizontally_tightest_fit(box, boxes):
+        if not boxes:
+            return
+        min_dis, min_i = 1000000, None
+        for i, b in enumerate(boxes):
+            if box.get("layoutno", "0") != b.get("layoutno", "0"):
+                continue
+            # layoutno is f"table-{index}" and resets per page, so the same string
+            # names a different table on every page. Only accept a candidate that
+            # shares vertical extent with the box; top/bottom are page-cumulative,
+            # so this rejects a same-layoutno column from another page whose x range
+            # happens to be closer.
+            if min(box["bottom"], b["bottom"]) <= max(box["top"], b["top"]):
+                continue
+            dis = min(abs(box["x0"] - b["x0"]), abs(box["x1"] - b["x1"]), abs(box["x0"] + box["x1"] - b["x1"] - b["x0"]) / 2)
+            if dis < min_dis:
+                min_i = i
+                min_dis = dis
+        return min_i
+
+    @staticmethod
+    def find_overlapped_with_threshold(box, boxes, thr=0.3):
+        if not boxes:
+            return
+        max_overlapped_i, max_overlapped, _max_overlapped = None, thr, 0
+        s, e = 0, len(boxes)
+        for i in range(s, e):
+            ov = Recognizer.overlapped_area(box, boxes[i])
+            _ov = Recognizer.overlapped_area(boxes[i], box)
+            if (ov, _ov) < (max_overlapped, _max_overlapped):
+                continue
+            max_overlapped_i = i
+            max_overlapped = ov
+            _max_overlapped = _ov
+
+        return max_overlapped_i
+
+    def preprocess(self, image_list):
+        inputs = []
+        if "scale_factor" in self.input_names:
+            preprocess_ops = []
+            for op_info in [
+                {"interp": 2, "keep_ratio": False, "target_size": [800, 608], "type": "LinearResize"},
+                {"is_scale": True, "mean": [0.485, 0.456, 0.406], "std": [0.229, 0.224, 0.225], "type": "StandardizeImage"},
+                {"type": "Permute"},
+                {"stride": 32, "type": "PadStride"},
+            ]:
+                new_op_info = op_info.copy()
+                op_type = new_op_info.pop("type")
+                preprocess_ops.append(getattr(operators, op_type)(**new_op_info))
+
+            for im_path in image_list:
+                im, im_info = preprocess(im_path, preprocess_ops)
+                inputs.append({"image": np.array((im,)).astype("float32"), "scale_factor": np.array((im_info["scale_factor"],)).astype("float32")})
+        else:
+            hh, ww = self.input_shape
+            for img in image_list:
+                h, w = img.shape[:2]
+                img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+                img = cv2.resize(np.array(img).astype("float32"), (ww, hh))
+                # Scale input pixel values to 0 to 1
+                img /= 255.0
+                img = img.transpose(2, 0, 1)
+                img = img[np.newaxis, :, :, :].astype(np.float32)
+                inputs.append({self.input_names[0]: img, "scale_factor": [w / ww, h / hh]})
+        return inputs
+
+    def postprocess(self, boxes, inputs, thr):
+        if "scale_factor" in self.input_names:
+            bb = []
+            for b in boxes:
+                clsid, bbox, score = int(b[0]), b[2:], b[1]
+                if score < thr:
+                    continue
+                if clsid >= len(self.label_list):
+                    continue
+                bb.append({"type": self.label_list[clsid].lower(), "bbox": [float(t) for t in bbox.tolist()], "score": float(score)})
+            return bb
+
+        def xywh2xyxy(x):
+            # [x, y, w, h] to [x1, y1, x2, y2]
+            y = np.copy(x)
+            y[:, 0] = x[:, 0] - x[:, 2] / 2
+            y[:, 1] = x[:, 1] - x[:, 3] / 2
+            y[:, 2] = x[:, 0] + x[:, 2] / 2
+            y[:, 3] = x[:, 1] + x[:, 3] / 2
+            return y
+
+        def compute_iou(box, boxes):
+            # Compute xmin, ymin, xmax, ymax for both boxes
+            xmin = np.maximum(box[0], boxes[:, 0])
+            ymin = np.maximum(box[1], boxes[:, 1])
+            xmax = np.minimum(box[2], boxes[:, 2])
+            ymax = np.minimum(box[3], boxes[:, 3])
+
+            # Compute intersection area
+            intersection_area = np.maximum(0, xmax - xmin) * np.maximum(0, ymax - ymin)
+
+            # Compute union area
+            box_area = (box[2] - box[0]) * (box[3] - box[1])
+            boxes_area = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
+            union_area = box_area + boxes_area - intersection_area
+
+            # Compute IoU
+            iou = intersection_area / union_area
+
+            return iou
+
+        def iou_filter(boxes, scores, iou_threshold):
+            sorted_indices = np.argsort(scores)[::-1]
+
+            keep_boxes = []
+            while sorted_indices.size > 0:
+                # Pick the last box
+                box_id = sorted_indices[0]
+                keep_boxes.append(box_id)
+
+                # Compute IoU of the picked box with the rest
+                ious = compute_iou(boxes[box_id, :], boxes[sorted_indices[1:], :])
+
+                # Remove boxes with IoU over the threshold
+                keep_indices = np.where(ious < iou_threshold)[0]
+
+                # print(keep_indices.shape, sorted_indices.shape)
+                sorted_indices = sorted_indices[keep_indices + 1]
+
+            return keep_boxes
+
+        boxes = np.squeeze(boxes).T
+        # Filter out object confidence scores below threshold
+        scores = np.max(boxes[:, 4:], axis=1)
+        boxes = boxes[scores > thr, :]
+        scores = scores[scores > thr]
+        if len(boxes) == 0:
+            return []
+
+        # Get the class with the highest confidence
+        class_ids = np.argmax(boxes[:, 4:], axis=1)
+        boxes = boxes[:, :4]
+        input_shape = np.array([inputs["scale_factor"][0], inputs["scale_factor"][1], inputs["scale_factor"][0], inputs["scale_factor"][1]])
+        boxes = np.multiply(boxes, input_shape, dtype=np.float32)
+        boxes = xywh2xyxy(boxes)
+
+        unique_class_ids = np.unique(class_ids)
+        indices = []
+        for class_id in unique_class_ids:
+            class_indices = np.where(class_ids == class_id)[0]
+            class_boxes = boxes[class_indices, :]
+            class_scores = scores[class_indices]
+            class_keep_boxes = iou_filter(class_boxes, class_scores, 0.2)
+            indices.extend(class_indices[class_keep_boxes])
+
+        return [{"type": self.label_list[class_ids[i]].lower(), "bbox": [float(t) for t in boxes[i].tolist()], "score": float(scores[i])} for i in indices]
+
+    def close(self):
+        logging.info("Close recognizer.")
+        if hasattr(self, "ort_sess"):
+            del self.ort_sess
+        gc.collect()
+
+    def __call__(self, image_list, thr=0.7, batch_size=16):
+        res = []
+        images = []
+        for i in range(len(image_list)):
+            if not isinstance(image_list[i], np.ndarray):
+                images.append(np.array(image_list[i]))
+            else:
+                images.append(image_list[i])
+
+        batch_loop_cnt = math.ceil(float(len(images)) / batch_size)
+        for i in range(batch_loop_cnt):
+            start_index = i * batch_size
+            end_index = min((i + 1) * batch_size, len(images))
+            batch_image_list = images[start_index:end_index]
+            inputs = self.preprocess(batch_image_list)
+            logging.debug("preprocess")
+            for ins in inputs:
+                bb = self.postprocess(self.ort_sess.run(None, {k: v for k, v in ins.items() if k in self.input_names}, self.run_options)[0], ins, thr)
+                res.append(bb)
+
+        # seeit.save_results(image_list, res, self.label_list, threshold=thr)
+
+        return res
+
+    def __del__(self):
+        self.close()

langparse/engines/pdf/deepdoc/rendering.py

@@ -0,0 +1,210 @@
+"""
+New code, not ported from upstream: deepdoc's TableStructureRecognizer emits
+tables as HTML (colspan/rowspan) or as Chinese natural-language sentences --
+neither matches langparse's cross-engine table shape. This module renders
+deepdoc's box list into ParsedPageResult, normalizing tables to
+{"rows": list[list[str]]} to match SimplePDFEngine/MinerUEngine and to keep
+services/fidelity.py's TEDS scoring working across engines.
+"""
+
+from __future__ import annotations
+
+from collections import defaultdict
+from html.parser import HTMLParser
+
+from langparse.types import ParsedElement, ParsedPageResult
+
+
+class _RawTableParser(HTMLParser):
+    def __init__(self) -> None:
+        super().__init__()
+        self.rows: list[list[dict]] = []
+        self._row: list[dict] | None = None
+        self._cell: list[str] | None = None
+        self._cell_attrs: dict[str, str] = {}
+
+    def handle_starttag(self, tag: str, attrs: list[tuple[str, str | None]]) -> None:
+        attrs_dict = dict(attrs)
+        if tag == "tr":
+            self._row = []
+        elif tag in ("td", "th"):
+            self._cell = []
+            self._cell_attrs = attrs_dict
+
+    def handle_data(self, data: str) -> None:
+        if self._cell is not None:
+            self._cell.append(data)
+
+    def handle_endtag(self, tag: str) -> None:
+        if tag in ("td", "th") and self._row is not None and self._cell is not None:
+            text = "".join(self._cell).strip()
+            colspan = int(self._cell_attrs.get("colspan", 1) or 1)
+            rowspan = int(self._cell_attrs.get("rowspan", 1) or 1)
+            self._row.append({"text": text, "colspan": colspan, "rowspan": rowspan})
+            self._cell = None
+        elif tag == "tr" and self._row is not None:
+            self.rows.append(self._row)
+            self._row = None
+
+
+def _extract_raw_rows(html: str) -> list[list[dict]]:
+    parser = _RawTableParser()
+    parser.feed(html)
+    return parser.rows
+
+
+def html_table_to_rows(html: str) -> list[list[str]]:
+    """Flatten an HTML table (with colspan/rowspan) into a plain row grid."""
+    raw_rows = _extract_raw_rows(html)
+    grid: list[list[str]] = []
+    carry: dict[int, tuple[str, int]] = {}  # column -> (text, additional rows remaining)
+
+    for raw_row in raw_rows:
+        row: list[str] = []
+        col = 0
+        cells = list(raw_row)
+        cell_index = 0
+        while cell_index < len(cells) or col in carry or any(c > col for c in carry):
+            if col in carry:
+                text, remaining = carry.pop(col)
+                row.append(text)
+                if remaining > 1:
+                    carry[col] = (text, remaining - 1)
+                col += 1
+                continue
+            if cell_index >= len(cells):
+                col += 1
+                continue
+            cell = cells[cell_index]
+            cell_index += 1
+            colspan = max(1, cell["colspan"])
+            rowspan = max(1, cell["rowspan"])
+            for _ in range(colspan):
+                row.append(cell["text"])
+                if rowspan > 1:
+                    carry[col] = (cell["text"], rowspan - 1)
+                col += 1
+        grid.append(row)
+    return grid
+
+
+def _bbox(box: dict) -> list[float]:
+    # box["top"]/box["bottom"] are in document-cumulative Y space (offset by
+    # the running sum of prior pages' heights, see _layouts_rec in
+    # pdf_parser.py) so boxes sort correctly across a multi-page document
+    # internally. box["positions"] -- a list of
+    # [page_number, left, right, top, bottom] tuples -- carries the same
+    # rectangle in page-local coordinates instead, which is what a
+    # ParsedElement.bbox must report (MinerUEngine's bbox is page-local too).
+    # Fall back to x0/top/x1/bottom when positions is absent, e.g. for
+    # hand-built test fixtures.
+    positions = box.get("positions")
+    if positions:
+        _page_number, left, right, top, bottom = positions[0]
+        return [float(left), float(top), float(right), float(bottom)]
+    return [
+        float(box.get("x0", 0.0)),
+        float(box.get("top", 0.0)),
+        float(box.get("x1", 0.0)),
+        float(box.get("bottom", 0.0)),
+    ]
+
+
+def _rows_to_markdown_table(rows: list[list[str]]) -> str:
+    if not rows:
+        return ""
+    lines = [f"| {' | '.join(rows[0])} |", f"| {' | '.join(['---'] * len(rows[0]))} |"]
+    for row in rows[1:]:
+        lines.append(f"| {' | '.join(row)} |")
+    return "\n".join(lines)
+
+
+def render_pages(
+    boxes: list[dict], ocr_pages: dict[int, bool] | None = None
+) -> list[ParsedPageResult]:
+    """Render deepdoc's flat box list (from RAGFlowPdfParser.parse_into_bboxes) into pages.
+
+    ocr_pages is an optional {page_number: bool} map (1-indexed, matching
+    box["page_number"]) saying whether a page's text should be credited to
+    OCR rather than the PDF's native text layer -- see
+    DeepDocEngine._classify_ocr_pages, which derives it via the same
+    needs_ocr() heuristic simple/ocr.py already uses. A page absent from the
+    map, or the map itself omitted, defaults to False/0.
+    """
+    boxes_by_page: dict[int, list[dict]] = defaultdict(list)
+    for box in boxes:
+        boxes_by_page[box["page_number"]].append(box)
+
+    ocr_pages = ocr_pages or {}
+    pages = []
+    for page_number in sorted(boxes_by_page):
+        markdown_parts: list[str] = []
+        plain_parts: list[str] = []
+        elements: list[ParsedElement] = []
+        tables: list[dict] = []
+        images: list[dict] = []
+
+        for box in boxes_by_page[page_number]:
+            layout_type = box.get("layout_type") or "text"
+            text = (box.get("text") or "").strip()
+            bbox = _bbox(box)
+
+            if layout_type == "table":
+                rows = html_table_to_rows(text)
+                tables.append({"rows": rows})
+                # Use original cells, not the span-expanded grid, so merged
+                # text appears once in search text and OCR character counts.
+                table_text = "\n".join(
+                    "\t".join(cell["text"] for cell in row if cell["text"])
+                    for row in _extract_raw_rows(text)
+                ).strip()
+                if table_text:
+                    plain_parts.append(table_text)
+                markdown_parts.append(_rows_to_markdown_table(rows))
+                elements.append(
+                    ParsedElement(
+                        kind="table", text=text, bbox=bbox, metadata={"layout_type": layout_type}
+                    )
+                )
+                continue
+
+            if layout_type == "figure":
+                images.append({"caption": text, "bbox": bbox})
+                if text:
+                    markdown_parts.append(f"*{text}*")
+                elements.append(
+                    ParsedElement(
+                        kind="figure", text=text, bbox=bbox, metadata={"layout_type": layout_type}
+                    )
+                )
+                continue
+
+            if not text:
+                continue
+
+            markdown_parts.append(f"# {text}" if layout_type == "title" else text)
+            plain_parts.append(text)
+            elements.append(
+                ParsedElement(
+                    kind=layout_type, text=text, bbox=bbox, metadata={"layout_type": layout_type}
+                )
+            )
+
+        plain_text = "\n".join(plain_parts)
+        page_ocr_applied = bool(ocr_pages.get(page_number, False))
+        pages.append(
+            ParsedPageResult(
+                page_number=page_number,
+                markdown_content="\n\n".join(part for part in markdown_parts if part),
+                plain_text=plain_text,
+                elements=elements,
+                tables=tables,
+                images=images,
+                metadata={
+                    "engine_name": "deepdoc",
+                    "ocr_applied": page_ocr_applied,
+                    "ocr_text_chars": len(plain_text) if page_ocr_applied else 0,
+                },
+            )
+        )
+    return pages