natural-pdf 0.1.6__py3-none-any.whl → 0.1.8__py3-none-any.whl

natural_pdf/core/element_manager.py

@@ -359,8 +359,10 @@ class ElementManager:
 
                 # Handle potential None confidence
                 raw_confidence = result.get("confidence")
-                confidence_value = float(raw_confidence) if raw_confidence is not None else None # Keep None if it was None
-                ocr_text = result.get("text") # Get text, will be None if detect_only
+                confidence_value = (
+                    float(raw_confidence) if raw_confidence is not None else None
+                )  # Keep None if it was None
+                ocr_text = result.get("text")  # Get text, will be None if detect_only
 
                 # Create the TextElement for the word
                 word_element_data = {
@@ -373,7 +375,7 @@ class ElementManager:
                     "height": pdf_height,
                     "object_type": "word",  # Treat OCR results as whole words
                     "source": "ocr",
-                    "confidence": confidence_value, # Use the handled confidence
+                    "confidence": confidence_value,  # Use the handled confidence
                     "fontname": "OCR",  # Use consistent OCR fontname
                     "size": (
                         round(pdf_height) if pdf_height > 0 else 10.0
@@ -391,7 +393,7 @@ class ElementManager:
                 ocr_char_dict.setdefault("adv", ocr_char_dict.get("width", 0))
 
                 # Add the char dict list to the word data before creating TextElement
-                word_element_data["_char_dicts"] = [ocr_char_dict] # Store itself as its only char
+                word_element_data["_char_dicts"] = [ocr_char_dict]  # Store itself as its only char
 
                 word_elem = TextElement(word_element_data, self._page)
                 added_word_elements.append(word_elem)
@@ -537,3 +539,70 @@ class ElementManager:
         """Get all region elements."""
         self.load_elements()
         return self._elements.get("regions", [])
+
+    def remove_ocr_elements(self):
+        """
+        Remove all elements with source="ocr" from the elements dictionary.
+        This should be called before adding new OCR elements if replacement is desired.
+        
+        Returns:
+            int: Number of OCR elements removed
+        """
+        # Load elements if not already loaded
+        self.load_elements()
+        
+        removed_count = 0
+        
+        # Filter out OCR elements from words
+        if "words" in self._elements:
+            original_len = len(self._elements["words"])
+            self._elements["words"] = [
+                word for word in self._elements["words"] 
+                if getattr(word, "source", None) != "ocr"
+            ]
+            removed_count += original_len - len(self._elements["words"])
+        
+        # Filter out OCR elements from chars
+        if "chars" in self._elements:
+            original_len = len(self._elements["chars"])
+            self._elements["chars"] = [
+                char for char in self._elements["chars"]
+                if (isinstance(char, dict) and char.get("source") != "ocr") or
+                   (not isinstance(char, dict) and getattr(char, "source", None) != "ocr")
+            ]
+            removed_count += original_len - len(self._elements["chars"])
+        
+        logger.info(f"Page {self._page.number}: Removed {removed_count} OCR elements.")
+        return removed_count
+
+    def remove_element(self, element, element_type="words"):
+        """
+        Remove a specific element from the managed elements.
+        
+        Args:
+            element: The element to remove
+            element_type: The type of element ('words', 'chars', etc.)
+            
+        Returns:
+            bool: True if removed successfully, False otherwise
+        """
+        # Load elements if not already loaded
+        self.load_elements()
+        
+        # Check if the collection exists
+        if element_type not in self._elements:
+            logger.warning(f"Cannot remove element: collection '{element_type}' does not exist")
+            return False
+            
+        # Try to remove the element
+        try:
+            if element in self._elements[element_type]:
+                self._elements[element_type].remove(element)
+                logger.debug(f"Removed element from {element_type}: {element}")
+                return True
+            else:
+                logger.debug(f"Element not found in {element_type}: {element}")
+                return False
+        except Exception as e:
+            logger.error(f"Error removing element from {element_type}: {e}", exc_info=True)
+            return False

natural_pdf/core/page.py

@@ -6,14 +6,19 @@ import logging
 import os
 import re
 import tempfile
+import time # Import time
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Union
+import concurrent.futures # Added import
+from tqdm.auto import tqdm # Added tqdm import
+import threading
 
 import pdfplumber
 from PIL import Image, ImageDraw
 
 from natural_pdf.elements.collections import ElementCollection
 from natural_pdf.elements.region import Region
+from natural_pdf.utils.locks import pdf_render_lock  # Import from utils instead
 
 if TYPE_CHECKING:
     import pdfplumber
@@ -46,10 +51,20 @@ from natural_pdf.widgets.viewer import _IPYWIDGETS_AVAILABLE, SimpleInteractiveV
 from natural_pdf.qa import DocumentQA, get_qa_engine
 from natural_pdf.ocr.utils import _apply_ocr_correction_to_elements
 
+# --- Classification Imports --- #
+from natural_pdf.classification.mixin import ClassificationMixin
+from natural_pdf.classification.manager import ClassificationManager # For type hint
+# --- End Classification Imports --- #
+
+from natural_pdf.utils.locks import pdf_render_lock # Import the lock
+from natural_pdf.elements.base import Element # Import base element
+from natural_pdf.classification.mixin import ClassificationMixin # Import classification mixin
+from natural_pdf.extraction.mixin import ExtractionMixin # Import extraction mixin
+
 logger = logging.getLogger(__name__)
 
 
-class Page:
+class Page(ClassificationMixin, ExtractionMixin):
     """
     Enhanced Page wrapper built on top of pdfplumber.Page.
 
@@ -73,14 +88,21 @@ class Page:
         self._text_styles = None  # Lazy-loaded text style analyzer results
         self._exclusions = []  # List to store exclusion functions/regions
 
+        # --- ADDED --- Metadata store for mixins
+        self.metadata: Dict[str, Any] = {}
+        # --- END ADDED ---
+
         # Region management
         self._regions = {
             "detected": [],  # Layout detection results
             "named": {},  # Named regions (name -> region)
         }
 
-        # Initialize ElementManager
-        self._element_mgr = ElementManager(self, font_attrs)
+        # Initialize ElementManager, passing font_attrs
+        self._element_mgr = ElementManager(self, font_attrs=font_attrs)
+        # self._highlighter = HighlightingService(self) # REMOVED - Use property accessor
+        # --- NEW --- Central registry for analysis results
+        self.analyses: Dict[str, Any] = {}
 
         # --- Get OCR Manager Instance ---
         if (
@@ -115,6 +137,8 @@ class Page:
         # Initialize the internal variable with a single underscore
         self._layout_analyzer = None
 
+        self._load_elements()
+
     @property
     def pdf(self) -> "PDF":
         """Provides public access to the parent PDF object."""
@@ -1233,7 +1257,7 @@ class Page:
         render_ocr: bool = False,
         resolution: Optional[float] = None,
         include_highlights: bool = True,
-        exclusions: Optional[str] = None, # New parameter
+        exclusions: Optional[str] = None,  # New parameter
         **kwargs,
     ) -> Optional[Image.Image]:
         """
@@ -1257,38 +1281,48 @@ class Page:
         """
         image = None
         render_resolution = resolution if resolution is not None else scale * 72
+        thread_id = threading.current_thread().name
+        logger.debug(f"[{thread_id}] Page {self.index}: Attempting to acquire pdf_render_lock for to_image...")
+        lock_wait_start = time.monotonic()
         try:
-            if include_highlights:
-                # Delegate rendering to the central service
-                image = self._highlighter.render_page(
-                    page_index=self.index,
-                    scale=scale, # Note: scale is used by highlighter internally for drawing
-                    labels=labels,
-                    legend_position=legend_position,
-                    render_ocr=render_ocr,
-                    resolution=render_resolution, # Pass the calculated resolution
-                    **kwargs,
-                )
-            else:
-                # Get the base page image directly from pdfplumber if no highlights needed
-                # Use the underlying pdfplumber page object
-                img_object = self._page.to_image(resolution=render_resolution, **kwargs)
-                # Access the PIL image directly (assuming pdfplumber structure)
-                image = (
-                    img_object.annotated
-                    if hasattr(img_object, "annotated")
-                    else img_object._repr_png_()
-                )
-                if isinstance(image, bytes):  # Handle cases where it returns bytes
-                    from io import BytesIO
+            # Acquire the global PDF rendering lock
+            with pdf_render_lock:
+                lock_acquired_time = time.monotonic()
+                logger.debug(f"[{thread_id}] Page {self.index}: Acquired pdf_render_lock (waited {lock_acquired_time - lock_wait_start:.2f}s). Starting render...")
+                if include_highlights:
+                    # Delegate rendering to the central service
+                    image = self._highlighter.render_page(
+                        page_index=self.index,
+                        scale=scale,  # Note: scale is used by highlighter internally for drawing
+                        labels=labels,
+                        legend_position=legend_position,
+                        render_ocr=render_ocr,
+                        resolution=render_resolution,  # Pass the calculated resolution
+                        **kwargs,
+                    )
+                else:
+                    # Get the base page image directly from pdfplumber if no highlights needed
+                    # Use the underlying pdfplumber page object
+                    img_object = self._page.to_image(resolution=render_resolution, **kwargs)
+                    # Access the PIL image directly (assuming pdfplumber structure)
+                    image = (
+                        img_object.annotated
+                        if hasattr(img_object, "annotated")
+                        else img_object._repr_png_()
+                    )
+                    if isinstance(image, bytes):  # Handle cases where it returns bytes
+                        from io import BytesIO
 
-                    image = Image.open(BytesIO(image)).convert(
-                        "RGB"
-                    )  # Convert to RGB for consistency
+                        image = Image.open(BytesIO(image)).convert(
+                            "RGB"
+                        )  # Convert to RGB for consistency
 
         except Exception as e:
             logger.error(f"Error rendering page {self.index}: {e}", exc_info=True)
             return None  # Return None on error
+        finally:
+            render_end_time = time.monotonic()
+            logger.debug(f"[{thread_id}] Page {self.index}: Released pdf_render_lock. Total render time (incl. lock wait): {render_end_time - lock_wait_start:.2f}s")
 
         if image is None:
             return None
@@ -1322,16 +1356,21 @@ class Page:
                             max(0, img_x0),
                             max(0, img_top),
                             min(image.width, img_x1),
-                            min(image.height, img_bottom)
+                            min(image.height, img_bottom),
                         )
                         if img_coords[0] < img_coords[2] and img_coords[1] < img_coords[3]:
-                           draw.rectangle(img_coords, fill="white")
+                            draw.rectangle(img_coords, fill="white")
                         else:
-                             logger.warning(f"Skipping invalid exclusion rect for masking: {img_coords}")
+                            logger.warning(
+                                f"Skipping invalid exclusion rect for masking: {img_coords}"
+                            )
 
-                    del draw # Release drawing context
+                    del draw  # Release drawing context
             except Exception as mask_error:
-                logger.error(f"Error applying exclusion mask to page {self.index}: {mask_error}", exc_info=True)
+                logger.error(
+                    f"Error applying exclusion mask to page {self.index}: {mask_error}",
+                    exc_info=True,
+                )
                 # Decide if you want to return None or continue without mask
                 # For now, continue without mask
 
@@ -1379,6 +1418,7 @@ class Page:
         resolution: Optional[int] = None,
         detect_only: bool = False,
         apply_exclusions: bool = True,
+        replace: bool = True,
     ) -> "Page":
         """
         Apply OCR to THIS page and add results to page elements via PDF.apply_ocr.
@@ -1392,13 +1432,21 @@ class Page:
             resolution: DPI resolution for rendering page image before OCR.
             apply_exclusions: If True (default), render page image for OCR
                               with excluded areas masked (whited out).
+            detect_only: If True, only detect text bounding boxes, don't perform OCR.
+            replace: If True (default), remove any existing OCR elements before
+                    adding new ones. If False, add new OCR elements to existing ones.
 
         Returns:
-            List of created TextElements derived from OCR results for this page.
+            Self for method chaining.
         """
         if not hasattr(self._parent, "apply_ocr"):
             logger.error(f"Page {self.number}: Parent PDF missing 'apply_ocr'. Cannot apply OCR.")
-            return [] # Return empty list for consistency
+            return self  # Return self for chaining
+
+        # Remove existing OCR elements if replace is True
+        if replace and hasattr(self, "_element_mgr"):
+            logger.info(f"Page {self.number}: Removing existing OCR elements before applying new OCR.")
+            self._element_mgr.remove_ocr_elements()
 
         logger.info(f"Page {self.number}: Delegating apply_ocr to PDF.apply_ocr.")
         try:
@@ -1414,18 +1462,13 @@ class Page:
                 resolution=resolution,
                 detect_only=detect_only,
                 apply_exclusions=apply_exclusions,
+                replace=replace,  # Pass the replace parameter to PDF.apply_ocr
             )
         except Exception as e:
             logger.error(f"Page {self.number}: Error during delegated OCR call: {e}", exc_info=True)
-            return []
+            return self  # Return self for chaining
 
-        # Return the OCR elements specifically added to this page
-        ocr_elements = [el for el in self.words if getattr(el, "source", None) == "ocr"]
-        logger.debug(
-            f"Page {self.number}: apply_ocr completed. Found {len(ocr_elements)} OCR elements."
-        )
-        # Note: The method is typed to return Page for chaining, but the log indicates
-        # finding elements. Let's stick to returning self for chaining consistency.
+        # Return self for chaining
         return self
 
     def extract_ocr_elements(
@@ -1459,30 +1502,32 @@ class Page:
             return []
 
         logger.info(f"Page {self.number}: Extracting OCR elements (extract only)...")
-        
+
         # Determine rendering resolution
-        final_resolution = resolution if resolution is not None else 150 # Default to 150 DPI
+        final_resolution = resolution if resolution is not None else 150  # Default to 150 DPI
         logger.debug(f"  Using rendering resolution: {final_resolution} DPI")
-        
+
         try:
             # Get base image without highlights using the determined resolution
-            image = self.to_image(resolution=final_resolution, include_highlights=False)
-            if not image:
-                logger.error(f"  Failed to render page {self.number} to image for OCR extraction.")
-                return []
-            logger.debug(f"  Rendered image size: {image.width}x{image.height}")
+            # Use the global PDF rendering lock
+            with pdf_render_lock:
+                image = self.to_image(resolution=final_resolution, include_highlights=False)
+                if not image:
+                    logger.error(f"  Failed to render page {self.number} to image for OCR extraction.")
+                    return []
+                logger.debug(f"  Rendered image size: {image.width}x{image.height}")
         except Exception as e:
             logger.error(f"  Failed to render page {self.number} to image: {e}", exc_info=True)
             return []
 
         # Prepare arguments for the OCR Manager call
         manager_args = {
-             "images": image, 
-             "engine": engine,
-             "languages": languages,
-             "min_confidence": min_confidence,
-             "device": device, 
-             "options": options
+            "images": image,
+            "engine": engine,
+            "languages": languages,
+            "min_confidence": min_confidence,
+            "device": device,
+            "options": options,
         }
         manager_args = {k: v for k, v in manager_args.items() if v is not None}
 
@@ -1514,7 +1559,7 @@ class Page:
         scale_x = self.width / image.width if image.width else 1
         scale_y = self.height / image.height if image.height else 1
         for result in results:
-            try: # Added try-except around result processing
+            try:  # Added try-except around result processing
                 x0, top, x1, bottom = [float(c) for c in result["bbox"]]
                 elem_data = {
                     "text": result["text"],
@@ -1525,15 +1570,17 @@ class Page:
                     "bottom": bottom * scale_y,
                     "width": (x1 - x0) * scale_x,
                     "height": (bottom - top) * scale_y,
-                    "object_type": "text", # Using text for temporary elements
+                    "object_type": "text",  # Using text for temporary elements
                     "source": "ocr",
-                    "fontname": "OCR-extract", # Different name for clarity
+                    "fontname": "OCR-extract",  # Different name for clarity
                     "size": 10.0,
                     "page_number": self.number,
                 }
                 temp_elements.append(TextElement(elem_data, self))
             except (KeyError, ValueError, TypeError) as convert_err:
-                 logger.warning(f"  Skipping invalid OCR result during conversion: {result}. Error: {convert_err}")
+                logger.warning(
+                    f"  Skipping invalid OCR result during conversion: {result}. Error: {convert_err}"
+                )
 
         logger.info(f"  Created {len(temp_elements)} TextElements from OCR (extract only).")
         return temp_elements
@@ -2020,41 +2067,166 @@ class Page:
     def correct_ocr(
         self,
         correction_callback: Callable[[Any], Optional[str]],
-    ) -> "Page": # Return self for chaining
+        max_workers: Optional[int] = None,
+        progress_callback: Optional[Callable[[], None]] = None, # Added progress callback
+    ) -> "Page":  # Return self for chaining
         """
         Applies corrections to OCR-generated text elements on this page
-        using a user-provided callback function.
+        using a user-provided callback function, potentially in parallel.
 
         Finds text elements on this page whose 'source' attribute starts
         with 'ocr' and calls the `correction_callback` for each, passing the
-        element itself.
-
-        The `correction_callback` should contain the logic to:
-        1. Determine if the element needs correction.
-        2. Perform the correction (e.g., call an LLM).
-        3. Return the new text (`str`) or `None`.
-
-        If the callback returns a string, the element's `.text` is updated.
-        Metadata updates (source, confidence, etc.) should happen within the callback.
+        element itself. Updates the element's text if the callback returns
+        a new string.
 
         Args:
             correction_callback: A function accepting an element and returning
                                  `Optional[str]` (new text or None).
+            max_workers: The maximum number of threads to use for parallel execution.
+                         If None or 0 or 1, runs sequentially.
+            progress_callback: Optional callback function to call after processing each element.
 
         Returns:
             Self for method chaining.
         """
-        logger.info(f"Page {self.number}: Starting OCR correction process using callback '{correction_callback.__name__}'")
+        logger.info(
+            f"Page {self.number}: Starting OCR correction with callback '{correction_callback.__name__}' (max_workers={max_workers})"
+        )
+
+        target_elements_collection = self.find_all(
+            selector="text[source=ocr]", apply_exclusions=False
+        )
+        target_elements = target_elements_collection.elements # Get the list
+
+        if not target_elements:
+            logger.info(f"Page {self.number}: No OCR elements found to correct.")
+            return self
 
-        # Find OCR elements specifically on this page
-        # Note: We typically want to correct even if the element falls in an excluded area
-        target_elements = self.find_all(selector="text[source^=ocr]", apply_exclusions=False)
+        processed_count = 0
+        updated_count = 0
+        error_count = 0
 
-        # Delegate to the utility function
-        _apply_ocr_correction_to_elements(
-            elements=target_elements, # Pass the ElementCollection directly
-            correction_callback=correction_callback,
-            caller_info=f"Page({self.number})", # Pass caller info
+        # Define the task to be run by the worker thread or sequentially
+        def _process_element_task(element):
+            try:
+                current_text = getattr(element, 'text', None)
+                # Call the user-provided callback
+                corrected_text = correction_callback(element)
+
+                # Validate result type
+                if corrected_text is not None and not isinstance(corrected_text, str):
+                    logger.warning(f"Page {self.number}: Correction callback for element '{getattr(element, 'text', '')[:20]}...' returned non-string, non-None type: {type(corrected_text)}. Skipping update.")
+                    return element, None, None # Treat as no correction
+
+                return element, corrected_text, None  # Return element, result, no error
+            except Exception as e:
+                logger.error(
+                    f"Page {self.number}: Error applying correction callback to element '{getattr(element, 'text', '')[:30]}...' ({element.bbox}): {e}",
+                    exc_info=False # Keep log concise
+                )
+                return element, None, e # Return element, no result, error
+            finally:
+                # --- Call progress callback here --- #
+                if progress_callback:
+                    try:
+                        progress_callback()
+                    except Exception as cb_e:
+                        # Log error in callback itself, but don't stop processing
+                        logger.error(f"Page {self.number}: Error executing progress_callback: {cb_e}", exc_info=False)
+
+        # Choose execution strategy based on max_workers
+        if max_workers is not None and max_workers > 1:
+            # --- Parallel execution --- #
+            logger.info(f"Page {self.number}: Running OCR correction in parallel with {max_workers} workers.")
+            futures = []
+            with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
+                # Submit all tasks
+                future_to_element = {executor.submit(_process_element_task, element): element for element in target_elements}
+
+                # Process results as they complete (progress_callback called by worker)
+                for future in concurrent.futures.as_completed(future_to_element):
+                    processed_count += 1
+                    try:
+                        element, corrected_text, error = future.result()
+                        if error:
+                            error_count += 1
+                            # Error already logged in worker
+                        elif corrected_text is not None:
+                            # Apply correction if text changed
+                            current_text = getattr(element, 'text', None)
+                            if corrected_text != current_text:
+                                element.text = corrected_text
+                                updated_count += 1
+                    except Exception as exc:
+                        # Catch errors from future.result() itself
+                        element = future_to_element[future] # Find original element
+                        logger.error(f"Page {self.number}: Internal error retrieving correction result for element {element.bbox}: {exc}", exc_info=True)
+                        error_count += 1
+                        # Note: progress_callback was already called in the worker's finally block
+
+        else:
+            # --- Sequential execution --- #
+            logger.info(f"Page {self.number}: Running OCR correction sequentially.")
+            for element in target_elements:
+                 # Call the task function directly (it handles progress_callback)
+                 processed_count += 1
+                 _element, corrected_text, error = _process_element_task(element)
+                 if error:
+                     error_count += 1
+                 elif corrected_text is not None:
+                     # Apply correction if text changed
+                     current_text = getattr(_element, 'text', None)
+                     if corrected_text != current_text:
+                         _element.text = corrected_text
+                         updated_count += 1
+
+        logger.info(
+             f"Page {self.number}: OCR correction finished. Processed: {processed_count}/{len(target_elements)}, Updated: {updated_count}, Errors: {error_count}."
         )
 
         return self # Return self for chaining
+
+    # --- Classification Mixin Implementation --- #
+    def _get_classification_manager(self) -> "ClassificationManager":
+        if not hasattr(self, 'pdf') or not hasattr(self.pdf, 'get_manager'):
+             raise AttributeError("ClassificationManager cannot be accessed: Parent PDF or get_manager method missing.")
+        try:
+             # Use the PDF's manager registry accessor
+             return self.pdf.get_manager('classification')
+        except (ValueError, RuntimeError, AttributeError) as e:
+            # Wrap potential errors from get_manager for clarity
+            raise AttributeError(f"Failed to get ClassificationManager from PDF: {e}") from e
+
+    def _get_classification_content(self, model_type: str, **kwargs) -> Union[str, "Image"]: # Use "Image" for lazy import
+        if model_type == 'text':
+            text_content = self.extract_text(layout=False, use_exclusions=False) # Simple join, ignore exclusions for classification
+            if not text_content or text_content.isspace():
+                raise ValueError("Cannot classify page with 'text' model: No text content found.")
+            return text_content
+        elif model_type == 'vision':
+            # Get resolution from manager/kwargs if possible, else default
+            manager = self._get_classification_manager()
+            default_resolution = 150
+            # Access kwargs passed to classify method if needed
+            resolution = kwargs.get('resolution', default_resolution) if 'kwargs' in locals() else default_resolution
+
+            # Use to_image, ensuring no highlights interfere
+            img = self.to_image(
+                resolution=resolution,
+                include_highlights=False,
+                labels=False,
+                exclusions=None # Don't mask exclusions for classification input image
+            )
+            if img is None:
+                raise ValueError("Cannot classify page with 'vision' model: Failed to render image.")
+            return img
+        else:
+            raise ValueError(f"Unsupported model_type for classification: {model_type}")
+
+    def _get_metadata_storage(self) -> Dict[str, Any]:
+        # Ensure metadata exists
+        if not hasattr(self, 'metadata') or self.metadata is None:
+            self.metadata = {}
+        return self.metadata
+
+    # --- Content Extraction ---