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

natural_pdf/ocr/ocr_factory.py

@@ -0,0 +1,114 @@
+import logging
+import importlib.util
+from typing import Dict, Any, Optional, Type, Union, List
+
+from .engine import OCREngine
+
+logger = logging.getLogger(__name__)
+
+
+class OCRFactory:
+    """Factory for creating and managing OCR engines with optional dependencies."""
+
+    @staticmethod
+    def create_engine(engine_type: str, **kwargs) -> OCREngine:
+        """Create and return an OCR engine instance.
+        
+        Args:
+            engine_type: One of 'surya', 'easyocr', 'paddle'
+            **kwargs: Arguments to pass to the engine constructor
+            
+        Returns:
+            An initialized OCR engine
+            
+        Raises:
+            ImportError: If the required dependencies aren't installed
+            ValueError: If the engine_type is unknown
+        """
+        if engine_type == "surya":
+            try:
+                from .engine_surya import SuryaOCREngine
+                return SuryaOCREngine(**kwargs)
+            except ImportError:
+                raise ImportError("Surya engine requires the 'surya' package. "
+                                 "Install with: pip install surya")
+        elif engine_type == "easyocr":
+            try:
+                from .engine_easyocr import EasyOCREngine
+                return EasyOCREngine(**kwargs)
+            except ImportError:
+                raise ImportError("EasyOCR engine requires the 'easyocr' package. "
+                                 "Install with: pip install easyocr")
+        elif engine_type == "paddle":
+            try:
+                from .engine_paddle import PaddleOCREngine
+                return PaddleOCREngine(**kwargs)
+            except ImportError:
+                raise ImportError("PaddleOCR engine requires 'paddleocr' and 'paddlepaddle'. "
+                                 "Install with: pip install paddleocr paddlepaddle")
+        else:
+            raise ValueError(f"Unknown engine type: {engine_type}")
+    
+    @staticmethod
+    def list_available_engines() -> Dict[str, bool]:
+        """Returns a dictionary of engine names and their availability status."""
+        engines = {}
+        
+        # Check Surya
+        try:
+            engines["surya"] = importlib.util.find_spec("surya") is not None
+        except ImportError:
+            engines["surya"] = False
+            
+        # Check EasyOCR
+        try:
+            engines["easyocr"] = importlib.util.find_spec("easyocr") is not None
+        except ImportError:
+            engines["easyocr"] = False
+            
+        # Check PaddleOCR
+        try:
+            paddle = importlib.util.find_spec("paddle") is not None or importlib.util.find_spec("paddlepaddle") is not None
+            paddleocr = importlib.util.find_spec("paddleocr") is not None
+            engines["paddle"] = paddle and paddleocr
+        except ImportError:
+            engines["paddle"] = False
+            
+        return engines
+    
+    @staticmethod
+    def get_recommended_engine(**kwargs) -> OCREngine:
+        """Returns the best available OCR engine based on what's installed.
+        
+        First tries engines in order of preference: EasyOCR, Paddle, Surya.
+        If none are available, raises ImportError with installation instructions.
+        
+        Args:
+            **kwargs: Arguments to pass to the engine constructor
+            
+        Returns:
+            The best available OCR engine instance
+            
+        Raises:
+            ImportError: If no engines are available
+        """
+        available = OCRFactory.list_available_engines()
+        
+        # Try engines in order of recommendation
+        if available.get("easyocr", False):
+            logger.info("Using EasyOCR engine (recommended)")
+            return OCRFactory.create_engine("easyocr", **kwargs)
+        elif available.get("paddle", False):
+            logger.info("Using PaddleOCR engine")
+            return OCRFactory.create_engine("paddle", **kwargs)
+        elif available.get("surya", False):
+            logger.info("Using Surya OCR engine")
+            return OCRFactory.create_engine("surya", **kwargs)
+        
+        # If we get here, no engines are available
+        raise ImportError(
+            "No OCR engines available. Please install at least one of: \n"
+            "- EasyOCR (recommended): pip install easyocr\n"
+            "- PaddleOCR: pip install paddleocr paddlepaddle\n"
+            "- Surya OCR: pip install surya"
+        ) 

natural_pdf/ocr/ocr_manager.py

@@ -9,8 +9,8 @@ from PIL import Image
 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 OCROptions  # <-- Import Surya Options
+from .engine_surya import SuryaOCREngine
+from .ocr_options import OCROptions
 from .ocr_options import BaseOCROptions, EasyOCROptions, PaddleOCROptions, SuryaOCROptions
 
 logger = logging.getLogger(__name__)
@@ -27,15 +27,6 @@ class OCRManager:
         # 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",
-        # 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
@@ -49,16 +40,16 @@ class OCRManager:
                 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
             if not engine_instance.is_available():
                 # Check availability before storing
