natural-pdf 0.1.4__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

@@ -1,191 +1,189 @@
 # ocr_manager.py
+import copy  # For deep copying options
 import logging
-from typing import Dict, List, Any, Optional, Union, Type
+from typing import Any, Dict, List, Optional, Type, Union
+
 from PIL import Image
-import copy # For deep copying options
 
 # Import engine classes and options
 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 (
-    BaseOCROptions, EasyOCROptions, PaddleOCROptions, SuryaOCROptions, 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__)
 
+
 class OCRManager:
     """Manages OCR engine selection, configuration, and execution."""
 
     # Registry mapping engine names to classes and default options
     ENGINE_REGISTRY: Dict[str, Dict[str, Any]] = {
-        'easyocr': {'class': EasyOCREngine, 'options_class': EasyOCROptions},
-        'paddle': {'class': PaddleOCREngine, 'options_class': PaddleOCROptions},
-        'surya': {'class': SuryaOCREngine, 'options_class': SuryaOCROptions}, # <-- Add Surya
+        "easyocr": {"class": EasyOCREngine, "options_class": EasyOCROptions},
+        "paddle": {"class": PaddleOCREngine, "options_class": PaddleOCROptions},
+        "surya": {"class": SuryaOCREngine, "options_class": SuryaOCROptions},  # <-- Add Surya
         # 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
+        self._engine_instances: Dict[str, OCREngine] = {}  # Cache for engine instances
         logger.info("OCRManager initialized.")
 
     def _get_engine_instance(self, engine_name: str) -> OCREngine:
         """Retrieves or creates an instance of the specified OCR engine."""
         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())}")
+            raise ValueError(
+                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
+            engine_class = self.ENGINE_REGISTRY[engine_name]["class"]
+            engine_instance = engine_class()  # Instantiate first
             if not engine_instance.is_available():
-                 # Check availability before storing
-                 raise RuntimeError(f"Engine '{engine_name}' is not available. Please check dependencies.")
-            self._engine_instances[engine_name] = engine_instance # Store if 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
 
         return self._engine_instances[engine_name]
 
     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}")
-
-
-        # --- Get Engine Instance and Process ---
+            raise TypeError("Input 'images' must be a PIL Image or a list of PIL Images.")
+
+        # --- 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 --- 
         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:
-                 # Ensure results is a list before trying to get lengths
-                 if isinstance(results, list):
-                     num_results_per_image = [len(res_list) if isinstance(res_list, list) else -1 for res_list in results] # Handle potential errors returning non-lists
-                     logger.info(f"Processing complete. Found results per image: {num_results_per_image}")
-                 else:
-                     logger.error(f"Processing complete but received unexpected result type for batch: {type(results)}")
+                # Ensure results is a list before trying to get lengths
+                if isinstance(results, list):
+                    num_results_per_image = [
+                        len(res_list) if isinstance(res_list, list) else -1 for res_list in results
+                    ]  # Handle potential errors returning non-lists
+                    logger.info(
+                        f"Processing complete. Found results per image: {num_results_per_image}"
+                    )
+                else:
+                    logger.error(
+                        f"Processing complete but received unexpected result type for batch: {type(results)}"
+                    )
             else:
-                 # Ensure results is a list
-                 if isinstance(results, list):
-                      logger.info(f"Processing complete. Found {len(results)} results.")
-                 else:
-                      logger.error(f"Processing complete but received unexpected result type for single image: {type(results)}")
-            return results # Return type matches input type due to engine logic
+                # Ensure results is a list
+                if isinstance(results, list):
+                    logger.info(f"Processing complete. Found {len(results)} results.")
+                else:
+                    logger.error(
+                        f"Processing complete but received unexpected result type for single image: {type(results)}"
+                    )
+            return results  # Return type matches input type due to engine logic
 
         except (ImportError, RuntimeError, ValueError, TypeError) as e:
-             logger.error(f"OCR processing failed for engine '{selected_engine_name}': {e}", exc_info=True)
-             raise # Re-raise expected errors
+            logger.error(
+                f"OCR processing failed for engine '{selected_engine_name}': {e}", exc_info=True
+            )
+            raise  # Re-raise expected errors
         except Exception as e:
-             logger.error(f"An unexpected error occurred during OCR processing: {e}", exc_info=True)
-             raise # Re-raise unexpected errors
-
+            logger.error(f"An unexpected error occurred during OCR processing: {e}", exc_info=True)
+            raise  # Re-raise unexpected errors
 
     def get_available_engines(self) -> List[str]:
         """Returns a list of registered engine names that are currently available."""
         available = []
         for name, registry_entry in self.ENGINE_REGISTRY.items():
-             try:
-                 # Temporarily instantiate to check availability without caching
-                 engine_class = registry_entry['class']
-                 if engine_class().is_available():
-                     available.append(name)
-             except Exception as e:
-                 logger.debug(f"Engine '{name}' check failed: {e}") # Log check failures at debug level
-                 pass # Ignore engines that fail to instantiate or check
+            try:
+                # Temporarily instantiate to check availability without caching
+                engine_class = registry_entry["class"]
+                if engine_class().is_available():
+                    available.append(name)
+            except Exception as e:
+                logger.debug(
+                    f"Engine '{name}' check failed: {e}"
+                )  # Log check failures at debug level
+                pass  # Ignore engines that fail to instantiate or check
         return available
-

natural_pdf/ocr/ocr_options.py

@@ -1,30 +1,29 @@
 # ocr_options.py
 import logging
 from dataclasses import dataclass, field
-from typing import List, Optional, Dict, Any, Tuple, Union
+from typing import Any, Dict, List, Optional, Tuple, Union
 
 # Configure logging
 # logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
 # logger = logging.getLogger(__name__)
 # Assume logger is configured elsewhere or remove if not needed globally
 
+
 # --- Base Options ---
 @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)
 
+
 # --- EasyOCR Specific Options ---
 @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'
-    detect_network: str = 'craft'
+    recog_network: str = "english_g2"
+    detect_network: str = "craft"
     download_enabled: bool = True
     detector: bool = True
     recognizer: bool = True
@@ -32,7 +31,7 @@ class EasyOCROptions(BaseOCROptions):
     quantize: bool = True
     cudnn_benchmark: bool = False
     detail: int = 1
-    decoder: str = 'greedy'
+    decoder: str = "greedy"
     beamWidth: int = 5
     batch_size: int = 1
     workers: int = 0
@@ -55,7 +54,7 @@ class EasyOCROptions(BaseOCROptions):
     y_ths: float = 0.5
     x_ths: float = 1.0
     add_margin: float = 0.1
-    output_format: str = 'standard'
+    output_format: str = "standard"
 
     # def __post_init__(self):
     #     logger.debug(f"Initialized EasyOCROptions: {self}")
@@ -71,7 +70,7 @@ class PaddleOCROptions(BaseOCROptions):
     ir_optim: bool = True
     use_tensorrt: bool = False
     min_subgraph_size: int = 15
-    precision: str = 'fp32'
+    precision: str = "fp32"
     enable_mkldnn: bool = False
     cpu_threads: int = 10
     use_fp16: bool = False
@@ -90,25 +89,22 @@ 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
 
 
 # --- Union type for type hinting ---
 OCROptions = Union[EasyOCROptions, PaddleOCROptions, SuryaOCROptions, BaseOCROptions]
-

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

natural_pdf/qa/__init__.py

@@ -1,3 +1,3 @@
 from natural_pdf.qa.document_qa import DocumentQA, get_qa_engine
 
-__all__ = ["DocumentQA", "get_qa_engine"]
+__all__ = ["DocumentQA", "get_qa_engine"]