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

natural_pdf/ocr/engine.py

@@ -11,35 +11,153 @@ from .ocr_options import BaseOCROptions
 logger = logging.getLogger(__name__)
 
 
+class TextRegion:
+    """Standard representation of an OCR text region."""
+
+    def __init__(
+        self,
+        bbox: Tuple[float, float, float, float],
+        text: str,
+        confidence: float,
+        source: str = "ocr",
+    ):
+        """
+        Initialize a text region.
+
+        Args:
+            bbox: Tuple of (x0, y0, x1, y1) coordinates
+            text: The recognized text
+            confidence: Confidence score (0.0-1.0)
+            source: Source of the text region (default: "ocr")
+        """
+        self.bbox = bbox
+        self.text = text
+        self.confidence = confidence
+        self.source = source
+
+    @classmethod
+    def from_polygon(cls, polygon: List[List[float]], text: str, confidence: float):
+        """Create from polygon coordinates [[x1,y1], [x2,y2], ...]"""
+        x_coords = [float(point[0]) for point in polygon]
+        y_coords = [float(point[1]) for point in polygon]
+        bbox = (min(x_coords), min(y_coords), max(x_coords), max(y_coords))
+        return cls(bbox, text, confidence)
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to dictionary representation for compatibility."""
+        return {
+            "bbox": self.bbox,
+            "text": self.text,
+            "confidence": self.confidence,
+            "source": self.source,
+        }
+
+
 class OCREngine(ABC):
     """Abstract Base Class for OCR engines."""
 
+    # Default values as class constants
+    DEFAULT_MIN_CONFIDENCE = 0.2
+    DEFAULT_LANGUAGES = ["en"]
+    DEFAULT_DEVICE = "cpu"
+
     def __init__(self):
         """Initializes the base OCR engine."""
         self.logger = logging.getLogger(f"{__name__}.{self.__class__.__name__}")
         self.logger.info(f"Initializing {self.__class__.__name__}")
+        self._model = None
+        self._initialized = False
         self._reader_cache = {}  # Cache for initialized models/readers
 
-    @abstractmethod
     def process_image(
         self,
-        images: Union[Image.Image, List[Image.Image]],  # Accept single or list
-        options: BaseOCROptions,
-    ) -> Union[List[Dict[str, Any]], List[List[Dict[str, Any]]]]:  # Return single or list of lists
+        images: Union[Image.Image, List[Image.Image]],
+        languages: Optional[List[str]] = None,
+        min_confidence: Optional[float] = None,
+        device: Optional[str] = None,
+        detect_only: bool = False,
+        options: Optional[BaseOCROptions] = None,
+    ) -> Union[List[Dict[str, Any]], List[List[Dict[str, Any]]]]:
         """
-        Processes a single image or a batch of images using the specific engine and options.
+        Process a single image or batch of images with OCR.
 
         Args:
-            images: A single PIL Image or a list of PIL Images.
-            options: An instance of a dataclass inheriting from BaseOCROptions
-                     containing configuration for this run.
+            images: A single PIL Image or a list of PIL Images
+            languages: List of languages to use (default: ['en'])
+            min_confidence: Minimum confidence threshold (default: 0.2)
+            device: Device to use for processing (default: 'cpu')
+            detect_only: Whether to only detect text regions without recognition
+            options: Engine-specific options
 
         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.
