vlm4ocr 0.1.0__py3-none-any.whl → 0.3.0__py3-none-any.whl

vlm4ocr/utils.py

@@ -1,49 +1,218 @@
+import abc
 import os
 import io
 import base64
-from typing import List
-from pdf2image import convert_from_path
+from typing import Dict, List, Tuple
+import json
+import json_repair
+import importlib.util
+from pdf2image import convert_from_path, pdfinfo_from_path
 from PIL import Image
+import asyncio
+import warnings
 
-def get_images_from_pdf(file_path: str) -> List[Image.Image]:
-    """ Extracts images from a PDF file. """
-    try:
-        images = convert_from_path(file_path)
-        if not images:
-            print(f"Warning: No images extracted from PDF: {file_path}")
-        return images
-    except Exception as e:
-        print(f"Error converting PDF to images: {e}")
-        raise ValueError(f"Failed to process PDF file '{os.path.basename(file_path)}'. Ensure poppler is installed and the file is valid.") from e
 
-def get_images_from_tiff(file_path: str) -> List[Image.Image]:
-    """ Extracts images from a TIFF file. """
-    images = []
-    try:
-        img = Image.open(file_path)
-        for i in range(img.n_frames):
-            img.seek(i)
-            images.append(img.copy())
-        if not images:
-            print(f"Warning: No images extracted from TIFF: {file_path}")
-        return images
-    except FileNotFoundError:
-        raise FileNotFoundError(f"TIFF file not found: {file_path}")
-    except Exception as e:
-        print(f"Error processing TIFF file: {e}")
-        raise ValueError(f"Failed to process TIFF file '{os.path.basename(file_path)}'. Ensure the file is a valid TIFF.") from e
+class DataLoader(abc.ABC):
+    def __init__(self, file_path: str):
+        self.file_path = file_path
+        if not os.path.exists(file_path):
+            raise FileNotFoundError(f"File not found: {file_path}")
 
+    @abc.abstractmethod
+    def get_all_pages(self) -> List[Image.Image]:
+        """ 
+        Abstract method to get all pages from the file. 
+        """
+        pass
 
