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

natural_pdf/ocr/engine_surya.py

@@ -1,15 +1,17 @@
 # ocr_engine_surya.py
-import logging
 import importlib.util
-from typing import Dict, List, Any, Optional, Tuple, Union
+import logging
+from typing import Any, Dict, List, Optional, Tuple, Union
+
 import numpy as np
 from PIL import Image
 
 from .engine import OCREngine
-from .ocr_options import SuryaOCROptions, BaseOCROptions
+from .ocr_options import BaseOCROptions, SuryaOCROptions
 
 logger = logging.getLogger(__name__)
 
+
 class SuryaOCREngine(OCREngine):
     """Surya OCR engine implementation."""
 
@@ -30,8 +32,9 @@ class SuryaOCREngine(OCREngine):
             raise ImportError("Surya OCR library is not installed or available.")
 
         try:
-            from surya.recognition import RecognitionPredictor
             from surya.detection import DetectionPredictor
+            from surya.recognition import RecognitionPredictor
+
             self._surya_recognition = RecognitionPredictor
             self._surya_detection = DetectionPredictor
             logger.info("Surya modules imported successfully.")
@@ -40,7 +43,7 @@ class SuryaOCREngine(OCREngine):
             # Add arguments from options if Surya supports them
             # Example: device = options.device or 'cuda' if torch.cuda.is_available() else 'cpu'
             # predictor_args = {'device': options.device} # If applicable
-            predictor_args = {} # Assuming parameterless init based on example
+            predictor_args = {}  # Assuming parameterless init based on example
 
             logger.info("Instantiating Surya DetectionPredictor...")
             self._detection_predictor = self._surya_detection(**predictor_args)
@@ -61,13 +64,17 @@ class SuryaOCREngine(OCREngine):
         """Check if the surya library is installed."""
         return importlib.util.find_spec("surya") is not None
 
-    def _standardize_results(self, raw_ocr_result: Any, options: SuryaOCROptions) -> List[Dict[str, Any]]:
+    def _standardize_results(
+        self, raw_ocr_result: Any, options: SuryaOCROptions
+    ) -> List[Dict[str, Any]]:
         """Standardizes raw results from a single image from Surya."""
         standardized_page = []
         min_confidence = options.min_confidence
 
         # Check if the result has the expected structure (OCRResult with text_lines)
-        if not hasattr(raw_ocr_result, 'text_lines') or not isinstance(raw_ocr_result.text_lines, list):
+        if not hasattr(raw_ocr_result, "text_lines") or not isinstance(
+            raw_ocr_result.text_lines, list
+        ):
             logger.warning(f"Unexpected Surya result format: {type(raw_ocr_result)}. Skipping.")
             return standardized_page
 
@@ -77,52 +84,54 @@ class SuryaOCREngine(OCREngine):
                 text = line.text
                 confidence = line.confidence
                 # Surya provides both polygon and bbox, bbox is already (x0, y0, x1, y1)
-                bbox_raw = line.bbox # Use bbox directly if available and correct format
+                bbox_raw = line.bbox  # Use bbox directly if available and correct format
 
                 if confidence >= min_confidence:
-                    bbox = self._standardize_bbox(bbox_raw) # Validate/convert format
+                    bbox = self._standardize_bbox(bbox_raw)  # Validate/convert format
                     if bbox:
-                        standardized_page.append({
-                            'bbox': bbox,
-                            'text': text,
-                            'confidence': confidence,
-                            'source': 'ocr'
-                        })
+                        standardized_page.append(
+                            {"bbox": bbox, "text": text, "confidence": confidence, "source": "ocr"}
+                        )
                     else:
                         # Try polygon if bbox failed standardization
                         bbox_poly = self._standardize_bbox(line.polygon)
                         if bbox_poly:
-                             standardized_page.append({
-                                'bbox': bbox_poly, 'text': text, 'confidence': confidence, 'source': 'ocr'
-                             })
+                            standardized_page.append(
+                                {
+                                    "bbox": bbox_poly,
+                                    "text": text,
+                                    "confidence": confidence,
+                                    "source": "ocr",
+                                }
+                            )
                         else:
-                             logger.warning(f"Skipping Surya line due to invalid bbox/polygon: {line}")
+                            logger.warning(
+                                f"Skipping Surya line due to invalid bbox/polygon: {line}"
+                            )
 
             except (AttributeError, ValueError, TypeError) as e:
-                 logger.warning(f"Skipping invalid Surya TextLine format: {line}. Error: {e}")
-                 continue
+                logger.warning(f"Skipping invalid Surya TextLine format: {line}. Error: {e}")
+                continue
         return standardized_page
 
     def process_image(
-        self,
-        images: Union[Image.Image, List[Image.Image]],
-        options: BaseOCROptions
+        self, images: Union[Image.Image, List[Image.Image]], options: BaseOCROptions
     ) -> Union[List[Dict[str, Any]], List[List[Dict[str, Any]]]]:
         """Processes a single image or a batch of images with Surya OCR."""
 
         if not isinstance(options, SuryaOCROptions):
-             logger.warning("Received BaseOCROptions, expected SuryaOCROptions. Using defaults.")
-             options = SuryaOCROptions(
-                 languages=options.languages,
-                 min_confidence=options.min_confidence,
-                 device=options.device,
-                 extra_args=options.extra_args
-             )
+            logger.warning("Received BaseOCROptions, expected SuryaOCROptions. Using defaults.")
+            options = SuryaOCROptions(
+                languages=options.languages,
+                min_confidence=options.min_confidence,
+                device=options.device,
+                extra_args=options.extra_args,
+            )
 
         # Ensure predictors are loaded/initialized
         self._lazy_load_predictors(options)
         if not self._recognition_predictor or not self._detection_predictor:
-             raise RuntimeError("Surya predictors could not be initialized.")
+            raise RuntimeError("Surya predictors could not be initialized.")
 
         # --- Prepare inputs for Surya ---
         is_batch = isinstance(images, list)
@@ -131,8 +140,8 @@ class SuryaOCREngine(OCREngine):
         input_langs: List[List[str]] = [options.languages for _ in input_images]
 
         if not input_images:
-             logger.warning("No images provided for Surya processing.")
-             return [] if not is_batch else [[]]
+            logger.warning("No images provided for Surya processing.")
+            return [] if not is_batch else [[]]
 
         # --- Run Surya Prediction ---
         try:
@@ -141,24 +150,26 @@ class SuryaOCREngine(OCREngine):
             # Call Surya's predictor
             # It returns a list of OCRResult objects, one per input image
             predictions = self._recognition_predictor(
-                images=input_images,
-                langs=input_langs,
-                det_predictor=self._detection_predictor
+                images=input_images, langs=input_langs, det_predictor=self._detection_predictor
             )
             logger.info(f"Surya prediction complete. Received {len(predictions)} results.")
 
             # --- Standardize Results ---
             if len(predictions) != len(input_images):
-                 logger.error(f"Surya result count ({len(predictions)}) does not match input count ({len(input_images)}). Returning empty results.")
-                 # Decide on error handling: raise error or return empty structure
-                 return [[] for _ in input_images] if is_batch else []
+                logger.error(
+                    f"Surya result count ({len(predictions)}) does not match input count ({len(input_images)}). Returning empty results."
+                )
+                # Decide on error handling: raise error or return empty structure
+                return [[] for _ in input_images] if is_batch else []
 
-            all_standardized_results = [self._standardize_results(res, options) for res in predictions]
+            all_standardized_results = [
+                self._standardize_results(res, options) for res in predictions
+            ]
 
             if is_batch:
-                return all_standardized_results # Return List[List[Dict]]
+                return all_standardized_results  # Return List[List[Dict]]
             else:
-                return all_standardized_results[0] # Return List[Dict] for single image
+                return all_standardized_results[0]  # Return List[Dict] for single image
 
         except Exception as e:
             logger.error(f"Error during Surya OCR processing: {e}", exc_info=True)
