natural-pdf 0.1.1__py3-none-any.whl → 0.1.3__py3-none-any.whl

natural_pdf/__init__.py

@@ -52,4 +52,36 @@ __version__ = "0.1.1"
 if HAS_QA:
     __all__ = ["PDF", "Page", "Region", "ElementCollection", "configure_logging", "DocumentQA", "get_qa_engine"]
 else:
-    __all__ = ["PDF", "Page", "Region", "ElementCollection", "configure_logging"]
+    __all__ = ["PDF", "Page", "Region", "ElementCollection", "configure_logging"]
+
+# Core classes
+from .core.pdf import PDF
+from .collections.pdf_collection import PDFCollection
+from .elements.region import Region
+
+# Search options (if extras installed)
+try:
+    from .search.search_options import TextSearchOptions, MultiModalSearchOptions, BaseSearchOptions
+except ImportError:
+    # Define dummy classes if extras not installed, so imports don't break
+    # but using them will raise the ImportError from check_haystack_availability
+    class TextSearchOptions:
+        def __init__(self, *args, **kwargs): pass
+    class MultiModalSearchOptions:
+        def __init__(self, *args, **kwargs): pass
+    class BaseSearchOptions:
+        def __init__(self, *args, **kwargs): pass
+
+# Expose logging setup? (Optional)
+# from . import logging_config
+# logging_config.setup_logging()
+
+# Explicitly define what gets imported with 'from natural_pdf import *'
+__all__ = [
+    'PDF',
+    'PDFCollection',
+    'Region',
+    'TextSearchOptions',       # Include search options
+    'MultiModalSearchOptions',
+    'BaseSearchOptions'
+]

natural_pdf/analyzers/layout/layout_analyzer.py

@@ -1,10 +1,11 @@
 import logging
 from typing import List, Dict, Any, Optional, Union
 from PIL import Image
+import copy
 
 from natural_pdf.elements.region import Region
 from natural_pdf.analyzers.layout.layout_manager import LayoutManager
-from natural_pdf.analyzers.layout.layout_options import LayoutOptions
+from natural_pdf.analyzers.layout.layout_options import LayoutOptions, TATRLayoutOptions, BaseLayoutOptions
 
 logger = logging.getLogger(__name__)
 
@@ -36,20 +37,25 @@ class LayoutAnalyzer:
         classes: Optional[List[str]] = None,
         exclude_classes: Optional[List[str]] = None,
         device: Optional[str] = None,
-        existing: str = "replace"
+        existing: str = "replace",
+        **kwargs
     ) -> List[Region]:
         """
         Analyze the page layout using the configured LayoutManager.
         
+        This method constructs the final options object, including internal context, 
+        and passes it to the LayoutManager.
+        
         Args:
-            engine: Name of the layout engine (e.g., 'yolo', 'tatr'). Uses manager's default if None.
-            options: Specific LayoutOptions object for advanced configuration.
+            engine: Name of the layout engine (e.g., 'yolo', 'tatr'). Uses manager's default if None and no options object given.
+            options: Specific LayoutOptions object for advanced configuration. If provided, simple args (confidence, etc.) are ignored.
             confidence: Minimum confidence threshold (simple mode).
             classes: Specific classes to detect (simple mode).
             exclude_classes: Classes to exclude (simple mode).
             device: Device for inference (simple mode).
             existing: How to handle existing detected regions: 'replace' (default) or 'append'.
-            
+            **kwargs: Additional engine-specific arguments (added to options.extra_args or used by constructor if options=None).
+
         Returns:
             List of created Region objects.
         """
@@ -57,72 +63,139 @@ class LayoutAnalyzer:
             logger.error(f"Page {self._page.number}: LayoutManager not available. Cannot analyze layout.")
             return []
 
-        logger.info(f"Page {self._page.number}: Analyzing layout (Engine: {engine or 'default'}, Options: {options is not None})...")
+        logger.info(f"Page {self._page.number}: Analyzing layout (Engine: {engine or 'default'}, Options provided: {options is not None})...")
 
