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

natural_pdf/__init__.py

@@ -12,17 +12,16 @@ logger = logging.getLogger("natural_pdf")
 logger.addHandler(logging.NullHandler())
 
 
-# Utility function for users to easily configure logging
 def configure_logging(level=logging.INFO, handler=None):
-    """Configure Natural PDF's logging.
+    """Configure logging for the natural_pdf package.
 
     Args:
-        level: The logging level (e.g., logging.INFO, logging.DEBUG)
-        handler: A custom handler, or None to use StreamHandler
+        level: Logging level (e.g., logging.INFO, logging.DEBUG)
+        handler: Optional custom handler. Defaults to a StreamHandler.
     """
-    # Remove NullHandler if present
-    if logger.handlers and isinstance(logger.handlers[0], logging.NullHandler):
-        logger.removeHandler(logger.handlers[0])
+    # Avoid adding duplicate handlers
+    if any(isinstance(h, logging.StreamHandler) for h in logger.handlers):
+        return
 
     if handler is None:
         handler = logging.StreamHandler()
@@ -32,11 +31,7 @@ def configure_logging(level=logging.INFO, handler=None):
     logger.addHandler(handler)
     logger.setLevel(level)
 
-    # Propagate level to all child loggers
-    for name in logging.root.manager.loggerDict:
-        if name.startswith("natural_pdf."):
-            logging.getLogger(name).setLevel(level)
-
+    logger.propagate = False
 
 from natural_pdf.core.page import Page
 from natural_pdf.core.pdf import PDF
@@ -53,18 +48,21 @@ except ImportError:
 
 __version__ = "0.1.1"
 
+__all__ = [
+    "PDF",
+    "PDFCollection",
+    "Page",
+    "Region",
+    "ElementCollection",
+    "TextSearchOptions",
+    "MultiModalSearchOptions",
+    "BaseSearchOptions",
+    "configure_logging",
+]
+
 if HAS_QA:
-    __all__ = [
-        "PDF",
-        "Page",
-        "Region",
-        "ElementCollection",
-        "configure_logging",
-        "DocumentQA",
-        "get_qa_engine",
-    ]
-else:
-    __all__ = ["PDF", "Page", "Region", "ElementCollection", "configure_logging"]
+    __all__.extend(["DocumentQA", "get_qa_engine"])
+
 
 from .collections.pdf_collection import PDFCollection
 

natural_pdf/analyzers/layout/gemini.py

@@ -0,0 +1,264 @@
+# layout_detector_gemini.py
+import importlib.util
+import logging
+import os
+from typing import Any, Dict, List, Optional
+import base64
+import io
+
+from pydantic import BaseModel, Field
+from PIL import Image
+
+# Use OpenAI library for interaction
+try:
+    from openai import OpenAI
+    from openai.types.chat import ChatCompletion
+    # Import OpenAIError for exception handling if needed
+except ImportError:
+    OpenAI = None
+    ChatCompletion = None
+
+try:
+    from .base import LayoutDetector
+    from .layout_options import BaseLayoutOptions, GeminiLayoutOptions
+except ImportError:
+    # Placeholders if run standalone or imports fail
+    class BaseLayoutOptions:
+        pass
+
+    class GeminiLayoutOptions(BaseLayoutOptions):
+        pass
+
+    class LayoutDetector:
+        def __init__(self):
+            self.logger = logging.getLogger()
+            self.supported_classes = set() # Will be dynamic based on user request
+
+        def _get_model(self, options):
+            raise NotImplementedError
+
+        def _normalize_class_name(self, n):
+            return n.lower().replace("_", "-").replace(" ", "-")
+
+        def validate_classes(self, c):
+            pass # Less strict validation needed for LLM
+
+    logging.basicConfig()
+
+logger = logging.getLogger(__name__)
+
+# Define Pydantic model for the expected output structure
+# 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)
+
+
+class GeminiLayoutDetector(LayoutDetector):
+    """Document layout detector using Google's Gemini models via OpenAI compatibility layer."""
+
+    # Base URL for the Gemini OpenAI-compatible endpoint
+    GEMINI_BASE_URL = "https://generativelanguage.googleapis.com/v1beta/openai/"
+
+    def __init__(self):
+        super().__init__()
+        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
+        if OpenAI is None:
+             logger.warning("openai package not found. Gemini detector (via OpenAI lib) will not be available.")
+             return False
+        return True
+
+    def _get_cache_key(self, options: GeminiLayoutOptions) -> str:
+        """Generate cache key based on model name."""
+        if not isinstance(options, GeminiLayoutOptions):
+            options = GeminiLayoutOptions() # Use defaults
+
+        model_key = options.model_name
+        # Prompt is built dynamically, so not part of cache key based on options
+        return f"{self.__class__.__name__}_{model_key}"
+
+    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"
+            )
+
+        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()
+        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."
+            )
+
+        # Ensure options are the correct type
+        if not isinstance(options, GeminiLayoutOptions):
+            self.logger.warning(
+                "Received BaseLayoutOptions, expected GeminiLayoutOptions. Using defaults."
+            )
+            options = GeminiLayoutOptions(
+                confidence=options.confidence,
+                classes=options.classes,
+                exclude_classes=options.exclude_classes,
+                device=options.device,
+                extra_args=options.extra_args,
+            )
+
+        model_name = self._get_model(options)
+        api_key = os.environ.get("GOOGLE_API_KEY")
+
+        detections = []
+        try:
+            # --- 1. Initialize OpenAI Client for Gemini ---
+            client = OpenAI(
+                api_key=api_key,
+                base_url=self.GEMINI_BASE_URL
+            )
+
+            # --- 2. Prepare Input for OpenAI API ---
+            if not options.classes:
+                logger.error("Gemini layout detection requires a list of classes to find.")
+                return []
+
+            width, height = image.size
+
+            # Convert image to base64
+            buffered = io.BytesIO()
+            image.save(buffered, format="PNG")
+            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)
+            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."
+            )
+
+            # Prepare messages for chat completions endpoint
+            messages = [
+                {
+                    "role": "user",
+                    "content": [
+                        {"type": "text", "text": prompt_text},
+                        {
+                            "type": "image_url",
+                            "image_url": {"url": image_url},
+                        },
+                    ],
+                }
+            ]
+
+            # --- 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}")
+
+            # 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
+            }
+            # Filter out None values
+            completion_kwargs = {k: v for k, v in completion_kwargs.items() if v is not None}
+
+            completion: ChatCompletion = client.beta.chat.completions.parse(
+                model=model_name,
+                messages=messages,
+                response_format=List[DetectedRegion], # Pass the Pydantic model list
+                **completion_kwargs
+            )
+
+            logger.debug(f"Gemini response received via OpenAI lib.")
+
+            # --- 4. Process Parsed Response --- 
+            if not completion.choices:
+                logger.error("Gemini response (via OpenAI lib) contained no choices.")
+                return []
+
+            # Get the parsed Pydantic objects
+            parsed_results = completion.choices[0].message.parsed
+            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)}")
+                 return []
+
+            # --- 5. Convert to Detections & Filter --- 
+            normalized_classes_req = {
+                self._normalize_class_name(c) for c in options.classes
+            }
+            normalized_classes_excl = {
+                self._normalize_class_name(c) for c in options.exclude_classes
+            } if options.exclude_classes else set()
+
+            for item in parsed_results:
+                # The item is already a validated DetectedRegion Pydantic object
+                # Access fields directly
+                label = item.label
+                bbox_raw = item.bbox
+                confidence_score = item.confidence
+
+                # Coordinates should already be floats, but ensure tuple format
+                xmin, ymin, xmax, ymax = tuple(bbox_raw)
+
+                # --- Apply Filtering --- 
+                normalized_class = self._normalize_class_name(label)
+
+                # Check against requested classes (Should be guaranteed by schema, but doesn't hurt)
+                if normalized_class not in normalized_classes_req:
+                    logger.warning(f"Gemini (via OpenAI) returned unexpected class '{label}' despite schema. Skipping.")
+                    continue
+
+                # Check against excluded classes
+                if normalized_class in normalized_classes_excl:
+                    logger.debug(f"Skipping excluded class '{label}' (normalized: {normalized_class}).")
+                    continue
+                
+                # Check against base confidence threshold from options
+                if confidence_score < options.confidence:
+                    logger.debug(f"Skipping item with confidence {confidence_score:.3f} below threshold {options.confidence}.")
+                    continue
+
+                # Add detection
+                detections.append({
+                    "bbox": (xmin, ymin, xmax, ymax),
+                    "class": label, # Use original label from LLM
+                    "confidence": confidence_score,
+                    "normalized_class": normalized_class,
+                    "source": "layout",
+                    "model": "gemini", # Keep model name generic as gemini
+                })
+
+            self.logger.info(
+                f"Gemini (via OpenAI lib) processed response. Detected {len(detections)} layout elements matching criteria."
+            )
+
+        except Exception as e:
+            # Catch potential OpenAI API errors or other issues
+            self.logger.error(f"Error during Gemini detection (via OpenAI lib): {e}", exc_info=True)
+            return []
+
+        return detections
+
+    def _normalize_class_name(self, name: str) -> str:
+        """Normalizes class names for filtering (lowercase, hyphenated)."""
+        return super()._normalize_class_name(name)
+
+    def validate_classes(self, classes: List[str]):
+        """Validation is less critical as we pass requested classes to the LLM."""
+        pass # Override base validation if needed, but likely not necessary
+

