rnsr 0.1.0__py3-none-any.whl

rnsr/ingestion/xy_cut.py

@@ -0,0 +1,555 @@
+"""
+Recursive XY-Cut Algorithm - Visual-Geometric Segmentation
+
+Implements the Recursive XY-Cut (RXYC) algorithm from Section 4.1.1:
+"A top-down page segmentation technique that is particularly effective 
+for discovering document structure without relying on text content."
+
+The algorithm:
+1. Treats document page as a binary image
+2. Calculates projection profiles (sum of black pixels) along X and Y axes
+3. Identifies "valleys" (whitespace gaps) as natural separators
+4. Recursively cuts at widest valleys to produce a tree of bounding boxes
+5. Larger boxes (detected early) = major structural elements
+6. Smaller, deeply nested boxes = paragraphs/cells
+
+Use this for:
+- Multi-column layouts
+- Complex L-shaped text wraps  
+- Documents with visual structure but no font variance
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import structlog
+
+logger = structlog.get_logger(__name__)
+
+
+@dataclass
+class BoundingRegion:
+    """A rectangular region on a page."""
+    
+    x0: float
+    y0: float
+    x1: float
+    y1: float
+    page_num: int = 0
+    
+    @property
+    def width(self) -> float:
+        return self.x1 - self.x0
+    
+    @property
+    def height(self) -> float:
+        return self.y1 - self.y0
+    
+    @property
+    def area(self) -> float:
+        return self.width * self.height
+    
+    def contains(self, other: "BoundingRegion") -> bool:
+        """Check if this region contains another."""
+        return (
+            self.x0 <= other.x0 and
+            self.y0 <= other.y0 and
+            self.x1 >= other.x1 and
+            self.y1 >= other.y1
+        )
+
+
+@dataclass
+class SegmentNode:
+    """A node in the XY-Cut segmentation tree."""
+    
+    region: BoundingRegion
+    children: list["SegmentNode"] = field(default_factory=list)
+    text: str = ""
+    node_type: str = "region"  # "region", "text_block", "header", "body"
+    depth: int = 0
+    
+    @property
+    def is_leaf(self) -> bool:
+        return len(self.children) == 0
+
+
+class RecursiveXYCutter:
+    """
+    Implements the Recursive XY-Cut algorithm for document segmentation.
+    
+    Per Section 6.2 of the research paper:
+    "A major failure mode of simple parsing is complex layouts 
+    (e.g., a figure spanning two columns, or an L-shaped text wrap). 
+    The Recursive XY-Cut handles this."
+    """
+    
+    def __init__(
+        self,
+        min_gap_ratio: float = 0.02,  # Minimum gap as ratio of page dimension
+        min_region_ratio: float = 0.01,  # Minimum region size ratio
+        max_depth: int = 10,  # Maximum recursion depth
+        valley_threshold: float = 0.1,  # Threshold for valley detection
+    ):
+        """
+        Initialize the XY-Cutter.
+        
+        Args:
+            min_gap_ratio: Minimum whitespace gap size as ratio of page size.
+            min_region_ratio: Minimum region size to consider.
+            max_depth: Maximum recursion depth.
+            valley_threshold: Threshold for detecting valleys in projection.
+        """
+        self.min_gap_ratio = min_gap_ratio
+        self.min_region_ratio = min_region_ratio
+        self.max_depth = max_depth
+        self.valley_threshold = valley_threshold
+    
+    def segment_pdf(self, pdf_path: Path | str) -> list[SegmentNode]:
+        """
+        Segment all pages of a PDF using XY-Cut.
+        
+        Args:
+            pdf_path: Path to the PDF file.
+            
+        Returns:
+            List of SegmentNode trees (one per page).
+        """
+        import fitz
+        
+        pdf_path = Path(pdf_path)
+        doc = fitz.open(pdf_path)
+        
+        page_trees = []
+        
+        for page_num in range(len(doc)):
+            page = doc[page_num]
+            logger.debug("segmenting_page", page=page_num)
+            tree = self.segment_page(page, page_num)
+            page_trees.append(tree)
+        
+        doc.close()
+        
+        logger.info(
+            "xy_cut_complete",
+            pages=len(page_trees),
+            total_regions=sum(self._count_nodes(t) for t in page_trees),
+        )
+        
+        return page_trees
+    
+    def segment_page(self, page: Any, page_num: int = 0) -> SegmentNode:
+        """
+        Segment a single page using XY-Cut.
+        
+        Args:
+            page: A fitz.Page object.
+            page_num: Page number for metadata.
+            
+        Returns:
+            Root SegmentNode with hierarchy of regions.
+        """
+        import fitz
+        
+        # Get page dimensions
+        rect = page.rect
+        page_width = rect.width
+        page_height = rect.height
+        
+        # Create initial region (full page)
+        root_region = BoundingRegion(
+            x0=0, y0=0,
+            x1=page_width, y1=page_height,
+            page_num=page_num,
+        )
+        
+        # Render page to pixmap for projection analysis
+        # Use lower resolution for speed (72 dpi = 1x)
+        mat = fitz.Matrix(1, 1)  # 72 dpi
+        pix = page.get_pixmap(matrix=mat, alpha=False)
+        
+        # Convert to numpy array (grayscale)
+        img = np.frombuffer(pix.samples, dtype=np.uint8)
+        img = img.reshape(pix.height, pix.width, 3)
+        gray = np.mean(img, axis=2)
+        
+        # Binarize (invert so text = 1, background = 0)
+        binary = (gray < 240).astype(np.float32)
+        
+        # Calculate minimum dimensions
+        min_gap_x = int(page_width * self.min_gap_ratio)
+        min_gap_y = int(page_height * self.min_gap_ratio)
+        min_region_w = int(page_width * self.min_region_ratio)
+        min_region_h = int(page_height * self.min_region_ratio)
+        
+        # Recursive cut
+        root = SegmentNode(region=root_region, depth=0)
+        self._recursive_cut(
+            binary, root,
+            0, 0, pix.width, pix.height,
+            min_gap_x, min_gap_y,
+            min_region_w, min_region_h,
+            page_width / pix.width,  # Scale factor
+            page_height / pix.height,
+        )
+        
+        return root
+    
+    def _recursive_cut(
+        self,
+        binary: np.ndarray,
+        parent: SegmentNode,
+        x0: int, y0: int, x1: int, y1: int,
+        min_gap_x: int, min_gap_y: int,
+        min_region_w: int, min_region_h: int,
+        scale_x: float, scale_y: float,
+    ) -> None:
+        """
+        Recursively cut a region.
+        
+        Per the research paper algorithm:
+        1. Calculate projection profiles
+        2. Find valleys (gaps of whitespace)
+        3. Split horizontally first (Y-cut), then vertically (X-cut)
+        4. Recurse on sub-regions
+        """
+        if parent.depth >= self.max_depth:
+            return
+        
+        width = x1 - x0
+        height = y1 - y0
+        
+        # Check minimum size
+        if width < min_region_w or height < min_region_h:
+            return
+        
+        # Extract region
+        region_pixels = binary[y0:y1, x0:x1]
+        
+        if region_pixels.size == 0:
+            return
+        
+        # Calculate projection profiles
+        y_proj = np.sum(region_pixels, axis=1)  # Horizontal projection
+        x_proj = np.sum(region_pixels, axis=0)  # Vertical projection
+        
+        # Try horizontal cut first (Y-cut - splits top/bottom)
+        y_valleys = self._find_valleys(y_proj, min_gap_y)
+        
+        if y_valleys:
+            # Split at the widest valley
+            best_valley = max(y_valleys, key=lambda v: v[1] - v[0])
+            cut_y = (best_valley[0] + best_valley[1]) // 2
+            
+            # Create two child regions
+            if cut_y - y0 > min_region_h:
+                top_region = BoundingRegion(
+                    x0=x0 * scale_x, y0=y0 * scale_y,
+                    x1=x1 * scale_x, y1=cut_y * scale_y,
+                    page_num=parent.region.page_num,
+                )
+                top_node = SegmentNode(region=top_region, depth=parent.depth + 1)
+                parent.children.append(top_node)
+                self._recursive_cut(
+                    binary, top_node,
+                    x0, y0, x1, cut_y,
+                    min_gap_x, min_gap_y,
+                    min_region_w, min_region_h,
+                    scale_x, scale_y,
+                )
+            
+            if y1 - cut_y > min_region_h:
+                bottom_region = BoundingRegion(
+                    x0=x0 * scale_x, y0=cut_y * scale_y,
+                    x1=x1 * scale_x, y1=y1 * scale_y,
+                    page_num=parent.region.page_num,
+                )
+                bottom_node = SegmentNode(region=bottom_region, depth=parent.depth + 1)
+                parent.children.append(bottom_node)
+                self._recursive_cut(
+                    binary, bottom_node,
+                    x0, cut_y, x1, y1,
+                    min_gap_x, min_gap_y,
+                    min_region_w, min_region_h,
+                    scale_x, scale_y,
+                )
+            return
+        
+        # No horizontal cut found - try vertical (X-cut - splits columns)
+        x_valleys = self._find_valleys(x_proj, min_gap_x)
+        
+        if x_valleys:
+            # Split at the widest valley
+            best_valley = max(x_valleys, key=lambda v: v[1] - v[0])
+            cut_x = (best_valley[0] + best_valley[1]) // 2
+            
+            # Create two child regions
+            if cut_x - x0 > min_region_w:
+                left_region = BoundingRegion(
+                    x0=x0 * scale_x, y0=y0 * scale_y,
+                    x1=cut_x * scale_x, y1=y1 * scale_y,
+                    page_num=parent.region.page_num,
+                )
+                left_node = SegmentNode(region=left_region, depth=parent.depth + 1)
+                parent.children.append(left_node)
+                self._recursive_cut(
+                    binary, left_node,
+                    x0, y0, cut_x, y1,
+                    min_gap_x, min_gap_y,
+                    min_region_w, min_region_h,
+                    scale_x, scale_y,
+                )
+            
+            if x1 - cut_x > min_region_w:
+                right_region = BoundingRegion(
+                    x0=cut_x * scale_x, y0=y0 * scale_y,
+                    x1=x1 * scale_x, y1=y1 * scale_y,
+                    page_num=parent.region.page_num,
+                )
+                right_node = SegmentNode(region=right_region, depth=parent.depth + 1)
+                parent.children.append(right_node)
+                self._recursive_cut(
+                    binary, right_node,
+                    cut_x, y0, x1, y1,
+                    min_gap_x, min_gap_y,
+                    min_region_w, min_region_h,
+                    scale_x, scale_y,
+                )
+            return
+        
+        # No cuts possible - this is a leaf (text block)
+        parent.node_type = "text_block"
+    
+    def _find_valleys(
+        self,
+        projection: np.ndarray,
+        min_gap: int,
+    ) -> list[tuple[int, int]]:
+        """
+        Find valleys (whitespace gaps) in a projection profile.
+        
+        A valley is a contiguous region where the projection is below threshold.
+        
+        Args:
+            projection: 1D array of projection values.
+            min_gap: Minimum gap size to consider.
+            
+        Returns:
+            List of (start, end) tuples for each valley.
+        """
+        if len(projection) == 0:
+            return []
+        
+        # Normalize projection
+        max_val = np.max(projection)
+        if max_val == 0:
+            return []
+        
+        normalized = projection / max_val
+        
+        # Find regions below threshold (valleys)
+        is_valley = normalized < self.valley_threshold
+        
+        valleys = []
+        in_valley = False
+        valley_start = 0
+        
+        for i, is_v in enumerate(is_valley):
+            if is_v and not in_valley:
+                # Start of valley
+                valley_start = i
+                in_valley = True
+            elif not is_v and in_valley:
+                # End of valley
+                if i - valley_start >= min_gap:
+                    valleys.append((valley_start, i))
+                in_valley = False
+        
+        # Handle valley at end
+        if in_valley and len(projection) - valley_start >= min_gap:
+            valleys.append((valley_start, len(projection)))
+        
+        return valleys
+    
+    def _count_nodes(self, node: SegmentNode) -> int:
+        """Count total nodes in a tree."""
+        return 1 + sum(self._count_nodes(c) for c in node.children)
+    
+    def extract_text_in_regions(
+        self,
+        page: Any,
+        root: SegmentNode,
+    ) -> None:
+        """
+        Extract text content for each leaf region.
+        
+        Args:
+            page: A fitz.Page object.
+            root: Root SegmentNode from segment_page().
+        """
+        self._extract_text_recursive(page, root)
+    
+    def _extract_text_recursive(
+        self,
+        page: Any,
+        node: SegmentNode,
+    ) -> None:
+        """Recursively extract text for leaf nodes."""
+        import fitz
+        
+        if node.is_leaf:
+            # Extract text from this region
+            rect = fitz.Rect(
+                node.region.x0,
+                node.region.y0,
+                node.region.x1,
+                node.region.y1,
+            )
+            node.text = page.get_text("text", clip=rect).strip()
+        else:
+            for child in node.children:
+                self._extract_text_recursive(page, child)
+
+
+def segment_pdf_with_xy_cut(pdf_path: Path | str) -> list[SegmentNode]:
+    """
+    Convenience function to segment a PDF using XY-Cut.
+    
+    Args:
+        pdf_path: Path to the PDF file.
+        
+    Returns:
+        List of SegmentNode trees (one per page).
+        
+    Example:
+        trees = segment_pdf_with_xy_cut("document.pdf")
+        for page_tree in trees:
+            for leaf in get_leaves(page_tree):
+                print(leaf.text)
+    """
+    cutter = RecursiveXYCutter()
+    return cutter.segment_pdf(pdf_path)
+
+
+def get_leaves(node: SegmentNode) -> list[SegmentNode]:
+    """Get all leaf nodes from a segment tree."""
+    if node.is_leaf:
+        return [node]
+    
+    leaves = []
+    for child in node.children:
+        leaves.extend(get_leaves(child))
+    return leaves
+
+
+def analyze_document_with_xycut(
+    pdf_path: Path | str,
+    use_layoutlm: bool = True,
+) -> Any:
+    """
+    Analyze document using XY-Cut + LayoutLM visual classification.
+    
+    Combines geometric segmentation (XY-Cut) with visual analysis (LayoutLM)
+    to create a hierarchical document tree.
+    
+    Args:
+        pdf_path: Path to PDF file.
+        use_layoutlm: Use LayoutLM to classify block types (Header/Body/Title).
+        
+    Returns:
+        DocumentTree with visually-detected structure.
+    """
+    from rnsr.models import DocumentNode, DocumentTree
+    import fitz
+    
+    pdf_path = Path(pdf_path)
+    
+    # Segment with XY-Cut
+    cutter = RecursiveXYCutter()
+    page_trees = cutter.segment_pdf(pdf_path)
+    
+    # Extract text for each region
+    doc = fitz.open(pdf_path)
+    for page_num, tree in enumerate(page_trees):
+        cutter.extract_text_in_regions(doc[page_num], tree)
+    
+    # Optionally classify blocks with LayoutLM
+    if use_layoutlm:
+        try:
+            from rnsr.ingestion.layout_model import classify_layout_blocks
+            from PIL import Image
+            
+            for page_num, tree in enumerate(page_trees):
+                # Render page as image
+                page = doc[page_num]
+                pix = page.get_pixmap(dpi=150)
+                img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
+                
+                # Get all leaf regions
+                leaves = get_leaves(tree)
+                bboxes = [(leaf.region.x0, leaf.region.y0, leaf.region.x1, leaf.region.y1) 
+                          for leaf in leaves]
+                texts = [leaf.text for leaf in leaves]
+                
+                # Classify with LayoutLM
+                if bboxes:
+                    results = classify_layout_blocks(img, bboxes, texts)
+                    
+                    # Update node types based on classification
+                    for leaf, result in zip(leaves, results):
+                        leaf.node_type = result["label"].lower()
+                        
+        except Exception as e:
+            logger.warning("layoutlm_classification_failed", error=str(e))
+    
+    doc.close()
+    
+    # Convert to DocumentTree
+    root = DocumentNode(id="root", level=0, header=pdf_path.stem)
+    
+    section_num = 0
+    for page_tree in page_trees:
+        for leaf in get_leaves(page_tree):
+            if not leaf.text.strip():
+                continue
+            
+            section_num += 1
+            
+            # Determine if it's a header based on LayoutLM classification
+            is_header = leaf.node_type in ("header", "title")
+            
+            if is_header:
+                # Create header node
+                section = DocumentNode(
+                    id=f"sec_{section_num:03d}",
+                    level=1,
+                    header=leaf.text.strip(),
+                    page_num=leaf.region.page_num,
+                )
+            else:
+                # Create body node with synthetic header
+                from rnsr.ingestion.semantic_fallback import _generate_synthetic_header
+                
+                section = DocumentNode(
+                    id=f"sec_{section_num:03d}",
+                    level=1,
+                    header=_generate_synthetic_header(leaf.text, section_num),
+                    content=leaf.text,
+                    page_num=leaf.region.page_num,
+                )
+            
+            root.children.append(section)
+    
+    return DocumentTree(
+        title=pdf_path.stem,
+        root=root,
+        total_nodes=section_num + 1,
+        ingestion_tier=1,
+        ingestion_method="xy_cut_layoutlm" if use_layoutlm else "xy_cut",
+    )
+
+