-        # --- Render Page Image ---
-        logger.debug(f"  Rendering page {self._page.number} to image for layout analysis...")
+        # --- Render Page Image (Standard Resolution) ---
+        logger.debug(f"  Rendering page {self._page.number} to image for initial layout detection...")
         try:
-            # Use a resolution suitable for layout analysis, potentially configurable
-            layout_scale = getattr(self._page._parent, '_config', {}).get('layout_image_scale', 1.5) # ~108 DPI default
+            layout_scale = getattr(self._page._parent, '_config', {}).get('layout_image_scale', 1.5)
             layout_resolution = layout_scale * 72
-            # Render without existing highlights to avoid interference
-            page_image = self._page.to_image(resolution=layout_resolution, include_highlights=False)
-            logger.debug(f"  Rendered image size: {page_image.width}x{page_image.height}")
+            std_res_page_image = self._page.to_image(resolution=layout_resolution, include_highlights=False)
+            if not std_res_page_image:
+                raise ValueError("Initial page rendering returned None")
+            logger.debug(f"  Initial rendered image size: {std_res_page_image.width}x{std_res_page_image.height}")
         except Exception as e:
-            logger.error(f"  Failed to render page {self._page.number} to image: {e}", exc_info=True)
+            logger.error(f"  Failed to render initial page image: {e}", exc_info=True)
             return []
+            
+        # --- Calculate Scaling Factors (Standard Res Image <-> PDF) ---
+        if std_res_page_image.width == 0 or std_res_page_image.height == 0:
+            logger.error(f"Page {self._page.number}: Invalid initial rendered image dimensions. Cannot scale results.")
+            return []
+        img_scale_x = self._page.width / std_res_page_image.width
+        img_scale_y = self._page.height / std_res_page_image.height
+        logger.debug(f"  StdRes Image -> PDF Scaling: x={img_scale_x:.4f}, y={img_scale_y:.4f}")
 
-        # --- Prepare Arguments for Layout Manager ---
-        manager_args = {'image': page_image, 'options': options, 'engine': engine}
-        if confidence is not None: manager_args['confidence'] = confidence
-        if classes is not None: manager_args['classes'] = classes
-        if exclude_classes is not None: manager_args['exclude_classes'] = exclude_classes
-        if device is not None: manager_args['device'] = device
-
-        # --- Call Layout Manager ---
-        logger.debug(f"  Calling Layout Manager...")
+        # --- Construct Final Options Object --- 
+        final_options: BaseLayoutOptions
+        
+        if options is not None:
+             # User provided a complete options object, use it directly
+             logger.debug("Using user-provided options object.")
+             final_options = copy.deepcopy(options) # Copy to avoid modifying original user object
+             if kwargs:
+                  logger.warning(f"Ignoring kwargs {list(kwargs.keys())} because a full options object was provided.")
+             # Infer engine from options type if engine arg wasn't provided
+             if engine is None:
+                  for name, registry_entry in self._layout_manager.ENGINE_REGISTRY.items():
+                       if isinstance(final_options, registry_entry['options_class']):
+                            engine = name
+                            logger.debug(f"Inferred engine '{engine}' from options type.")
+                            break
+                  if engine is None:
+                       logger.warning("Could not infer engine from provided options object.")
+        else:
+             # Construct options from simple args (engine, confidence, classes, etc.)
+             logger.debug("Constructing options from simple arguments.")
+             selected_engine = engine or self._layout_manager.get_available_engines()[0] # Use provided or first available
+             engine_lower = selected_engine.lower()
+             registry = self._layout_manager.ENGINE_REGISTRY
+             
+             if engine_lower not in registry:
+                  raise ValueError(f"Unknown or unavailable engine: '{selected_engine}'. Available: {list(registry.keys())}")
+                  
+             options_class = registry[engine_lower]['options_class']
+             
+             # Get base defaults
+             base_defaults = BaseLayoutOptions()
+             
+             # Prepare args for constructor, prioritizing explicit args over defaults
+             constructor_args = {
+                 'confidence': confidence if confidence is not None else base_defaults.confidence,
+                 'classes': classes, # Pass None if not provided
+                 'exclude_classes': exclude_classes, # Pass None if not provided
+                 'device': device if device is not None else base_defaults.device,
+                 'extra_args': kwargs # Pass other kwargs here
+             }
+             # Remove None values unless they are valid defaults (like classes=None)
+             # We can pass all to the dataclass constructor; it handles defaults
+             
+             try:
+                  final_options = options_class(**constructor_args)
+                  logger.debug(f"Constructed options: {final_options}")
+             except TypeError as e:
+                  logger.error(f"Failed to construct options object {options_class.__name__} with args {constructor_args}: {e}")
+                  # Filter kwargs to only include fields defined in the specific options class? Complex.
+                  # Re-raise for now, indicates programming error or invalid kwarg.
+                  raise e
+
+        # --- Add Internal Context to extra_args (ALWAYS) --- 
+        if not hasattr(final_options, 'extra_args') or final_options.extra_args is None:
+             final_options.extra_args = {}
+        final_options.extra_args['_page_ref'] = self._page
+        final_options.extra_args['_img_scale_x'] = img_scale_x
+        final_options.extra_args['_img_scale_y'] = img_scale_y
+        logger.debug(f"Added internal context to final_options.extra_args: {final_options.extra_args}")
+
+        # --- Call Layout Manager with the Final Options --- 
+        logger.debug(f"Calling Layout Manager with final options object.")
         try:
