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

natural_pdf/analyzers/layout/docling.py

@@ -1,24 +1,38 @@
 # layout_detector_docling.py
-import logging
 import importlib.util
+import logging
 import os
 import tempfile
-from typing import List, Dict, Any, Optional
+from typing import Any, Dict, List, Optional
+
 from PIL import Image
 
 # Assuming base class and options are importable
 try:
     from .base import LayoutDetector
-    from .layout_options import DoclingLayoutOptions, BaseLayoutOptions
+    from .layout_options import BaseLayoutOptions, DoclingLayoutOptions
 except ImportError:
     # Placeholders if run standalone or imports fail
-    class BaseLayoutOptions: pass
-    class DoclingLayoutOptions(BaseLayoutOptions): pass
+    class BaseLayoutOptions:
+        pass
+
+    class DoclingLayoutOptions(BaseLayoutOptions):
+        pass
+
     class LayoutDetector:
-         def __init__(self): self.logger=logging.getLogger(); self.supported_classes=set()
-         def _get_model(self, options): raise NotImplementedError
-         def _normalize_class_name(self, n): return n
-         def validate_classes(self, c): pass
+        def __init__(self):
+            self.logger = logging.getLogger()
+            self.supported_classes = set()
+
+        def _get_model(self, options):
+            raise NotImplementedError
+
+        def _normalize_class_name(self, n):
+            return n
+
+        def validate_classes(self, c):
+            pass
+
     logging.basicConfig()
 
 logger = logging.getLogger(__name__)
@@ -42,11 +56,27 @@ class DoclingLayoutDetector(LayoutDetector):
         super().__init__()
         # Docling classes are dynamic/hierarchical, define common ones
         self.supported_classes = {
-            'Header', 'Footer', 'Paragraph', 'Heading', 'List', 'ListItem',
-            'Table', 'Figure', 'Caption', 'Footnote', 'PageNumber', 'Equation',
-            'Code', 'Title', 'Author', 'Abstract', 'Section', 'Unknown', 'Metadata' # Add more as needed
+            "Header",
+            "Footer",
+            "Paragraph",
+            "Heading",
+            "List",
+            "ListItem",
+            "Table",
+            "Figure",
+            "Caption",
+            "Footnote",
+            "PageNumber",
+            "Equation",
+            "Code",
+            "Title",
+            "Author",
+            "Abstract",
+            "Section",
+            "Unknown",
+            "Metadata",  # Add more as needed
         }
-        self._docling_document_cache = {} # Cache the output doc per image/options if needed
+        self._docling_document_cache = {}  # Cache the output doc per image/options if needed
 
     def is_available(self) -> bool:
         """Check if docling is installed."""
@@ -55,9 +85,9 @@ class DoclingLayoutDetector(LayoutDetector):
     def _get_cache_key(self, options: BaseLayoutOptions) -> str:
         """Generate cache key based on device and potentially converter args."""
         if not isinstance(options, DoclingLayoutOptions):
-             options = DoclingLayoutOptions(device=options.device, extra_args=options.extra_args)
+            options = DoclingLayoutOptions(device=options.device, extra_args=options.extra_args)
 
-        device_key = str(options.device).lower() if options.device else 'default_device'
+        device_key = str(options.device).lower() if options.device else "default_device"
         # Include hash of extra_args if they affect model loading/converter init
         extra_args_key = hash(frozenset(options.extra_args.items()))
         return f"{self.__class__.__name__}_{device_key}_{extra_args_key}"
@@ -88,12 +118,17 @@ class DoclingLayoutDetector(LayoutDetector):
             raise RuntimeError("Docling dependency not installed.")
 
         if not isinstance(options, DoclingLayoutOptions):
-             self.logger.warning("Received BaseLayoutOptions, expected DoclingLayoutOptions. Using defaults.")
-             options = DoclingLayoutOptions(
-                 confidence=options.confidence, classes=options.classes,
-                 exclude_classes=options.exclude_classes, device=options.device,
-                 extra_args=options.extra_args, verbose=options.extra_args.get('verbose', False)
-             )
+            self.logger.warning(
+                "Received BaseLayoutOptions, expected DoclingLayoutOptions. Using defaults."
+            )
+            options = DoclingLayoutOptions(
+                confidence=options.confidence,
+                classes=options.classes,
+                exclude_classes=options.exclude_classes,
+                device=options.device,
+                extra_args=options.extra_args,
+                verbose=options.extra_args.get("verbose", False),
+            )
 
         # Validate classes before proceeding (note: Docling classes are case-sensitive)
         # self.validate_classes(options.classes or []) # Validation might be tricky due to case sensitivity
@@ -105,18 +140,20 @@ class DoclingLayoutDetector(LayoutDetector):
 
         # Docling convert method requires an image path. Save temp file.
         detections = []
-        docling_doc = None # To store the result
+        docling_doc = None  # To store the result
         with tempfile.TemporaryDirectory() as temp_dir:
             temp_image_path = os.path.join(temp_dir, f"docling_input_{os.getpid()}.png")
             try:
