natural-pdf 0.1.7__py3-none-any.whl → 0.1.9__py3-none-any.whl

natural_pdf/__init__.py

@@ -3,6 +3,9 @@ Natural PDF - A more intuitive interface for working with PDFs.
 """
 
 import logging
+import os
+
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
 
 # Create library logger
 logger = logging.getLogger("natural_pdf")

natural_pdf/analyzers/layout/base.py

@@ -98,7 +98,7 @@ class LayoutDetector(ABC):
                 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  # Re-raise exception after logging
+                raise
         else:
             self.logger.debug(f"Using cached model for key: {cache_key}")
         return self._model_cache[cache_key]
@@ -135,7 +135,6 @@ class LayoutDetector(ABC):
             return
 
         if classes:
-            # Normalize both requested and supported classes for comparison
             normalized_supported = {self._normalize_class_name(c) for c in self.supported_classes}
             normalized_requested = {self._normalize_class_name(c) for c in classes}
             unsupported_normalized = normalized_requested - normalized_supported
@@ -153,7 +152,4 @@ class LayoutDetector(ABC):
     def __del__(self):
         """Cleanup resources."""
         self.logger.info(f"Cleaning up {self.__class__.__name__} resources.")
-        # Clear model cache to free up memory/GPU resources if models are large
-        # Consider implications if models are shared or expensive to reload
-        # del self._model_cache # Optional: uncomment if models should be released aggressively
         self._model_cache.clear()

natural_pdf/analyzers/layout/gemini.py

@@ -1,13 +1,13 @@
 # layout_detector_gemini.py
+import base64
 import importlib.util
+import io
 import logging
 import os
 from typing import Any, Dict, List, Optional
-import base64
-import io
 
-from pydantic import BaseModel, Field
 from PIL import Image
+from pydantic import BaseModel, Field
 
 # Use OpenAI library for interaction
 try:
@@ -53,10 +53,8 @@ logger = logging.getLogger(__name__)
 # This is used by the openai library's `response_format`
 class DetectedRegion(BaseModel):
     label: str = Field(description="The identified class name.")
-    bbox: List[float] = Field(
-        description="Bounding box coordinates [xmin, ymin, xmax, ymax].", min_items=4, max_items=4
-    )
-    confidence: float = Field(description="Confidence score [0.0, 1.0].", ge=0.0, le=1.0)
+    bbox: List[float] = Field(description="Bounding box coordinates [xmin, ymin, xmax, ymax].")
+    confidence: float = Field(description="Confidence score [0.0, 1.0].")
 
 
 class GeminiLayoutDetector(LayoutDetector):
@@ -70,16 +68,10 @@ class GeminiLayoutDetector(LayoutDetector):
         self.supported_classes = set()  # Indicate dynamic nature
 
     def is_available(self) -> bool:
-        """Check if openai library is installed and GOOGLE_API_KEY is available."""
-        api_key = os.environ.get("GOOGLE_API_KEY")
-        if not api_key:
-            logger.warning(
-                "GOOGLE_API_KEY environment variable not set. Gemini detector (via OpenAI lib) will not be available."
-            )
-            return False
+        """Check if openai library is installed."""
         if OpenAI is None:
             logger.warning(
-                "openai package not found. Gemini detector (via OpenAI lib) will not be available."
+                "openai package not found. Gemini detector (via OpenAI lib) will not be available. Run: pip install openai"
             )
             return False
         return True
@@ -96,44 +88,65 @@ class GeminiLayoutDetector(LayoutDetector):
     def _load_model_from_options(self, options: GeminiLayoutOptions) -> Any:
         """Validate options and return the model name."""
         if not self.is_available():
-            raise RuntimeError(
-                "OpenAI library not installed or GOOGLE_API_KEY not set. Please run: pip install openai"
-            )
+            raise RuntimeError("OpenAI library not installed. Please run: pip install openai")
 
         if not isinstance(options, GeminiLayoutOptions):
             raise TypeError("Incorrect options type provided for Gemini model loading.")
 
-        # Simply return the model name, client is created in detect()
+        # Model loading is deferred to detect() based on whether a client is provided
         return options.model_name
 
     def detect(self, image: Image.Image, options: BaseLayoutOptions) -> List[Dict[str, Any]]:
         """Detect layout elements in an image using Gemini via OpenAI library."""
         if not self.is_available():
-            raise RuntimeError("OpenAI library not installed or GOOGLE_API_KEY not set.")
+            # The is_available check now only confirms library presence
+            raise RuntimeError("OpenAI library not installed. Please run: pip install openai")
 
         # Ensure options are the correct type
-        if not isinstance(options, GeminiLayoutOptions):
+        final_options: GeminiLayoutOptions
+        if isinstance(options, GeminiLayoutOptions):
+            final_options = options
+        else:
+            # If base options are passed, try to convert, keeping extra_args
+            # Note: This won't transfer a 'client' if it was somehow attached to BaseLayoutOptions
             self.logger.warning(
-                "Received BaseLayoutOptions, expected GeminiLayoutOptions. Using defaults."
+                "Received BaseLayoutOptions, expected GeminiLayoutOptions. Converting and using defaults."
             )
-            options = GeminiLayoutOptions(
+            final_options = GeminiLayoutOptions(
                 confidence=options.confidence,
                 classes=options.classes,
                 exclude_classes=options.exclude_classes,
-                device=options.device,
+                device=options.device,  # device is not used by Gemini detector currently
                 extra_args=options.extra_args,
+                # client will be None here, forcing default client creation below
             )
 
-        model_name = self._get_model(options)
-        api_key = os.environ.get("GOOGLE_API_KEY")
-
+        model_name = self._get_model(final_options)
         detections = []
+
         try:
-            # --- 1. Initialize OpenAI Client for Gemini ---
-            client = OpenAI(api_key=api_key, base_url=self.GEMINI_BASE_URL)
+            # --- 1. Initialize OpenAI Client ---
+            client: Optional[OpenAI] = None
+            # Use the provided client instance
+            if hasattr(final_options.client, "beta") and hasattr(
+                final_options.client.beta.chat.completions, "parse"
+            ):
+                client = final_options.client
+                logger.debug("Using provided client instance.")
+            else:
+                logger.error(
+                    "Provided client does not seem compatible (missing beta.chat.completions.parse)."
+                )
+                raise TypeError(
+                    "Provided client is not compatible with the expected OpenAI interface."
+                )
+
+            if not client:
+                # This should not happen if logic above is correct, but as a safeguard
+                raise RuntimeError("Failed to obtain a valid client for Gemini detection.")
 
             # --- 2. Prepare Input for OpenAI API ---
-            if not options.classes:
+            if not final_options.classes:
                 logger.error("Gemini layout detection requires a list of classes to find.")
                 return []
 
@@ -145,15 +158,13 @@ class GeminiLayoutDetector(LayoutDetector):
             img_base64 = base64.b64encode(buffered.getvalue()).decode("utf-8")
             image_url = f"data:image/png;base64,{img_base64}"
 
-            # Construct the prompt text
-            class_list_str = ", ".join(f"`{c}`" for c in options.classes)
+            class_list_str = ", ".join(f"`{c}`" for c in final_options.classes)
             prompt_text = (
                 f"Analyze the provided image of a document page ({width}x{height}). "
                 f"Identify all regions corresponding to the following types: {class_list_str}. "
-                f"Return ONLY the structured data requested."
+                f"Return ONLY the structured data requested as formatted JSON."
             )
 
-            # Prepare messages for chat completions endpoint
             messages = [
                 {
                     "role": "user",
@@ -167,27 +178,26 @@ class GeminiLayoutDetector(LayoutDetector):
                 }
             ]
 
-            # --- 3. Call OpenAI API using .parse for structured output ---
             logger.debug(
-                f"Running Gemini detection via OpenAI lib (Model: {model_name}). Asking for classes: {options.classes}"
+                f"Running Gemini detection via OpenAI lib (Model: {model_name}). Asking for classes: {final_options.classes}"
             )
 
-            # Extract relevant generation parameters from extra_args if provided
-            # Mapping common names: temperature, top_p, max_tokens
             completion_kwargs = {
-                "temperature": options.extra_args.get("temperature", 0.2),  # Default to low temp
-                "top_p": options.extra_args.get("top_p"),
-                "max_tokens": options.extra_args.get(
-                    "max_tokens", 4096
-                ),  # Map from max_output_tokens
+                "temperature": final_options.extra_args.get(
+                    "temperature", 0.0
+                ),  # Default to low temp
+                "max_tokens": final_options.extra_args.get("max_tokens", 4096),
             }
-            # Filter out None values
+
             completion_kwargs = {k: v for k, v in completion_kwargs.items() if v is not None}
 
+            class ImageContents(BaseModel):
+                regions: List[DetectedRegion]
+
             completion: ChatCompletion = client.beta.chat.completions.parse(
                 model=model_name,
                 messages=messages,
-                response_format=List[DetectedRegion],  # Pass the Pydantic model list
+                response_format=ImageContents,
                 **completion_kwargs,
             )
 
@@ -199,7 +209,7 @@ class GeminiLayoutDetector(LayoutDetector):
                 return []
 
             # Get the parsed Pydantic objects
-            parsed_results = completion.choices[0].message.parsed
+            parsed_results = completion.choices[0].message.parsed.regions
             if not parsed_results or not isinstance(parsed_results, list):
                 logger.error(
                     f"Gemini response (via OpenAI lib) did not contain a valid list of parsed regions. Found: {type(parsed_results)}"
@@ -207,10 +217,10 @@ class GeminiLayoutDetector(LayoutDetector):
                 return []
 
             # --- 5. Convert to Detections & Filter ---
-            normalized_classes_req = {self._normalize_class_name(c) for c in options.classes}
+            normalized_classes_req = {self._normalize_class_name(c) for c in final_options.classes}
             normalized_classes_excl = (
-                {self._normalize_class_name(c) for c in options.exclude_classes}
-                if options.exclude_classes
+                {self._normalize_class_name(c) for c in final_options.exclude_classes}
+                if final_options.exclude_classes
                 else set()
             )
 
@@ -242,9 +252,9 @@ class GeminiLayoutDetector(LayoutDetector):
                     continue
 
                 # Check against base confidence threshold from options
-                if confidence_score < options.confidence:
+                if confidence_score < final_options.confidence:
                     logger.debug(
-                        f"Skipping item with confidence {confidence_score:.3f} below threshold {options.confidence}."
+                        f"Skipping item with confidence {confidence_score:.3f} below threshold {final_options.confidence}."
                     )
                     continue
 

natural_pdf/analyzers/layout/layout_analyzer.py

@@ -7,6 +7,7 @@ from PIL import Image
 from natural_pdf.analyzers.layout.layout_manager import LayoutManager
 from natural_pdf.analyzers.layout.layout_options import (
     BaseLayoutOptions,
+    GeminiLayoutOptions,
     LayoutOptions,
     TATRLayoutOptions,
 )
@@ -82,10 +83,10 @@ class LayoutAnalyzer:
             f"  Rendering page {self._page.number} to image for initial layout detection..."
         )
         try:
-            layout_scale = getattr(self._page._parent, "_config", {}).get("layout_image_scale", 1.5)
+            layout_scale = getattr(self._page._parent, "_config", {}).get("layout_image_scale", 1.0)
             layout_resolution = layout_scale * 72
             std_res_page_image = self._page.to_image(
-                resolution=layout_resolution, include_highlights=False
+                resolution=layout_resolution, include_highlights=False, scale=1.0
             )
             if not std_res_page_image:
                 raise ValueError("Initial page rendering returned None")
@@ -110,12 +111,11 @@ class LayoutAnalyzer:
         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."
+                    f"Ignoring simple mode keyword arguments {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:
@@ -145,16 +145,39 @@ class LayoutAnalyzer:
             # Get base defaults
             base_defaults = BaseLayoutOptions()
 
+            # Separate client from other kwargs
+            client_instance = kwargs.pop("client", None)  # Get client, remove from kwargs
+
+            # Separate model_name if provided for Gemini
+            model_name_kwarg = None
+            if issubclass(options_class, GeminiLayoutOptions):
+                model_name_kwarg = kwargs.pop("model_name", None)
+
             # 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
+                # Pass client explicitly if constructing Gemini options
+                # Note: We check issubclass *before* calling constructor
+                **(
+                    {"client": client_instance}
+                    if client_instance and issubclass(options_class, GeminiLayoutOptions)
+                    else {}
+                ),
+                # Pass model_name explicitly if constructing Gemini options and it was provided
+                **(
+                    {"model_name": model_name_kwarg}
+                    if model_name_kwarg and issubclass(options_class, GeminiLayoutOptions)
+                    else {}
+                ),
+                "extra_args": kwargs,  # Pass REMAINING kwargs here
             }
             # Remove None values unless they are valid defaults (like classes=None)
             # We can pass all to the dataclass constructor; it handles defaults
+            # **Filter constructor_args to remove None values that aren't defaults?**
+            # For simplicity, let dataclass handle it for now.
 
             try:
                 final_options = options_class(**constructor_args)
@@ -167,24 +190,30 @@ class LayoutAnalyzer:
                 # Re-raise for now, indicates programming error or invalid kwarg.
                 raise e
 
-        # --- Add Internal Context to extra_args (ALWAYS) ---
+        # --- Add Internal Context to extra_args (Applies to the final_options object) ---
         if not hasattr(final_options, "extra_args") or final_options.extra_args is None:
+            # Ensure extra_args exists, potentially overwriting if needed
+            final_options.extra_args = {}
+        elif not isinstance(final_options.extra_args, dict):
+            logger.warning(
+                f"final_options.extra_args was not a dict ({type(final_options.extra_args)}), replacing with internal context."
+            )
             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}"
+            f"Added/updated internal context in final_options.extra_args: {final_options.extra_args}"
         )
 
-        # --- Call Layout Manager with the Final Options ---
+        # --- Call Layout Manager (ALWAYS with options object) ---
         logger.debug(f"Calling Layout Manager with final options object.")
         try:
-            # Pass only image and the constructed options object
+            # ALWAYS pass the constructed/modified options object
             detections = self._layout_manager.analyze_layout(
                 image=std_res_page_image,
-                options=final_options,
-                # No engine, confidence, classes etc. passed here directly
+                options=final_options,  # Pass the final object with internal context
             )
             logger.info(f"  Layout Manager returned {len(detections)} detections.")
         # Specifically let errors about unknown/unavailable engines propagate

natural_pdf/analyzers/layout/layout_manager.py

@@ -96,9 +96,6 @@ class LayoutManager:
             "options_class": GeminiLayoutOptions,
         }
 
-    # Define the limited set of kwargs allowed for the simple analyze_layout call
-    SIMPLE_MODE_ALLOWED_KWARGS = {"engine", "confidence", "classes", "exclude_classes", "device"}
-
     def __init__(self):
         """Initializes the Layout Manager."""
         # Cache for detector instances (different from model cache inside detector)
@@ -145,109 +142,54 @@ class LayoutManager:
     def analyze_layout(
         self,
         image: Image.Image,
-        engine: Optional[str] = None,  # Default engine handled below
-        options: Optional[LayoutOptions] = None,
-        **kwargs,
+        options: LayoutOptions,
     ) -> List[Dict[str, Any]]:
         """