+                # Construct helpful error message with install hint
+                install_hint = f"pip install 'natural-pdf[{engine_name}]'"
                 raise RuntimeError(
-                    f"Engine '{engine_name}' is not available. Please check dependencies."
+                    f"Engine '{engine_name}' is not available. Please install the required dependencies: {install_hint}"
                 )
             self._engine_instances[engine_name] = engine_instance  # Store if available
 
@@ -66,106 +57,87 @@ class OCRManager:
 
     def apply_ocr(
         self,
-        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
+        images: Union[Image.Image, List[Image.Image]],
+        # --- Explicit Common Parameters ---
+        engine: Optional[str] = None,
+        languages: Optional[List[str]] = None,
+        min_confidence: Optional[float] = None,
+        device: Optional[str] = None,
+        detect_only: bool = False,
+        # --- Engine-Specific Options ---
+        options: Optional[Any] = None, # e.g. EasyOCROptions(), PaddleOCROptions()
+    ) -> Union[List[Dict[str, Any]], List[List[Dict[str, Any]]]]:
         """
-        Applies OCR to a single image or a batch of images using either simple
-        keyword arguments or an options object.
+        Applies OCR to a single image or a batch of images.
 
         Args:
             images: A single PIL Image or a list of PIL Images to process.
-            engine: Name of the engine to use (e.g., 'easyocr', 'paddle', 'surya').
-                    Ignored if 'options' object is provided. Defaults to 'easyocr'.
-            options: An instance of EasyOCROptions, PaddleOCROptions, or SuryaOCROptions
-                     for detailed configuration. If provided, simple kwargs (languages, etc.)
-                     and the 'engine' arg are ignored.
-            **kwargs: For simple mode, accepts: 'languages', 'min_confidence', 'device'.
-                      Other kwargs will raise a TypeError unless 'options' is provided.
+            engine: Name of the engine (e.g., 'easyocr', 'paddle', 'surya').
+                    Defaults to 'easyocr' if not specified.
+            languages: List of language codes (e.g., ['en', 'fr'], ['en', 'german']).
+                       **Passed directly to the engine.** Must be codes understood
+                       by the specific engine. No mapping is performed by the manager.
+            min_confidence: Minimum confidence threshold (0.0-1.0).
+                            Passed directly to the engine.
+            device: Device string (e.g., 'cpu', 'cuda').
+                    Passed directly to the engine.
+            detect_only: If True, only detect text regions, do not perform OCR.
+            options: An engine-specific options object (e.g., EasyOCROptions) or dict
+                     containing additional parameters specific to the chosen engine.
+                     Passed directly to the engine.
 
         Returns:
             If input is a single image: List of result dictionaries.
-            If input is a list of images: List of lists of result dictionaries,
-                                          corresponding to each input image.
+            If input is a list of images: List of lists of result dictionaries.
 
         Raises:
             ValueError: If the engine name is invalid.
-            TypeError: If unexpected keyword arguments are provided in simple mode,
-                       or if input 'images' is not a PIL Image or list of PIL Images.
-            RuntimeError: If the selected engine is not available.
+            TypeError: If input 'images' is not valid or options type is incompatible.
+            RuntimeError: If the selected engine is not available or processing fails.
         """
-        final_options: BaseOCROptions
-        selected_engine_name: str
-
         # --- 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.")
-        # 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
-            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"]):
-                    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."
-                )
-            if kwargs:
-                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}"
-            )
-
-            if selected_engine_name not in self.ENGINE_REGISTRY:
-                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."
-                )
-
-            # Get the *correct* options class for the selected engine
-            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"),
-                # 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}")
+        # --- Determine Engine --- 
+        selected_engine_name = (engine or "easyocr").lower()
+        if selected_engine_name not in self.ENGINE_REGISTRY:
+            raise ValueError(
+                f"Unknown OCR engine: '{selected_engine_name}'. Available: {list(self.ENGINE_REGISTRY.keys())}"
+            )
+        logger.debug(f"Selected engine: '{selected_engine_name}'")
+
+        # --- Prepare Options --- 
+        final_options = copy.deepcopy(options) if options is not None else None
+        
+        # Type check options object if provided
+        if final_options is not None:
+            options_class = self.ENGINE_REGISTRY[selected_engine_name].get("options_class", BaseOCROptions)
+            if not isinstance(final_options, options_class):
+                 # Allow dicts to be passed directly too, assuming engine handles them
+                if not isinstance(final_options, dict):
+                     raise TypeError(
+                        f"Provided options type '{type(final_options).__name__}' is not compatible with engine '{selected_engine_name}'. Expected '{options_class.__name__}' or dict."
+                    )
 
-        # --- Get Engine Instance and Process ---
+        # --- Get Engine Instance and Process --- 
         try:
             engine_instance = self._get_engine_instance(selected_engine_name)
             processing_mode = "batch" if is_batch else "single image"
             logger.info(f"Processing {processing_mode} with engine '{selected_engine_name}'...")
