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

natural_pdf/ocr/engine_easyocr.py

@@ -1,22 +1,24 @@
 # ocr_engine_easyocr.py
-import logging
 import importlib.util
-from typing import Dict, List, Any, Optional, Tuple, Union
+import inspect  # Used for dynamic parameter passing
+import logging
+from typing import Any, Dict, List, Optional, Tuple, Union
+
 import numpy as np
 from PIL import Image
-import inspect # Used for dynamic parameter passing
 
 from .engine import OCREngine
-from .ocr_options import EasyOCROptions, BaseOCROptions
+from .ocr_options import BaseOCROptions, EasyOCROptions
 
 logger = logging.getLogger(__name__)
 
+
 class EasyOCREngine(OCREngine):
     """EasyOCR engine implementation."""
 
     def __init__(self):
         super().__init__()
-        self._easyocr = None # Lazy load easyocr module
+        self._easyocr = None  # Lazy load easyocr module
 
     def _lazy_import_easyocr(self):
         """Imports easyocr only when needed."""
@@ -25,6 +27,7 @@ class EasyOCREngine(OCREngine):
                 raise ImportError("EasyOCR is not installed or available.")
             try:
                 import easyocr
+
                 self._easyocr = easyocr
                 logger.info("EasyOCR module imported successfully.")
             except ImportError as e:
@@ -56,15 +59,18 @@ class EasyOCREngine(OCREngine):
 
         constructor_sig = inspect.signature(easyocr.Reader.__init__)
         constructor_args = {}
-        constructor_args['lang_list'] = options.languages
-        constructor_args['gpu'] = 'cuda' in str(options.device).lower() or 'mps' in str(options.device).lower()
+        constructor_args["lang_list"] = options.languages
+        constructor_args["gpu"] = (
+            "cuda" in str(options.device).lower() or "mps" in str(options.device).lower()
+        )
 
         for field_name, param in constructor_sig.parameters.items():
-            if field_name in ['self', 'lang_list', 'gpu']: continue
+            if field_name in ["self", "lang_list", "gpu"]:
+                continue
             if hasattr(options, field_name):
-                 constructor_args[field_name] = getattr(options, field_name)
+                constructor_args[field_name] = getattr(options, field_name)
             elif field_name in options.extra_args:
-                 constructor_args[field_name] = options.extra_args[field_name]
+                constructor_args[field_name] = options.extra_args[field_name]
 
         logger.debug(f"EasyOCR Reader constructor args: {constructor_args}")
         try:
@@ -81,22 +87,29 @@ class EasyOCREngine(OCREngine):
         readtext_sig = inspect.signature(reader.readtext)
         readtext_args = {}
         for field_name, param in readtext_sig.parameters.items():
-             if field_name == 'image': continue
-             if hasattr(options, field_name):
-                 readtext_args[field_name] = getattr(options, field_name)
-             elif field_name in options.extra_args:
-                 readtext_args[field_name] = options.extra_args[field_name]
+            if field_name == "image":
+                continue
+            if hasattr(options, field_name):
+                readtext_args[field_name] = getattr(options, field_name)
+            elif field_name in options.extra_args:
+                readtext_args[field_name] = options.extra_args[field_name]
         logger.debug(f"EasyOCR readtext args: {readtext_args}")
         return readtext_args
 
-    def _standardize_results(self, raw_results: List[Any], options: EasyOCROptions) -> List[Dict[str, Any]]:
+    def _standardize_results(
+        self, raw_results: List[Any], options: EasyOCROptions
+    ) -> List[Dict[str, Any]]:
         """Standardizes raw results from EasyOCR's readtext."""
         standardized_results = []
         min_confidence = options.min_confidence
 
         for detection in raw_results:
             try:
-                if options.detail == 1 and isinstance(detection, (list, tuple)) and len(detection) >= 3:
+                if (
+                    options.detail == 1
+                    and isinstance(detection, (list, tuple))
+                    and len(detection) >= 3
+                ):
                     bbox_raw = detection[0]
                     text = str(detection[1])
                     confidence = float(detection[2])