-def get_image_from_file(file_path: str) -> Image.Image:
-    """ Loads a single image file. """
-    try:
-        image = Image.open(file_path)
-        image.load()
-        return image
-    except FileNotFoundError:
-        raise FileNotFoundError(f"Image file not found: {file_path}")
-    except Exception as e:
-        raise ValueError(f"Failed to load image file '{os.path.basename(file_path)}': {e}") from e
+    @abc.abstractmethod
+    def get_page(self, page_index:int) -> Image.Image:
+        """ 
+        Abstract method to get pages from the file. 
+        
+        Parameters:
+        ----------
+        page_index : int
+            Index of the page to retrieve. 
+        """
+        pass
+
+    @abc.abstractmethod
+    async def get_page_async(self, page_index:int) -> Image.Image:
+        """ 
+        Abstract method to get pages from the file. 
+        
+        Parameters:
+        ----------
+        page_index : int
+            Index of the page to retrieve.
+        """
+        pass
+
+    @abc.abstractmethod
+    def get_page_count(self) -> int:
+        """ Returns the number of pages in the PDF file. """
+        pass
+
+
+class PDFDataLoader(DataLoader):
+    def __init__(self, file_path: str):
+        super().__init__(file_path)
+        self.info = pdfinfo_from_path(self.file_path, userpw=None, poppler_path=None)
+
+    def get_all_pages(self) -> List[Image.Image]:
+        """ 
+        Extracts pages from a PDF file. 
+        """
+        try:
+            return convert_from_path(self.file_path)
+
+        except Exception as e:
+            print(f"Error converting PDF to images: {e}")
+            raise ValueError(f"Failed to process PDF file '{os.path.basename(self.file_path)}'. Ensure poppler is installed and the file is valid.") from e
+
+    def get_page(self, page_index:int) -> Image.Image:
+        """
+        Extracts a page from a PDF file.
+
+        Parameters:
+        ----------
+        page_index : int
+            Index of the page to retrieve.
+        """
+        try:
+            return convert_from_path(self.file_path, first_page=page_index + 1, last_page=page_index + 1)[0]
+        except Exception as e:
+            print(f"Error converting PDF to images: {e}")
+            raise ValueError(f"Failed to process PDF file '{os.path.basename(self.file_path)}'. Ensure poppler is installed and the file is valid.") from e
+
+
+    async def get_page_async(self, page_index:int) -> Image.Image:
+        """ 
+        Asynchronously extracts a page from a PDF file. 
+        
+        Parameters:
+        ----------
+        page_index : int
+            Index of the page to retrieve. 
+        """
+        loop = asyncio.get_running_loop()
+        return await loop.run_in_executor(None, self.get_page, page_index)
+
+
+    def get_page_count(self) -> int:
+        """ Returns the number of pages in the PDF file. """
+        return self.info['Pages'] if 'Pages' in self.info else 0
+    
+
+class TIFFDataLoader(DataLoader):
+    def __init__(self, file_path: str):
+        super().__init__(file_path)
+
+    def get_all_pages(self) -> List[Image.Image]:
+        """ 
+        Extracts images from a TIFF file. 
+        """
+        try:
+            img = Image.open(self.file_path)
+            images = []
+            for i in range(img.n_frames):
+                img.seek(i)
+                images.append(img.copy())
+            return images
+        except Exception as e:
+            print(f"Error extracting images from TIFF: {e}")
+            raise ValueError(f"Failed to process TIFF file '{os.path.basename(self.file_path)}'. Ensure the file is valid.") from e
+        
+
+    def get_page(self, page_index:int) -> Image.Image:
+        """
+        Extracts a page from a TIFF file.
+
+        Parameters:
+        ----------
+        page_index : int
+            Index of the page to retrieve. 
+        """
+        try:
+            img = Image.open(self.file_path)
+            img.seek(page_index)
+            return img.copy()
+        except IndexError:
+            raise ValueError(f"Page index {page_index} out of range for TIFF file '{os.path.basename(self.file_path)}'.") from None
+        except Exception as e:
+            print(f"Error extracting page {page_index} from TIFF: {e}")
+            raise ValueError(f"Failed to process TIFF file '{os.path.basename(self.file_path)}'. Ensure the file is valid.") from e
+
+    async def get_page_async(self, page_index:int) -> Image.Image:
+        """ 
+        Asynchronously extracts images from a TIFF file. 
+        
+        Parameters:
+        ----------
+        page_index : int
+            Index of the page to retrieve.
+        """
+        loop = asyncio.get_running_loop()
+        return await loop.run_in_executor(None, self.get_page, page_index)
+
+    def get_page_count(self) -> int:
+        """ Returns the number of images (pages) in the TIFF file. """
+        try:
+            img = Image.open(self.file_path)
+            return img.n_frames 
+        except Exception as e:
+            print(f"Error getting page count from TIFF: {e}")
+            raise ValueError(f"Failed to process TIFF file '{os.path.basename(self.file_path)}'. Ensure the file is valid.") from e
+
+
+class ImageDataLoader(DataLoader):
+    def get_all_pages(self) -> List[Image.Image]:
+        """ 
+        Loads a single image file. 
+        """
+        try:
+            image = Image.open(self.file_path)
+            image.load()
+            return [image]
+        except FileNotFoundError:
+            raise FileNotFoundError(f"Image file not found: {self.file_path}")
+        except Exception as e:
+            raise ValueError(f"Failed to load image file '{os.path.basename(self.file_path)}': {e}") from e
+        
+    def get_page(self, page_index:int) -> Image.Image:
+        """ 
+        Loads a single image file. 
+        
+        Parameters:
+        ----------
+        page_index : int
+            Index of the page to retrieve. Not applicable for single image files.
+        """
+        try:
+            image = Image.open(self.file_path)
+            image.load()
+            return image
+        except FileNotFoundError:
+            raise FileNotFoundError(f"Image file not found: {self.file_path}")
+        except Exception as e:
+            raise ValueError(f"Failed to load image file '{os.path.basename(self.file_path)}': {e}") from e
+        
+    async def get_page_async(self, page_index:int) -> Image.Image:
+        """ 
+        Asynchronously loads a single image file. 
+        
+        Parameters:
+        ----------
+        page_index : int
+            Index of the page to retrieve. Not applicable for single image files.
+        """
+        loop = asyncio.get_running_loop()
+        return await loop.run_in_executor(None, self.get_page, page_index)
+
+    def get_page_count(self) -> int:
+        """ Returns 1 as there is only one image in a single image file. """
+        return 1
 
 
 def image_to_base64(image:Image.Image, format:str="png") -> str:
@@ -61,4 +230,182 @@ def image_to_base64(image:Image.Image, format:str="png") -> str:
     
 def clean_markdown(text:str) -> str:
     cleaned_text = text.replace("```markdown", "").replace("```", "")