@@ -168,4 +179,3 @@ class SuryaOCREngine(OCREngine):
     # Note: Caching is handled differently for Surya as predictors are stateful
     # and initialized once. The base class _reader_cache is not used here.
     # If predictors could be configured per-run, caching would need rethinking.
-

natural_pdf/ocr/ocr_manager.py

@@ -1,68 +1,76 @@
 # ocr_manager.py
+import copy  # For deep copying options
 import logging
-from typing import Dict, List, Any, Optional, Union, Type
+from typing import Any, Dict, List, Optional, Type, Union
+
 from PIL import Image
-import copy # For deep copying options
 
 # Import engine classes and options
 from .engine import OCREngine
 from .engine_easyocr import EasyOCREngine
 from .engine_paddle import PaddleOCREngine
-from .engine_surya import SuryaOCREngine # <-- Import Surya Engine
-from .ocr_options import (
-    BaseOCROptions, EasyOCROptions, PaddleOCROptions, SuryaOCROptions, OCROptions # <-- Import Surya Options
-)
+from .engine_surya import SuryaOCREngine  # <-- Import Surya Engine
+from .ocr_options import OCROptions  # <-- Import Surya Options
+from .ocr_options import BaseOCROptions, EasyOCROptions, PaddleOCROptions, SuryaOCROptions
 
 logger = logging.getLogger(__name__)
 
+
 class OCRManager:
     """Manages OCR engine selection, configuration, and execution."""
 
     # Registry mapping engine names to classes and default options
     ENGINE_REGISTRY: Dict[str, Dict[str, Any]] = {
-        'easyocr': {'class': EasyOCREngine, 'options_class': EasyOCROptions},
-        'paddle': {'class': PaddleOCREngine, 'options_class': PaddleOCROptions},
-        'surya': {'class': SuryaOCREngine, 'options_class': SuryaOCROptions}, # <-- Add Surya
+        "easyocr": {"class": EasyOCREngine, "options_class": EasyOCROptions},
+        "paddle": {"class": PaddleOCREngine, "options_class": PaddleOCROptions},
+        "surya": {"class": SuryaOCREngine, "options_class": SuryaOCROptions},  # <-- Add Surya
         # Add other engines here
     }
 
     # Define the limited set of kwargs allowed for the simple apply_ocr call
     SIMPLE_MODE_ALLOWED_KWARGS = {
-        'engine', 'languages', 'min_confidence', 'device'
+        "engine",
+        "languages",
+        "min_confidence",
+        "device",
         # Add image pre-processing args like 'resolution', 'width' if handled here
     }
 
     def __init__(self):
         """Initializes the OCR Manager."""
-        self._engine_instances: Dict[str, OCREngine] = {} # Cache for engine instances
+        self._engine_instances: Dict[str, OCREngine] = {}  # Cache for engine instances
         logger.info("OCRManager initialized.")
 
     def _get_engine_instance(self, engine_name: str) -> OCREngine:
         """Retrieves or creates an instance of the specified OCR engine."""
         engine_name = engine_name.lower()
         if engine_name not in self.ENGINE_REGISTRY:
-            raise ValueError(f"Unknown OCR engine: '{engine_name}'. Available: {list(self.ENGINE_REGISTRY.keys())}")
+            raise ValueError(
+                f"Unknown OCR engine: '{engine_name}'. Available: {list(self.ENGINE_REGISTRY.keys())}"
+            )
 
         # Surya engine might manage its own predictor state, consider if caching instance is always right
         # For now, we cache the engine instance itself.
         if engine_name not in self._engine_instances:
             logger.info(f"Creating instance of engine: {engine_name}")
-            engine_class = self.ENGINE_REGISTRY[engine_name]['class']
-            engine_instance = engine_class() # Instantiate first
+            engine_class = self.ENGINE_REGISTRY[engine_name]["class"]
+            engine_instance = engine_class()  # Instantiate first
             if not engine_instance.is_available():