-
-            # Call the engine's process_image, passing single image or list
-            results = engine_instance.process_image(images, final_options)
+            logger.debug(f"  Engine Args: languages={languages}, min_confidence={min_confidence}, device={device}, options={final_options}")
+
+            # Call the engine's process_image, passing common args and options object
+            # **ASSUMPTION**: Engine process_image signatures are updated to accept these common args.
+            results = engine_instance.process_image(
+                images=images, 
+                languages=languages,
+                min_confidence=min_confidence,
+                device=device,
+                detect_only=detect_only,
+                options=final_options
+            )
 
             # Log result summary based on mode
             if is_batch:

natural_pdf/ocr/ocr_options.py

@@ -13,10 +13,6 @@ from typing import Any, Dict, List, Optional, Tuple, Union
 @dataclass
 class BaseOCROptions:
     """Base class for OCR engine options."""
-
-    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.
     extra_args: Dict[str, Any] = field(default_factory=dict)
 
 
@@ -24,7 +20,6 @@ class BaseOCROptions:
 @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"
@@ -69,7 +64,6 @@ 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
@@ -95,24 +89,20 @@ class PaddleOCROptions(BaseOCROptions):
     cls: Optional[bool] = None
 
     def __post_init__(self):
-        if self.use_gpu is None:
-            if self.device and "cuda" in self.device.lower():
-                self.use_gpu = True
-            else:
-                self.use_gpu = False
-        # logger.debug(f"Initialized PaddleOCROptions: {self}")
+        pass
+    #     if self.use_gpu is None:
+    #         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).
-    # For now, it primarily uses the base options like 'languages' and 'min_confidence'.
-    # Configuration like batch sizes are often set via environment variables for Surya.
     pass
 
 

natural_pdf/ocr/utils.py

@@ -0,0 +1,98 @@
+import io
+import base64
+import logging
+from typing import TYPE_CHECKING, Callable, Iterable, Optional, Any
+from natural_pdf.elements.text import TextElement
+from tqdm.auto import tqdm
+
+if TYPE_CHECKING:
+    from natural_pdf.elements.base import Element
+
+logger = logging.getLogger(__name__)
+
+def _apply_ocr_correction_to_elements(
+    elements: Iterable["Element"],
+    correction_callback: Callable[[Any], Optional[str]],
+) -> None:
+    """
+    Applies correction callback to a list of elements in place,
+    showing a progress bar.
+
+    Iterates through elements, calls the callback, and updates
+    element.text if a new string is returned.
+    
+    Args:
+        elements: An iterable of Element objects.
+        correction_callback: A function accepting an element and returning
+                             Optional[str] (new text or None).
+    """
+    corrections_applied = 0
+    elements_checked = 0
+
+    # Prepare the iterable with tqdm
+    element_iterable = tqdm(elements, desc=f"Correcting OCR", unit="element")
+
+    for element in element_iterable:
+        # Check if the element is likely from OCR and has text attribute
+        element_source = getattr(element, 'source', None)
+        if isinstance(element_source, str) and element_source.startswith('ocr') and hasattr(element, 'text'):
+            elements_checked += 1
+            current_text = getattr(element, 'text') 
+
+            new_text = correction_callback(element)
+
+            if new_text is not None:
+                if new_text != current_text:
+                    element.text = new_text 
+                    corrections_applied += 1
+
+    logger.info(f"OCR correction finished. Checked: {elements_checked}, Applied: {corrections_applied}")
+
+
+def direct_ocr_llm(element,
+                   client,
+                   model="",
+                   resolution=150,
+                   prompt="OCR this image. Return only the exact text from the image. Include misspellings, punctuation, etc.",
+                   padding=2) -> str:
+    """Convenience method to directly OCR a region of the page."""
+
+    if isinstance(element, TextElement):
+        region = element.expand(left=padding, right=padding, top=padding, bottom=padding)
+    else:
+        region = element
+
+    buffered = io.BytesIO()
+    region_img = region.to_image(resolution=resolution, include_highlights=False)
+    region_img.save(buffered, format="PNG")
+    base64_image = base64.b64encode(buffered.getvalue()).decode('utf-8')
+
+    response = client.chat.completions.create(
+        model=model,
+        messages=[
+            {
+                "role": "system",
+                "content": "You are an expert OCR engineer. You will be given an image of a region of a page. You will return the exact text from the image."
+            },
+            {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "text",
+                        "text": prompt
+                    },
+                    {
+                        "type": "image_url",
+                        "image_url": {
+                            "url": f"data:image/png;base64,{base64_image}"
+                        }
+                    }
+                ]
+            }
+        ]
+    )
+    
+    corrected = response.choices[0].message.content
+    logger.debug(f"Corrected {region.extract_text()} to {corrected}")
+
+    return corrected