-                                          An empty list indicates failure for that image.
+            For a single image: List of text region dictionaries
+            For a batch: List of lists of text region dictionaries
         """
+        # Convert single image to batch format
+        single_image = not isinstance(images, list)
+        image_batch = [images] if single_image else images
+
+        # Use default values where parameters are not provided
+        effective_languages = languages or self.DEFAULT_LANGUAGES
+        effective_confidence = (
+            min_confidence if min_confidence is not None else self.DEFAULT_MIN_CONFIDENCE
+        )
+        effective_device = device or self.DEFAULT_DEVICE
+
+        # Ensure the model is initialized
+        self._ensure_initialized(effective_languages, effective_device, options)
+
+        # Process each image in the batch
+        results = []
+        for img in image_batch:
+            # Preprocess the image for the specific engine
+            processed_img = self._preprocess_image(img)
+
+            # Process the image with the engine-specific implementation
+            raw_results = self._process_single_image(processed_img, detect_only, options)
+
+            # Convert results to standardized format
+            text_regions = self._standardize_results(raw_results, effective_confidence, detect_only)
+
+            # Convert TextRegion objects to dictionaries for backward compatibility
+            region_dicts = [region.to_dict() for region in text_regions]
+            results.append(region_dicts)
+
+        # Return results in the appropriate format
+        return results[0] if single_image else results
+
+    def _ensure_initialized(
+        self, languages: List[str], device: str, options: Optional[BaseOCROptions]
+    ):
+        """Ensure the model is initialized with the correct parameters."""
+        if not self._initialized:
+            self._initialize_model(languages, device, options)
+            self._initialized = True
+
+    @abstractmethod
+    def _initialize_model(
+        self, languages: List[str], device: str, options: Optional[BaseOCROptions]
+    ):
+        """Initialize the OCR model with the given parameters."""
+        raise NotImplementedError("Subclasses must implement this method")
+
+    @abstractmethod
+    def _preprocess_image(self, image: Image.Image) -> Any:
+        """Convert PIL Image to engine-specific format."""
+        raise NotImplementedError("Subclasses must implement this method")
+
+    @abstractmethod
+    def _process_single_image(
+        self, image: Any, detect_only: bool, options: Optional[BaseOCROptions]
+    ) -> Any:
+        """Process a single image with the initialized model."""
+        raise NotImplementedError("Subclasses must implement this method")
+
+    @abstractmethod
+    def _standardize_results(
+        self, raw_results: Any, min_confidence: float, detect_only: bool
+    ) -> List[TextRegion]:
+        """Convert engine-specific results to standardized TextRegion objects."""
         raise NotImplementedError("Subclasses must implement this method")
 
     @abstractmethod
@@ -63,48 +181,41 @@ class OCREngine(ABC):
         Returns:
             A string cache key.
         """
-        # Basic key includes languages and device
-        lang_key = "-".join(sorted(options.languages))
-        device_key = str(options.device).lower()
+        lang_key = "-".join(sorted(getattr(options, "languages", self.DEFAULT_LANGUAGES)))
+        device_key = str(getattr(options, "device", self.DEFAULT_DEVICE)).lower()
         return f"{self.__class__.__name__}_{lang_key}_{device_key}"
 
-    def _standardize_bbox(self, bbox: Any) -> Optional[Tuple[float, float, float, float]]:
-        """
-        Helper to standardize bounding boxes to (x0, y0, x1, y1) format.
-
-        Args:
-            bbox: The bounding box in the engine's native format.
-                  Expected formats:
-                  - List/Tuple of 4 numbers: (x0, y0, x1, y1)
-                  - List of points: [[x1,y1],[x2,y2],[x3,y3],[x4,y4]] (polygon)
-
-        Returns:
-            Tuple[float, float, float, float] or None if conversion fails.
-        """
-        try:
-            if (
-                isinstance(bbox, (list, tuple))
-                and len(bbox) == 4
-                and all(isinstance(n, (int, float)) for n in bbox)
-            ):
-                # Already in (x0, y0, x1, y1) format (or similar)
+    def _standardize_bbox(self, bbox: Any) -> Tuple[float, float, float, float]:
+        """Standardizes bounding boxes to (x0, y0, x1, y1) format. Raises ValueError if standardization fails."""
+        # Check if it's already in the correct tuple/list format
+        if (
+            isinstance(bbox, (list, tuple))
+            and len(bbox) == 4
+            and all(isinstance(n, (int, float)) for n in bbox)
+        ):
+            try:
                 return tuple(float(c) for c in bbox[:4])
-            elif (
-                isinstance(bbox, (list, tuple))
-                and len(bbox) > 0
-                and isinstance(bbox[0], (list, tuple))
-            ):
-                # Polygon format [[x1,y1],[x2,y2],...]
+            except (ValueError, TypeError) as e:
+                raise ValueError(f"Invalid number format in bbox: {bbox}") from e
+
+        # Check if it's in polygon format [[x1,y1],[x2,y2],...]
+        elif (
+            isinstance(bbox, (list, tuple))
+            and len(bbox) > 0
+            and isinstance(bbox[0], (list, tuple))
+            and len(bbox[0]) == 2  # Ensure points are pairs
+        ):
+            try:
                 x_coords = [float(point[0]) for point in bbox]
                 y_coords = [float(point[1]) for point in bbox]