-                 # Check availability before storing
-                 raise RuntimeError(f"Engine '{engine_name}' is not available. Please check dependencies.")
-            self._engine_instances[engine_name] = engine_instance # Store if available
+                # Check availability before storing
+                raise RuntimeError(
+                    f"Engine '{engine_name}' is not available. Please check dependencies."
+                )
+            self._engine_instances[engine_name] = engine_instance  # Store if available
 
         return self._engine_instances[engine_name]
 
     def apply_ocr(
         self,
-        images: Union[Image.Image, List[Image.Image]], # Accept single or list
-        engine: Optional[str] = 'easyocr', # Default engine
+        images: Union[Image.Image, List[Image.Image]],  # Accept single or list
+        engine: Optional[str] = "easyocr",  # Default engine
         options: Optional[OCROptions] = None,
-        **kwargs
-    ) -> Union[List[Dict[str, Any]], List[List[Dict[str, Any]]]]: # Return single or list of lists
+        **kwargs,
+    ) -> Union[List[Dict[str, Any]], List[List[Dict[str, Any]]]]:  # Return single or list of lists
         """
         Applies OCR to a single image or a batch of images using either simple
         keyword arguments or an options object.
@@ -94,54 +102,62 @@ class OCRManager:
         # --- Validate input type ---
         is_batch = isinstance(images, list)
         if not is_batch and not isinstance(images, Image.Image):
-             raise TypeError("Input 'images' must be a PIL Image or a list of PIL Images.")
+            raise TypeError("Input 'images' must be a PIL Image or a list of PIL Images.")
         # Allow engines to handle non-PIL images in list if they support it/log warnings
         # if is_batch and not all(isinstance(img, Image.Image) for img in images):
         #     logger.warning("Batch may contain items that are not PIL Images.")
 
-
         # --- Determine Options and Engine ---
         if options is not None:
             # Advanced Mode
             logger.debug(f"Using advanced mode with options object: {type(options).__name__}")
-            final_options = copy.deepcopy(options) # Prevent modification of original
+            final_options = copy.deepcopy(options)  # Prevent modification of original
             found_engine = False
             for name, registry_entry in self.ENGINE_REGISTRY.items():
                 # Check if options object is an instance of the registered options class
-                if isinstance(options, registry_entry['options_class']):
+                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 engine options.")
+                raise TypeError(
+                    f"Provided options object type '{type(options).__name__}' does not match any registered engine options."
+                )
             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 'options' and will be ignored."
+                )
         else:
             # Simple Mode
-            selected_engine_name = engine.lower() if engine else 'easyocr' # Fallback default
-            logger.debug(f"Using simple mode with engine: '{selected_engine_name}' and kwargs: {kwargs}")
+            selected_engine_name = engine.lower() if engine else "easyocr"  # Fallback default
+            logger.debug(
+                f"Using simple mode with engine: '{selected_engine_name}' and kwargs: {kwargs}"
+            )
 
             if selected_engine_name not in self.ENGINE_REGISTRY:
-                 raise ValueError(f"Unknown OCR engine: '{selected_engine_name}'. Available: {list(self.ENGINE_REGISTRY.keys())}")
+                raise ValueError(
+                    f"Unknown OCR engine: '{selected_engine_name}'. Available: {list(self.ENGINE_REGISTRY.keys())}"
+                )
 
             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.")
+                raise TypeError(
+                    f"Got unexpected keyword arguments in simple mode: {list(unexpected_kwargs)}. Use the 'options' parameter for detailed configuration."
+                )
 
             # Get the *correct* options class for the selected engine
-            options_class = self.ENGINE_REGISTRY[selected_engine_name]['options_class']
+            options_class = self.ENGINE_REGISTRY[selected_engine_name]["options_class"]
 
             # Create options instance using provided simple kwargs or defaults
             simple_args = {
-                'languages': kwargs.get('languages', ['en']),
-                'min_confidence': kwargs.get('min_confidence', 0.5),
-                'device': kwargs.get('device', 'cpu')
+                "languages": kwargs.get("languages", ["en"]),
+                "min_confidence": kwargs.get("min_confidence", 0.5),
+                "device": kwargs.get("device", "cpu"),
                 # Note: 'extra_args' isn't populated in simple mode
             }
             final_options = options_class(**simple_args)
             logger.debug(f"Constructed options for simple mode: {final_options}")
 
-
         # --- Get Engine Instance and Process ---
         try:
             engine_instance = self._get_engine_instance(selected_engine_name)
@@ -153,39 +169,49 @@ class OCRManager:
 
             # Log result summary based on mode
             if is_batch:
-                 # Ensure results is a list before trying to get lengths
-                 if isinstance(results, list):
-                     num_results_per_image = [len(res_list) if isinstance(res_list, list) else -1 for res_list in results] # Handle potential errors returning non-lists
-                     logger.info(f"Processing complete. Found results per image: {num_results_per_image}")
-                 else:
-                     logger.error(f"Processing complete but received unexpected result type for batch: {type(results)}")
+                # Ensure results is a list before trying to get lengths
+                if isinstance(results, list):
+                    num_results_per_image = [
+                        len(res_list) if isinstance(res_list, list) else -1 for res_list in results
+                    ]  # Handle potential errors returning non-lists
+                    logger.info(
+                        f"Processing complete. Found results per image: {num_results_per_image}"
+                    )
+                else:
+                    logger.error(
+                        f"Processing complete but received unexpected result type for batch: {type(results)}"
+                    )
             else:
-                 # Ensure results is a list
-                 if isinstance(results, list):
-                      logger.info(f"Processing complete. Found {len(results)} results.")
-                 else:
-                      logger.error(f"Processing complete but received unexpected result type for single image: {type(results)}")
-            return results # Return type matches input type due to engine logic
+                # Ensure results is a list
+                if isinstance(results, list):
+                    logger.info(f"Processing complete. Found {len(results)} results.")
+                else:
+                    logger.error(
+                        f"Processing complete but received unexpected result type for single image: {type(results)}"
+                    )
+            return results  # Return type matches input type due to engine logic
 
         except (ImportError, RuntimeError, ValueError, TypeError) as e:
-             logger.error(f"OCR processing failed for engine '{selected_engine_name}': {e}", exc_info=True)
-             raise # Re-raise expected errors
+            logger.error(
+                f"OCR processing failed for engine '{selected_engine_name}': {e}", exc_info=True
+            )
+            raise  # Re-raise expected errors
         except Exception as e:
-             logger.error(f"An unexpected error occurred during OCR processing: {e}", exc_info=True)
-             raise # Re-raise unexpected errors
-
+            logger.error(f"An unexpected error occurred during OCR processing: {e}", exc_info=True)
+            raise  # Re-raise unexpected errors
 
     def get_available_engines(self) -> List[str]:
         """Returns a list of registered engine names that are currently available."""
         available = []
         for name, registry_entry in self.ENGINE_REGISTRY.items():
-             try:
-                 # Temporarily instantiate to check availability without caching
-                 engine_class = registry_entry['class']
-                 if engine_class().is_available():
-                     available.append(name)
-             except Exception as e:
-                 logger.debug(f"Engine '{name}' check failed: {e}") # Log check failures at debug level
-                 pass # Ignore engines that fail to instantiate or check
+            try:
+                # Temporarily instantiate to check availability without caching
+                engine_class = registry_entry["class"]
+                if engine_class().is_available():
+                    available.append(name)
+            except Exception as e:
+                logger.debug(
+                    f"Engine '{name}' check failed: {e}"
+                )  # Log check failures at debug level
+                pass  # Ignore engines that fail to instantiate or check
         return available
-

natural_pdf/ocr/ocr_options.py

@@ -1,30 +1,34 @@
 # ocr_options.py
 import logging
 from dataclasses import dataclass, field
-from typing import List, Optional, Dict, Any, Tuple, Union
+from typing import Any, Dict, List, Optional, Tuple, Union
 
 # Configure logging
 # logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
 # logger = logging.getLogger(__name__)
 # Assume logger is configured elsewhere or remove if not needed globally
 
+
 # --- Base Options ---
 @dataclass
 class BaseOCROptions:
     """Base class for OCR engine options."""
-    languages: List[str] = field(default_factory=lambda: ['en'])
+
+    languages: List[str] = field(default_factory=lambda: ["en"])
     min_confidence: float = 0.5
-    device: Optional[str] = 'cpu' # Suggestion, actual device usage depends on engine impl.
+    device: Optional[str] = "cpu"  # Suggestion, actual device usage depends on engine impl.
     extra_args: Dict[str, Any] = field(default_factory=dict)
 
+
 # --- EasyOCR Specific Options ---
 @dataclass
 class EasyOCROptions(BaseOCROptions):
     """Specific options for the EasyOCR engine."""
+
     model_storage_directory: Optional[str] = None
     user_network_directory: Optional[str] = None
-    recog_network: str = 'english_g2'
-    detect_network: str = 'craft'
+    recog_network: str = "english_g2"
+    detect_network: str = "craft"
     download_enabled: bool = True
     detector: bool = True
     recognizer: bool = True
@@ -32,7 +36,7 @@ class EasyOCROptions(BaseOCROptions):
     quantize: bool = True
     cudnn_benchmark: bool = False
     detail: int = 1
-    decoder: str = 'greedy'
+    decoder: str = "greedy"
     beamWidth: int = 5
     batch_size: int = 1
     workers: int = 0
@@ -55,7 +59,7 @@ class EasyOCROptions(BaseOCROptions):
     y_ths: float = 0.5
     x_ths: float = 1.0
     add_margin: float = 0.1
-    output_format: str = 'standard'
+    output_format: str = "standard"
 
     # def __post_init__(self):
     #     logger.debug(f"Initialized EasyOCROptions: {self}")
@@ -65,13 +69,14 @@ class EasyOCROptions(BaseOCROptions):
 @dataclass
 class PaddleOCROptions(BaseOCROptions):
     """Specific options for the PaddleOCR engine."""
+
     use_angle_cls: bool = True
     use_gpu: Optional[bool] = None
     gpu_mem: int = 500
     ir_optim: bool = True
     use_tensorrt: bool = False
     min_subgraph_size: int = 15
-    precision: str = 'fp32'
+    precision: str = "fp32"
     enable_mkldnn: bool = False
     cpu_threads: int = 10
     use_fp16: bool = False
@@ -91,16 +96,18 @@ class PaddleOCROptions(BaseOCROptions):
 
     def __post_init__(self):
         if self.use_gpu is None:
-            if self.device and 'cuda' in self.device.lower():
+            if self.device and "cuda" in self.device.lower():
                 self.use_gpu = True
             else:
                 self.use_gpu = False
         # logger.debug(f"Initialized PaddleOCROptions: {self}")
 
+
 # --- Surya Specific Options ---
 @dataclass
 class SuryaOCROptions(BaseOCROptions):
     """Specific options for the Surya OCR engine."""
+
     # Currently, Surya example shows languages passed at prediction time.
     # Add fields here if Surya's RecognitionPredictor or DetectionPredictor
     # constructors accept relevant arguments (e.g., model paths, device settings).
@@ -111,4 +118,3 @@ class SuryaOCROptions(BaseOCROptions):
 
 # --- Union type for type hinting ---
 OCROptions = Union[EasyOCROptions, PaddleOCROptions, SuryaOCROptions, BaseOCROptions]
-

natural_pdf/qa/__init__.py

@@ -1,3 +1,3 @@
 from natural_pdf.qa.document_qa import DocumentQA, get_qa_engine
 
-__all__ = ["DocumentQA", "get_qa_engine"]
+__all__ = ["DocumentQA", "get_qa_engine"]