natural_pdf/analyzers/layout/layout_manager.py

@@ -37,9 +37,15 @@ try:
 except ImportError:
     DoclingLayoutDetector = None
 
+try:
+    from .gemini import GeminiLayoutDetector
+except ImportError:
+    GeminiLayoutDetector = None
+
 from .layout_options import (
     BaseLayoutOptions,
     DoclingLayoutOptions,
+    GeminiLayoutOptions,
     LayoutOptions,
     PaddleLayoutOptions,
     SuryaLayoutOptions,
@@ -83,6 +89,13 @@ class LayoutManager:
             "options_class": DoclingLayoutOptions,
         }
 
+    # Add Gemini entry if available
+    if GeminiLayoutDetector:
+        ENGINE_REGISTRY["gemini"] = {
+            "class": GeminiLayoutDetector,
+            "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"}
 
@@ -108,8 +121,22 @@ class LayoutManager:
             detector_instance = engine_class()  # Instantiate
             if not detector_instance.is_available():
                 # Check availability before storing
+                # Construct helpful error message with install hint
+                install_hint = ""
+                if engine_name == "yolo":
+                    install_hint = "pip install 'natural-pdf[layout_yolo]'"
+                elif engine_name == "tatr":
+                    install_hint = "pip install 'natural-pdf[core-ml]'"
+                elif engine_name == "paddle":
+                    install_hint = "pip install 'natural-pdf[paddle]'"
+                elif engine_name == "surya":
+                    install_hint = "pip install 'natural-pdf[surya]'"
+                # Add other engines like docling if they become optional extras
+                else:
+                    install_hint = f"(Check installation requirements for {engine_name})"
+
                 raise RuntimeError(
-                    f"Layout engine '{engine_name}' is not available. Please check dependencies."
+                    f"Layout engine '{engine_name}' is not available. Please install the required dependencies: {install_hint}"
                 )
             self._detector_instances[engine_name] = detector_instance  # Store if available
 