-            detections = self._layout_manager.analyze_layout(**manager_args)
+            # Pass only image and the constructed options object
+            detections = self._layout_manager.analyze_layout(
+                image=std_res_page_image, 
+                options=final_options
+                # No engine, confidence, classes etc. passed here directly
+            )
             logger.info(f"  Layout Manager returned {len(detections)} detections.")
         except Exception as e:
             logger.error(f"  Layout analysis failed: {e}", exc_info=True)
             return []
 
-        # --- Process Detections (Convert to Regions, Scale Coords) ---
-        # Calculate scale factor to convert from image back to PDF coordinates
-        if page_image.width == 0 or page_image.height == 0:
-            logger.error(f"Page {self._page.number}: Invalid rendered image dimensions ({page_image.width}x{page_image.height}). Cannot scale layout results.")
-            return []
-        scale_x = self._page.width / page_image.width
-        scale_y = self._page.height / page_image.height
-        logger.debug(f"  Scaling factors: x={scale_x:.4f}, y={scale_y:.4f}")
-
+        # --- Process Detections (Convert to Regions, Scale Coords from Image to PDF) ---
         layout_regions = []
         docling_id_to_region = {} # For hierarchy if using Docling
 
         for detection in detections:
             try:
+                # bbox is relative to std_res_page_image
                 x_min, y_min, x_max, y_max = detection['bbox']
 
                 # Convert coordinates from image to PDF space
-                pdf_x0 = x_min * scale_x
-                pdf_y0 = y_min * scale_y
-                pdf_x1 = x_max * scale_x
-                pdf_y1 = y_max * scale_y
-
-                # Create a Region object
+                pdf_x0 = x_min * img_scale_x
+                pdf_y0 = y_min * img_scale_y
+                pdf_x1 = x_max * img_scale_x
+                pdf_y1 = y_max * img_scale_y
+                
+                # Ensure PDF coords are valid
+                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)
+
+                # Create a Region object with PDF coordinates
                 region = Region(self._page, (pdf_x0, pdf_y0, pdf_x1, pdf_y1))
-                region.region_type = detection.get('class', 'unknown') # Original class name
-                region.normalized_type = detection.get('normalized_class', 'unknown') # Hyphenated name
+                region.region_type = detection.get('class', 'unknown')
+                region.normalized_type = detection.get('normalized_class', 'unknown')
                 region.confidence = detection.get('confidence', 0.0)
-                region.model = detection.get('model', engine or 'unknown') # Store model name
+                region.model = detection.get('model', engine or 'unknown')
                 region.source = 'detected'
-
+                
                 # Add extra info if available
                 if 'text' in detection: region.text_content = detection['text']
                 if 'docling_id' in detection: region.docling_id = detection['docling_id']
                 if 'parent_id' in detection: region.parent_id = detection['parent_id']
