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

natural_pdf/ocr/ocr_manager.py

@@ -2,6 +2,8 @@
 import copy  # For deep copying options
 import logging
 from typing import Any, Dict, List, Optional, Type, Union
+import threading # Import threading for lock
+import time # Import time for timing
 
 from PIL import Image
 
@@ -30,30 +32,68 @@ class OCRManager:
     def __init__(self):
         """Initializes the OCR Manager."""
         self._engine_instances: Dict[str, OCREngine] = {}  # Cache for engine instances
+        self._engine_locks: Dict[str, threading.Lock] = {} # Lock per engine type for initialization
+        self._engine_inference_locks: Dict[str, threading.Lock] = {} # Lock per engine type for inference
         logger.info("OCRManager initialized.")
 
     def _get_engine_instance(self, engine_name: str) -> OCREngine:
-        """Retrieves or creates an instance of the specified OCR engine."""
+        """Retrieves or creates an instance of the specified OCR engine, ensuring thread-safe initialization."""
         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())}"
             )
 
-        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 install the required dependencies: {install_hint}"
-                )
-            self._engine_instances[engine_name] = engine_instance  # Store if available
+        # Quick check if instance already exists (avoid lock contention)
+        if engine_name in self._engine_instances:
+            return self._engine_instances[engine_name]
 
