first commit

Co-authored-by: Cursor <cursoragent@cursor.com>
2026-07-16 11:12:17 +08:00
commit 4003624b8c
80 changed files with 9990 additions and 0 deletions
@@ -0,0 +1,91 @@
+"""Image/table region cropping with OCR and descriptions."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import fitz
+
+from rag_cut.models import Block, BlockType
+from rag_cut.parsers.pdf.layout import LayoutRegion, PageLayout
+from rag_cut.parsers.pdf.ocr import describe_visual, ocr_image_file, ocr_pixmap
+from rag_cut.parsers.pdf.table_extract import build_table_block
+
+
+def _crop_region(page: fitz.Page, region: LayoutRegion, zoom: float = 2.0) -> fitz.Pixmap:
+    clip = fitz.Rect(region.x0, region.y0, region.x1, region.y1)
+    return page.get_pixmap(matrix=fitz.Matrix(zoom, zoom), clip=clip, alpha=False)
+
+
+def _save_pixmap(pix: fitz.Pixmap, path: Path) -> None:
+    if pix.n - pix.alpha > 3:
+        pix = fitz.Pixmap(fitz.csRGB, pix)
+    pix.save(str(path))
+
+
+def extract_visual_blocks(
+    page: fitz.Page,
+    layout: PageLayout,
+    doc: fitz.Document,
+    assets_dir: Path,
+    chapter_title: str | None,
+    img_counter: int,
+) -> tuple[list[Block], int]:
+    blocks: list[Block] = []
+    page_no = layout.page_index + 1
+
+    for region in layout.regions:
+        if region.kind == "table":
+            table_block = build_table_block(page, region, layout, assets_dir, chapter_title)
+            if table_block:
+                blocks.append(table_block)
+            continue
+
+        if region.kind != "image":
+            continue
+
+        img_counter += 1
+        xref = region.data.get("xref")
+        img_id = f"page{page_no}_img{img_counter}.png"
+        img_path = assets_dir / img_id
+
+        try:
+            # Prefer on-page crop so scaled/clipped placements OCR the visible region.
+            crop_pix = _crop_region(page, region)
+            _save_pixmap(crop_pix, img_path)
+        except Exception:
+            try:
+                pix = fitz.Pixmap(doc, xref)
+                if pix.n - pix.alpha > 3:
+                    pix = fitz.Pixmap(fitz.csRGB, pix)
+                pix.save(str(img_path))
+            except Exception:
+                continue
+
+        ocr_text = ocr_image_file(img_path)
+        if not ocr_text:
+            try:
+                ocr_text = ocr_pixmap(_crop_region(page, region))
+            except Exception:
+                ocr_text = ""
+
+        meta = {
+            "page": page_no,
+            "bbox": [region.x0, region.y0, region.x1, region.y1],
+            "xref": xref,
+        }
+        if chapter_title:
+            meta["chapter"] = chapter_title
+
+        blocks.append(
+            Block(
+                type=BlockType.IMAGE,
+                image_id=img_id,
+                image_path=str(img_path),
+                ocr_text=ocr_text,
+                text=describe_visual(ocr_text, "图片", img_id),
+                meta=meta,
+            )
+        )
+
+    return blocks, img_counter