-                # Add other fields like polygon, position, row/col index if needed
 
                 layout_regions.append(region)
 
@@ -163,4 +236,20 @@ class LayoutAnalyzer:
         self._page.detected_layout_regions = self._page._regions['detected']
         logger.info(f"Layout analysis complete for page {self._page.number}.")
         
+        # --- Auto-create cells if requested by TATR options ---
+        if isinstance(final_options, TATRLayoutOptions) and final_options.create_cells:
+            logger.info(f"  Option create_cells=True detected for TATR. Attempting cell creation...")
+            created_cell_count = 0
+            for region in layout_regions:
+                # Only attempt on regions identified as tables by the TATR model
+                if region.model == 'tatr' and region.region_type == 'table':
+                    try:
+                        # create_cells now modifies the page elements directly and returns self
+                        region.create_cells()
+                        # We could potentially count cells created here if needed, 
+                        # but the method logs its own count.
+                    except Exception as cell_error:
+                        logger.warning(f"    Error calling create_cells for table region {region.bbox}: {cell_error}")
+            logger.info(f"  Finished cell creation process triggered by options.")
+        
         return layout_regions 

natural_pdf/analyzers/layout/layout_manager.py

@@ -120,9 +120,10 @@ class LayoutManager:
 
         # --- Determine Options and Engine ---
         if options is not None:
-            # Advanced Mode
-            logger.debug(f"LayoutManager: Using advanced mode with options object: {type(options).__name__}")
-            final_options = copy.deepcopy(options) # Use copy
+            # Advanced Mode: An options object was provided directly (or constructed by LayoutAnalyzer)
+            # Use this object directly, do not deep copy or reconstruct.
+            logger.debug(f"LayoutManager: Using provided options object: {type(options).__name__}")
+            final_options = options # Use the provided object directly
             found_engine = False
             for name, registry_entry in self.ENGINE_REGISTRY.items():
                 if isinstance(options, registry_entry['options_class']):
@@ -131,12 +132,14 @@ class LayoutManager:
                     break
             if not found_engine:
                  raise TypeError(f"Provided options object type '{type(options).__name__}' does not match any registered layout engine options.")
+            # Ignore simple kwargs if options object is present
             if kwargs:
-                logger.warning(f"Keyword arguments {list(kwargs.keys())} were provided alongside 'options' and will be ignored.")
+                logger.warning(f"Keyword arguments {list(kwargs.keys())} were provided alongside an 'options' object and will be ignored.")
         else:
-            # Simple Mode
+            # Simple Mode: No options object provided initially. 
+            # Determine engine from kwargs or default, then construct options.
             selected_engine_name = default_engine.lower()
-            logger.debug(f"LayoutManager: Using simple mode with engine: '{selected_engine_name}' and kwargs: {kwargs}")
+            logger.debug(f"LayoutManager: Using simple mode. Engine: '{selected_engine_name}', kwargs: {kwargs}")
 
             if selected_engine_name not in self.ENGINE_REGISTRY:
                  raise ValueError(f"Unknown or unavailable layout engine: '{selected_engine_name}'. Available: {available_engines}")

natural_pdf/analyzers/layout/layout_options.py

@@ -34,7 +34,7 @@ class TATRLayoutOptions(BaseLayoutOptions):
     max_detection_size: int = 800
     max_structure_size: int = 1000
     # Whether to create cell regions (can be slow)
-    create_cells: bool = False # Keep the flag for cell creation control
+    create_cells: bool = True
 
 # --- Paddle Specific Options ---
 @dataclass
@@ -51,10 +51,8 @@ class PaddleLayoutOptions(BaseLayoutOptions):
 @dataclass
 class SuryaLayoutOptions(BaseLayoutOptions):
     """Options specific to Surya layout detection."""
-    # Surya doesn't seem to have many config options based on the example,
-    # but we can add placeholders if needed. Device is handled by BaseLayoutOptions.
     model_name: str = "default" # Placeholder if different models become available
-    verbose: bool = False # Verbose logging for the detector class
+    recognize_table_structure: bool = True # Automatically run table structure recognition?
 
 # --- Docling Specific Options ---
 @dataclass