-        Analyzes layout of a single image using simple args or an options object.
+        Analyzes layout of a single image using a specific options object.
 
         Args:
             image: The PIL Image to analyze.
-            engine: Name of the engine (e.g., 'yolo', 'tatr'). Ignored if 'options' provided.
-                    Defaults to the first available engine if None.
-            options: Specific LayoutOptions object for advanced configuration.
-            **kwargs: For simple mode, accepts: 'confidence', 'classes',
-                      'exclude_classes', 'device'.
+            options: Specific LayoutOptions object containing configuration and context.
+                     This object MUST be provided.
 
         Returns:
             A list of standardized detection dictionaries.
         """
-        final_options: BaseLayoutOptions
-        selected_engine_name: str
-
-        if not isinstance(image, Image.Image):
-            raise TypeError("Input 'image' must be a PIL Image.")
-
-        available_engines = self.get_available_engines()
-        if not available_engines:
-            raise RuntimeError("No layout engines are available. Please check dependencies.")
-
-        # Determine default engine if not specified
-        default_engine = engine if engine else available_engines[0]
-
-        # --- Determine Options and Engine ---
-        if options is not None:
-            # 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"]):
-                    selected_engine_name = name
-                    found_engine = True
-                    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 an 'options' object and will be ignored."
-                )
-        else:
-            # 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. Engine: '{selected_engine_name}', kwargs: {kwargs}"
+        selected_engine_name: Optional[str] = None
+        found_engine = False
+        for name, registry_entry in self.ENGINE_REGISTRY.items():
+            if isinstance(options, registry_entry["options_class"]):
+                selected_engine_name = name
+                found_engine = True
+                break
+        if not found_engine or selected_engine_name is None:
+            available_options_types = [
+                reg["options_class"].__name__ for reg in self.ENGINE_REGISTRY.values()
+            ]
+            raise TypeError(
+                f"Provided options object type '{type(options).__name__}' does not match any registered layout engine options: {available_options_types}"
             )
 
-            if selected_engine_name not in self.ENGINE_REGISTRY:
-                raise ValueError(
-                    f"Unknown or unavailable layout engine: '{selected_engine_name}'. Available: {available_engines}"
-                )
-
-            unexpected_kwargs = set(kwargs.keys()) - self.SIMPLE_MODE_ALLOWED_KWARGS
-            if unexpected_kwargs:
-                raise TypeError(
-                    f"Got unexpected keyword arguments in simple mode: {list(unexpected_kwargs)}. Use the 'options' parameter for detailed configuration."
-                )
-
-            options_class = self.ENGINE_REGISTRY[selected_engine_name]["options_class"]
-            # Use BaseLayoutOptions defaults unless overridden by kwargs
-            base_defaults = BaseLayoutOptions()
-            simple_args = {
-                "confidence": kwargs.get("confidence", base_defaults.confidence),
-                "classes": kwargs.get("classes"),
-                "exclude_classes": kwargs.get("exclude_classes"),
-                "device": kwargs.get("device", base_defaults.device),
-            }
-            # Filter out None values before passing to constructor
-            simple_args_filtered = {k: v for k, v in simple_args.items() if v is not None}
-            final_options = options_class(**simple_args_filtered)
-            logger.debug(f"LayoutManager: Constructed options for simple mode: {final_options}")
-
-        # --- Get Engine Instance and Process ---
         try:
             engine_instance = self._get_engine_instance(selected_engine_name)
             logger.info(f"Analyzing layout with engine '{selected_engine_name}'...")
 
-            # Call the engine's detect method
-            detections = engine_instance.detect(image, final_options)
+            detections = engine_instance.detect(image, options)  # Pass options directly
 
             logger.info(f"Layout analysis complete. Found {len(detections)} regions.")
             return detections
 
         except (ImportError, RuntimeError, ValueError, TypeError) as e:
-            logger.error(
-                f"Layout analysis failed for engine '{selected_engine_name}': {e}", exc_info=True
-            )
+            # Add engine name to error message if possible
+            engine_context = f" for engine '{selected_engine_name}'" if selected_engine_name else ""
+            logger.error(f"Layout analysis failed{engine_context}: {e}", exc_info=True)
             raise  # Re-raise expected errors
         except Exception as e:
-            logger.error(f"An unexpected error occurred during layout analysis: {e}", exc_info=True)
+            engine_context = f" for engine '{selected_engine_name}'" if selected_engine_name else ""
+            logger.error(
+                f"An unexpected error occurred during layout analysis{engine_context}: {e}",
+                exc_info=True,
+            )
             raise  # Re-raise unexpected errors
 
     def get_available_engines(self) -> List[str]:

natural_pdf/analyzers/layout/layout_options.py

@@ -43,6 +43,12 @@ class TATRLayoutOptions(BaseLayoutOptions):
     max_structure_size: int = 1000
     # Whether to create cell regions (can be slow)
     create_cells: bool = True
+    # Image enhancement options
+    enhance_contrast: float = 1.5  # Contrast enhancement factor (1.0 = no change)
+    # Special thresholds for specific elements
+    column_threshold: Optional[float] = (
+        None  # Lower threshold for columns (default: confidence * 0.8)
+    )
 
 
 # --- Paddle Specific Options ---
@@ -86,6 +92,7 @@ class GeminiLayoutOptions(BaseLayoutOptions):
     """Options specific to Gemini-based layout detection (using OpenAI compatibility)."""
 
     model_name: str = "gemini-2.0-flash"
+    client: Optional[Any] = None  # Allow passing a pre-configured client
     # Removed: prompt_template, temperature, top_p, max_output_tokens
     # These are typically passed directly to the chat completion call or via extra_args