-                self.logger.debug(f"Saving temporary image for Docling detector to: {temp_image_path}")
-                image.convert("RGB").save(temp_image_path) # Ensure RGB
+                self.logger.debug(
+                    f"Saving temporary image for Docling detector to: {temp_image_path}"
+                )
+                image.convert("RGB").save(temp_image_path)  # Ensure RGB
 
                 # Convert the document using Docling's DocumentConverter
                 self.logger.debug("Running Docling conversion...")
                 # Docling convert returns a Result object with a 'document' attribute
                 result = converter.convert(temp_image_path)
-                docling_doc = result.document # Store the DoclingDocument
+                docling_doc = result.document  # Store the DoclingDocument
                 self.logger.info(f"Docling conversion complete.")
 
                 # Convert Docling document to our detection format
@@ -124,12 +161,14 @@ class DoclingLayoutDetector(LayoutDetector):
 
             except Exception as e:
                 self.logger.error(f"Error during Docling detection: {e}", exc_info=True)
-                raise # Re-raise the exception
+                raise  # Re-raise the exception
             finally:
-                 # Ensure temp file is removed
-                 if os.path.exists(temp_image_path):
-                      try: os.remove(temp_image_path)
-                      except OSError as e_rm: self.logger.warning(f"Could not remove temp file {temp_image_path}: {e_rm}")
+                # Ensure temp file is removed
+                if os.path.exists(temp_image_path):
+                    try:
+                        os.remove(temp_image_path)
+                    except OSError as e_rm:
+                        self.logger.warning(f"Could not remove temp file {temp_image_path}: {e_rm}")
 
         # Cache the docling document if needed elsewhere (maybe associate with page?)
         # self._docling_document_cache[image_hash] = docling_doc # Needs a way to key this
@@ -137,26 +176,37 @@ class DoclingLayoutDetector(LayoutDetector):
         self.logger.info(f"Docling detected {len(detections)} layout elements matching criteria.")
         return detections
 
-    def _convert_docling_to_detections(self, doc, options: DoclingLayoutOptions) -> List[Dict[str, Any]]:
+    def _convert_docling_to_detections(
+        self, doc, options: DoclingLayoutOptions
+    ) -> List[Dict[str, Any]]:
         """Convert a Docling document to our standard detection format."""
-        if not doc or not hasattr(doc, 'pages') or not doc.pages:
+        if not doc or not hasattr(doc, "pages") or not doc.pages:
             self.logger.warning("Invalid or empty Docling document for conversion.")
             return []
 
         detections = []
-        id_to_detection_index = {} # Map Docling ID to index in detections list
+        id_to_detection_index = {}  # Map Docling ID to index in detections list
 
         # Prepare normalized class filters once
-        normalized_classes_req = {self._normalize_class_name(c) for c in options.classes} if options.classes else None
-        normalized_classes_excl = {self._normalize_class_name(c) for c in options.exclude_classes} if options.exclude_classes else set()
+        normalized_classes_req = (
+            {self._normalize_class_name(c) for c in options.classes} if options.classes else None
+        )
+        normalized_classes_excl = (
+            {self._normalize_class_name(c) for c in options.exclude_classes}
+            if options.exclude_classes
+            else set()
+        )
 
         # --- Iterate through elements using Docling's structure ---
         # This requires traversing the hierarchy (e.g., doc.body.children)
         # or iterating through specific lists like doc.texts, doc.tables etc.
         elements_to_process = []
-        if hasattr(doc, 'texts'): elements_to_process.extend(doc.texts)
-        if hasattr(doc, 'tables'): elements_to_process.extend(doc.tables)
-        if hasattr(doc, 'pictures'): elements_to_process.extend(doc.pictures)
+        if hasattr(doc, "texts"):
+            elements_to_process.extend(doc.texts)
+        if hasattr(doc, "tables"):
+            elements_to_process.extend(doc.tables)
+        if hasattr(doc, "pictures"):
+            elements_to_process.extend(doc.pictures)
         # Add other element types from DoclingDocument as needed
 
         self.logger.debug(f"Converting {len(elements_to_process)} Docling elements...")
@@ -164,16 +214,19 @@ class DoclingLayoutDetector(LayoutDetector):
         for elem in elements_to_process:
             try:
                 # Get Provenance (bbox and page number)
-                if not hasattr(elem, 'prov') or not elem.prov: continue
-                prov = elem.prov[0] # Use first provenance
-                if not hasattr(prov, 'bbox') or not prov.bbox: continue
+                if not hasattr(elem, "prov") or not elem.prov:
+                    continue
+                prov = elem.prov[0]  # Use first provenance
+                if not hasattr(prov, "bbox") or not prov.bbox:
+                    continue
                 bbox = prov.bbox
                 page_no = prov.page_no
 
                 # Get Page Dimensions (crucial for coordinate conversion)