@@ -104,38 +117,40 @@ class EasyOCREngine(OCREngine):
                     if confidence >= min_confidence:
                         bbox = self._standardize_bbox(bbox_raw)
                         if bbox:
-                            standardized_results.append({
-                                'bbox': bbox, 'text': text, 'confidence': confidence, 'source': 'ocr'
-                            })
+                            standardized_results.append(
+                                {
+                                    "bbox": bbox,
+                                    "text": text,
+                                    "confidence": confidence,
+                                    "source": "ocr",
+                                }
+                            )
                         else:
-                             logger.warning(f"Skipping result due to invalid bbox: {bbox_raw}")
+                            logger.warning(f"Skipping result due to invalid bbox: {bbox_raw}")
 
                 elif options.detail == 0 and isinstance(detection, str):
-                     standardized_results.append({
-                         'bbox': None, 'text': detection, 'confidence': 1.0, 'source': 'ocr'
-                     })
+                    standardized_results.append(
+                        {"bbox": None, "text": detection, "confidence": 1.0, "source": "ocr"}
+                    )
             except (IndexError, ValueError, TypeError) as e:
-                 logger.warning(f"Skipping invalid detection format: {detection}. Error: {e}")
-                 continue
+                logger.warning(f"Skipping invalid detection format: {detection}. Error: {e}")
+                continue
         return standardized_results
 
