natural-pdf 0.2.18__py3-none-any.whl → 0.2.19__py3-none-any.whl

natural_pdf/__init__.py

@@ -66,6 +66,7 @@ class Options:
         self.layout = ConfigSection(
             directional_offset=0.01,  # Offset in points when using directional methods
             auto_multipage=False,  # Whether directional methods span pages by default
+            directional_within=None,  # Region to constrain directional operations to
         )
 
 
@@ -126,6 +127,9 @@ from natural_pdf.elements.region import Region
 from natural_pdf.flows.flow import Flow
 from natural_pdf.flows.region import FlowRegion
 
+# Judge for visual classification
+from natural_pdf.judge import Decision, Judge, JudgeError, PickResult
+
 # Search options (if extras installed)
 try:
     from natural_pdf.search.search_options import (
@@ -165,6 +169,10 @@ __all__ = [
     "Flow",
     "FlowRegion",
     "Guides",
+    "Judge",
+    "Decision",
+    "PickResult",
+    "JudgeError",
     "TextSearchOptions",
     "MultiModalSearchOptions",
     "BaseSearchOptions",

natural_pdf/analyzers/checkbox/__init__.py

@@ -0,0 +1,6 @@
+"""Checkbox detection analyzers for natural-pdf."""
+
+from .checkbox_manager import CheckboxManager
+from .checkbox_options import CheckboxOptions, RTDETRCheckboxOptions
+
+__all__ = ["CheckboxManager", "CheckboxOptions", "RTDETRCheckboxOptions"]

natural_pdf/analyzers/checkbox/base.py

@@ -0,0 +1,265 @@
+"""Base class for checkbox detection engines."""
+
+import logging
+from abc import ABC, abstractmethod
+from typing import Any, Dict, List, Set
+
+from PIL import Image
+
+from .checkbox_options import CheckboxOptions
+
+logger = logging.getLogger(__name__)
+
+
+class CheckboxDetector(ABC):
+    """Abstract base class for checkbox detection engines.
+
+    This class defines the standard interface that all checkbox detection engines
+    must implement in natural-pdf. Checkbox detectors analyze document images to
+    identify checkboxes and their states (checked/unchecked).
+
+    Subclasses must implement:
+    - detect(): Core checkbox detection for a single image
+    - is_available(): Check if engine dependencies are installed
+    - _load_model_from_options(): Load and configure the detection model
+    - _get_cache_key(): Generate cache keys for model instances
+
+    Attributes:
+        logger: Logger instance for the specific detector.
+        _model_cache: Dictionary cache for loaded model instances.
+    """
+
+    def __init__(self):
+        """Initialize the base checkbox detector."""
+        self.logger = logging.getLogger(f"{__name__}.{self.__class__.__name__}")
+        self.logger.info(f"Initializing {self.__class__.__name__}")
+        self._model_cache: Dict[str, Any] = {}  # Cache for initialized models
+
+    @abstractmethod
+    def detect(self, image: Image.Image, options: CheckboxOptions) -> List[Dict[str, Any]]:
+        """
+        Detect checkboxes in a given PIL Image.
+
+        Args:
+            image: PIL Image of the page/region to analyze.
+            options: Instance of CheckboxOptions with configuration.
+
+        Returns:
+            List of detection dictionaries with:
+            - 'bbox': Tuple[float, float, float, float] - (x0, y0, x1, y1) relative to image
+            - 'class': str - Original class name from model (e.g., 'checkbox', 'checked_checkbox')
+            - 'normalized_class': str - Always 'checkbox'
+            - 'is_checked': bool - Whether checkbox is checked
+            - 'checkbox_state': str - 'checked' or 'unchecked'
+            - 'confidence': float - Confidence score (0.0-1.0)
+            - 'model': str - Name of the model used
+            - 'source': str - Always 'checkbox'
+        """
+        raise NotImplementedError("Subclasses must implement this method")
+
+    @classmethod
+    @abstractmethod
+    def is_available(cls) -> bool:
+        """
+        Check if the detector's dependencies are installed and usable.
+
+        Returns:
+            True if the detector is available, False otherwise.
+        """
+        raise NotImplementedError("Subclasses must implement this method")
+
+    def _get_cache_key(self, options: CheckboxOptions) -> str:
+        """
+        Generate a cache key for model loading based on relevant options.
+
+        Args:
+            options: The options dataclass instance.
+
+        Returns:
+            A string cache key.
+        """
+        # Base key includes device, subclasses should add model specifics
+        device_key = str(options.device).lower()
+        return f"{self.__class__.__name__}_{device_key}"
+
+    def _get_model(self, options: CheckboxOptions) -> Any:
+        """
+        Get or initialize the underlying model based on options, using caching.
+        """
+        cache_key = self._get_cache_key(options)
+        if cache_key not in self._model_cache:
+            self.logger.info(f"Loading model for cache key: {cache_key}")
+            try:
+                # Ensure dependencies are met before loading
+                if not self.is_available():
+                    raise RuntimeError(f"{self.__class__.__name__} dependencies are not met.")
+                self._model_cache[cache_key] = self._load_model_from_options(options)
+                self.logger.info(f"Model loaded successfully for key: {cache_key}")
+            except Exception as e:
+                self.logger.error(f"Failed to load model for key {cache_key}: {e}", exc_info=True)
+                # Remove potentially corrupted cache entry
+                self._model_cache.pop(cache_key, None)
+                raise
+        else:
+            self.logger.debug(f"Using cached model for key: {cache_key}")
+        return self._model_cache[cache_key]
+
+    @abstractmethod
+    def _load_model_from_options(self, options: CheckboxOptions) -> Any:
+        """
+        Load and configure the detection model based on provided options.
+
+        Args:
+            options: The options dataclass instance.
+
+        Returns:
+            The loaded model object(s).
+        """
+        raise NotImplementedError("Subclasses must implement _load_model_from_options")
+
+    def _map_label_to_state(self, label: str, options: CheckboxOptions) -> tuple[bool, str]:
+        """
+        Map model output label to checkbox state.
+
+        Args:
+            label: Raw label from model (e.g., 'checked_checkbox', '1')
+            options: Options containing label mapping
+
+        Returns:
+            Tuple of (is_checked: bool, state: str)
+        """
+        # Normalize label
+        normalized_label = str(label).lower().strip()
+
+        # Check mapping
+        if normalized_label in options.label_mapping:
+            state = options.label_mapping[normalized_label]
+            is_checked = state == "checked"
+            return is_checked, state
+
+        # Default heuristic if not in mapping
+        if any(term in normalized_label for term in ["checked", "tick", "filled", "1"]):
+            return True, "checked"
+        else:
+            return False, "unchecked"
+
+    def _apply_nms(
+        self, detections: List[Dict[str, Any]], iou_threshold: float
+    ) -> List[Dict[str, Any]]:
+        """
+        Apply non-maximum suppression to remove overlapping detections.
+        For checkboxes, we reject ANY meaningful overlap.
+
+        Args:
+            detections: List of detection dictionaries
+            iou_threshold: IoU threshold for suppression (ignored for checkboxes - we use stricter rules)
+
+        Returns:
+            Filtered list of detections
+        """
+        if not detections:
+            return detections
+
+        # Sort by confidence (descending), then by area (ascending) to prefer smaller boxes
+        def sort_key(det):
+            bbox = det["bbox"]
+            area = (bbox[2] - bbox[0]) * (bbox[3] - bbox[1])
+            return (-det["confidence"], area)
+
+        sorted_detections = sorted(detections, key=sort_key)
+
+        keep = []
+        for i, det in enumerate(sorted_detections):
+            should_keep = True
+            det_bbox = det["bbox"]
+
+            for kept_det in keep:
+                kept_bbox = kept_det["bbox"]
+
+                # Check for ANY overlap at all
+                if self._boxes_overlap(det_bbox, kept_bbox):
+                    should_keep = False
+                    logger.debug(f"Rejecting box {det_bbox} due to overlap with {kept_bbox}")
+                    break
+
+            if should_keep:
+                keep.append(det)
+                logger.debug(f"Keeping box {det_bbox} with confidence {det['confidence']}")
+
+        logger.info(f"NMS: Reduced {len(detections)} detections to {len(keep)}")
+        return keep
+
+    def _boxes_overlap(self, box1: tuple, box2: tuple) -> bool:
+        """Check if two boxes have any overlap at all."""
+        x1_min, y1_min, x1_max, y1_max = box1
+        x2_min, y2_min, x2_max, y2_max = box2
+
+        # Check if boxes are separated
+        if x1_max <= x2_min or x2_max <= x1_min:
+            return False
+        if y1_max <= y2_min or y2_max <= y1_min:
+            return False
+
+        # If we get here, boxes overlap
+        return True
+
+    def _compute_intersection_ratio(self, box1: tuple, box2: tuple) -> float:
+        """
+        Compute intersection ratio relative to the smaller box.
+        This is more aggressive than IoU for checkbox detection.
+        """
+        x1_min, y1_min, x1_max, y1_max = box1
+        x2_min, y2_min, x2_max, y2_max = box2
+
+        # Intersection
+        inter_xmin = max(x1_min, x2_min)
+        inter_ymin = max(y1_min, y2_min)
+        inter_xmax = min(x1_max, x2_max)
+        inter_ymax = min(y1_max, y2_max)
+
+        if inter_xmax < inter_xmin or inter_ymax < inter_ymin:
+            return 0.0
+
+        inter_area = (inter_xmax - inter_xmin) * (inter_ymax - inter_ymin)
+
+        # Areas of both boxes
+        area1 = (x1_max - x1_min) * (y1_max - y1_min)
+        area2 = (x2_max - x2_min) * (y2_max - y2_min)
+
+        # Ratio relative to smaller box
+        smaller_area = min(area1, area2)
+        if smaller_area == 0:
+            return 0.0
+
+        return inter_area / smaller_area
+
+    def _compute_iou(self, box1: tuple, box2: tuple) -> float:
+        """Compute IoU between two boxes."""
+        x1_min, y1_min, x1_max, y1_max = box1
+        x2_min, y2_min, x2_max, y2_max = box2
+
+        # Intersection
+        inter_xmin = max(x1_min, x2_min)
+        inter_ymin = max(y1_min, y2_min)
+        inter_xmax = min(x1_max, x2_max)
+        inter_ymax = min(y1_max, y2_max)
+
+        if inter_xmax < inter_xmin or inter_ymax < inter_ymin:
+            return 0.0
+
+        inter_area = (inter_xmax - inter_xmin) * (inter_ymax - inter_ymin)
+
+        # Union
+        area1 = (x1_max - x1_min) * (y1_max - y1_min)
+        area2 = (x2_max - x2_min) * (y2_max - y2_min)
+        union_area = area1 + area2 - inter_area
+
+        if union_area == 0:
+            return 0.0
+
+        return inter_area / union_area
+
+    def __del__(self):
+        """Cleanup resources."""
+        self.logger.info(f"Cleaning up {self.__class__.__name__} resources.")
+        self._model_cache.clear()

natural_pdf/analyzers/checkbox/checkbox_analyzer.py

@@ -0,0 +1,329 @@
+"""Checkbox analyzer for PDF pages and regions."""
+
+import logging
+from typing import Any, Dict, List, Optional, Union
+
+from PIL import Image
+
+from natural_pdf.elements.region import Region
+
+from .checkbox_manager import CheckboxManager
+from .checkbox_options import CheckboxOptions
+
+logger = logging.getLogger(__name__)
+
+
+class CheckboxAnalyzer:
+    """
+    Handles checkbox analysis for PDF pages and regions, including image rendering,
+    coordinate scaling, region creation, and result storage.
+    """
+
+    def __init__(self, element, checkbox_manager: Optional[CheckboxManager] = None):
+        """
+        Initialize the checkbox analyzer.
+
+        Args:
+            element: The Page or Region object to analyze
+            checkbox_manager: Optional CheckboxManager instance. If None, creates a new one.
+        """
+        self._element = element
+        self._checkbox_manager = checkbox_manager or CheckboxManager()
+
+        # Determine if element is a page or region
+        self._is_page = hasattr(element, "number") and hasattr(element, "_parent")
+        self._is_region = hasattr(element, "bbox") and hasattr(element, "_page")
+
+        if self._is_region:
+            self._page = element._page
+        else:
+            self._page = element
+
+    def detect_checkboxes(
+        self,
+        engine: Optional[str] = None,
+        options: Optional[Union[CheckboxOptions, Dict[str, Any]]] = None,
+        confidence: Optional[float] = None,
+        resolution: Optional[int] = None,
+        device: Optional[str] = None,
+        existing: str = "replace",
+        limit: Optional[int] = None,
+        **kwargs,
+    ) -> List[Region]:
+        """
+        Detect checkboxes in the page or region.
+
+        Args:
+            engine: Name of the detection engine (default: 'rtdetr')
+            options: CheckboxOptions instance or dict of options
+            confidence: Minimum confidence threshold
+            resolution: DPI for rendering (default: 150)
+            device: Device for inference
+            existing: How to handle existing checkbox regions: 'replace' (default) or 'append'
+            limit: Maximum number of checkboxes to detect
+            **kwargs: Additional engine-specific arguments
+
+        Returns:
+            List of created Region objects representing checkboxes
+        """
+        logger.info(
+            f"Detecting checkboxes (Engine: {engine or 'default'}, "
+            f"Element type: {'region' if self._is_region else 'page'})"
+        )
+
+        # Prepare options
+        if options is None:
+            # Build options from simple arguments
+            option_kwargs = {}
+            if confidence is not None:
+                option_kwargs["confidence"] = confidence
+            if resolution is not None:
+                option_kwargs["resolution"] = resolution
+            if device is not None:
+                option_kwargs["device"] = device
+            option_kwargs.update(kwargs)
+
+            # Let manager create appropriate options
+            final_options = None
+            final_kwargs = option_kwargs
+        else:
+            # Use provided options
+            final_options = options
+            # Apply any overrides
+            final_kwargs = {}
+            if confidence is not None:
+                final_kwargs["confidence"] = confidence
+            if resolution is not None:
+                final_kwargs["resolution"] = resolution
+            if device is not None:
+                final_kwargs["device"] = device
+            final_kwargs.update(kwargs)
+
+        # Render image
+        try:
+            resolution_val = (
+                resolution
+                or (
+                    final_options.resolution
+                    if final_options and hasattr(final_options, "resolution")
+                    else None
+                )
+                or 150
+            )
+
+            if self._is_region:
+                # For regions, crop the page image to just the region bounds
+                page_image = self._page.render(resolution=resolution_val)
+                if not page_image:
+                    raise ValueError("Page rendering returned None")
+
+                # Calculate region bounds in image coordinates
+                img_scale_x = page_image.width / self._page.width
+                img_scale_y = page_image.height / self._page.height
+
+                x0, y0, x1, y1 = self._element.bbox
+                img_x0 = int(x0 * img_scale_x)
+                img_y0 = int(y0 * img_scale_y)
+                img_x1 = int(x1 * img_scale_x)
+                img_y1 = int(y1 * img_scale_y)
+
+                # Crop to region
+                image = page_image.crop((img_x0, img_y0, img_x1, img_y1))
+
+                # Store crop offset for coordinate transformation
+                crop_offset = (img_x0, img_y0)
+
+            else:
+                # For pages, use the full image
+                image = self._page.render(resolution=resolution_val)
+                if not image:
+                    raise ValueError("Page rendering returned None")
+                crop_offset = (0, 0)
+
+            logger.debug(f"Rendered image size: {image.width}x{image.height}")
+
+        except Exception as e:
+            logger.error(f"Failed to render image: {e}", exc_info=True)
+            return []
+
+        # Calculate scaling factors
+        if self._is_region:
+            # For regions, scale is relative to the cropped image
+            scale_x = (self._element.bbox[2] - self._element.bbox[0]) / image.width
+            scale_y = (self._element.bbox[3] - self._element.bbox[1]) / image.height
+            pdf_offset = (self._element.bbox[0], self._element.bbox[1])
+        else:
+            # For pages, scale is from image to PDF coordinates
+            scale_x = self._page.width / image.width
+            scale_y = self._page.height / image.height
+            pdf_offset = (0, 0)
+
+        # Run detection
+        try:
+            detections = self._checkbox_manager.detect_checkboxes(
+                image=image, engine=engine, options=final_options, **final_kwargs
+            )
+            logger.info(f"Detected {len(detections)} checkboxes")
+        except Exception as e:
+            logger.error(f"Checkbox detection failed: {e}", exc_info=True)
+            return []
+
+        # Process detections into regions
+        checkbox_regions = []
+
+        for detection in detections:
+            try:
+                # Get image coordinates
+                img_x0, img_y0, img_x1, img_y1 = detection["bbox"]
+
+                if self._is_region:
+                    # For regions, add crop offset and scale to page image coords
+                    page_img_x0 = img_x0 + crop_offset[0]
+                    page_img_y0 = img_y0 + crop_offset[1]
+                    page_img_x1 = img_x1 + crop_offset[0]
+                    page_img_y1 = img_y1 + crop_offset[1]
+
+                    # Then scale to PDF coords
+                    pdf_x0 = page_img_x0 * (
+                        self._page.width / (self._page.render(resolution=resolution_val).width)
+                    )
+                    pdf_y0 = page_img_y0 * (
+                        self._page.height / (self._page.render(resolution=resolution_val).height)
+                    )
+                    pdf_x1 = page_img_x1 * (
+                        self._page.width / (self._page.render(resolution=resolution_val).width)
+                    )
+                    pdf_y1 = page_img_y1 * (
+                        self._page.height / (self._page.render(resolution=resolution_val).height)
+                    )
+                else:
+                    # For pages, directly scale to PDF coordinates
+                    pdf_x0 = img_x0 * scale_x + pdf_offset[0]
+                    pdf_y0 = img_y0 * scale_y + pdf_offset[1]
+                    pdf_x1 = img_x1 * scale_x + pdf_offset[0]
+                    pdf_y1 = img_y1 * scale_y + pdf_offset[1]
+
+                # Ensure valid bounds
+                pdf_x0, pdf_x1 = min(pdf_x0, pdf_x1), max(pdf_x0, pdf_x1)
+                pdf_y0, pdf_y1 = min(pdf_y0, pdf_y1), max(pdf_y0, pdf_y1)
+                pdf_x0 = max(0, pdf_x0)
+                pdf_y0 = max(0, pdf_y0)
+                pdf_x1 = min(self._page.width, pdf_x1)
+                pdf_y1 = min(self._page.height, pdf_y1)
+
+                # For region detection, skip checkboxes outside the region bounds
+                if self._is_region:
+                    region_x0, region_y0, region_x1, region_y1 = self._element.bbox
+                    # Check if checkbox center is within region
+                    cb_center_x = (pdf_x0 + pdf_x1) / 2
+                    cb_center_y = (pdf_y0 + pdf_y1) / 2
+                    if not (
+                        region_x0 <= cb_center_x <= region_x1
+                        and region_y0 <= cb_center_y <= region_y1
+                    ):
+                        continue  # Skip this checkbox
+
+                # Create region
+                region = Region(self._page, (pdf_x0, pdf_y0, pdf_x1, pdf_y1))
+                region.region_type = "checkbox"
+                region.normalized_type = "checkbox"
+                region.is_checked = detection.get("is_checked", False)
+                region.checkbox_state = detection.get("checkbox_state", "unchecked")
+                region.confidence = detection.get("confidence", 0.0)
+                region.model = detection.get("model", "checkbox_detector")
+                region.source = "checkbox"
+
+                # Store original class for debugging
+                region.original_class = detection.get("class", "unknown")
+
+                # Check if region contains text - if so, it's probably not a checkbox
+                # Get reject_with_text setting from options or kwargs, default to True
+                reject_with_text = True
+                if final_options:
+                    reject_with_text = getattr(final_options, "reject_with_text", True)
+                else:
+                    reject_with_text = kwargs.get("reject_with_text", True)
+
+                if reject_with_text:
+                    text_in_region = region.extract_text().strip()
+                    if text_in_region:
+                        # Allow only single characters that might be check marks
+                        if len(text_in_region) > 1 or text_in_region.isalnum():
+                            logger.debug(
+                                f"Rejecting checkbox at {region.bbox} - contains text: '{text_in_region}'"
+                            )
+                            continue
+
+                checkbox_regions.append(region)
+
+            except Exception as e:
+                logger.warning(f"Could not process checkbox detection: {detection}. Error: {e}")
+                continue
+
+        # Apply limit if specified
+        if limit is not None and len(checkbox_regions) > limit:
+            # Sort by confidence (highest first) and take top N
+            checkbox_regions = sorted(checkbox_regions, key=lambda r: r.confidence, reverse=True)[
+                :limit
+            ]
+
+        # Final cleanup - ensure no overlapping boxes (this shouldn't be needed if NMS worked)
+        cleaned_regions = []
+        for region in checkbox_regions:
+            overlaps = False
+            for kept_region in cleaned_regions:
+                # Check if bboxes overlap
+                r1 = region.bbox
+                r2 = kept_region.bbox
+                if not (r1[2] <= r2[0] or r2[2] <= r1[0] or r1[3] <= r2[1] or r2[3] <= r1[1]):
+                    overlaps = True
+                    logger.warning(
+                        f"Found overlapping checkbox regions after NMS: {r1} overlaps {r2}"
+                    )
+                    break
+            if not overlaps:
+                cleaned_regions.append(region)
+
+        if len(cleaned_regions) < len(checkbox_regions):
+            logger.warning(
+                f"Removed {len(checkbox_regions) - len(cleaned_regions)} overlapping checkboxes in final cleanup"
+            )
+            checkbox_regions = cleaned_regions
+
+        # Store results
+        logger.debug(f"Storing {len(checkbox_regions)} checkbox regions (mode: {existing})")
+
+        # Initialize storage if needed
+        if not hasattr(self._page, "_regions"):
+            self._page._regions = {}
+
+        # Handle existing regions
+        if existing.lower() == "append":
+            if "checkbox" not in self._page._regions:
+                self._page._regions["checkbox"] = []
+            self._page._regions["checkbox"].extend(checkbox_regions)
+        else:  # replace
+            # Remove old checkbox regions from element manager
+            if "checkbox" in self._page._regions:
+                old_checkboxes = self._page._regions["checkbox"]
+                if (
+                    hasattr(self._page._element_mgr, "_elements")
+                    and self._page._element_mgr._elements
+                ):
+                    current_regions = self._page._element_mgr._elements.get("regions", [])
+                    # Remove old checkbox regions
+                    self._page._element_mgr._elements["regions"] = [
+                        r for r in current_regions if r not in old_checkboxes
+                    ]
+            self._page._regions["checkbox"] = checkbox_regions
+
+        # Add to element manager
+        for region in checkbox_regions:
+            self._page._element_mgr.add_region(region)
+
+        # Store for easy access
+        self._page.detected_checkbox_regions = self._page._regions.get("checkbox", [])
+
+        logger.info(f"Checkbox detection complete. Found {len(checkbox_regions)} checkboxes.")
+
+        return checkbox_regions