-                x0 = min(x_coords)
-                y0 = min(y_coords)
-                x1 = max(x_coords)
-                y1 = max(y_coords)
-                return (x0, y0, x1, y1)
-        except Exception as e:
-            self.logger.warning(f"Could not standardize bounding box: {bbox}. Error: {e}")
-        return None
+                if not x_coords or not y_coords:  # Handle empty polygon case
+                    raise ValueError("Empty polygon provided")
+                return (min(x_coords), min(y_coords), max(x_coords), max(y_coords))
+            except (ValueError, TypeError, IndexError) as e:
+                raise ValueError(f"Invalid polygon format or values: {bbox}") from e
+
+        # If it's neither format, raise an error
+        raise ValueError(f"Could not standardize bounding box from unexpected format: {bbox}")
 
     def __del__(self):
         """Cleanup resources when the engine is deleted."""

natural_pdf/ocr/engine_easyocr.py

@@ -1,13 +1,12 @@
 # ocr_engine_easyocr.py
 import importlib.util
-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
 
-from .engine import OCREngine
+from .engine import OCREngine, TextRegion
 from .ocr_options import BaseOCROptions, EasyOCROptions
 
 logger = logging.getLogger(__name__)
@@ -18,178 +17,200 @@ class EasyOCREngine(OCREngine):
 
     def __init__(self):
         super().__init__()
-        self._easyocr = None  # Lazy load easyocr module
-
-    def _lazy_import_easyocr(self):
-        """Imports easyocr only when needed."""
-        if self._easyocr is None:
-            if not self.is_available():
-                raise ImportError("EasyOCR is not installed or available.")
-            try:
-                import easyocr
-
-                self._easyocr = easyocr
-                logger.info("EasyOCR module imported successfully.")
-            except ImportError as e:
-                logger.error(f"Failed to import EasyOCR: {e}")
-                raise
-        return self._easyocr
+        # No longer need _easyocr attribute
+        # self._easyocr = None
 
     def is_available(self) -> bool:
         """Check if EasyOCR is installed."""
         return importlib.util.find_spec("easyocr") is not None
 
-    def _get_cache_key(self, options: EasyOCROptions) -> str:
-        """Generate a more specific cache key for EasyOCR."""
-        base_key = super()._get_cache_key(options)
-        recog_key = options.recog_network
-        detect_key = options.detect_network
-        quantize_key = str(options.quantize)
-        return f"{base_key}_{recog_key}_{detect_key}_{quantize_key}"
-
-    def _get_reader(self, options: EasyOCROptions):
-        """Get or initialize an EasyOCR reader based on options."""
-        cache_key = self._get_cache_key(options)
-        if cache_key in self._reader_cache:
-            logger.debug(f"Using cached EasyOCR reader for key: {cache_key}")
-            return self._reader_cache[cache_key]
-
-        logger.info(f"Creating new EasyOCR reader for key: {cache_key}")
-        easyocr = self._lazy_import_easyocr()
-
-        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()
-        )
-
-        for field_name, param in constructor_sig.parameters.items():
-            if field_name in ["self", "lang_list", "gpu"]:
-                continue
-            if hasattr(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]
-
-        logger.debug(f"EasyOCR Reader constructor args: {constructor_args}")
+    def _initialize_model(
+        self, languages: List[str], device: str, options: Optional[BaseOCROptions]
+    ):
+        """Initialize the EasyOCR model."""
+        # Import directly here
+        try:
+            import easyocr
+
+            self.logger.info("EasyOCR module imported successfully.")
+        except ImportError as e:
+            self.logger.error(f"Failed to import EasyOCR: {e}")
+            raise
+
+        # Cast to EasyOCROptions if possible, otherwise use default
+        easy_options = options if isinstance(options, EasyOCROptions) else EasyOCROptions()
+
+        # Prepare constructor arguments
+        use_gpu = "cuda" in device.lower() or "mps" in device.lower()
+
+        constructor_args = {
+            "lang_list": languages,
+            "gpu": use_gpu,
+            # Explicitly map relevant options
+            "model_storage_directory": easy_options.model_storage_directory,
+            "user_network_directory": easy_options.user_network_directory,
+            "recog_network": easy_options.recog_network,
+            "detect_network": easy_options.detect_network,
+            "download_enabled": easy_options.download_enabled,
+            "detector": easy_options.detector,
+            "recognizer": easy_options.recognizer,
+            "verbose": easy_options.verbose,
+            "quantize": easy_options.quantize,
+            "cudnn_benchmark": easy_options.cudnn_benchmark,
+        }
+
+        # Filter out None values, as EasyOCR expects non-None or default behaviour
+        constructor_args = {k: v for k, v in constructor_args.items() if v is not None}
+
+        self.logger.debug(f"EasyOCR Reader constructor args: {constructor_args}")
+
+        # Create the reader
         try:
-            reader = easyocr.Reader(**constructor_args)
-            self._reader_cache[cache_key] = reader
-            logger.info("EasyOCR reader created successfully.")
-            return reader
+            self._model = easyocr.Reader(**constructor_args)
+            self.logger.info("EasyOCR reader created successfully")
         except Exception as e:
-            logger.error(f"Failed to create EasyOCR reader: {e}", exc_info=True)
+            self.logger.error(f"Failed to create EasyOCR reader: {e}")
             raise
 
-    def _prepare_readtext_args(self, options: EasyOCROptions, reader) -> Dict[str, Any]:
-        """Helper to prepare arguments for the readtext method."""
-        readtext_sig = inspect.signature(reader.readtext)
+    def _preprocess_image(self, image: Image.Image) -> np.ndarray:
+        """Convert PIL Image to numpy array for EasyOCR."""
+        return np.array(image)
+
+    def _process_single_image(
+        self, image: np.ndarray, detect_only: bool, options: Optional[EasyOCROptions]
+    ) -> Any:
+        """Process a single image with EasyOCR."""
+        if self._model is None:
+            raise RuntimeError("EasyOCR model not initialized")
+
+        # Cast options to proper type if provided
+        easy_options = options if isinstance(options, EasyOCROptions) else EasyOCROptions()
+
+        # Prepare readtext arguments (only needed if not detect_only)
         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]
-        logger.debug(f"EasyOCR readtext args: {readtext_args}")
-        return readtext_args
+        if not detect_only:
+            for param in [
+                "detail",
+                "paragraph",
+                "min_size",
+                "contrast_ths",
+                "adjust_contrast",
+                "filter_ths",
+                "text_threshold",
+                "low_text",
+                "link_threshold",
+                "canvas_size",
+                "mag_ratio",
+                "slope_ths",
+                "ycenter_ths",
+                "height_ths",
+                "width_ths",
+                "y_ths",
+                "x_ths",
+                "add_margin",
+                "output_format",
+            ]:
+                if hasattr(easy_options, param):
+                    val = getattr(easy_options, param)
+                    if val is not None:
+                        readtext_args[param] = val
+
+        # Process differently based on detect_only flag
+        if detect_only:
+            # Returns tuple (horizontal_list, free_list)
+            # horizontal_list is a list containing one item: the list of boxes
+            # Each box is [[x1, y1], [x2, y1], [x2, y2], [x1, y2]]
+            bboxes_tuple = self._model.detect(
+                image, **readtext_args
+            )  # Pass args here too? Check EasyOCR docs if needed.
+            if (
+                bboxes_tuple
+                and isinstance(bboxes_tuple, tuple)
+                and len(bboxes_tuple) > 0
+                and isinstance(bboxes_tuple[0], list)
+            ):
+                return bboxes_tuple[0]  # Return the list of polygons directly
+            else:
+                self.logger.warning(f"EasyOCR detect returned unexpected format: {bboxes_tuple}")
+                return []  # Return empty list on unexpected format
+        else:
+            return self._model.readtext(image, **readtext_args)
 
     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