-
     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 EasyOCR."""
 
         if not isinstance(options, EasyOCROptions):
-             logger.warning("Received BaseOCROptions, expected EasyOCROptions. Using defaults.")
-             # Create default EasyOCR options if base was passed, preserving base settings
-             options = EasyOCROptions(
-                 languages=options.languages,
-                 min_confidence=options.min_confidence,
-                 device=options.device,
-                 extra_args=options.extra_args # Pass along any extra args
-             )
+            logger.warning("Received BaseOCROptions, expected EasyOCROptions. Using defaults.")
+            # Create default EasyOCR options if base was passed, preserving base settings
+            options = EasyOCROptions(
+                languages=options.languages,
+                min_confidence=options.min_confidence,
+                device=options.device,
+                extra_args=options.extra_args,  # Pass along any extra args
+            )
 
         reader = self._get_reader(options)
         readtext_args = self._prepare_readtext_args(options, reader)
@@ -147,9 +162,9 @@ class EasyOCREngine(OCREngine):
             logger.info(f"Processing batch of {len(images)} images with EasyOCR (iteratively)...")
             for i, img in enumerate(images):
                 if not isinstance(img, Image.Image):
-                     logger.warning(f"Item at index {i} in batch is not a PIL Image. Skipping.")
-                     all_results.append([])
-                     continue
+                    logger.warning(f"Item at index {i} in batch is not a PIL Image. Skipping.")
+                    all_results.append([])
+                    continue
                 img_array = np.array(img)
                 try:
                     logger.debug(f"Processing image {i+1}/{len(images)} in batch.")
@@ -157,10 +172,12 @@ class EasyOCREngine(OCREngine):
                     standardized = self._standardize_results(raw_results, options)
                     all_results.append(standardized)
                 except Exception as e:
-                    logger.error(f"Error processing image {i+1} in EasyOCR batch: {e}", exc_info=True)
-                    all_results.append([]) # Append empty list for failed image
+                    logger.error(
+                        f"Error processing image {i+1} in EasyOCR batch: {e}", exc_info=True
+                    )
+                    all_results.append([])  # Append empty list for failed image
             logger.info(f"Finished processing batch with EasyOCR.")
-            return all_results # Return List[List[Dict]]
+            return all_results  # Return List[List[Dict]]
 
         elif isinstance(images, Image.Image):
             # --- Single Image Processing ---
@@ -170,10 +187,9 @@ class EasyOCREngine(OCREngine):
                 raw_results = reader.readtext(img_array, **readtext_args)
                 standardized = self._standardize_results(raw_results, options)
                 logger.info(f"Finished processing single image. Found {len(standardized)} results.")
-                return standardized # Return List[Dict]
+                return standardized  # Return List[Dict]
             except Exception as e:
                 logger.error(f"Error processing single image with EasyOCR: {e}", exc_info=True)
-                return [] # Return empty list on failure
+                return []  # Return empty list on failure
         else:
             raise TypeError("Input 'images' must be a PIL Image or a list of PIL Images.")
-

natural_pdf/ocr/engine_paddle.py

@@ -1,27 +1,49 @@
 # ocr_engine_paddleocr.py
-import logging
 import importlib.util
-from typing import Dict, List, Any, Optional, Tuple, Union
+import inspect  # Used for dynamic parameter passing
+import logging
+from typing import Any, Dict, List, Optional, Tuple, Union
+
 import numpy as np
 from PIL import Image
-import inspect # Used for dynamic parameter passing
 
 from .engine import OCREngine
-from .ocr_options import PaddleOCROptions, BaseOCROptions
+from .ocr_options import BaseOCROptions, PaddleOCROptions
 
 logger = logging.getLogger(__name__)
 
+
 class PaddleOCREngine(OCREngine):
     """PaddleOCR engine implementation."""
 
-    LANGUAGE_MAP = { 
-        'en': 'en', 'zh': 'ch', 'zh-cn': 'ch', 'zh-tw': 'chinese_cht',
-        'ja': 'japan', 'ko': 'korean', 'th': 'thai', 'fr': 'french',
-        'de': 'german', 'ru': 'russian', 'ar': 'arabic', 'hi': 'hindi',
-        'vi': 'vietnam', 'fa': 'cyrillic', 'ur': 'cyrillic', 'rs': 'serbian',
-        'oc': 'latin', 'rsc': 'cyrillic', 'bg': 'bulgarian', 'uk': 'cyrillic',
-        'be': 'cyrillic', 'te': 'telugu', 'kn': 'kannada', 'ta': 'tamil',
-        'latin': 'latin', 'cyrillic': 'cyrillic', 'devanagari': 'devanagari',
+    LANGUAGE_MAP = {
+        "en": "en",
+        "zh": "ch",
+        "zh-cn": "ch",
+        "zh-tw": "chinese_cht",
+        "ja": "japan",
+        "ko": "korean",
+        "th": "thai",
+        "fr": "french",
+        "de": "german",
+        "ru": "russian",
+        "ar": "arabic",
+        "hi": "hindi",
+        "vi": "vietnam",
+        "fa": "cyrillic",
+        "ur": "cyrillic",
+        "rs": "serbian",
+        "oc": "latin",
+        "rsc": "cyrillic",
+        "bg": "bulgarian",
+        "uk": "cyrillic",
+        "be": "cyrillic",
+        "te": "telugu",
+        "kn": "kannada",
+        "ta": "tamil",
+        "latin": "latin",
+        "cyrillic": "cyrillic",
+        "devanagari": "devanagari",
     }
 
     def __init__(self):
@@ -36,6 +58,7 @@ class PaddleOCREngine(OCREngine):
             try:
                 import paddle
                 import paddleocr
+
                 self._paddleocr = paddleocr
                 logger.info("PaddleOCR module imported successfully.")
             except ImportError as e:
@@ -45,19 +68,21 @@ class PaddleOCREngine(OCREngine):
 
     def is_available(self) -> bool:
         """Check if PaddleOCR and paddlepaddle are installed."""
-        paddle_installed = importlib.util.find_spec("paddle") is not None or \
-                           importlib.util.find_spec("paddlepaddle") is not None
+        paddle_installed = (
+            importlib.util.find_spec("paddle") is not None
+            or importlib.util.find_spec("paddlepaddle") is not None
+        )
         paddleocr_installed = importlib.util.find_spec("paddleocr") is not None
         return paddle_installed and paddleocr_installed
 
     def _map_language(self, iso_lang: str) -> str:
         """Map ISO language code to PaddleOCR language code."""
-        return self.LANGUAGE_MAP.get(iso_lang.lower(), 'en')
+        return self.LANGUAGE_MAP.get(iso_lang.lower(), "en")
 
     def _get_cache_key(self, options: PaddleOCROptions) -> str:
         """Generate a more specific cache key for PaddleOCR."""
         base_key = super()._get_cache_key(options)
-        primary_lang = self._map_language(options.languages[0]) if options.languages else 'en'
+        primary_lang = self._map_language(options.languages[0]) if options.languages else "en"
         angle_cls_key = str(options.use_angle_cls)
         precision_key = options.precision
         return f"{base_key}_{primary_lang}_{angle_cls_key}_{precision_key}"
@@ -74,31 +99,34 @@ class PaddleOCREngine(OCREngine):
 
         constructor_sig = inspect.signature(paddleocr.PaddleOCR.__init__)
         constructor_args = {}
-        constructor_args['lang'] = self._map_language(options.languages[0]) if options.languages else 'en'
+        constructor_args["lang"] = (
+            self._map_language(options.languages[0]) if options.languages else "en"
+        )
 
         for field_name, param in constructor_sig.parameters.items():
-            if field_name in ['self', 'lang']: continue
-            if field_name == 'use_gpu':
-                 constructor_args['use_gpu'] = options.use_gpu
-                 continue
+            if field_name in ["self", "lang"]:
+                continue
+            if field_name == "use_gpu":
+                constructor_args["use_gpu"] = options.use_gpu
+                continue
             if hasattr(options, field_name):
-                 constructor_args[field_name] = getattr(options, field_name)
+                constructor_args[field_name] = getattr(options, field_name)
             elif field_name in options.extra_args:
-                 constructor_args[field_name] = options.extra_args[field_name]
+                constructor_args[field_name] = options.extra_args[field_name]
 
-        constructor_args.pop('device', None)
+        constructor_args.pop("device", None)
         logger.debug(f"PaddleOCR constructor args: {constructor_args}")
 
         try:
-            show_log = constructor_args.get('show_log', False)
-            original_log_level = logging.getLogger('ppocr').level
+            show_log = constructor_args.get("show_log", False)
+            original_log_level = logging.getLogger("ppocr").level
             if not show_log:
-                 logging.getLogger('ppocr').setLevel(logging.ERROR)
+                logging.getLogger("ppocr").setLevel(logging.ERROR)
 
             reader = paddleocr.PaddleOCR(**constructor_args)
 
             if not show_log:
-                 logging.getLogger('ppocr').setLevel(original_log_level)
+                logging.getLogger("ppocr").setLevel(original_log_level)
 
             self._reader_cache[cache_key] = reader
             logger.info("PaddleOCR reader created successfully.")
@@ -108,32 +136,36 @@ class PaddleOCREngine(OCREngine):
             raise
 
     def _prepare_ocr_args(self, options: PaddleOCROptions) -> Dict[str, Any]:
-         """Helper to prepare arguments for the ocr method (excluding image)."""
-         ocr_args = {}
-         # Determine 'cls' value based on options precedence
-         ocr_args['cls'] = options.cls if options.cls is not None else options.use_angle_cls
-         ocr_args['det'] = options.det
-         ocr_args['rec'] = options.rec
-         # Add extra args if needed (less common for ocr method itself)
-         # for field_name in options.extra_args:
-         #      if field_name in ['cls', 'det', 'rec']: # Check against known ocr args
-         #          ocr_args[field_name] = options.extra_args[field_name]
-         logger.debug(f"PaddleOCR ocr args (excluding image): {ocr_args}")
-         return ocr_args
-
-    def _standardize_results(self, raw_page_results: Optional[List[Any]], options: PaddleOCROptions) -> List[Dict[str, Any]]:
+        """Helper to prepare arguments for the ocr method (excluding image)."""
+        ocr_args = {}
+        # Determine 'cls' value based on options precedence
+        ocr_args["cls"] = options.cls if options.cls is not None else options.use_angle_cls
+        ocr_args["det"] = options.det
+        ocr_args["rec"] = options.rec
+        # Add extra args if needed (less common for ocr method itself)
+        # for field_name in options.extra_args:
+        #      if field_name in ['cls', 'det', 'rec']: # Check against known ocr args
+        #          ocr_args[field_name] = options.extra_args[field_name]
+        logger.debug(f"PaddleOCR ocr args (excluding image): {ocr_args}")
+        return ocr_args
+
+    def _standardize_results(
+        self, raw_page_results: Optional[List[Any]], options: PaddleOCROptions
+    ) -> List[Dict[str, Any]]:
         """Standardizes raw results from a single page/image from PaddleOCR."""
         standardized_page = []
-        if not raw_page_results: # Handle None or empty list
+        if not raw_page_results:  # Handle None or empty list
             return standardized_page
 
         min_confidence = options.min_confidence
         for detection in raw_page_results:
             try:
-                if not isinstance(detection, (list, tuple)) or len(detection) < 2: continue
+                if not isinstance(detection, (list, tuple)) or len(detection) < 2:
+                    continue
                 bbox_raw = detection[0]
                 text_confidence = detection[1]
-                if not isinstance(text_confidence, tuple) or len(text_confidence) < 2: continue
+                if not isinstance(text_confidence, tuple) or len(text_confidence) < 2:
+                    continue
 
                 text = str(text_confidence[0])
                 confidence = float(text_confidence[1])
@@ -141,55 +173,52 @@ class PaddleOCREngine(OCREngine):
                 if confidence >= min_confidence:
                     bbox = self._standardize_bbox(bbox_raw)
                     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:
                         logger.warning(f"Skipping result due to invalid bbox: {bbox_raw}")
             except (IndexError, ValueError, TypeError) as e:
-                 logger.warning(f"Skipping invalid detection format: {detection}. Error: {e}")
-                 continue
+                logger.warning(f"Skipping invalid detection format: {detection}. Error: {e}")
+                continue
         return standardized_page
 
     def _pil_to_bgr(self, image: Image.Image) -> np.ndarray:
         """Converts PIL Image to BGR numpy array."""
-        if image.mode == 'BGR': # Already BGR
-             return np.array(image)
-        img_rgb = image.convert('RGB')
+        if image.mode == "BGR":  # Already BGR
+            return np.array(image)
+        img_rgb = image.convert("RGB")
         img_array_rgb = np.array(img_rgb)
-        img_array_bgr = img_array_rgb[:, :, ::-1] # Convert RGB to BGR
+        img_array_bgr = img_array_rgb[:, :, ::-1]  # Convert RGB to BGR
         return img_array_bgr
 
-
     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 PaddleOCR."""
 
         if not isinstance(options, PaddleOCROptions):
-             logger.warning("Received BaseOCROptions, expected PaddleOCROptions. Using defaults.")
-             options = PaddleOCROptions(
-                 languages=options.languages,
-                 min_confidence=options.min_confidence,
-                 device=options.device,
-                 extra_args=options.extra_args
-             )
+            logger.warning("Received BaseOCROptions, expected PaddleOCROptions. Using defaults.")
+            options = PaddleOCROptions(
+                languages=options.languages,
+                min_confidence=options.min_confidence,
+                device=options.device,
+                extra_args=options.extra_args,
+            )
 
         reader = self._get_reader(options)
         ocr_args = self._prepare_ocr_args(options)
-        
+
         # Helper function to process one image
         def process_one(img):
             try:
                 img_array_bgr = self._pil_to_bgr(img)
                 raw_results = reader.ocr(img_array_bgr, **ocr_args)
-                
+
                 page_results = []
                 if raw_results and isinstance(raw_results, list) and len(raw_results) > 0:
                     page_results = raw_results[0]
-                
+
                 return self._standardize_results(page_results, options)
             except Exception as e:
                 logger.error(f"Error processing image with PaddleOCR: {e}")