natural_pdf/analyzers/layout/layout_options.py

@@ -80,6 +80,16 @@ class DoclingLayoutOptions(BaseLayoutOptions):
     # Other kwargs like 'device', 'batch_size' can go in extra_args
 
 
+# --- Gemini Specific Options ---
+@dataclass
+class GeminiLayoutOptions(BaseLayoutOptions):
+    """Options specific to Gemini-based layout detection (using OpenAI compatibility)."""
+
+    model_name: str = "gemini-2.0-flash"
+    # Removed: prompt_template, temperature, top_p, max_output_tokens
+    # These are typically passed directly to the chat completion call or via extra_args
+
+
 # --- Union Type ---
 LayoutOptions = Union[
     YOLOLayoutOptions,
@@ -87,5 +97,6 @@ LayoutOptions = Union[
     PaddleLayoutOptions,
     SuryaLayoutOptions,
     DoclingLayoutOptions,
+    GeminiLayoutOptions,
     BaseLayoutOptions,  # Include base for typing flexibility
 ]

natural_pdf/analyzers/layout/yolo.py

@@ -91,7 +91,9 @@ class YOLODocLayoutDetector(LayoutDetector):
     def _load_model_from_options(self, options: YOLOLayoutOptions) -> Any:
         """Load the YOLOv10 model based on options."""
         if not self.is_available():
-            raise RuntimeError("YOLO dependencies (doclayout_yolo, huggingface_hub) not installed.")
+            raise RuntimeError(
+                "YOLO dependencies not installed. Please run: pip install 'natural-pdf[layout_yolo]'"
+            )
         self.logger.info(f"Loading YOLO model: {options.model_repo}/{options.model_file}")
         try:
             model_path = hf_hub_download(repo_id=options.model_repo, filename=options.model_file)
@@ -105,7 +107,9 @@ class YOLODocLayoutDetector(LayoutDetector):
     def detect(self, image: Image.Image, options: BaseLayoutOptions) -> List[Dict[str, Any]]:
         """Detect layout elements in an image using YOLO."""
         if not self.is_available():
-            raise RuntimeError("YOLO dependencies (doclayout_yolo, huggingface_hub) not installed.")
+            raise RuntimeError(
+                "YOLO dependencies not installed. Please run: pip install 'natural-pdf[layout_yolo]'"
+            )
 
         # Ensure options are the correct type, falling back to defaults if base type passed
         if not isinstance(options, YOLOLayoutOptions):

natural_pdf/collections/pdf_collection.py

@@ -267,6 +267,27 @@ class PDFCollection(SearchableMixin):  # Inherit from the mixin
         # Implementation requires integrating with classification models or logic
         raise NotImplementedError("categorize requires classification implementation.")
 
+    def export_ocr_correction_task(self, output_zip_path: str, **kwargs):
+        """
+        Exports OCR results from all PDFs in this collection into a single
+        correction task package (zip file).
+
+        Args:
+            output_zip_path: The path to save the output zip file.
+            **kwargs: Additional arguments passed to create_correction_task_package
+                      (e.g., image_render_scale, overwrite).
+        """
+        try:
+            from natural_pdf.utils.packaging import create_correction_task_package
+            # Pass the collection itself (self) as the source
+            create_correction_task_package(source=self, output_zip_path=output_zip_path, **kwargs)
+        except ImportError:
+            logger.error("Failed to import 'create_correction_task_package'. Packaging utility might be missing.")
+            # Or raise
+        except Exception as e:
+            logger.error(f"Failed to export correction task for collection: {e}", exc_info=True)
+            raise # Re-raise the exception from the utility function
+
     # --- Mixin Required Implementation ---
     def get_indexable_items(self) -> Iterable[Indexable]:
         """Yields Page objects from the collection, conforming to Indexable."""

natural_pdf/core/element_manager.py

@@ -312,6 +312,7 @@ class ElementManager:
 
         Args:
             ocr_results: List of OCR results dictionaries with 'text', 'bbox', 'confidence'.
+                         Confidence can be None for detection-only results.
             scale_x: Factor to convert image x-coordinates to PDF coordinates.
             scale_y: Factor to convert image y-coordinates to PDF coordinates.
 
@@ -356,9 +357,14 @@ class ElementManager:
                 pdf_bottom = bottom_img * scale_y
                 pdf_height = (bottom_img - top_img) * scale_y
 
+                # 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
+
                 # Create the TextElement for the word
                 word_element_data = {
-                    "text": result["text"],
+                    "text": ocr_text,
                     "x0": pdf_x0,
                     "top": pdf_top,
                     "x1": pdf_x1,
@@ -367,7 +373,7 @@ class ElementManager:
                     "height": pdf_height,
                     "object_type": "word",  # Treat OCR results as whole words
                     "source": "ocr",
-                    "confidence": float(result.get("confidence", 0.0)),
+                    "confidence": confidence_value, # Use the handled confidence
                     "fontname": "OCR",  # Use consistent OCR fontname
                     "size": (
                         round(pdf_height) if pdf_height > 0 else 10.0
@@ -385,7 +391,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]
+                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)
@@ -393,16 +399,13 @@ class ElementManager:
                 # Append the word element to the manager's list
                 self._elements["words"].append(word_elem)
 
-                # Also create and append a representative character dictionary
-                # for consistency if someone iterates through manager.chars later.
-                # This char dict represents the entire OCR word as a single 'char'.
-                char_dict_data = ocr_char_dict  # Use the one we already created
-                char_dict_data["object_type"] = "char"  # Mark as char type
-                # pdfplumber char dicts don't typically have width/height/doctop,
-                # but keeping them won't hurt WordExtractor if it encounters them.
-                char_dict_data.setdefault("adv", char_dict_data.get("width", 0))
-
-                self._elements["chars"].append(char_dict_data)  # Append the dictionary
+                # Only add a representative char dict if text actually exists
+                if ocr_text is not None:
+                    # This char dict represents the entire OCR word as a single 'char'.
+                    char_dict_data = ocr_char_dict  # Use the one we already created
+                    char_dict_data["object_type"] = "char"  # Mark as char type
+                    char_dict_data.setdefault("adv", char_dict_data.get("width", 0))
+                    self._elements["chars"].append(char_dict_data)  # Append the dictionary
 
             except (KeyError, ValueError, TypeError) as e:
                 logger.error(f"Failed to process OCR result: {result}. Error: {e}", exc_info=True)