-
+        self, raw_results: Any, min_confidence: float, detect_only: bool
+    ) -> List[TextRegion]:
+        """Convert EasyOCR results to standardized TextRegion objects."""
+        standardized_regions = []
+
+        if detect_only:
+            # In detect_only mode, raw_results is already a list of bounding boxes
+            # Each bbox is in [x_min, x_max, y_min, y_max] format
+            if isinstance(raw_results, list):
+                for detection in raw_results:
+                    try:
+                        if isinstance(detection, (list, tuple)) and len(detection) == 4:
+                            x_min, x_max, y_min, y_max = detection
+                            # Convert to standardized (x0, y0, x1, y1) format
+                            try:
+                                bbox = (float(x_min), float(y_min), float(x_max), float(y_max))
+                                standardized_regions.append(
+                                    TextRegion(bbox, text=None, confidence=None)
+                                )
+                            except (ValueError, TypeError) as e:
+                                raise ValueError(
+                                    f"Invalid number format in EasyOCR detect bbox: {detection}"
+                                ) from e
+                        else:
+                            raise ValueError(f"Invalid detection format from EasyOCR: {detection}")
+                    except ValueError as e:
+                        # Re-raise any value errors from standardization or format checks
+                        raise e
+                    except Exception as e:
+                        # Catch other potential processing errors
+                        raise ValueError(
+                            f"Error processing EasyOCR detection item: {detection}"
+                        ) from e
+            else:
+                raise ValueError(
+                    f"Expected list of bounding boxes in detect_only mode, got: {raw_results}"
+                )
+
+            return standardized_regions
+
+        # Full OCR mode (readtext results)
         for detection in raw_results:
             try:
-                if (
-                    options.detail == 1
-                    and isinstance(detection, (list, tuple))
-                    and len(detection) >= 3
-                ):
-                    bbox_raw = detection[0]
+                # Detail mode (list/tuple result)
+                if isinstance(detection, (list, tuple)) and len(detection) >= 3:
+                    bbox_raw = detection[0]  # This is usually a polygon [[x1,y1],...]
                     text = str(detection[1])
                     confidence = float(detection[2])
 
                     if confidence >= min_confidence:
-                        bbox = self._standardize_bbox(bbox_raw)
-                        if bbox:
-                            standardized_results.append(
-                                {
-                                    "bbox": bbox,
-                                    "text": text,
-                                    "confidence": confidence,
-                                    "source": "ocr",
-                                }
-                            )
-                        else:
-                            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"}
+                        try:
+                            # Use the standard helper for polygons
+                            bbox = self._standardize_bbox(bbox_raw)
+                            standardized_regions.append(TextRegion(bbox, text, confidence))
+                        except ValueError as e:
+                            raise ValueError(
+                                f"Could not standardize bounding box from EasyOCR readtext: {bbox_raw}"
+                            ) from e
+
+                # Simple mode (string result)
+                elif isinstance(detection, str):
+                    if 0.0 >= min_confidence:  # Always include if min_confidence is 0
+                        standardized_regions.append(TextRegion((0, 0, 0, 0), detection, 1.0))
+                else:
+                    # Handle unexpected format in OCR mode
+                    raise ValueError(
+                        f"Invalid OCR detection format from EasyOCR readtext: {detection}"
                     )
-            except (IndexError, ValueError, TypeError) as e:
-                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
-    ) -> 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
-            )
-
-        reader = self._get_reader(options)
-        readtext_args = self._prepare_readtext_args(options, reader)
-
-        # --- Handle single image or batch ---
-        if isinstance(images, list):
-            # --- Batch Processing (Iterative for EasyOCR) ---
-            all_results = []
-            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
-                img_array = np.array(img)
-                try:
-                    logger.debug(f"Processing image {i+1}/{len(images)} in batch.")
-                    raw_results = reader.readtext(img_array, **readtext_args)
-                    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.info(f"Finished processing batch with EasyOCR.")
-            return all_results  # Return List[List[Dict]]
-
-        elif isinstance(images, Image.Image):
-            # --- Single Image Processing ---
-            logger.info("Processing single image with EasyOCR...")
-            img_array = np.array(images)
-            try:
-                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]
+
+            except ValueError as e:
+                # Re-raise any value errors from standardization or format checks
+                raise e
             except Exception as e:
-                logger.error(f"Error processing single image with EasyOCR: {e}", exc_info=True)
-                return []  # Return empty list on failure
-        else:
-            raise TypeError("Input 'images' must be a PIL Image or a list of PIL Images.")
+                # Catch other potential processing errors
+                raise ValueError(f"Error processing EasyOCR detection item: {detection}") from e
+
+        return standardized_regions