Delete table_extraction.py

table_extraction.py

@@ -1,565 +0,0 @@
-"""
-table_extraction.py
-====================
-
-Right-hand branch of the diagram: when a matched region's class is "table",
-its cropped image goes through TableFormerONNX -> OTSL string -> markdown.
-
-`TableFormerONNX`, `ocr_anchor_cells`, `seq_to_otsl`, and `table_image_to_otsl`
-are carried over from docling_.ipynb essentially unchanged (decoupled from
-that notebook's specific OCR-backend import style; backends now come from
-ocr_backends.py). `otsl_to_markdown` is new — it's the "otsl -> markdown" step
-in the diagram, which in the original notebook was done ad hoc in later
-cells (build a <doctag>, run DocTagsDocument/DoclingDocument). Folding it
-into one function makes it reusable per-table inside the full-page pipeline.
-It returns `(markdown, dataframe)` — the dataframe is built from
-`TableItem.export_to_dataframe()` with `collapse_spanned_columns` applied,
-since a raw rectangular DataFrame otherwise duplicates a colspan'd cell's
-value into every column it visually spans (see that function's docstring
-for why, and why `pd.read_html(export_to_html())` has the identical issue).
-"""
-
-from __future__ import annotations
-
-import json
-import logging
-import time
-from pathlib import Path
-from typing import Optional, Union
-from urllib.parse import urlparse
-
-import cv2
-import numpy as np
-import onnxruntime as ort
-import pandas as pd
-from PIL import Image
-
-from .ocr_backends import OCRBackend
-
-logger = logging.getLogger(__name__)
-
-ANCHOR_TAGS = {"fcel", "ecel", "ched", "rhed", "srow"}  # tags that occupy their own grid slot
-EMPTY_OTSL_PLACEHOLDER = "?"  # inserted when a non-empty-tagged cell OCRs to nothing,
-                              # so the docling_core parser never sees an empty content cell
-
-MEAN = np.array([0.94247851, 0.94254675, 0.94292611], dtype=np.float32)
-STD = np.array([0.17910956, 0.17940403, 0.17931663], dtype=np.float32)
-IMAGE_SIZE = 448
-
-
-# --------------------------------------------------------------------------- #
-# bbox merge helpers (lcel chains)
-# --------------------------------------------------------------------------- #
-def merge_bboxes(box1: np.ndarray, box2: np.ndarray) -> np.ndarray:
-    """Union the cxcywh boxes of an lcel chain's stub anchor and the real cell that closes it."""
-    new_w = (box2[0] + box2[2] / 2) - (box1[0] - box1[2] / 2)
-    new_h = (box2[1] + box2[3] / 2) - (box1[1] - box1[3] / 2)
-    new_left = box1[0] - box1[2] / 2
-    new_top = min((box2[1] - box2[3] / 2), (box1[1] - box1[3] / 2))
-    return np.array(
-        [new_left + new_w / 2, new_top + new_h / 2, new_w, new_h], dtype=np.float32
-    )
-
-
-def apply_bbox_merge(coords: np.ndarray, bboxes_to_merge: dict) -> np.ndarray:
-    coords = [np.asarray(c, dtype=np.float32) for c in coords]
-    merged, skip = [], set()
-    for i, box1 in enumerate(coords):
-        if i in bboxes_to_merge:
-            j = bboxes_to_merge[i]
-            if j >= 0:
-                skip.add(j)
-                merged.append(merge_bboxes(box1, coords[j]))
-        elif i not in skip:
-            merged.append(box1)
-    return np.stack(merged) if merged else np.empty((0, 4), dtype=np.float32)
-
-
-def make_session(path: str, threads: int = 2) -> ort.InferenceSession:
-    opts = ort.SessionOptions()
-    opts.inter_op_num_threads = 2
-    opts.intra_op_num_threads = 4
-    opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
-    opts.enable_mem_pattern = True
-    opts.enable_mem_reuse = True
-    opts.enable_profiling = False
-    return ort.InferenceSession(path, sess_options=opts, providers=["CPUExecutionProvider"])
-
-
-# --------------------------------------------------------------------------- #
-# TableFormer ONNX runner
-# --------------------------------------------------------------------------- #
-class TableFormerONNX:
-    """
-    encoder + decoder-step (autoregressive loop, KV-cache) + bbox_decoder,
-    all ONNX Runtime. Produces a structure sequence (OTSL tags) plus a
-    normalized cxcywh bbox per anchor cell.
-
-    Usage:
-        runner = TableFormerONNX(artifact_root="tableformerv1", variant="accurate")
-    """
-
-    def __init__(self, artifact_root: str = "tableformerv1", variant: str = "accurate", threads: int = 2):
-        root = Path(artifact_root)
-        onnx_dir = root / "onnx" / variant
-        self.enc = make_session(str(onnx_dir / f"tableformer_{variant}_encoder.onnx"), threads)
-        self.dec = make_session(str(onnx_dir / f"tableformer_{variant}_decoder_step.onnx"), threads)
-        self.bbox = make_session(str(onnx_dir / f"tableformer_{variant}_bbox_decoder.onnx"), threads)
-
-        with open(root / "tm_config.json", encoding="utf-8") as f:
-            self.tag_map = json.load(f)["dataset_wordmap"]["word_map_tag"]
-        self.rev_tag = {v: k for k, v in self.tag_map.items()}
-
-        cache_input = next(i for i in self.dec.get_inputs() if i.name == "cache")
-        self.num_layers = int(cache_input.shape[0])
-        self.embed_dim = int(cache_input.shape[3])
-
-    @staticmethod
-    def is_valid_url(url_string: str) -> bool:
-        try:
-            result = urlparse(url_string)
-            return all([result.scheme in ("http", "https"), result.netloc])
-        except ValueError:
-            return False
-
-    def image_preprocess(self, img) -> np.ndarray:
-        """Accepts a path, URL, PIL.Image, or HWC RGB ndarray."""
-        if isinstance(img, str) and self.is_valid_url(img):
-            import requests
-
-            img = Image.open(requests.get(img, stream=True, timeout=30).raw)
-        if isinstance(img, (str, Path)):
-            img = Image.open(img)
-        if isinstance(img, Image.Image):
-            img = np.array(img.convert("RGB"))
-        elif isinstance(img, np.ndarray) and img.ndim == 3 and img.shape[2] == 3:
-            pass
-        else:
-            raise TypeError(f"Unsupported image type: {type(img)}")
-
-        img = (img.astype(np.float32) - 255.0 * MEAN) / STD
-        img = cv2.resize(img, (IMAGE_SIZE, IMAGE_SIZE), interpolation=cv2.INTER_CUBIC)
-        img = img.transpose(2, 1, 0) / 255.0
-        return np.expand_dims(img.astype(np.float32), axis=0)
-
-    def draw_bbox(self, img_path: str, bboxs: np.ndarray, out_path: str = "draw_p.png") -> None:
-        """Debug helper — draws predicted cxcywh boxes on the ORIGINAL-resolution image."""
-        from PIL import ImageDraw
-
-        pil_image = Image.open(img_path)
-        width, height = pil_image.size
-        draw = ImageDraw.Draw(pil_image)
-        for cx, cy, w, h in np.asarray(bboxs):
-            xyxy = (
-                int((cx - 0.5 * w) * width), int((cy - 0.5 * h) * height),
-                int((cx + 0.5 * w) * width), int((cy + 0.5 * h) * height),
-            )
-            draw.rectangle(xyxy, outline="red", width=2)
-        pil_image.save(out_path)
-
-    def predict(self, img: np.ndarray, max_steps: int = 1024) -> dict:
-        """img: float32 [1, 3, 448, 448] (preprocessed). Returns seq / outputs_class /
-        outputs_coord (cxcywh, normalized [0,1]) / bboxes_to_merge / timings."""
-        t0 = time.time()
-        enc_out, memory = self.enc.run(None, {"image": img})
-        t_enc = time.time() - t0
-
-        wm = self.tag_map
-        decoded_tags = np.array([[wm["<start>"]]], dtype=np.int64)
-        cache = np.zeros((self.num_layers, 0, 1, self.embed_dim), dtype=np.float32)
-        skip_next_tag = True
-        prev_tag_ucel = False
-        line_num = 0
-        first_lcel = True
-        bbox_ind = 0
-        cur_bbox_ind = -1
-        bboxes_to_merge: dict = {}
-        tag_H_buf = []
-
-        t1 = time.time()
-        for _ in range(max_steps):
-            logits, last_hidden, cache = self.dec.run(
-                None, {"decoded_tags": decoded_tags, "memory": memory, "cache": cache}
-            )
-            new_tag = int(np.argmax(logits, axis=1)[0])
-
-            if line_num == 0 and new_tag == wm["xcel"]:
-                new_tag = wm["lcel"]
-            if prev_tag_ucel and new_tag == wm["lcel"]:
-                new_tag = wm["fcel"]
-
-            if new_tag == wm["<end>"]:
-                decoded_tags = np.concatenate(
-                    [decoded_tags, np.array([[new_tag]], dtype=np.int64)], axis=0
-                )
-                break
-
-            if not skip_next_tag:
-                if new_tag in (wm["fcel"], wm["ecel"], wm["ched"], wm["rhed"],
-                               wm["srow"], wm["nl"], wm["ucel"]):
-                    tag_H_buf.append(last_hidden[:, 0, :].copy())
-                    if first_lcel is not True:
-                        bboxes_to_merge[cur_bbox_ind] = bbox_ind
-                    bbox_ind += 1
-
-            if new_tag != wm["lcel"]:
-                first_lcel = True
-            else:
-                if first_lcel:
-                    tag_H_buf.append(last_hidden[:, 0, :].copy())
-                    first_lcel = False
-                    cur_bbox_ind = bbox_ind
-                    bboxes_to_merge[cur_bbox_ind] = -1
-                    bbox_ind += 1
-
-            skip_next_tag = new_tag in (wm["nl"], wm["ucel"], wm["xcel"])
-            prev_tag_ucel = new_tag == wm["ucel"]
-            if new_tag == wm["nl"]:
-                line_num += 1
-
-            decoded_tags = np.concatenate(
-                [decoded_tags, np.array([[new_tag]], dtype=np.int64)], axis=0
-            )
-        t_dec = time.time() - t1
-
-        seq = decoded_tags.squeeze().tolist()
-
-        if tag_H_buf:
-            tag_H_stacked = np.stack(
-                [h[None, ...] if h.ndim == 1 else h for h in tag_H_buf], axis=0
-            ).reshape(-1, 1, self.embed_dim).astype(np.float32)
-            t2 = time.time()
-            cls_logits, coord = self.bbox.run(None, {"enc_out": enc_out, "tag_H_stacked": tag_H_stacked})
-            t_bbox = time.time() - t2
-            coord = apply_bbox_merge(coord, bboxes_to_merge)
-        else:
-            cls_logits = np.empty((0, 3), dtype=np.float32)
-            coord = np.empty((0, 4), dtype=np.float32)
-            t_bbox = 0.0
-
-        return {
-            "seq": seq,
-            "outputs_class": cls_logits,
-            "outputs_coord": coord,
-            "bboxes_to_merge": bboxes_to_merge,
-            "timings": {"encoder": t_enc, "decoder": t_dec, "bbox": t_bbox,
-                        "total": t_enc + t_dec + t_bbox},
-        }
-
-
-# --------------------------------------------------------------------------- #
-# OTSL assembly: structure tokens + per-cell OCR -> docling-ready OTSL string
-# --------------------------------------------------------------------------- #
-def compute_ioa(ocr_bbox: list[float], cell_bbox: list[float]) -> float:
-    """Intersection over OCR-Token Area (IoA): maps a token to a cell if most
-    of the token sits inside it."""
-    ox1, oy1, ox2, oy2 = ocr_bbox
-    cx1, cy1, cx2, cy2 = cell_bbox
-
-    ix1 = max(ox1, cx1)
-    iy1 = max(oy1, cy1)
-    ix2 = min(ox2, cx2)
-    iy2 = min(oy2, cy2)
-
-    if ix2 <= ix1 or iy2 <= iy1:
-        return 0.0
-
-    inter_area = (ix2 - ix1) * (iy2 - iy1)
-    ocr_area = (ox2 - ox1) * (oy2 - oy1)
-    return inter_area / max(ocr_area, 1e-6)
-
-
-def ocr_anchor_cells(
-    orig_image_path: str,
-    seq: list,
-    box_o: np.ndarray,
-    rev_tag: dict,
-    backend: Optional[OCRBackend] = None,
-    tokens: Optional[list[dict]] = None,
-    pad: int = 4,
-    narrow_width: float = 0.08,
-) -> list:
-    """Maps text tokens into TableFormer cell boundaries via spatial intersection.
-
-    Either `backend` (run OCR on the table crop) or pre-extracted `tokens`
-    (e.g. pdfplumber words already expressed in the crop's local pixel space —
-    see pipeline.py) must be supplied. `tokens` takes priority: when a PDF
-    has a real text layer, reusing those exact words is both cheaper and more
-    accurate than re-OCRing a table crop pdfplumber already read correctly.
-    """
-    if tokens is None and backend is None:
-        raise ValueError("ocr_anchor_cells needs either `tokens` or `backend`")
-
-    orig = Image.open(orig_image_path).convert("RGB")
-    orig_w, orig_h = orig.size
-
-    anchor_tags = [rev_tag[t] for t in seq[1:-1] if rev_tag[t] in ANCHOR_TAGS]
-    assert len(anchor_tags) == box_o.shape[0], (
-        f"{len(anchor_tags)} anchor tags vs {box_o.shape[0]} boxes — structural alignment issue."
-    )
-
-    ocr_tokens = tokens if tokens is not None else backend.get_text_boxes(orig)
-    source = "pdfplumber tokens" if tokens is not None else type(backend).__name__
-    logger.debug(
-        "ocr_anchor_cells: %d anchor cells, %d tokens from %s",
-        len(anchor_tags),
-        len(ocr_tokens),
-        source,
-    )
-    texts = []
-
-    for i, (tag, (cx, cy, w, h)) in enumerate(zip(anchor_tags, box_o)):
-        if tag == "ecel":
-            texts.append("")
-            continue
-
-        x1 = max(0, (cx - w / 2) * orig_w - pad)
-        y1 = max(0, (cy - h / 2) * orig_h - pad)
-        x2 = min(orig_w, (cx + w / 2) * orig_w + pad)
-        y2 = min(orig_h, (cy + h / 2) * orig_h + pad)
-        cell_bbox = [x1, y1, x2, y2]
-
-        matched_tokens = [t for t in ocr_tokens if compute_ioa(t["bbox"], cell_bbox) >= 0.45]
-
-        if not matched_tokens:
-            texts.append("")
-            logger.debug("  cell[%d] tag=%s (empty)", i, tag)
-            continue
-
-        matched_tokens.sort(key=lambda t: (int(t["bbox"][1] / 6), t["bbox"][0]))
-        cell_string = " ".join(t["text"] for t in matched_tokens)
-        texts.append(cell_string.strip())
-        logger.debug("  cell[%d] tag=%s text=%r", i, tag, cell_string.strip()[:60])
-
-    return texts
-
-
-def _otsl_cell_text(tag: str, text: str) -> str | None:
-    if tag == "ecel":
-        return None
-    cleaned = (text or "").strip()
-    return cleaned if cleaned else EMPTY_OTSL_PLACEHOLDER
-
-
-def seq_to_otsl(seq: list, rev_tag: dict, texts: list) -> str:
-    """Re-walk seq, inserting OCR text right after each anchor tag."""
-    out, ti = [], 0
-    for tok in seq[1:-1]:
-        name = rev_tag[tok]
-        if name == "nl":
-            out.append("<nl>")
-            continue
-        out.append(f"<{name}>")
-        if name in ANCHOR_TAGS:
-            cell_text = _otsl_cell_text(name, texts[ti])
-            if cell_text is not None:
-                out.append(cell_text)
-            ti += 1
-    return "".join(out)
-
-
-def table_image_to_otsl(
-    runner: TableFormerONNX,
-    orig_image_path: str,
-    backend: Optional[OCRBackend] = None,
-    tokens: Optional[list[dict]] = None,
-) -> str:
-    """End-to-end: image -> ONNX TableFormer -> anchors filled from `tokens`
-    (preferred, e.g. pdfplumber words) or OCR'd via `backend` -> OTSL string."""
-    logger.info("TableFormer: processing %s", orig_image_path)
-    preproc = runner.image_preprocess(orig_image_path)
-    out = runner.predict(preproc)
-    logger.debug(
-        "TableFormer: seq length=%d, anchor boxes=%d",
-        len(out["seq"]),
-        out["outputs_coord"].shape[0],
-    )
-    texts = ocr_anchor_cells(
-        orig_image_path, out["seq"], out["outputs_coord"], runner.rev_tag,
-        backend=backend, tokens=tokens,
-    )
-    otsl = seq_to_otsl(out["seq"], runner.rev_tag, texts)
-    logger.debug("TableFormer OTSL preview: %s", otsl[:200])
-    return otsl
-
-
-def collapse_spanned_columns(df: "pd.DataFrame") -> "pd.DataFrame":
-    """
-    Collapse adjacent columns that are exact duplicates (same header text,
-    identical value in every row) back into a single column.
-
-    A plain rectangular DataFrame has no native way to represent a cell that
-    spans multiple columns (e.g. an OTSL `lcel` chain / HTML `colspan`).
-    Both `TableItem.export_to_dataframe()` and the `export_to_html()` +
-    `pd.read_html()` round-trip handle this the same way: they repeat the
-    spanned cell's value into every column it covers, so a 2-column-wide
-    "Description" header ends up as two separate "Description" columns
-    with identical content. This collapses those back into one.
-
-    Note: this is a heuristic — if two genuinely distinct adjacent columns
-    happen to share both the same header text and identical values in
-    every row of this specific table, they'll also get collapsed. That's
-    expected to be rare in practice (real tables don't usually have two
-    different columns with the same name and identical data), but worth
-    knowing if a table doesn't collapse cleanly.
-    """
-    cols = list(df.columns)
-    if len(cols) <= 1:
-        return df
-
-    keep = [0]
-    for i in range(1, len(cols)):
-        prev = keep[-1]
-        # pandas suffixes repeated header names with ".1", ".2", ... on
-        # read_html; export_to_dataframe doesn't, so strip defensively.
-        prev_header = str(cols[prev]).split(".")[0]
-        cur_header = str(cols[i]).split(".")[0]
-        same_header = prev_header == cur_header
-        same_values = df.iloc[:, prev].astype(str).equals(df.iloc[:, i].astype(str))
-        if same_header and same_values:
-            continue  # duplicate of the column we're keeping — drop it
-        keep.append(i)
-
-    out = df.iloc[:, keep].copy()
-    out.columns = [str(c).split(".")[0] for c in out.columns]
-    return out
-
-
-def _parse_md_table_row(line: str) -> list[str]:
-    """Split one markdown table row into its cell strings, stripping whitespace."""
-    return [c.strip() for c in line.strip().strip("|").split("|")]
-
-
-def _is_separator_row(cells: list[str]) -> bool:
-    """Return True if every cell is a markdown alignment/separator marker (e.g. `---`, `:---:`)."""
-    return all(set(c.replace(":", "").replace("-", "")) == set() and len(c) > 0 for c in cells)
-
-
-def collapse_markdown_table(md: str) -> str:
-    """
-    Apply the same duplicate-adjacent-column collapse to every markdown table
-    inside `md`, in-place. This is necessary because markdown has no colspan
-    syntax — docling_core's `export_to_markdown()` represents a colspan-N cell
-    by repeating its value in N consecutive columns, which looks like
-    `| Description | Description | Total |` for a 2-column-wide "Description"
-    header. This function detects and removes those duplicate columns from the
-    markdown string so the rendered table matches what the source document
-    actually meant.
-
-    The collapse criterion is identical to `collapse_spanned_columns`: two
-    adjacent columns are merged when they share the same header text (after
-    stripping pandas `.1`/`.2` dedup suffixes, though docling markdown doesn't
-    add those) AND every data row has the same value in both columns.
-
-    Non-table lines (headings, paragraphs, etc.) are passed through unchanged.
-    """
-    lines = md.splitlines()
-    out_lines: list[str] = []
-    i = 0
-
-    while i < len(lines):
-        line = lines[i]
-
-        # Detect the start of a markdown table: a line containing at least one |
-        # that isn't just punctuation.
-        if "|" not in line:
-            out_lines.append(line)
-            i += 1
-            continue
-
-        # Collect the full table block (contiguous lines containing |).
-        table_block: list[str] = []
-        while i < len(lines) and "|" in lines[i]:
-            table_block.append(lines[i])
-            i += 1
-
-        if len(table_block) < 2:
-            # Not a real table (need at least header + separator).
-            out_lines.extend(table_block)
-            continue
-
-        # Parse into header, optional separator, and data rows.
-        header_cells = _parse_md_table_row(table_block[0])
-        sep_idx = 1 if _is_separator_row(_parse_md_table_row(table_block[1])) else None
-
-        data_rows: list[list[str]] = []
-        for row_line in table_block[(2 if sep_idx is not None else 1):]:
-            data_rows.append(_parse_md_table_row(row_line))
-
-        if not data_rows:
-            out_lines.extend(table_block)
-            continue
-
-        # Pad short rows so indexing is safe.
-        n_cols = len(header_cells)
-        data_rows = [r + [""] * max(0, n_cols - len(r)) for r in data_rows]
-
-        # Determine which columns to keep — same logic as collapse_spanned_columns.
-        keep: list[int] = [0]
-        for j in range(1, n_cols):
-            prev = keep[-1]
-            prev_header = header_cells[prev].split(".")[0]
-            cur_header = header_cells[j].split(".")[0]
-            same_header = prev_header == cur_header
-            same_values = all(
-                (row[prev] if prev < len(row) else "") == (row[j] if j < len(row) else "")
-                for row in data_rows
-            )
-            if same_header and same_values:
-                continue
-            keep.append(j)
-
-        if len(keep) == n_cols:
-            # Nothing to collapse — pass through unchanged.
-            out_lines.extend(table_block)
-            continue
-
-        # Rebuild the table with only the kept columns.
-        def _fmt_row(cells: list[str], indices: list[int]) -> str:
-            return "| " + " | ".join(cells[k] for k in indices) + " |"
-
-        def _fmt_sep(indices: list[int]) -> str:
-            return "| " + " | ".join("---" for _ in indices) + " |"
-
-        out_lines.append(_fmt_row(header_cells, keep))
-        if sep_idx is not None:
-            out_lines.append(_fmt_sep(keep))
-        for row in data_rows:
-            out_lines.append(_fmt_row(row, keep))
-
-    return "\n".join(out_lines)
-
-
-def otsl_to_markdown(
-    otsl: str, dummy_image_size: tuple[int, int] = (512, 512)
-) -> tuple[str, "pd.DataFrame"]:
-    """
-    Wrap a raw OTSL tag string as <doctag><otsl>...</otsl></doctag>, load it
-    through docling_core's DocTagsDocument/DoclingDocument, and export both
-    markdown and a DataFrame for the (first) detected table.
-
-    Returns `(markdown, dataframe)` — both have duplicate columns from
-    colspan cells collapsed:
-    - markdown: via `collapse_markdown_table` (markdown has no colspan syntax
-      so docling repeats the spanned cell's text into every column it covers)
-    - dataframe: via `TableItem.export_to_dataframe()` + `collapse_spanned_columns`
-      (same duplication issue at the DataFrame level, same fix)
-    """
-    from docling_core.types.doc.document import DocTagsDocument, DoclingDocument
-
-    otsl_doc = f"<doctag><otsl>{otsl}</otsl></doctag>"
-    dummy_image = Image.new("RGB", dummy_image_size, color="black")
-    doctags_doc = DocTagsDocument.from_doctags_and_image_pairs([otsl_doc], [dummy_image])
-    doc = DoclingDocument.load_from_doctags(doctags_doc)
-
-    md = collapse_markdown_table(doc.export_to_markdown())
-
-    if doc.tables:
-        df = doc.tables[0].export_to_dataframe(doc=doc)
-        df = collapse_spanned_columns(df)
-        logger.info("OTSL -> markdown: %d rows x %d cols", df.shape[0], df.shape[1])
-    else:
-        df = pd.DataFrame()
-        logger.warning("OTSL -> markdown: no tables detected in OTSL output")
-
-    return md, df