-    return cleaned_text
+    return cleaned_text
+
+def _find_dict_strings( text: str) -> List[str]:
+    """
+    Extracts balanced JSON-like dictionaries from a string, even if nested.
+
+    Parameters:
+    -----------
+    text : str
+        the input text containing JSON-like structures.
+
+    Returns : List[str]
+        A list of valid JSON-like strings representing dictionaries.
+    """
+    open_brace = 0
+    start = -1
+    json_objects = []
+
+    for i, char in enumerate(text):
+        if char == '{':
+            if open_brace == 0:
+                # start of a new JSON object
+                start = i 
+            open_brace += 1
+        elif char == '}':
+            open_brace -= 1
+            if open_brace == 0 and start != -1:
+                json_objects.append(text[start:i + 1])
+                start = -1
+
+    return json_objects
+
+def extract_json(gen_text:str) -> List[Dict[str, str]]:
+    """ 
+    This method inputs a generated text and output a JSON of information tuples
+    """
+    out = []
+    dict_str_list = _find_dict_strings(gen_text)
+    for dict_str in dict_str_list:
+        try:
+            dict_obj = json.loads(dict_str)
+            out.append(dict_obj)
+        except json.JSONDecodeError:
+            dict_obj = json_repair.repair_json(dict_str, skip_json_loads=True, return_objects=True)
+            if dict_obj:
+                warnings.warn(f'JSONDecodeError detected, fixed with repair_json:\n{dict_str}', RuntimeWarning)
+                out.append(dict_obj)
+            else:
+                warnings.warn(f'JSONDecodeError could not be fixed:\n{dict_str}', RuntimeWarning)
+    return out
+
+def get_default_page_delimiter(output_mode:str) -> str:
+    """ 
+    Returns the default page delimiter based on the environment variable.
+
+    Parameters:
+    ----------
+    output_mode : str
+        The output mode, which can be "markdown", "HTML", or "text".
+    
+    Returns:
+    -------
+    str
+        The default page delimiter.
+    """
+    if output_mode not in ["markdown", "HTML", "text", "JSON"]:
+        raise ValueError("output_mode must be 'markdown', 'HTML', 'text', or 'JSON'")
+    
+    if output_mode == "markdown":
+        return "\n\n---\n\n"
+    elif output_mode == "HTML":
+        return "<br><br>"
+    elif output_mode == "text":
+        return "\n\n---\n\n"
+    elif output_mode == "JSON":
+        return "\n\n---\n\n"
+
+
+class ImageProcessor:
+    def __init__(self):
+        self.has_tesseract = importlib.util.find_spec("pytesseract") is not None
+
+    def rotate_correction(self, image: Image.Image) -> Tuple[Image.Image, int]:
+        """ 
+        This method use Tesseract OSD to correct the rotation of the image. 
+        
+        Parameters:
+        ----------
+        image : Image.Image
+            The image to be corrected.
+
+        Returns:
+        -------
+        Image.Image
+            The corrected image.
+        int
+            The rotation angle in degrees.
+        """
+        if importlib.util.find_spec("pytesseract") is None:
+            raise ImportError("pytesseract is not installed. Please install it to use this feature.")
+        
+        import pytesseract
+
+        try:
+            osd = pytesseract.image_to_osd(image, output_type=pytesseract.Output.DICT)
+            rotation_angle = osd['rotate']
+            if rotation_angle != 0:
+                return image.rotate(rotation_angle, expand=True), rotation_angle
+            
+            return image, 0
+        except Exception as e:
+            print(f"Error correcting image rotation: {e}")
+            raise ValueError(f"Failed to correct image rotation: {e}") from e
+
+    async def rotate_correction_async(self, image: Image.Image) -> Tuple[Image.Image, int]:
+        """ 
+        Asynchronous version of rotate_correction method.
+
+        Parameters:
+        ----------
+        image : Image.Image
+            The image to be corrected.
+
+        Returns:
+        -------
+        Image.Image
+            The corrected image.
+        int
+            The rotation angle in degrees.
+        """
+        loop = asyncio.get_running_loop()
+        return await loop.run_in_executor(None, self.rotate_correction, image)
+
+    def resize(self, image: Image.Image, max_dimension_pixels:int=4000) -> Tuple[Image.Image, bool]:
+        """ 
+        Resizes the image to fit within the specified maximum dimension while maintaining aspect ratio.
+        
+        Parameters:
+        ----------
+        max_dimension_pixels : int
+            The maximum dimension (width or height) in pixels.
+
+        Returns:
+        -------
+        Image.Image
+            The resized image.
+        bool
+            True if the image was resized, False otherwise.
+        """
+        width, height = image.size
+        if width > max_dimension_pixels or height > max_dimension_pixels:
+            if width > height:
+                new_width = max_dimension_pixels
+                new_height = int((max_dimension_pixels / width) * height)
+            else:
+                new_height = max_dimension_pixels
+                new_width = int((max_dimension_pixels / height) * width)
+            return image.resize((new_width, new_height), resample=Image.Resampling.LANCZOS), True  # Resizing was done
+        
+        return image, False  # No resizing needed
+
+    async def resize_async(self, image: Image.Image, max_dimension_pixels:int=4000) -> Tuple[Image.Image, bool]:
+        """ 
+        Asynchronous version of resize method.
+        
+        Parameters:
+        ----------
+        max_dimension_pixels : int
+            The maximum dimension (width or height) in pixels.
+
+        Returns:
+        -------
+        Image.Image
+            The resized image.
+        bool
+            True if the image was resized, False otherwise.
+        """
+        loop = asyncio.get_running_loop()
+        return await loop.run_in_executor(None, self.resize, image, max_dimension_pixels)