natural-pdf 0.1.7__py3-none-any.whl → 0.1.9__py3-none-any.whl

natural_pdf/ocr/ocr_options.py

@@ -67,9 +67,10 @@ class EasyOCROptions(BaseOCROptions):
 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
+    gpu_mem: int = 8000  # Default from Paddle documentation
     ir_optim: bool = True
     use_tensorrt: bool = False
     min_subgraph_size: int = 15
@@ -77,19 +78,73 @@ 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"
+    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"
+    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
+    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
+    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
+
+    # 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
@@ -111,5 +166,33 @@ class SuryaOCROptions(BaseOCROptions):
     pass
 
 
+# --- Doctr Specific Options ---
+@dataclass
+class DoctrOCROptions(BaseOCROptions):
+    """Specific options for the doctr engine."""
+
+    # OCR predictor options
+    det_arch: str = "db_resnet50"
+    reco_arch: str = "crnn_vgg16_bn"
+    pretrained: bool = True
+    assume_straight_pages: bool = True  # Faster if pages are straight
+    export_as_straight_boxes: bool = False  # Output straight boxes even if rotated text is detected
+
+    # Additional options from standalone predictors
+    # Detection predictor options
+    symmetric_pad: bool = True
+    preserve_aspect_ratio: bool = True
+    batch_size: int = 1
+
+    # Postprocessing parameters
+    bin_thresh: Optional[float] = None  # Default is usually 0.3
+    box_thresh: Optional[float] = None  # Default is usually 0.1
+
+    # Options for orientation predictors
+    use_orientation_predictor: bool = False  # Whether to use page orientation predictor
+
+
 # --- Union type for type hinting ---
-OCROptions = Union[EasyOCROptions, PaddleOCROptions, SuryaOCROptions, BaseOCROptions]
+OCROptions = Union[
+    EasyOCROptions, PaddleOCROptions, SuryaOCROptions, DoctrOCROptions, BaseOCROptions
+]

natural_pdf/ocr/utils.py

@@ -1,13 +1,18 @@
-import io
 import base64
+import io
 import logging
-from typing import TYPE_CHECKING, Callable, Iterable, Optional, Any
-from natural_pdf.elements.text import TextElement
+from typing import TYPE_CHECKING, Any, Callable, Iterable, Optional
+
 from tqdm.auto import tqdm
 
+from natural_pdf.elements.text import TextElement
+
 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__)
 
 
@@ -72,7 +77,7 @@ def direct_ocr_llm(
     client,
     model="",
     resolution=150,
-    prompt="OCR this image. Return only the exact text from the image. Include misspellings, punctuation, etc.",
+    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."""
@@ -83,7 +88,15 @@ def direct_ocr_llm(
         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")
 
@@ -107,7 +120,7 @@ def direct_ocr_llm(
         ],
     )
 
-    corrected = response.choices[0].message.content
+    corrected = response.choices[0].message.content.strip()
     logger.debug(f"Corrected {region.extract_text()} to {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

natural_pdf/search/__init__.py

@@ -24,63 +24,54 @@ from .search_service_protocol import Indexable, IndexConfigurationError, SearchS
 
 logger = logging.getLogger(__name__)
 
-# --- Factory Function ---
-
 
+# Factory Function
 def get_search_service(
-    collection_name: str,  # Add collection_name as a required argument
-    persist: bool = False,  # Default to In-Memory
-    # Configuration for the service itself
-    default_persist_path: Optional[str] = None,
+    collection_name: str,
+    persist: bool = False,
+    uri: Optional[str] = None,
     default_embedding_model: Optional[str] = None,
-    # Potential future args: cache_services=True? service_type='haystack'?
 ) -> SearchServiceProtocol:
     """
     Factory function to get an instance of the configured search service.
 
-    A service instance is tied to a specific collection name.
+    A service instance is tied to a specific index name (collection/table).
 
     Currently, only returns HaystackSearchService but is structured for future extension.
 
     Args:
-        collection_name: The name of the collection this service instance will manage.
+        collection_name: The logical name for the index this service instance manages
+                         (used as table_name for LanceDB).
         persist: If True, creates a service instance configured for persistent
-                 storage (ChromaDB). If False (default), uses In-Memory.
-        default_persist_path: Override the default path for persistent storage.
+                 storage (currently LanceDB). If False (default), uses InMemory.
+        uri: Override the default path/URI for persistent storage.
         default_embedding_model: Override the default embedding model used by the service.
         **kwargs: Reserved for future configuration options.
 
     Returns:
-        An instance conforming to the SearchServiceProtocol for the specified collection.
+        An instance conforming to the SearchServiceProtocol for the specified collection/table.
     """
     logger.debug(
-        f"Calling get_search_service factory for collection '{collection_name}' (persist={persist})..."
+        f"Calling get_search_service factory for index '{collection_name}' (persist={persist}, uri={uri})..."
     )
 
-    # For now, we only have one implementation
     # Collect arguments relevant to HaystackSearchService.__init__
     service_args = {}
-    service_args["collection_name"] = collection_name  # Pass collection_name
-    service_args["persist"] = persist  # Pass persist flag to service constructor
-    if default_persist_path is not None:
-        service_args["default_persist_path"] = default_persist_path
+    service_args["table_name"] = collection_name
+    service_args["persist"] = persist
+    if uri is not None:
+        service_args["uri"] = uri
     if default_embedding_model is not None:
-        service_args["default_embedding_model"] = default_embedding_model
+        service_args["embedding_model"] = default_embedding_model
 
-    # TODO: Implement caching/registry if needed to return the same instance
-    # for the same configuration instead of always creating a new one.
-    # cache_key = tuple(sorted(service_args.items()))
-    # if cache_key in _service_instance_cache:
-    #    return _service_instance_cache[cache_key]
+    # Cache logic commented out as before
 
     try:
         service_instance = HaystackSearchService(**service_args)
-        # _service_instance_cache[cache_key] = service_instance
-        logger.info(
-            f"Created new HaystackSearchService instance for collection '{collection_name}'."
-        )
+        logger.info(f"Created new HaystackSearchService instance for index '{collection_name}'.")
         return service_instance
     except ImportError as e:
+        # Error message remains valid
         logger.error(
             f"Failed to instantiate Search Service due to missing dependencies: {e}", exc_info=True
         )
@@ -92,9 +83,4 @@ def get_search_service(
         raise RuntimeError("Could not create Search Service instance.") from e
 
 
-# --- Optional: Define a default instance for extreme ease of use? ---
-# try:
-#     default_search_service = get_search_service()
-# except Exception:
-#     default_search_service = None
-#     logger.warning("Could not create default search service instance on import.")
+# Default instance commented out as before