-        return self._engine_instances[engine_name]
+        # Get or create the lock for this engine type
+        if engine_name not in self._engine_locks:
+            self._engine_locks[engine_name] = threading.Lock()
+        
+        engine_init_lock = self._engine_locks[engine_name]
+
+        # Acquire lock to safely check and potentially initialize the engine
+        with engine_init_lock:
+            # Double-check if another thread initialized it while we waited for the lock
+            if engine_name in self._engine_instances:
+                return self._engine_instances[engine_name]
+
+            # If still not initialized, create it now under the lock
+            logger.info(f"[{threading.current_thread().name}] Creating shared instance of engine: {engine_name}")
+            engine_class = self.ENGINE_REGISTRY[engine_name]["class"]
+            start_time = time.monotonic() # Optional: time initialization
+            try:
+                engine_instance = engine_class()  # Instantiate first
+                if not engine_instance.is_available():
+                    # Check availability before storing
+                    install_hint = f"pip install 'natural-pdf[{engine_name}]'"
+                    raise RuntimeError(
+                        f"Engine '{engine_name}' is not available. Please install the required dependencies: {install_hint}"
+                    )
+                # Store the shared instance
+                self._engine_instances[engine_name] = engine_instance
+                end_time = time.monotonic()
+                logger.info(f"[{threading.current_thread().name}] Shared instance of {engine_name} created successfully (Duration: {end_time - start_time:.2f}s).")
+                return engine_instance
+            except Exception as e:
+                 # Ensure we don't leave a partial state if init fails
+                 logger.error(f"[{threading.current_thread().name}] Failed to create shared instance of {engine_name}: {e}", exc_info=True)
+                 # Remove potentially partial entry if exists
+                 if engine_name in self._engine_instances: del self._engine_instances[engine_name]
+                 raise # Re-raise the exception after logging
+
+    def _get_engine_inference_lock(self, engine_name: str) -> threading.Lock:
+        """Gets or creates the inference lock for a given engine type."""
+        engine_name = engine_name.lower()
+        # Assume engine_name is valid as it's checked before this would be called
+        if engine_name not in self._engine_inference_locks:
+            # Create lock if it doesn't exist (basic thread safety for dict access)
+            # A more robust approach might lock around this check/creation too,
+            # but contention here is less critical than for engine init or inference itself.
+            self._engine_inference_locks[engine_name] = threading.Lock()
+        return self._engine_inference_locks[engine_name]
 
     def apply_ocr(
         self,
@@ -65,7 +105,7 @@ class OCRManager:
         device: Optional[str] = None,
         detect_only: bool = False,
         # --- Engine-Specific Options ---
-        options: Optional[Any] = None, # e.g. EasyOCROptions(), PaddleOCROptions()
+        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.
@@ -100,7 +140,7 @@ class OCRManager:
         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.")
 
-        # --- Determine Engine --- 
+        # --- Determine Engine ---
         selected_engine_name = (engine or "easyocr").lower()
         if selected_engine_name not in self.ENGINE_REGISTRY:
             raise ValueError(
@@ -108,37 +148,61 @@ class OCRManager:
             )
         logger.debug(f"Selected engine: '{selected_engine_name}'")
 
-        # --- Prepare Options --- 
+        # --- 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)
+            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
+                # Allow dicts to be passed directly too, assuming engine handles them
                 if not isinstance(final_options, dict):
-                     raise TypeError(
+                    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}'...")
-            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 thread name for clarity during parallel calls
+            thread_id = threading.current_thread().name
+            logger.info(f"[{thread_id}] Processing {processing_mode} using shared engine instance '{selected_engine_name}'...")
+            logger.debug(
+                f"  Engine Args: languages={languages}, min_confidence={min_confidence}, device={device}, options={final_options}"
             )
 
+            # Log image dimensions before processing
+            if is_batch:
+                image_dims = [f"{img.width}x{img.height}" for img in images if hasattr(img, 'width') and hasattr(img, 'height')]
+                logger.debug(f"[{thread_id}] Processing batch of {len(images)} images with dimensions: {image_dims}")
+            elif hasattr(images, 'width') and hasattr(images, 'height'):
+                logger.debug(f"[{thread_id}] Processing single image with dimensions: {images.width}x{images.height}")
+            else:
+                logger.warning(f"[{thread_id}] Could not determine dimensions of input image(s).")
+
+            # Acquire lock specifically for the inference call
+            inference_lock = self._get_engine_inference_lock(selected_engine_name)
+            logger.debug(f"[{thread_id}] Attempting to acquire inference lock for {selected_engine_name}...")
+            inference_wait_start = time.monotonic()
+            with inference_lock:
+                inference_acquired_time = time.monotonic()
+                logger.debug(f"[{thread_id}] Acquired inference lock for {selected_engine_name} (waited {inference_acquired_time - inference_wait_start:.2f}s). Calling process_image...")
+                inference_start_time = time.monotonic()
+
+                results = engine_instance.process_image(
+                    images=images,
+                    languages=languages,
+                    min_confidence=min_confidence,
+                    device=device,
+                    detect_only=detect_only,
+                    options=final_options,
+                )
+                inference_end_time = time.monotonic()
+                logger.debug(f"[{thread_id}] process_image call finished for {selected_engine_name} (Duration: {inference_end_time - inference_start_time:.2f}s). Releasing lock.")
+
             # Log result summary based on mode
             if is_batch:
                 # Ensure results is a list before trying to get lengths

natural_pdf/ocr/ocr_options.py

@@ -13,6 +13,7 @@ from typing import Any, Dict, List, Optional, Tuple, Union
 @dataclass
 class BaseOCROptions:
     """Base class for OCR engine options."""
+
     extra_args: Dict[str, Any] = field(default_factory=dict)
 
 
@@ -20,6 +21,7 @@ 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"
@@ -64,9 +66,10 @@ class EasyOCROptions(BaseOCROptions):
 @dataclass
 class PaddleOCROptions(BaseOCROptions):
     """Specific options for the PaddleOCR engine."""
-    use_angle_cls: bool = True
+
+    # General
     use_gpu: Optional[bool] = None
-    gpu_mem: int = 500
+    gpu_mem: int = 8000 # Default from Paddle documentation
     ir_optim: bool = True
     use_tensorrt: bool = False
     min_subgraph_size: int = 15
@@ -74,22 +77,46 @@ class PaddleOCROptions(BaseOCROptions):
     enable_mkldnn: bool = False
     cpu_threads: int = 10
     use_fp16: bool = False
+    show_log: bool = False
+    use_onnx: bool = False
+    use_zero_copy_run: bool = False
+
+    # Detection
+    det: bool = True
+    det_algorithm: str = "DB"
     det_model_dir: Optional[str] = None
+    det_limit_side_len: int = 960 # Corresponds to det_max_side_len
+    # DB specific
+    det_db_thresh: float = 0.3
+    det_db_box_thresh: float = 0.5
+    det_db_unclip_ratio: float = 2.0
+    # EAST specific
+    det_east_score_thresh: float = 0.8
+    det_east_cover_thresh: float = 0.1
+    det_east_nms_thresh: float = 0.2
+
+    # Recognition
+    rec: bool = True
+    rec_algorithm: str = "CRNN"
     rec_model_dir: Optional[str] = None
-    cls_model_dir: Optional[str] = None
-    det_limit_side_len: int = 960
-    rec_batch_num: int = 6
+    rec_image_shape: str = "3, 32, 320" # Kept as string per Paddle examples
+    rec_batch_num: int = 30 # Default from Paddle documentation
     max_text_length: int = 25
+    rec_char_dict_path: Optional[str] = None # Path to char dictionary file
     use_space_char: bool = True
     drop_score: float = 0.5
-    show_log: bool = False
-    use_onnx: bool = False
-    det: bool = True
-    rec: bool = True
-    cls: Optional[bool] = None
+
+    # Classification
+    cls: Optional[bool] = None # Often inferred from use_angle_cls
+    use_angle_cls: bool = False # Default from Paddle documentation
+    cls_model_dir: Optional[str] = None
+    cls_image_shape: str = "3, 48, 192" # Kept as string per Paddle examples
+    label_list: List[str] = field(default_factory=lambda: ['0', '180']) # Default from Paddle doc
+    cls_batch_num: int = 30
 
     def __post_init__(self):
         pass
+
     #     if self.use_gpu is None:
     #         if self.device and "cuda" in self.device.lower():
     #             self.use_gpu = True
@@ -102,6 +129,7 @@ class PaddleOCROptions(BaseOCROptions):
 @dataclass
 class SuryaOCROptions(BaseOCROptions):
     """Specific options for the Surya OCR engine."""
+
     # Currently, Surya example shows languages passed at prediction time.
     pass
 

natural_pdf/ocr/utils.py

@@ -8,53 +8,76 @@ from tqdm.auto import tqdm
 if TYPE_CHECKING:
     from natural_pdf.elements.base import Element
 
+# Import the global PDF render lock from dedicated locks module
+from natural_pdf.utils.locks import pdf_render_lock
+
 logger = logging.getLogger(__name__)
 
+
 def _apply_ocr_correction_to_elements(
     elements: Iterable["Element"],
     correction_callback: Callable[[Any], Optional[str]],
+    caller_info: str = "Utility",
 ) -> None:
     """
-    Applies correction callback to a list of elements in place,
+    Applies OCR 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.
-    
+    Iterates through elements, checks if source starts with 'ocr', 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).
+        caller_info: String identifying the calling context for logs.
     """
+    if not callable(correction_callback):
+        # Raise error here so individual methods don't need to repeat the check
+        raise TypeError("`correction_callback` must be a callable function.")
+
+    if not elements:
+        logger.warning(f"{caller_info}: No elements provided for correction.")
+        return
+
     corrections_applied = 0
     elements_checked = 0
 
     # Prepare the iterable with tqdm
-    element_iterable = tqdm(elements, desc=f"Correcting OCR", unit="element")
+    element_iterable = tqdm(elements, desc=f"Correcting OCR ({caller_info})", 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'):
+        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') 
+            current_text = getattr(element, "text")  # Already checked hasattr
 
             new_text = correction_callback(element)
 
             if new_text is not None:
                 if new_text != current_text:
-                    element.text = new_text 
+                    element.text = new_text  # Update in place
                     corrections_applied += 1
 
-    logger.info(f"OCR correction finished. Checked: {elements_checked}, Applied: {corrections_applied}")
+    logger.info(
+        f"{caller_info}: OCR correction finished. Checked: {elements_checked}, Applied: {corrections_applied}"
+    )
+    # No return value needed, modifies elements in place
 
 
-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:
+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. If you cannot see any text, return an empty string.",
+    padding=2,
+) -> str:
     """Convenience method to directly OCR a region of the page."""
 
     if isinstance(element, TextElement):
@@ -63,36 +86,39 @@ def direct_ocr_llm(element,
         region = element
 
     buffered = io.BytesIO()
-    region_img = region.to_image(resolution=resolution, include_highlights=False)
+    # Use the global PDF render lock when rendering images
+    with pdf_render_lock:
+        region_img = region.to_image(resolution=resolution, include_highlights=False)
+
+    # Handle cases where image creation might fail (e.g., zero-dim region)
+    if region_img is None:
+        logger.warning(f"Could not generate image for region {region.bbox}, skipping OCR.")
+        return "" # Return empty string if image creation failed
+
     region_img.save(buffered, format="PNG")
-    base64_image = base64.b64encode(buffered.getvalue()).decode('utf-8')
+    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."
+                "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": "text", "text": prompt},
                     {
                         "type": "image_url",
-                        "image_url": {
-                            "url": f"data:image/png;base64,{base64_image}"
-                        }
-                    }
-                ]
-            }
-        ]
+                        "image_url": {"url": f"data:image/png;base64,{base64_image}"},
+                    },
+                ],
+            },
+        ],
     )
-    
-    corrected = response.choices[0].message.content
+
+    corrected = response.choices[0].message.content.strip()
     logger.debug(f"Corrected {region.extract_text()} to {corrected}")
 
-    return corrected
+    return corrected

natural_pdf/qa/document_qa.py

@@ -58,10 +58,6 @@ class DocumentQA:
             import torch
             from transformers import pipeline
 
-            # Determine device
-            if device is None:
-                device = "cuda" if torch.cuda.is_available() else "cpu"
-
             logger.info(f"Initializing DocumentQA with model {model_name} on {device}")
 
             # Initialize the pipeline