-                if not hasattr(doc.pages.get(page_no), 'size'): continue
+                if not hasattr(doc.pages.get(page_no), "size"):
+                    continue
                 page_height = doc.pages[page_no].size.height
-                page_width = doc.pages[page_no].size.width # Needed? Bbox seems absolute
+                page_width = doc.pages[page_no].size.width  # Needed? Bbox seems absolute
 
                 # Convert coordinates from Docling's system (often bottom-left origin)
                 # to standard top-left origin (0,0 at top-left)
@@ -182,46 +235,51 @@ class DoclingLayoutDetector(LayoutDetector):
                 x1 = float(bbox.r)
                 # Convert y: top_y = page_height - bottom_left_t
                 #            bottom_y = page_height - bottom_left_b
-                y0 = float(page_height - bbox.t) # Top y
-                y1 = float(page_height - bbox.b) # Bottom y
+                y0 = float(page_height - bbox.t)  # Top y
+                y1 = float(page_height - bbox.b)  # Bottom y
 
                 # Ensure y0 < y1
-                if y0 > y1: y0, y1 = y1, y0
+                if y0 > y1:
+                    y0, y1 = y1, y0
                 # Ensure x0 < x1
-                if x0 > x1: x0, x1 = x1, x0
+                if x0 > x1:
+                    x0, x1 = x1, x0
 
                 # Get Class Label
-                label_orig = str(getattr(elem, 'label', 'Unknown')) # Default if no label
+                label_orig = str(getattr(elem, "label", "Unknown"))  # Default if no label
                 normalized_label = self._normalize_class_name(label_orig)
 
                 # Apply Class Filtering
-                if normalized_classes_req and normalized_label not in normalized_classes_req: continue
-                if normalized_label in normalized_classes_excl: continue
+                if normalized_classes_req and normalized_label not in normalized_classes_req:
+                    continue
+                if normalized_label in normalized_classes_excl:
+                    continue
 
                 # Get Confidence (Docling often doesn't provide per-element confidence)
-                confidence = getattr(elem, 'confidence', 0.95) # Assign default confidence
-                if confidence < options.confidence: continue # Apply confidence threshold
+                confidence = getattr(elem, "confidence", 0.95)  # Assign default confidence
+                if confidence < options.confidence:
+                    continue  # Apply confidence threshold
 
                 # Get Text Content
-                text_content = getattr(elem, 'text', None)
+                text_content = getattr(elem, "text", None)
 
                 # Get IDs for hierarchy
-                docling_id = getattr(elem, 'self_ref', None)
-                parent_id_obj = getattr(elem, 'parent', None)
-                parent_id = getattr(parent_id_obj, 'self_ref', None) if parent_id_obj else None
+                docling_id = getattr(elem, "self_ref", None)
+                parent_id_obj = getattr(elem, "parent", None)
+                parent_id = getattr(parent_id_obj, "self_ref", None) if parent_id_obj else None
 
                 # Create Detection Dictionary
                 detection = {
-                    'bbox': (x0, y0, x1, y1),
-                    'class': label_orig,
-                    'normalized_class': normalized_label,
-                    'confidence': confidence,
-                    'text': text_content,
-                    'docling_id': docling_id,
-                    'parent_id': parent_id,
-                    'page_number': page_no, # Add page number if useful
-                    'source': 'layout',
-                    'model': 'docling'
+                    "bbox": (x0, y0, x1, y1),
+                    "class": label_orig,
+                    "normalized_class": normalized_label,
+                    "confidence": confidence,
+                    "text": text_content,
+                    "docling_id": docling_id,
+                    "parent_id": parent_id,
+                    "page_number": page_no,  # Add page number if useful
+                    "source": "layout",
+                    "model": "docling",
                 }
                 detections.append(detection)
 
@@ -229,8 +287,8 @@ class DoclingLayoutDetector(LayoutDetector):
                 # if docling_id: id_to_detection_index[docling_id] = len(detections) - 1
 
             except Exception as conv_e:
-                 self.logger.warning(f"Could not convert Docling element: {elem}. Error: {conv_e}")
-                 continue
+                self.logger.warning(f"Could not convert Docling element: {elem}. Error: {conv_e}")
+                continue
 
         return detections
 
@@ -241,7 +299,8 @@ class DoclingLayoutDetector(LayoutDetector):
         """
         # This requires caching the doc based on image/options or re-running.
         # For simplicity, let's just re-run detect if needed.
-        self.logger.warning("get_docling_document: Re-running detection to ensure document is generated.")
-        self.detect(image, options) # Run detect to populate internal doc
-        return getattr(self, '_docling_document', None) # Return the stored doc
-
+        self.logger.warning(
+            "get_docling_document: Re-running detection to ensure document is generated."
+        )
+        self.detect(image, options)  # Run detect to populate internal doc
+        return getattr(self, "_docling_document", None)  # Return the stored doc

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
+