Scanipy stands for "scan it with Python"—it's your smart Python library for scanning and parsing complex PDF files like books, reports, articles, and academic papers. Utilizing cutting-edge Deep Learning algorithms, Scanipy transforms your PDFs into a treasure trove of extractable information: tables, images, equations, and text.
from pyarxiv import query, download_entries
from pyarxiv.arxiv_categories import ArxivCategory, arxiv_category_map
import requests
import os
import pandas as pd

def choose_topic(topic, max_results):
    # Perform the arXiv API query with the 'max_results' parameter
    entries = query(title=topic, max_results=max_results)
    
    # Pull title, author, date, and link to PDF of paper from "entries"
    # and put each in its own list

    titles = map(lambda x: x['title'], entries)
    authors = map(lambda x: x['author'], entries)
    updated = map(lambda x: x['updated'], entries)
    links = map(lambda x: x['link'], entries)

    # Create an empty dataframe called "papers"
    papers = pd.DataFrame()

    # Insert columns into "papers" from the previously created lists

    papers['Title'] = pd.Series(titles)
    papers['Author'] = pd.Series(authors)
    papers['Updated'] = pd.Series(updated)
    papers['Link'] = pd.Series(links)

    # Slice HH:MM:SS off of each row in the date column

    papers['Updated'] = papers['Updated'].str.slice(stop = 10)

    # Reformat URL string to take the user to the PDF of the paper

    papers['Link'] = papers['Link'].str.replace("abs", "pdf", case = True)

    # Sort the dataframe in descending order by date

    papers = papers.sort_values(by = 'Updated', ascending = False).reset_index(drop = True).head(max_results)
    
    # Calculate the number of rows in the "papers" DataFrame
    num_rows = len(papers)
    print("Number of papers:", num_rows)
    
    return papers

dataset/access_arxiv.py
SVJLucas/Scanipy
import requests
import os
import tarfile

def identify_main_tex_file(tar_ref):
    main_tex_file = None

    for file_info in tar_ref:
        if file_info.name.endswith(".tex"):
            # Read the content of the TeX file
            with tar_ref.extractfile(file_info) as tex_file:
                tex_content = tex_file.read().decode("utf-8")

            # Check if the TeX file contains criteria indicating it's the main file
            if is_main_tex_file(tex_content):
                main_tex_file = file_info
                break

    # If no main TeX file is found, fall back to the first TeX file encountered
    if main_tex_file is None and len(tar_ref.getnames()) > 0:
        main_tex_file = tar_ref.getnames()[0]

    return main_tex_file

def is_main_tex_file(tex_content):
    # Implement your logic here to determine if the TeX file is the main file
    # You can check for specific content, keywords, or patterns in the TeX file
    # For example, you might look for \begin{document} or other markers.

    # Replace this example logic with your own criteria
    return "\\begin{document}" in tex_content

def downloading_tar_pdf(df):
    # Iterate through the DataFrame and download PDFs
    for index, row in df.iterrows():
        link = row['Link']

        # Split the link by '/'
        parts = link.split('/')

        # Get the last part of the URL
        arxiv_id = parts[-1]

        # Define the directory where you want to save the files
        save_directory = "arxiv_files"
        save_tar = os.path.join(save_directory, "tars", arxiv_id)
        save_pdf = os.path.join(save_directory, "pdf")

        # Download the TeX file in a tar archive
        tex_tar_url = f"https://arxiv.org/e-print/{arxiv_id}"
        tex_tar_response = requests.get(tex_tar_url)

        # Check if the response is a tar archive
        if tex_tar_response.headers.get('content-type') == 'application/x-eprint-tar':
            # Save the tar file
            # Create the directory if it doesn't exist
            if not os.path.exists(save_tar):
                os.makedirs(save_tar)
            tar_file_path = os.path.join(save_tar, f"{arxiv_id}.tar")
            with open(tar_file_path, 'wb') as tar_file:
                tar_file.write(tex_tar_response.content)

            # Extract tar only if there is a TeX file
            tex_file_found = False
            with tarfile.open(tar_file_path, 'r') as tar_ref:
                main_tex_file = identify_main_tex_file(tar_ref)

                for file_info in tar_ref:
                    if not tex_file_found:
                        if file_info == main_tex_file:
                            tex_file_found = True

                        # Extract the file
                        tar_ref.extract(file_info, save_tar)
                        # Get the extracted file's name
                        extracted_file_name = file_info.name

                        if tex_file_found and extracted_file_name.endswith(".tex"):
                            # Rename the main TeX file to {arxiv_id}.tex
                            new_tex_file_name = f"{arxiv_id}.tex"
                            os.rename(os.path.join(save_tar, extracted_file_name),
                                      os.path.join(save_tar, new_tex_file_name))
                            #print(f"Main TeX file saved as {new_tex_file_name}")
                        else:
                            # Keep the names of auxiliary TeX files
                            os.rename(os.path.join(save_tar, extracted_file_name),
                                      os.path.join(save_tar, extracted_file_name))
                            #print(f"Auxiliary TeX file saved as {extracted_file_name}")

            # Remove the tar file after extraction
            os.remove(tar_file_path)
            #print(f"TAR files saved to {tar_file_path}")

            # If a TeX file was found, download the PDF
            if tex_file_found:
                pdf_url = f"https://arxiv.org/pdf/{arxiv_id}.pdf"
                # Create the directory if it doesn't exist
                if not os.path.exists(save_pdf):
                    os.makedirs(save_pdf)
                pdf_file_path = os.path.join(save_pdf, f"{arxiv_id}.pdf")
                pdf_response = requests.get(pdf_url)
                with open(pdf_file_path, 'wb') as pdf_file:
                    pdf_file.write(pdf_response.content)
                #print(f"PDF file saved to {pdf_file_path}")
        else:
            a=1
            ##print(f"TeX file tar archive not found for ArXiv ID: {arxiv_id}")

dataset/get_tar_pdf.py
SVJLucas/Scanipy
import scanipy

def main():
    parser = scanipy.Parser()
    document = parser.extract("test.pdf")
    document.to_markdown('output')

if __name__ == "__main__":
    main()

run.py
SVJLucas/Scanipy
from scanipy.parser import Parser
scanipy/__init__.py
SVJLucas/Scanipy
import os

# Model's path
ROOT = 'scanipy'
WEIGHTS = os.path.join(ROOT, 'deeplearning', 'weights')

scanipy/constants/paths.py
SVJLucas/Scanipy
from .tablestructureanalyzer import TableStructureAnalyzer
from .tablefinder import TableFinder
from .layoutdetector import LayoutDetector
from .equationfinder import EquationFinder
from .equationtolatex import EquationToLatex
from .textocr import TextOCR
scanipy/deeplearning/models/__init__.py
SVJLucas/Scanipy
import PIL
from typing import Union, List
from cnstd import LayoutAnalyzer
from scanipy.elements import EquationElement

class EquationFinder:
    """
    Locate equations within images using a fine-tuned Layout Analyzer model for equation detection.

    This is a chinese Mathematical Formula Detection (MFD) YoloV7 trained in the IBEM dataset (english) and the CnMFD_Dataset (chinese).
    
    Source: https://github.com/breezedeus/cnstd

    Attributes:
        model (object): The loaded YoloV7 for table detection.

    Example:
        >>> finder = EquationFinder()
        >>> detection_result = finder(image)
    """

    def __init__(self, device):
        """
        Initialize the EquationFinder class by loading the pre-trained model for equation detection.
        Args:
            device: The device on which the Yolov7 model will run.
        """

        # Verify the type of the device argument
        if not isinstance(device, str):
            raise TypeError("Device must be a string.")
          
        # Load the pre-trained Layout Analyzer from CNSTD
        self.model = LayoutAnalyzer(model_name='mfd',
                                    device=device)

    def __str__(self):
        """
        Returns the official string representation of the EquationFinder object.
        
        Returns:
            str: A string that can be used to recreate the EquationFinder object.
        """

        # Create a string representation of the EquationFinder object
        return self.__repr__()

    def __repr__(self):
        """
        Returns the official string representation of the EquationFinder object.
        
        Returns:
            str: A string that can be used to recreate the EquationFinder object.
        """

        # Create the official string representation of the EquationFinder object
        return f"EquationFinder(device='{self.model.device}')"


    def __call__(self, image: PIL.Image, pipeline_step: Union[int, None] = None) -> List[EquationElement]:
        """
        Detects equation elements in the given image and returns them as a list.
        
        Args:
            image (PIL.Image): The image in which to detect layout elements.
            pipeline_step (int or None): An optional integer representing the step in a pipeline. Defaults to None.
        
        Returns:
            empty_elements (List[EquationElement]): A list of detected layout elements, but only with its positions (no extracted content yet).
        """
        # Use the loaded model to detect equations in the given image
        equations = self.model(image)
        empty_elements = []

        # Getting image width and height 
        width, height = image.size
        
        for equation in equations:
          
            # Getting position values
            x_min, y_min = equation['box'][0]
            x_max, y_max = equation['box'][2]

            # Verify if x_min is less than x_max and y_min is less than y_max
            if x_min >= x_max or y_min >= y_max:
                raise ValueError("Invalid coordinates: x_min should be less than x_max and y_min should be less than y_max")


            # Normalizing coordinates to be in range [0,1]
            x_min = x_min/width
            x_max = x_max/width
            y_min = y_min/height
            y_max = y_max/height

            # Verifying if the equation is in the middle of  text
            is_inside_text = bool(equation['type']=='embedding')

            # Creating element
            element = EquationElement(x_min, y_min, x_max, y_max, pipeline_step, is_inside_text=is_inside_text)

            # Adding EquationElement in the list
            empty_elements.append(element)
          
        return empty_elements
        

scanipy/deeplearning/models/equationfinder.py
SVJLucas/Scanipy
import pickle
from pix2tex.cli import LatexOCR


class EquationToLatex:
    """
    Convert images of mathematical formulas to LaTeX code using the pix2tex model.

    This model utilizes a Vision Transformer (ViT) encoder with a ResNet backbone and a Transformer decoder.
    It is specifically designed to convert images of mathematical formulas into their corresponding LaTeX code.

    Source: https://github.com/lukas-blecher/LaTeX-OCR

    Attributes:
        model (object): The loaded pix2tex model for LaTeX OCR.

    Example:
        >>> equation_to_latex = EquationToLatex()
        >>> image = Image.open('equation_image.png')
        >>> latex_code = equation_to_latex(image)
    """

    def __init__(self):
        """
        Initialize the EquationToLatex class by loading the pre-trained pix2tex model.
        """
        # Load the pre-trained pix2tex model from a pickle file
        self.model = LatexOCR()

    def __call__(self, image):
        """
        Convert an image containing a mathematical formula to LaTeX code using the pre-trained pix2tex model.

        Args:
            image (PIL.Image): The image that contains a mathematical formula.

        Returns:
            str: The LaTeX code corresponding to the mathematical formula in the image.
        """
        # Use the loaded model to convert the image into LaTeX code
        return self.model(image)
    def __repr__(self):
        """
        Returns the official string representation of the EquationToLatex object.
        
        Returns:
            str: A string that can be used to recreate the EquationToLatex object.
        """
        return "EquationToLatex()"

    def __str__(self):
        """
        Returns a string representation of the EquationToLatex object, which is the same as its official representation.
        
        Returns:
            str: A string that can be used to recreate the EquationToLatex object.
        """
        return self.__repr__()

scanipy/deeplearning/models/equationtolatex.py
SVJLucas/Scanipy
import PIL
import PIL.Image
from PIL import Image, ImageDraw
import fitz
from typing import Union, List
from layoutparser.models import Detectron2LayoutModel
from scanipy.elements import TextElement,TitleElement,ImageElement,TableElement,EquationElement


class LayoutDetector:
    """
    This class is responsible for parsing the layout of a document using a pre-trained Detectron2 model
    from LayoutParser python package.

    Source: https://github.com/Layout-Parser/layout-parser

    Attributes:
        model: The Detectron2 model for layout detection.
    """

    def __init__(self, device='cpu'):
        """
        Initializes the LayoutDetector class with a given device.
        
        Args:
            device: The device on which the Detectron2 model will run.
        """
        # Verify the type of the device argument
        if not isinstance(device, str):
            raise TypeError("Device must be a string.")
        
        # Initialize the Detectron2 model with specific configurations and label mapping
        self.model = Detectron2LayoutModel(
            'lp://PubLayNet/mask_rcnn_X_101_32x8d_FPN_3x/config',
            extra_config=["MODEL.ROI_HEADS.SCORE_THRESH_TEST", 0.8],
            label_map={0: "Text", 1: "Title", 2: "List", 3: "Table", 4: "Figure"},
            device=device
        )
    def __repr__(self):
        """
        Returns the official string representation of the LayoutDetector object.
        
        Returns:
            str: A string that can be used to recreate the LayoutDetector object.
        """
        return "LayoutDetector()"

    def __str__(self):
        """
        Returns a string representation of the LayoutDetector object, which is the same as its official representation.
        
        Returns:
            str: A string that can be used to recreate the LayoutDetector object.
        """
        return self.__repr__()

    def __call__(self, image: PIL.Image.Image, page_number: Union[int, None] = None) -> List[Union[TextElement, TitleElement, ImageElement, TableElement]]:
        """
        Detects layout elements in the given image and returns them as a list.
        
        Args:
            image: The image in which to detect layout elements.
            page_number: An optional integer representing the page number. Defaults to None.
        
        Returns:
            empty_elements: A list of detected layout elements, but only with its positions (no extracted content yet).
        """
        # Verify the type of the image argument
        if not isinstance(image, PIL.Image.Image):
            raise TypeError("Image must be a PIL.Image object.")
        
        # Verify the type of the page_number argument
        if not (isinstance(page_number, int) or page_number is None):
            raise TypeError("page_number must be an integer or None.")
        
        # Perform layout detection on the image and convert the result to a DataFrame
        layout = self.model.detect(image).to_dataframe()
        
        # Initialize an empty list to store detected elements
        elements = []

        # Getting image width and height 
        width, height = image.size
        
        # Iterate through each block in the DataFrame
        for _, block in layout.iterrows():
            # Extract coordinates of the bounding box
            x_min, y_min, x_max, y_max = block.x_1, block.y_1, block.x_2, block.y_2

            # Normalizing coordinates to be in range [0,1]
            x_min = x_min/width
            x_max = x_max/width
            y_min = y_min/height
            y_max = y_max/height
            
            # Use pattern matching to identify the type of the block and create the corresponding element
            element = None # Make sure that a new element is assigned
            match block.type:
                case "List": #TODO
                    element = TextElement(x_min, y_min, x_max, y_max, page_number=page_number)
                case "Text":
                    element = TextElement(x_min, y_min, x_max, y_max, page_number=page_number)
                case "Title":
                    element = TitleElement(x_min, y_min, x_max, y_max, page_number=page_number)
                case "Figure":
                    element = ImageElement(x_min, y_min, x_max, y_max, page_number=page_number)
                case "Table":
                    element = TableElement(x_min, y_min, x_max, y_max, page_number=page_number)
            # Append the detected element to the list
            elements.append(element)
        
        # Return the list of detected elements
        return elements

    # Debug Function
    def _draw_rectangle(self, image: Image.Image, x1: float, y1: float, x2: float, y2: float, 
                        outline_color: tuple[int, int, int] = (255, 0, 0), thickness: int = 2) -> Image.Image:
        """
        Draw a rectangle on a PIL Image.

        Args:
            image (PIL.Image.Image): The input PIL Image.
            x1 (float): The x-coordinate of the top-left corner.
            y1 (float): The y-coordinate of the top-left corner.
            x2 (float): The x-coordinate of the bottom-right corner.
            y2 (float): The y-coordinate of the bottom-right corner.
            outline_color (tuple): The outline color as an RGB tuple (default is red).
            thickness (int): The outline thickness (default is 2).

        Returns:
            PIL.Image.Image: A new PIL Image with the rectangle drawn.
        """
        # Calculate PDF coordinates based on text element's normalized coordinates
        x1 = int(x1 * image.width)
        y1 = int(y1 * image.height)
        x2 = int(x2 * image.width)
        y2 = int(y2 * image.height)

        # Create a copy of the input image to avoid modifying the original
        image_copy = image.copy()
        
        # Create a drawing context on the copy of the image
        draw = ImageDraw.Draw(image_copy)

        # Draw the rectangle
        draw.rectangle((x1, y1, x2, y2), outline=outline_color, width=thickness)

        return image_copy

scanipy/deeplearning/models/layoutdetector.py
SVJLucas/Scanipy
'''
This file contains the implementation of TableStructureAnalyzer class

Created by:
    Lucas José (lucas.jose.veloso.de.souza@gmail.com)
    Date: September 11, 2023
'''

import pickle
from transformers import pipeline


class TableFinder:
    """
    Locate tables within images using a fine-tuned Table Transformer model for table detection.

    This model is a Table Transformer (DETR) variant trained by Microsoft on the PubTables1M dataset.
    It is designed to detect tables in images.

    Source: https://github.com/microsoft/table-transformer

    Attributes:
        model (object): The loaded fine-tuned Table Transformer model for table detection.

    Example:
        >>> finder = TableFinder()
        >>> detection_result = finder(image)
    """

    def __init__(self):
        """
        Initialize the TableFinder class by loading the pre-trained model for table detection.
        """
        # Load the pre-trained Table Transformer model from a pickle file for table detection
        self.model = pipeline("object-detection", model="microsoft/table-transformer-detection")

    def __call__(self, image):
        """
        Detect tables in a given image using the pre-trained Table Transformer model.

        Args:
            image: The image in which tables are to be detected.

        Returns:
            object: The detected table information including bounding boxes and other attributes.
        """
        # Use the loaded model to detect tables in the given image
        return self.model(image)

scanipy/deeplearning/models/tablefinder.py
SVJLucas/Scanipy
import pickle
from transformers import pipeline


class TableStructureAnalyzer:
    """
    Analyze the structure of tables in images using a fine-tuned Table Transformer model.

    The model is based on the Table Transformer (DETR) by Microsoft, trained on PubTables1M.
    It can detect structural elements like rows and columns within tables in images.

    Source: https://github.com/microsoft/table-transformer

    Attributes:
        model (object): The loaded fine-tuned Table Transformer model.

    Example:
        >>> analyzer = TableStructureAnalyzer()
        >>> result = analyzer(image)
    """

    def __init__(self):
        """
        Initialize the TableStructureAnalyzer by loading the pre-trained model.
        """
        # Load the fine-tuned Table Transformer model from a pickle file
        self.model = pipeline("object-detection", model="microsoft/table-transformer-structure-recognition")

    def __repr__(self):
        """
        Returns the official string representation of the EquationToLatex object.
        
        Returns:
            str: A string that can be used to recreate the EquationToLatex object.
        """
        return "TableStructureAnalyzer()"

    def __str__(self):
        """
        Returns a string representation of the EquationToLatex object, which is the same as its official representation.
        
        Returns:
            str: A string that can be used to recreate the EquationToLatex object.
        """
        return self.__repr__()

    def __call__(self, image):
        """
        Analyze the table structure in a given image.

        Args:
            image: The image containing the table to be analyzed.

        Returns:
            object: The analyzed table structure, including elements like rows and columns.
        """
        # Use the loaded model to analyze the table structure in the given image
        return self.model(image)

scanipy/deeplearning/models/tablestructureanalyzer.py
SVJLucas/Scanipy
import easyocr
import pytesseract
import numpy as np
from PIL import Image
from typing import Union

class TextOCR:
    def __init__(self, lang: str, device: str):
        """
        Initialize the TextOCR object.

        Args:
            lang (str): The language for OCR.
            device (str): The device to use for OCR ('cpu' or other).

        Raises:
            TypeError: If the types of the arguments are not as expected.
        """
        # Verify the input variable types
        if not isinstance(lang, str):
            raise TypeError("lang must be a string")
        if not isinstance(device, str):
            raise TypeError("device must be a string")

        # Initialize instance variables
        self.device = device
        if device == 'cpu':
            if lang != 'en':
                self.lang = 'en+' + lang
            else:
                self.lang = 'en'
            self.model = pytesseract.image_to_string
        else:
            self.lang = lang
            if self.lang != 'en':
                self.model = easyocr.Reader([self.lang, 'en']).readtext
            else:
                self.model = easyocr.Reader(['en']).readtext

    def __call__(self, image: Image.Image) -> str:
        """
        Perform OCR on the given image.

        Args:
            image (Image): The image to perform OCR on.

        Returns:
            str: The extracted text.
        """
        # Verify the input image type
        if not isinstance(image, Image.Image):
            raise TypeError("image must be a PIL.Image object")

        # Perform OCR based on the device
        if self.device == 'cpu':
            if self.lang == 'en':
                text = self.model(image)
            else:
                text = self.model(image, lang=self.lang) # type: ignore
        else:
            np_image = np.array(image)
            detections = self.model.recognize(np_image)
            text = ' '.join(detection[1] for detection in detections)

        # Remove newline characters from the extracted text
        text = text.replace('\n', '')

        return text

    def __repr__(self) -> str:
        """
        Returns a official string representation of the TextOCR object.

        Returns:
            str: A string representation of the object.
        """
        return f"TextOCR(lang={self.lang}, device={self.device})"

    def __str__(self) -> str:
        """
        Returns a string representation of the TextOCR object, which is the same as its official representation.
        
        Returns:
            str: A string that can be used to recreate the EquationToLatex object.
        """
        return self.__repr__()

scanipy/deeplearning/models/textocr.py
SVJLucas/Scanipy
import os
import layoutparser as lp
import numpy as np
from scanipy.elements import TableElement, TextElement, ImageElement, EquationElement
import matplotlib.pyplot as plt
import matplotlib.patches as patches



class Document:
    """
    Represents a document containing various elements, such as images.

    Attributes:
        elements (list): A list of elements in the document.
    """

    def __init__(self):
        self.elements = {}
        self.images = []
        self.layouts = []
        self.table_extractor_data = []

    def to_markdown(self, output_folder, filename='output.md'):
        """
        Generate a Markdown document from the elements and save it to a file.

        :param output_folder: The folder where the Markdown file will be saved.
        """
        if not os.path.exists(output_folder):
            os.makedirs(output_folder)
        output = ""
        sorted_pages = sorted(list(self.elements.keys()))
        for page in sorted_pages:
            sorted_elements = self.get_ordered_elements(page)
            for element in sorted_elements:
                element_output = element.generate_markdown(output_folder)
                output += element_output

        output_path = os.path.join(output_folder, filename)
        with open(output_path, 'w') as f:
            f.write(output)

    def get_ordered_elements(self, page):
        return sorted(self.elements[page])

    def add_element(self, page, element):
        if self.elements.get(page) is None:
            self.elements[page] = []
        self.elements[page].append(element)

    # def visualize_pipeline(self, page=0, step=0):
    #     if step == 0:
    #         return lp.draw_box(self.images[page], self.layouts[page], box_width=5, box_alpha=0.2)
    #     if step == 1:
    #         self._visualize_tables(self.table_extractor_data[page]['image'],
    #                                self.table_extractor_data[page]['detected_tables'])
    #         return
    #     raise ValueError(f'step {step} not recognized')

    def visualize_block(self, page=0, block=0):
        block = self.layouts[page][block]
        segment_image = (block
                         .pad(left=5, right=15, top=5, bottom=5)
                         .crop_image(np.asarray(self.images[page])))
        return block

    def _visualize_tables(self, image, detected_tables):
        # Create a matplotlib figure and axis for visualization
        fig, ax = plt.subplots(1)
        # Display the RGB image
        ax.imshow(image)

        # Loop through each detected table to draw its bounding box
        for table in detected_tables:
            xmin, ymin, xmax, ymax = table['box'].values()

            # Create a red rectangle around the table
            rect = patches.Rectangle((xmin, ymin), xmax - xmin, ymax - ymin, linewidth=1, edgecolor='r', facecolor='none')

            # Add the rectangle to the plot
            ax.add_patch(rect)

            # Add label and confidence score
            label = f"{table['label']} ({table['score']:.2f})"
            plt.text(xmin, ymin, label, color='white', fontsize=12, bbox=dict(facecolor='red', alpha=0.5))

        # Hide axes and show the plot
        plt.axis('off')
        plt.show()

    def store_page(self, image, layout):
        self.layouts.append(layout)
        self.images.append(image)

    def save_tables(self, image, detected_tables):
        self.table_extractor_data.append({'image': image, 'detected_tables': detected_tables})





scanipy/document.py
SVJLucas/Scanipy
from .equation_element import EquationElement
from .image_element import ImageElement
from .table_element import TableElement
from .title_element import TitleElement
from .text_element import TextElement
from .element import Element


scanipy/elements/__init__.py
SVJLucas/Scanipy
from typing import Union

# Define the Element class
class Element:
    def __init__(self, x_min: float, y_min: float, x_max: float, y_max: float,
                 pipeline_step:Union[int, None]=None, page_number: Union[int, None] = None,
                 intersection_percentage_threshold = 90):
        """
        Initialize an Element object with normalized coordinates, an optional pipeline step, and an intersection percentage threshold.

        Args:
            x_min (float): The minimum x-coordinate of the element, normalized to the image width (range: 0 to 1).
            y_min (float): The minimum y-coordinate of the element, normalized to the image height (range: 0 to 1).
            x_max (float): The maximum x-coordinate of the element, normalized to the image width (range: 0 to 1).
            y_max (float): The maximum y-coordinate of the element, normalized to the image height (range: 0 to 1).
            pipeline_step (Union[int, None], optional): The processing step in the pipeline to which this element belongs. Defaults to None.
            page_number (int): Specifies the page number on which the element is located.
            intersection_percentage_threshold (int, optional): The minimum percentage of intersection required for two elements to be considered overlapping. Defaults to 90 (%).

        Raises:
            ValueError:
                - If x_min >= x_max or y_min >= y_max.
                - If any coordinate is not in the range [0, 1].
            TypeError:
                - If the coordinates are not floats.
                - If pipeline_step is neither an integer nor None.
        """
        # Verify the input variable types
        if not isinstance(x_min, float) or not isinstance(y_min, float) or not isinstance(x_max, float) or not isinstance(y_max, float):
            raise TypeError("Coordinates must be floats")

        # Verify if coordinates are in the range 0 to 1
        if not (0 <= x_min <= 1 and 0 <= y_min <= 1 and 0 <= x_max <= 1 and 0 <= y_max <= 1):
            raise ValueError("Coordinates must be in the range from 0.0 to 1.0")

        # Verify if x_min is less than x_max and y_min is less than y_max
        if x_min >= x_max or y_min >= y_max:
            raise ValueError("Invalid coordinates: x_min should be less than x_max and y_min should be less than y_max")

        # Verify the pipeline_step type
        if pipeline_step is not None and not isinstance(pipeline_step, int):
            raise TypeError("pipeline_step must be an integer or None")
            
        # Verify the page_number type
        if pipeline_step is not None and not isinstance(page_number, int):
            raise TypeError("page_number must be an integer or None")

        # Initialize instance variables
        self._x_min = x_min
        self._y_min = y_min
        self._x_max = x_max
        self._y_max = y_max
        self._pipeline_step = pipeline_step
        self._page_number = page_number
        self._intersection_percentage_threshold = intersection_percentage_threshold

        # Calculate the center coordinates and width
        self.x_center = (self.x_min + self.x_max) / 2
        self.y_center = (self.y_min + self.y_max) / 2
        self.width = x_max - x_min
        self.height = y_max - y_min

    def __repr__(self) -> str:
        """
        Provides a human-readable representation of the Element object.

        Returns:
            str: A string representation of the Element object.
        """
        return f"Element(x_min={self._x_min}, y_min={self._y_min}, x_max={self._x_max}, y_max={self._y_max}, pipeline_step={self._pipeline_step}, page_number={self._page_number})"

    def __str__(self) -> str:
        """
        Returns a string representation of the object, which is the same as its official representation.
        
        Returns:
            str: A string that can be used to recreate the object.
        """
        return self.__repr__()

    def __lt__(self, other)->bool:
        """
        Less than operator for Element objects.

        Args:
            other (Element): Another Element object.

        Returns:
            bool: True if this Element is "less than" the other, False otherwise.

        Raises:
            TypeError: If 'other' is not an instance or subclass of Element.
        """

        if not isinstance(other, Element):
            raise TypeError("other must be an instance or subclass of Element")

        return self._column_before(other) or (self._same_column(other) and (self.y_min < other.y_min))

    def _column_before(self, other)->bool:
        """
        Check if this Element is in a column before the other Element.

        Args:
            other (Element): Another Element object.

        Returns:
            bool: True if in a column before, False otherwise.
        """

        if not isinstance(other, Element):
            raise TypeError("other must be an instance or subclass of Element")

        max_width = max(self.width, other.width)
        return self.x_min < other.x_min - max_width / 2

    def _same_column(self, other)->bool:
        """
        Check if this Element is in the same column as the other Element.

        Args:
            other (Element): Another Element object.

        Returns:
            bool: True if in the same column, False otherwise.
        """

        if not isinstance(other, Element):
            raise TypeError("other must be an instance or subclass of Element")

        max_width = max(self.width, other.width)
        return abs(self.x_min - other.x_min) < max_width / 2

    def is_in(self, other)->bool:
        """
        Check if this Element is inside the other Element.

        Args:
            other (Element): Another Element object.

        Returns:
            bool: True if inside, False otherwise.
        """
        if not isinstance(other, Element):
            raise TypeError("other must be an instance or subclass of Element")
        return self._intersection_percentage(other) > self._intersection_percentage_threshold

    def _intersection_area(self, other)->int:
        """
        Calculate the area of intersection between this Element and another.

        Args:
            other (Element): Another Element object.

        Returns:
            int: The area of intersection.
        """

        if not isinstance(other, Element):
            raise TypeError("other must be an instance or subclass of Element")

        x_min = max(self.x_min, other.x_min)
        x_max = min(self.x_max, other.x_max)
        y_min = max(self.y_min, other.y_min)
        y_max = min(self.y_max, other.y_max)

        if x_max < x_min or y_max < y_min:
            return 0

        return (x_max - x_min) * (y_max - y_min)

    def _intersection_percentage(self, other):
        """
        Calculate the percentage of intersection between this Element and another.

        Args:
            other (Element): Another Element object.

        Returns:
            float: The percentage of intersection.
        """
        if not isinstance(other, Element):
            raise TypeError("other must be an instance or subclass of Element")

        intersection_area = self._intersection_area(other)
        self_area = (self.x_max - self.x_min) * (self.y_max - self.y_min)

        return intersection_area / self_area * 100

    @property
    def x_min(self) -> float:
        """
        Gets the minimum x-coordinate of the element.

        Returns:
            float: The minimum x-coordinate, normalized to the image width.
        """
        return self._x_min

    @x_min.setter
    def x_min(self, x_min: float):
        """
        Sets the minimum x-coordinate of the element.

        Args:
            x_min (float): The new minimum x-coordinate, normalized to the image width (range: 0 to 1).

        Raises:
            ValueError: If x_min is not in the range [0, 1].
            TypeError: If x_min is not a float.
        """
        if not isinstance(x_min, float):
            raise TypeError("x_min must be a float.")
        if not 0 <= x_min <= 1:
            raise ValueError("x_min must be in the range [0, 1].")
        self._x_min = x_min

    @property
    def y_min(self) -> float:
        """
        Gets the minimum y-coordinate of the element.

        Returns:
            float: The minimum y-coordinate, normalized to the image height.
        """
        return self._y_min

    @y_min.setter
    def y_min(self, y_min: float):
        """
        Sets the minimum y-coordinate of the element.

        Args:
            y_min (float): The new minimum y-coordinate, normalized to the image height (range: 0 to 1).

        Raises:
            ValueError: If y_min is not in the range [0, 1].
            TypeError: If y_min is not a float.
        """
        if not isinstance(y_min, float):
            raise TypeError("y_min must be a float.")
        if not 0 <= y_min <= 1:
            raise ValueError("y_min must be in the range [0, 1].")
        self._y_min = y_min

    @property
    def x_max(self) -> float:
        """
        Gets the maximum x-coordinate of the element.

        Returns:
            float: The maximum x-coordinate, normalized to the image width.
        """
        return self._x_max

    @x_max.setter
    def x_max(self, x_max: float):
        """
        Sets the maximum x-coordinate of the element.

        Args:
            x_max (float): The new maximum x-coordinate, normalized to the image width (range: 0 to 1).

        Raises:
            ValueError: If x_max is not in the range [0, 1].
            TypeError: If x_max is not a float.
        """
        if not isinstance(x_max, float):
            raise TypeError("x_max must be a float.")
        if not 0 <= x_max <= 1:
            raise ValueError("x_max must be in the range [0, 1].")
        self._x_max = x_max

    @property
    def y_max(self) -> float:
        """
        Gets the maximum y-coordinate of the element.

        Returns:
            float: The maximum y-coordinate, normalized to the image height.
        """
        return self._y_max

    @y_max.setter
    def y_max(self, y_max: float):
        """
        Sets the maximum y-coordinate of the element.

        Args:
            y_max (float): The new maximum y-coordinate, normalized to the image height (range: 0 to 1).

        Raises:
            ValueError: If y_max is not in the range [0, 1].
            TypeError: If y_max is not a float.
        """
        if not isinstance(y_max, float):
            raise TypeError("y_max must be a float.")
        if not 0 <= y_max <= 1:
            raise ValueError("y_max must be in the range [0, 1].")
        self._y_max = y_max
    @property
    def pipeline_step(self) -> Union[int, None]:
        """
        Gets the pipeline step associated with the element.

        Returns:
            Union[int, None]: The pipeline step, which can be an integer or None.
        """
        return self._pipeline_step

    @pipeline_step.setter
    def pipeline_step(self, pipeline_step: Union[int, None]):
        """
        Sets the pipeline step associated with the element.

        Args:
            pipeline_step (Union[int, None]): The new pipeline step, which can be an integer or None.

        Raises:
            TypeError: If pipeline_step is neither an integer nor None.
        """
        if pipeline_step is not None and not isinstance(pipeline_step, int):
            raise TypeError("pipeline_step must be either an integer or None.")
        self._pipeline_step = pipeline_step

    @property
    def page_number(self) -> Union[int, None]:
        """
        Gets the page number associated with the element.

        Returns:
            Union[int, None]: The page number, which can be an integer or None.
        """
        return self._page_number

    @page_number.setter
    def page_number(self, page_number: Union[int, None]):
        """
        Sets the page number associated with the element.

        Args:
            page_number (Union[int, None]): The new page number, which can be an integer or None.

        Raises:
            TypeError: If page_number is neither an integer nor None.
        """
        if page_number is not None and not isinstance(page_number, int):
            raise TypeError("pipeline_step must be either an integer or None.")
        self._page_number = page_number

scanipy/elements/element.py
SVJLucas/Scanipy
import logging
from typing import Union
from .element import Element

# Define the EquationElement class
class EquationElement(Element):
    """
    Represents an equation element within a document.

    Attributes:
        latex_content (Union[str, None]): The LaTeX content of the equation element.
        is_inside_text (bool): Flag indicating if the equation is part of a text block.
    """

    def __init__(self, x_min: float, y_min: float, x_max: float, y_max: float, 
                 pipeline_step: Union[int, None] = None, page_number: Union[int, None] = None,
                 is_inside_text: bool = False):
        """
        Initialize an EquationElement object.

        Args:
            x_min (float): The minimum x-coordinate of the element, normalized to the image width (range: 0 to 1).
            y_min (float): The minimum y-coordinate of the element, normalized to the image height (range: 0 to 1).
            x_max (float): The maximum x-coordinate of the element, normalized to the image width (range: 0 to 1).
            y_max (float): The maximum y-coordinate of the element, normalized to the image height (range: 0 to 1).
            pipeline_step (Union[int, None], optional): The pipeline step, can be None.
            page_number (int): Specifies the page number on which the element is located.

        Raises:
            TypeError: If the types of the arguments are not as expected.
        """

        # Initialize instance variables by calling the parent class constructor
        super().__init__(x_min, y_min, x_max, y_max, pipeline_step, page_number)

        # Initialize additional instance variables specific to EquationElement
        self._latex_content = None
        self._is_inside_text = is_inside_text

    # Setter and Getter for latex_content
    @property
    def latex_content(self) -> Union[str, None]:
        """
        Gets the LaTeX content of the equation element.

        Returns:
            Union[str, None]: The LaTeX content of the equation element.
        """
        return self._latex_content

    @latex_content.setter
    def latex_content(self, value: Union[str, None]):
        """
        Sets the LaTeX content of the equation element.

        Args:
            value (Union[str, None]): The new LaTeX content for the equation element.

        Raises:
            TypeError: If the provided value is not a string or None.
        """
        if value is not None and not isinstance(value, str):
            raise TypeError("latex_content must be a string or None")
        self._latex_content = value

    # Setter and Getter for is_inside_text
    @property
    def is_inside_text(self) -> bool:
        """
        Gets the flag indicating if the equation is part of a text block.

        Returns:
            bool: True if the equation is part of a text block, False otherwise.
        """
        return self._is_inside_text

    @is_inside_text.setter
    def is_inside_text(self, value: bool):
        """
        Sets the flag indicating if the equation is part of a text block.

        Args:
            value (bool): The new flag value.

        Raises:
            TypeError: If the provided value is not a boolean.
        """
        if not isinstance(value, bool):
            raise TypeError("is_inside_text must be a boolean")
        self._is_inside_text = value

    def __repr__(self) -> str:
        """
        Returns a string representation of the EquationElement object.

        Returns:
            str: A string representation of the object.
        """
        return f"EquationElement(x_min={self.x_min}, y_min={self.y_min}, x_max={self.x_max}, y_max={self.y_max}, pipeline_step={self.pipeline_step}, latex_content={self.latex_content}, is_inside_text={self.is_inside_text})"

    def __str__(self)->str:
        """
        Returns a string representation of the object, which is the same as its official representation.
        
        Returns:
            str: A string that can be used to recreate the object.
        """
        return self.__repr__()

    def generate_markdown(self, output_directory: str) -> str:
        """
        Generate the Markdown representation of the equation element.

        Args:
            output_directory (str): The directory where the equation will be saved.

        Returns:
            str: Markdown representation of the equation element.
        """
        if self.latex_content == None:
            logging.warning('Tried to write a NoneType object')
            return '\n\n'

        # Concatenate the LaTeX content with newlines for Markdown formatting
        formatted_equation = self.latex_content + '\n\n'
        return formatted_equation

scanipy/elements/equation_element.py
SVJLucas/Scanipy
import os
import logging
from .element import Element
from .equation_element import EquationElement
from typing import Union, Any
import matplotlib.pyplot as plt

# Define the ImageElement class, which inherits from the Element class
class ImageElement(Element):
    """
    Represents an image element within a document.

    Attributes:
        unique_key (str): A unique identifier for the image element.
        image_content (Any): The actual content of the image.
        image_extension (str): The file extension for the image (e.g., 'jpg', 'png').
        has_equation_inside (bool): If there's some equation inside the image.
    """

    def __init__(self, x_min: float, y_min: float, x_max: float, y_max: float,
                 pipeline_step: Union[int, None] = None, page_number: Union[int, None] = None):
        """
        Initialize an ImageElement object.

        Args:
            x_min (float): The minimum x-coordinate of the element, normalized to the image width (range: 0 to 1).
            y_min (float): The minimum y-coordinate of the element, normalized to the image height (range: 0 to 1).
            x_max (float): The maximum x-coordinate of the element, normalized to the image width (range: 0 to 1).
            y_max (float): The maximum y-coordinate of the element, normalized to the image height (range: 0 to 1).
            pipeline_step (Union[int, None], optional): The pipeline step, can be None.
            page_number (int): Specifies the page number on which the element is located.

        Raises:
            ValueError: If x_min >= x_max or y_min >= y_max.
            TypeError: If the types of the arguments are not as expected.
        """

        # Initialize instance variables by calling the parent class constructor
        super().__init__(x_min, y_min, x_max, y_max, pipeline_step, page_number)

        # Initialize additional instance variables specific to ImageElement
        self._unique_key = None
        self._image_content = None
        self._image_extension = None
        self._has_equation_inside = False
        self._equation_inside = None
        
    @property
    def equation_inside(self) -> Union[EquationElement,None]:
        """
        Get the equation inside the element, if any.
    
        Returns:
            EquationElement if the element contains equations, otherwise None.
        """
        return self._equation_inside

    @equation_inside.setter
    def equation_inside(self, value: EquationElement):
        """
        Sets the equation inside the element.

        Args:
            value (EquationElement): The new EquationElement value.

        Raises:
            TypeError: If the provided value is not a EquationElement.
        """
        if not isinstance(value, EquationElement):
            raise TypeError("equation_inside must be a EquationElement")
        self._equation_inside = value

    @property
    def has_equation_inside(self) -> bool:
        """
        Determines if the element contains equations.
    
        This method checks if the element has any equations inside it and returns a boolean value accordingly.
    
        Returns:
            bool: True if the element contains equations, otherwise False.
        """
        return self._has_equation_inside

    @has_equation_inside.setter
    def has_equation_inside(self, value: bool):
        """
        Sets a boolean that checks if the element has any equations inside it.

        Args:
            value (bool): The new boolean value.

        Raises:
            TypeError: If the provided value is not a bool.
        """
        if not isinstance(value, bool):
            raise TypeError("has_equation_inside must be a bool")
        self._has_equation_inside = value

    @property
    def unique_key(self) -> str | None:
        """
        Gets the unique key of the image element.

        Returns:
            str: The unique key of the image element.
        """
        return self._unique_key

    @unique_key.setter
    def unique_key(self, value: str):
        """
        Sets the unique key of the image element.

        Args:
            value (str): The new unique key.

        Raises:
            TypeError: If the provided value is not a string.
        """
        if not isinstance(value, str):
            raise TypeError("unique_key must be a string")
        self._unique_key = value

    @property
    def image_content(self) -> Any:
        """
        Gets the content of the image element.

        Returns:
            Any: The content of the image element.
        """
        return self._image_content

    @image_content.setter
    def image_content(self, value: Any):
        """
        Sets the content of the image element.

        Args:
            value (Any): The new content for the image element.
        """
        # Type verification can be more specific depending on what 'Any' encompasses
        self._image_content = value

    @property
    def image_extension(self) -> str | None:
        """
        Gets the file extension of the image element.

        Returns:
            str: The file extension of the image element.
        """
        return self._image_extension

    @image_extension.setter
    def image_extension(self, value: str):
        """
        Sets the file extension of the image element.

        Args:
            value (str): The new file extension for the image element.

        Raises:
            TypeError: If the provided value is not a string.
        """
        if not isinstance(value, str):
            raise TypeError("image_extension must be a string")
        self._image_extension = value

    def __repr__(self) -> str:
        """
        Returns a string representation of the ImageElement object.
    
        Returns:
            str: A string representation of the object.
        """
        return f"ImageElement(unique_key={self.unique_key}, x_min={self.x_min}, y_min={self.y_min}, x_max={self.x_max}, y_max={self.y_max}, pipeline_step={self.pipeline_step}, image_extension={self.image_extension}, has_equation_inside={self.has_equation_inside})"

    def __str__(self):
        """
        Returns a string representation of the object, which is the same as its official representation.
        
        Returns:
            str: A string that can be used to recreate the object.
        """
        return self.__repr__()
        
    def display_image(self):
        """
        Display the image content using Matplotlib.

        Raises:
            ValueError: If image_content is None or not properly set.
        """
        if self.image_content is None:
            raise ValueError("image_content is not set")

        # Assuming image_content has format compatible with imshow
        plt.imshow(self.image_content)
        plt.title(f"Image: {self.unique_key}")
        plt.axis('off')  # Turn off axis numbers and ticks
        plt.show()

    def generate_markdown(self, output_directory: str) -> str:
        """
        Generate the Markdown representation of the image element and save the image to a file.

        Args:
            output_directory (str): The directory where the image file will be saved.

        Returns:
            str: Markdown representation of the image element.
        """
        # If the image is an equation, don't print it
        if self.has_equation_inside:
            return '\n'


        if self.image_content == None:
            logging.warning('Tried to write a NoneType object')
            return '\n\n'
        
        # Create the filename for the image
        filename = f"{self.unique_key}.{self.image_extension}"

        # Create the full path where the image will be saved
        output_path = os.path.join(output_directory, filename)

        # Save the image content to the file
        self.image_content.save(open(output_path, "wb"))

        # Return the Markdown representation of the image
        return f"\n![image]({filename})\n\n"


scanipy/elements/image_element.py
SVJLucas/Scanipy
import logging
from typing import Union
import pandas as pd
from .element import Element
from .equation_element import EquationElement

# Define the TableElement class
class TableElement(Element):
    """
    Represents a table element within a document for displaying tabular data.

    Attributes:
        table_data (pandas.DataFrame): The DataFrame containing the tabular data.
        has_equation_inside (bool): If there's some equation inside the table.
    """

    def __init__(self, x_min: float, y_min: float, x_max: float, y_max: float, 
                 pipeline_step: Union[int, None] = None, page_number: Union[int, None] = None):
        """
        Initialize a TableElement object.

        Args:
            x_min (float): The minimum x-coordinate of the element, normalized to the image width (range: 0 to 1).
            y_min (float): The minimum y-coordinate of the element, normalized to the image height (range: 0 to 1).
            x_max (float): The maximum x-coordinate of the element, normalized to the image width (range: 0 to 1).
            y_max (float): The maximum y-coordinate of the element, normalized to the image height (range: 0 to 1).
            pipeline_step (Union[int, None], optional): The pipeline step, can be None.
            page_number (int): Specifies the page number on which the element is located.

        Raises:
            TypeError: If the types of the arguments are not as expected.
        """
                     
        # Initialize instance variables by calling the parent class constructor
        super().__init__(x_min, y_min, x_max, y_max, pipeline_step, page_number)

        # Initialize additional instance variable specific to TableElement
        self._table_data = None
        self._has_equation_inside = False
        self._equation_inside = None
        
    @property
    def equation_inside(self) -> Union[EquationElement,None]:
        """
        Get the equation inside the element, if any.
    
        Returns:
            EquationElement if the element contains equations, otherwise None.
        """
        return self._equation_inside

    @equation_inside.setter
    def equation_inside(self, value: EquationElement):
        """
        Sets the equation inside the element.

        Args:
            value (EquationElement): The new EquationElement value.

        Raises:
            TypeError: If the provided value is not a EquationElement.
        """
        if not isinstance(value, EquationElement):
            raise TypeError("equation_inside must be a EquationElement")
        self._equation_inside = value

    @property
    def has_equation_inside(self) -> bool:
        """
        Determines if the element contains equations.
    
        This method checks if the element has any equations inside it and returns a boolean value accordingly.
    
        Returns:
            bool: True if the element contains equations, otherwise False.
        """
        return self._has_equation_inside

    @has_equation_inside.setter
    def has_equation_inside(self, value: bool):
        """
        Sets a boolean that checks if the element has any equations inside it.

        Args:
            value (bool): The new boolean value.

        Raises:
            TypeError: If the provided value is not a bool.
        """
        if not isinstance(value, bool):
            raise TypeError("has_equation_inside must be a bool")
        self._has_equation_inside = value
        
    # Setter and Getter for table_data
    @property
    def table_data(self) -> Union[pd.DataFrame, None]:
        """
        Gets the DataFrame containing the tabular data.

        Returns:
            Union[pd.DataFrame, None]: The DataFrame containing the tabular data.
        """
        return self._table_data

    @table_data.setter
    def table_data(self, value: Union[pd.DataFrame, None]):
        """
        Sets the DataFrame containing the tabular data.

        Args:
            value (Union[pd.DataFrame, None]): The new DataFrame for the table element.

        Raises:
            TypeError: If the provided value is not a DataFrame or None.
        """
        if value is not None and not isinstance(value, pd.DataFrame):
            raise TypeError("table_data must be a pandas DataFrame or None")
        self._table_data = value

    def __repr__(self) -> str:
        """
        Returns a string representation of the TableElement object.

        Returns:
            str: A string representation of the object.
        """
        return f"TableElement(x_min={self.x_min}, y_min={self.y_min}, x_max={self.x_max}, y_max={self.y_max}, pipeline_step={self.pipeline_step}, table_data={self.table_data}, has_equation_inside={self.has_equation_inside})"

    def __str__(self):
        """
        Returns a string representation of the object, which is the same as its official representation.
        
        Returns:
            str: A string that can be used to recreate the object.
        """
        return self.__repr__()

    def generate_markdown(self, output_directory: str) -> str:
        """
        Generate the Markdown representation of the table element.

        Args:
            output_directory (str): The directory where the table will be saved.

        Returns:
            str: Markdown representation of the table element.
        """
        if self.table_data is None:
            logging.warning('Tried to write a NoneType object')
            return '\n\n'

        if self.table_data.empty:
            logging.warning('Tried to write an empty DataFrame')
            return '\n\n'

        # Convert the DataFrame to Markdown format
        formatted_table = self.table_data.to_markdown() + '\n\n'
        return formatted_table

scanipy/elements/table_element.py
SVJLucas/Scanipy
import logging
from typing import Union
from .element import Element
from .equation_element import EquationElement
# Define the TextElement class, which inherits from the Element class
class TextElement(Element):
    """
    Represents a text element within a document.

    Attributes:
        text_content (str): The content of the text element.
        has_equation_inside (bool): If there's some equation inside the text.
        
    """

    def __init__(self, x_min: float, y_min: float, x_max: float, y_max: float,
                 pipeline_step: Union[int, None] = None, page_number: Union[int, None] = None):
        """
        Initialize a TextElement object.

        Args:
            x_min (float): The minimum x-coordinate of the element, normalized to the image width (range: 0 to 1).
            y_min (float): The minimum y-coordinate of the element, normalized to the image height (range: 0 to 1).
            x_max (float): The maximum x-coordinate of the element, normalized to the image width (range: 0 to 1).
            y_max (float): The maximum y-coordinate of the element, normalized to the image height (range: 0 to 1).
            pipeline_step (Union[int, None], optional): The pipeline step, can be None.
            page_number (int): Specifies the page number on which the element is located.

        Raises:
            TypeError: If the types of the arguments are not as expected.
        """

        # Initialize instance variables by calling the parent class constructor
        super().__init__(x_min, y_min, x_max, y_max, pipeline_step, page_number)

        # Initialize additional instance variable specific to TextElement
        self._text_content = None
        self._has_equation_inside = False
        self._equation_inside = None
        
    @property
    def equation_inside(self) -> Union[EquationElement,None]:
        """
        Get the equation inside the element, if any.
    
        Returns:
            EquationElement if the element contains equations, otherwise None.
        """
        return self._equation_inside

    @equation_inside.setter
    def equation_inside(self, value: EquationElement):
        """
        Sets the equation inside the element.

        Args:
            value (EquationElement): The new EquationElement value.

        Raises:
            TypeError: If the provided value is not a EquationElement.
        """
        if not isinstance(value, EquationElement):
            raise TypeError("equation_inside must be a EquationElement")
        self._equation_inside = value
    @property
    def has_equation_inside(self) -> bool:
        """
        Determines if the element contains equations.
    
        This method checks if the element has any equations inside it and returns a boolean value accordingly.
    
        Returns:
            bool: True if the element contains equations, otherwise False.
        """
        return self._has_equation_inside

    @has_equation_inside.setter
    def has_equation_inside(self, value: bool):
        """
        Sets a boolean that checks if the element has any equations inside it.

        Args:
            value (bool): The new boolean value.

        Raises:
            TypeError: If the provided value is not a bool.
        """
        if not isinstance(value, bool):
            raise TypeError("has_equation_inside must be a bool")
        self._has_equation_inside = value
        
    # Setter and Getter for text_content
    @property
    def text_content(self) -> Union[str, None]:
        """
        Gets the content of the text element.

        Returns:
            Union[str, None]: The content of the text element.
        """
        return self._text_content

    @text_content.setter
    def text_content(self, value: Union[str, None]):
        """
        Sets the content of the text element.

        Args:
            value (Union[str, None]): The new content for the text element.

        Raises:
            TypeError: If the provided value is not a string or None.
        """
        if value is not None and not isinstance(value, str):
            raise TypeError("text_content must be a string or None")
        self._text_content = value

    def __repr__(self) -> str:
        """
        Returns a string representation of the TextElement object.

        Returns:
            str: A string representation of the object.
        """
        return f"TextElement(x_min={self.x_min}, y_min={self.y_min}, x_max={self.x_max}, y_max={self.y_max}, pipeline_step={self.pipeline_step}, text_content={self.text_content}, has_equation_inside={self.has_equation_inside})"

    def __str__(self)->str:
        """
        Returns a string representation of the object, which is the same as its official representation.
        
        Returns:
            str: A string that can be used to recreate the object.
        """
        return self.__repr__()

    def generate_markdown(self, output_directory: str) -> str:
        """
        Generate the Markdown representation of the text element.

        Args:
            output_directory (str): The directory where the text file will be saved.

        Returns:
            str: Markdown representation of the text element.
        """
        if self.text_content == None:
            logging.warning('Tried to write a NoneType object')
            return '\n\n'

        # Concatenate the text content with newlines for Markdown formatting
        formatted_text = self.text_content + '\n\n'
        return formatted_text

scanipy/elements/text_element.py
SVJLucas/Scanipy
from typing import Union
import logging
from .element import Element
from .equation_element import EquationElement


# Define the TitleElement class
class TitleElement(Element):
    """
    Represents a title element within a document.

    Attributes:
        title_content (str): The content of the title element.
        has_equation_inside (bool): If there's some equation inside the title.
    """

    def __init__(self, x_min: float, y_min: float, x_max: float, y_max: float, 
                 pipeline_step: Union[int, None] = None, page_number: Union[int, None] = None):
        """
        Initialize a TitleElement object.

        Args:
            x_min (float): The minimum x-coordinate of the element, normalized to the image width (range: 0 to 1).
            y_min (float): The minimum y-coordinate of the element, normalized to the image height (range: 0 to 1).
            x_max (float): The maximum x-coordinate of the element, normalized to the image width (range: 0 to 1).
            y_max (float): The maximum y-coordinate of the element, normalized to the image height (range: 0 to 1).
            pipeline_step (Union[int, None], optional): The pipeline step, can be None.
            page_number (int): Specifies the page number on which the element is located.

        Raises:
            TypeError: If the types of the arguments are not as expected.
        """

        # Initialize instance variables by calling the parent class constructor
        super().__init__(x_min, y_min, x_max, y_max, pipeline_step, page_number)

        # Initialize additional instance variable specific to TitleElement
        self._title_content = None
        self._has_equation_inside = False
        self._equation_inside = None
        
    @property
    def equation_inside(self) -> Union[EquationElement,None]:
        """
        Get the equation inside the element, if any.
    
        Returns:
            EquationElement if the element contains equations, otherwise None.
        """
        return self._equation_inside

    @equation_inside.setter
    def equation_inside(self, value: EquationElement):
        """
        Sets the equation inside the element.

        Args:
            value (EquationElement): The new EquationElement value.

        Raises:
            TypeError: If the provided value is not a EquationElement.
        """
        if not isinstance(value, EquationElement):
            raise TypeError("equation_inside must be a EquationElement")
        self._equation_inside = value
    @property
    def has_equation_inside(self) -> bool:
        """
        Determines if the element contains equations.
    
        This method checks if the element has any equations inside it and returns a boolean value accordingly.
    
        Returns:
            bool: True if the element contains equations, otherwise False.
        """
        return self._has_equation_inside

    @has_equation_inside.setter
    def has_equation_inside(self, value: bool):
        """
        Sets a boolean that checks if the element has any equations inside it.

        Args:
            value (bool): The new boolean value.

        Raises:
            TypeError: If the provided value is not a bool.
        """
        if not isinstance(value, bool):
            raise TypeError("has_equation_inside must be a bool")
        self._has_equation_inside = value
        
    # Setter and Getter for title_content
    @property
    def title_content(self) -> Union[str, None]:
        """
        Gets the content of the title element.

        Returns:
            Union[str, None]: The content of the title element.
        """
        return self._title_content

    @title_content.setter
    def title_content(self, value: Union[str, None]):
        """
        Sets the content of the title element.

        Args:
            value (Union[str, None]): The new content for the title element.

        Raises:
            TypeError: If the provided value is not a string or None.
        """
        if value is not None and not isinstance(value, str):
            raise TypeError("title_content must be a string or None")
        self._title_content = value

    def __repr__(self) -> str:
        """
        Returns a string representation of the TitleElement object.

        Returns:
            str: A string representation of the object.
        """
        return f"TitleElement(x_min={self.x_min}, y_min={self.y_min}, x_max={self.x_max}, y_max={self.y_max}, pipeline_step={self.pipeline_step}, title_content={self.title_content}, has_equation_inside={self.has_equation_inside})"

    def __str__(self)-> str:
        """
        Returns a string representation of the object, which is the same as its official representation.
        
        Returns:
            str: A string that can be used to recreate the object.
        """
        return self.__repr__()
    
    def generate_markdown(self, output_directory: str) -> str:
        """
        Generate the Markdown representation of the title element.

        Args:
            output_directory (str): The directory where the text file will be saved.

        Returns:
            str: Markdown representation of the title element.
        """
        if self.title_content == None:
            logging.warning('Tried to write a NoneType object')
            return '\n\n'

        # Add Markdown header formatting to the title content
        formatted_title = '## ' + self.title_content + '\n\n'
        return formatted_title

scanipy/elements/title_element.py
SVJLucas/Scanipy
from .table import TableDataExtractor
from .title import TitleExtractor
from .text import TextExtractor
from .equation import EquationExtractor
from .image import ImageExtractor

scanipy/extractors/__init__.py
SVJLucas/Scanipy
from typing import Union
from PIL import Image
from fitz import Page
from scanipy.deeplearning.models import EquationToLatex
from scanipy.pdfhandler import PDFPage
from .extractor import Extractor
from scanipy.elements import EquationElement

# Define the EquationExtractor class
class EquationExtractor(Extractor):
    """
    Represents an equation extractor for extracting equations from a document.

    Attributes:
        latex_ocr (str): Deep Learning Model to extract equations from images.
    """

    def __init__(self):
        """
        Initialize an EquationExtractor object.
        """
        # Initialize the OCR model for converting equation images to LaTeX
        self.latex_ocr = EquationToLatex()

    def extract(self, page: PDFPage, equation_element: EquationElement) -> EquationElement:
        """
        Extracts an equation from a given page image based on the coordinates in the equation element.

        Args:
            page_image (PIL.Image): The page image from which to extract the equation.
            equation_element (EquationElement): The equation element containing the coordinates for extraction.

        Returns:
            EquationElement: The updated equation element with the extracted LaTeX content.

        Raises:
            TypeError: If the types of the arguments are not as expected.
        """
        # Separate PIL and pdfplumber elements from the page
        pdf_page = page.get_pdf()
        page_image = page.get_image()

        # Verify the input variable types
        if not isinstance(page_image, Image.Image):
            raise TypeError("page_image must be a PIL.Image object")
        if not isinstance(equation_element, EquationElement):
            raise TypeError("equation_element must be an EquationElement object")

        # Extract the coordinates from the equation element
        left = int(equation_element.x_min * page_image.width)
        upper = int(equation_element.y_min * page_image.height)
        right = int(equation_element.x_max * page_image.width)
        lower = int(equation_element.y_max * page_image.height)

        # Crop the image based on the coordinates
        equation_image = page_image.crop((left, upper, right, lower))

        # Convert the cropped equation image to LaTeX using the OCR model
        latex = self.latex_ocr(equation_image)

        # Update the equation element with the extracted LaTeX content
        equation_element.latex_content = latex

        return equation_element

    def __str__(self) -> str:
        """
        Returns a string representation of the EquationExtractor object.

        Returns:
            str: A string representation of the object.
        """
        return f"EquationExtractor(latex_ocr={self.latex_ocr})"

scanipy/extractors/equation.py
SVJLucas/Scanipy
from scanipy.pdfhandler import PDFPage

class Extractor:
    def extract(self, page: PDFPage):
        raise NotImplementedError

scanipy/extractors/extractor.py
SVJLucas/Scanipy
import random
import string
from typing import Union
from PIL import Image
from pdfplumber.page import Page
from .extractor import Extractor
from scanipy.elements import ImageElement
from scanipy.pdfhandler import PDFPage

# Define the ImageExtractor class
class ImageExtractor(Extractor):
    """
    Represents an image extractor for extracting images from a document.

    Attributes:
        unique_keys (set): A set to store unique keys for each extracted image.
    """

    def __init__(self):
        """
        Initialize an ImageExtractor object.
        """
        # Initialize a set to store unique keys for each extracted image
        self.unique_keys = set()

    def generate_random_string(self, length: int = 32) -> str:
        """
        Generate a random string of a given length consisting of characters a-z, A-Z, and 0-9.

        Args:
            length (int): The length of the random string to be generated. Default is 32.

        Returns:
            str: The generated random string.
        """
        # Verify the input variable type
        if not isinstance(length, int):
            raise TypeError("length must be an integer")

        # Define the characters that can be used in the random string
        characters = string.ascii_letters + string.digits  # a-z, A-Z, 0-9

        # Generate the random string using random.choices()
        random_string = ''.join(random.choices(characters, k=length))

        return random_string

    def extract(self, page: PDFPage, image_element: ImageElement, unique_key: Union[str, None] = None, image_extension: str = 'png' ) -> ImageElement:
        """
        Extracts an image from a given page image based on the coordinates in the image element.

        Args:
            page_image (PIL.Image): The page image from which to extract the image.
            image_element (ImageElement): The image element containing the coordinates for extraction.
            unique_key (Union[str, None], optional): A unique key for the image. If None, a random unique key will be generated.
            image_extension (str): The extension of the image. Defauls to be 'png'.

        Returns:
            ImageElement: The updated image element with the extracted image content.

        Raises:
            TypeError: If the unique_key is not unique.
        """
        # Separate PIL and pdfplumber elements from the page
        pdf_page = page.get_pdf()
        page_image = page.get_image()

        # Verify the input variable types
        if not isinstance(page_image, Image.Image):
            raise TypeError("page_image must be a PIL.Image object")
        if not isinstance(image_element, ImageElement):
            raise TypeError("image_element must be an ImageElement object")
        if not isinstance(unique_key, str) and unique_key is not None:
            raise TypeError("unique_key must be an string or None")
        if not isinstance(image_extension, str) and image_extension is not None:
            raise TypeError("image_extension must be an string or None")
        
        # Extract the coordinates from the image element
        left = int(image_element.x_min * page_image.width)
        upper = int(image_element.y_min * page_image.height)
        right = int(image_element.x_max * page_image.width)
        lower = int(image_element.y_max * page_image.height)

        # Crop the image based on the coordinates
        image_content = page_image.crop((left, upper, right, lower))

        # Generate a unique key if not provided
        if unique_key is None:
            unique_key = self.generate_random_string()
            while unique_key in self.unique_keys:
                unique_key = self.generate_random_string()
        elif unique_key in self.unique_keys:
            raise TypeError("The parameters unique_key must be unique.")

        # Update the set of used keys
        self.unique_keys.add(unique_key)

        # Update the image element with the extracted image content and unique key
        image_element.image_content = image_content
        image_element.unique_key = unique_key
        image_element.image_extension  = image_extension

        return image_element

    def __str__(self) -> str:
        """
        Returns a string representation of the ImageExtractor object.

        Returns:
            str: A string representation of the object.
        """
        return f"ImageExtractor(unique_keys={self.unique_keys})"

scanipy/extractors/image.py
SVJLucas/Scanipy
from typing import Union
from PIL import Image
from fitz import Page
from scanipy.deeplearning.models import TableStructureAnalyzer
from scanipy.elements import TableElement
from scanipy.pdfhandler import PDFPage
from .extractor import Extractor

from math import ceil, floor
import numpy as np
import pandas as pd
import cv2
import pytesseract #TODO use TextOCR or 

# Define the TableDataExtractor class
class TableDataExtractor(Extractor):
    """
    Represents an table extractor for extracting tables from a document.

    Attributes:
        latex_ocr (str): Deep Learning Model to extract tables from images.
    """

    def __init__(self, table_expansion_margin=10, threshold_percentage=0.10):
        """
        Initialize an TableDataExtractor object.
        """
        # Initialize the model for identifying table structures
        self.model = TableStructureAnalyzer()

        # Expand the bounding box slightly for better cropping
        self._table_expansion_margin = table_expansion_margin

        # Use a percentage (e.g., 10%) of the average height as the threshold for a new row
        self._threshold_percentage = threshold_percentage
        self.test = []

    def _get_cell_coordinates(self, page_image, table_element):
        """
        Obtains the coordinates of cells based on the analyzed table structure.

        Args:
            analyzed_structure (List[Dict]): Analyzed table structure.

        Returns:
            List[Dict]: List of cell coordinates.
        """
        # Convert PIL image to OpenCV format
        image_cv = cv2.cvtColor(np.array(page_image), cv2.COLOR_RGB2BGR)

        # Get image dimensions
        img_height, img_width = image_cv.shape[:2]

        # Expand the bounding box slightly for better cropping
        xmin, ymin = max(0, table_element.x_min - self._table_expansion_margin), max(0, table_element.y_min - self._table_expansion_margin)
        xmax, ymax = min(img_width, table_element.x_max + self._table_expansion_margin), min(img_height, table_element.y_max + self._table_expansion_margin)

        # Crop the image based on the coordinates
        table_image = page_image.crop((xmin, ymin, xmax, ymax))

        # Extract the structure from the cropped table image
        table_structure = self.model(table_image)

        # Verify if table structure is empty
        if table_structure is None:
            raise ValueError('Found a null table structure') #TODO make a logic for this case
        
        # For some reason, this improves the cell boxes #FIXME
        for box in table_structure: # type: ignore
            box['box']['ymax'] =  box['box']['ymax'] + 5 # type: ignore

        # Extract row and column bounding boxes from the table structure
        rows = [box['box'] for box in table_structure if box['label'] == 'table row'] # type: ignore
        columns = [box['box'] for box in table_structure if box['label'] == 'table column'] # type: ignore


        # Initialize a list to store cell positions
        cells = []

        # Loop through each row and column to find overlapping regions as cells
        #TODO: adapt for different table layouts
        for row in rows:
            for column in columns:
                cell_xmin = max(row['xmin'], column['xmin']) # type: ignore
                cell_ymin = max(row['ymin'], column['ymin']) # type: ignore
                cell_xmax = min(row['xmax'], column['xmax']) # type: ignore
                cell_ymax = min(row['ymax'], column['ymax']) # type: ignore

                # Add the cell only if it has a non-zero area
                if cell_xmin < cell_xmax and cell_ymin < cell_ymax:
                    cells.append({
                        'xmin': xmin + cell_xmin,
                        'ymin': ymin + cell_ymin,
                        'xmax': xmin + cell_xmax,
                        'ymax': ymin + cell_ymax
                    })
        # print(cells)
        return cells
    
    def _get_dataframe(self, rows, page_image):
        # Initialize an empty DataFrame
        df = pd.DataFrame()

        # Populate the DataFrame
        for row_idx, row_boxes in enumerate(rows):
            row_data = []
            for col_idx, box in enumerate(row_boxes):
                image_cv = cv2.cvtColor(np.array(page_image), cv2.COLOR_RGB2BGR)
                text = self.get_text_from_image(image_cv, box)
                row_data.append(text)

            # If it's the first row, set the DataFrame columns
            if row_idx == 0:
                df = pd.DataFrame(columns=row_data)
            else:
                df.loc[row_idx - 1] = row_data # type: ignore
        
        return df
    
    def _pad_image_with_white(self, image_cv, pad):

        # Get the dimensions of the original image
        height, width, _ = image_cv.shape

        # Calculate the dimensions of the new padded image
        padded_height = height + pad + pad
        padded_width = width + pad + pad

        # Create white image
        padded_image = np.ones((padded_height, padded_width, 3), dtype=np.uint8) * 255

        # Paste the original image onto the canvas with padding 
        padded_image[pad:pad + height, pad:pad + width] = image_cv

        return padded_image
    
    def _remove_table_lines(self, image_cv):
        # Convert to grayscale
        image_cv_gray = cv2.cvtColor(image_cv, cv2.COLOR_BGR2GRAY)

        # Use Hough Line Transform to detect lines
        edges = cv2.bitwise_not(image_cv_gray).astype('uint8')
        lines = cv2.HoughLines(edges, 1, np.pi / 180, 150)

        # Draw the detected lines on a copy of the original image
        cleaned_image = image_cv.copy()
        if lines is not None:
            for line in lines:
                rho, theta = line[0]
                a = np.cos(theta)
                b = np.sin(theta)
                x0 = a * rho
                y0 = b * rho
                x1 = int(x0 + 1000 * (-b))
                y1 = int(y0 + 1000 * (a))
                x2 = int(x0 - 1000 * (-b))
                y2 = int(y0 - 1000 * (a))
                cv2.line(cleaned_image, (x1, y1), (x2, y2), (255, 255, 255), 6)

        return cleaned_image


    def _separate_rows(self, cell_coordinates):
        # Sort boxes by ymin to separate rows
        cell_coordinates = sorted(cell_coordinates, key=lambda x: x['ymin'])

        # Calculate the average height of all boxes
        avg_height = np.mean([box['ymax'] - box['ymin'] for box in cell_coordinates])

        # Use a percentage (e.g., 10%) of the average height as the threshold for a new row
        threshold = self._threshold_percentage * avg_height

        # Identify unique rows by their y-coordinates
        rows = []
        last_ymin = None
        for box in cell_coordinates:
            ymin = box['ymin']
            if last_ymin is None or abs(ymin - last_ymin) > threshold:  # Adjust the threshold as needed
                rows.append([])
            rows[-1].append(box)
            last_ymin = ymin

        # Sort each row by xmin to arrange columns
        for row in rows:
            row.sort(key=lambda x: x['xmin'])

        return rows

    def get_text_from_image(self, image_cv, box):
        xmin, ymin, xmax, ymax = box['xmin'], box['ymin'], box['xmax'], box['ymax']

        # Extract each cell image
        # cell_image = image_cv[ymin:ymax, xmin:xmax]
        cell_image = image_cv[floor(ymin):ceil(ymax), floor(xmin):ceil(xmax)]

        text = ''
        for pad in [3,8,13,20]: # don't ask me why, but some numbers can only be read with a SPECIFIC padding #FIXME #despair
            # Pad image with white to improve OCR
            cell_image_padded = self._pad_image_with_white(cell_image, pad)
            
            # Clean the table lines to improve OCR performance
            cell_image_padded = self._remove_table_lines(cell_image_padded)

            # Convert to PIL Image
            cell_image_padded_pil = Image.fromarray(cv2.cvtColor(cell_image_padded, cv2.COLOR_BGR2RGB))

            # Perform OCR to get text
            text = pytesseract.image_to_string(cell_image_padded_pil).strip().replace('\n', '')

            if text != '':
                break

        return text

    def extract(self, page: PDFPage, table_element: TableElement) -> TableElement:
        """
        Extracts an table from a given page image based on the coordinates in the table element.

        Args:
            page_image (PIL.Image): The page image from which to extract the table.
            table_element (TableElement): The table element containing the coordinates for extraction.

        Returns:
            TableElement: The updated table element with the extracted LaTeX content.

        Raises:
            TypeError: If the types of the arguments are not as expected.
        """
        # Separate PIL and pdfplumber elements from the page
        pdf_page = page.get_pdf()
        page_image = page.get_image()

        # Verify the input variable types
        if not isinstance(page_image, Image.Image):
            raise TypeError("page_image must be a PIL.Image object")
        if not isinstance(table_element, TableElement):
            raise TypeError("table_element must be an TableElement object") #TODO

        # Gather the individual cells from the table
        cell_coordinates = self._get_cell_coordinates(page_image, table_element)

        rows = self._separate_rows(cell_coordinates)

        dataframe = self._get_dataframe(rows, page_image)

        # Update the table element with the extracted LaTeX content
        table_element._table_data = dataframe

        return table_element

    def __str__(self) -> str:
        """
        Returns a string representation of the TableDataExtractor object.

        Returns:
            str: A string representation of the object.
        """
        return f"TableDataExtractor(latex_ocr={self.model})"

scanipy/extractors/table.py
SVJLucas/Scanipy
import torch
import pix2text
from PIL import Image
from pdfplumber.page import Page
from .extractor import Extractor
from scanipy.elements import TextElement 
from scanipy.deeplearning.models import TextOCR
from scanipy.pdfhandler import PDFPage
from typing import Union, Tuple, List, Dict

# Define the TextExtractor class
class TextExtractor(Extractor):
    """
    Represents a text extractor for extracting text from a document.
    """

    def __init__(self, use_ocr: bool, lang: str = 'en', tolerance: float = 1.5):
      """
      Initialize a TextExtractor object.

      Args:
          use_ocr (bool): Whether to use OCR for text extraction or not.
          lang (str): The language of the text. Defaults to english.
          tolerance (float): The tolerance level for text extraction. Default is 1.5.

      Raises:
          TypeError: If the types of the arguments are not as expected.
      """

      # Verify the input variable types
      if not isinstance(use_ocr, bool):
          raise TypeError("use_ocr must be a boolean")

      # Initialize OCR settings and OCR model
      self.use_ocr = use_ocr

      # Check if CUDA is available and set the device accordingly
      self.device = 'cuda' if torch.cuda.is_available() else 'cpu'

      # Initialize OCR model for text and equations
      self.text_equations_ocr = pix2text.Pix2Text(device=self.device)

      # Set the language for OCR
      self.lang = lang

      # Initialize the TextOCR object
      self.text_ocr = TextOCR(self.lang, self.device)

      # Set the tolerance level for text extraction
      self.tolerance = tolerance

    def _process_text_image(self, text_element: TextElement, text_image: Image.Image, pdf_page: Page) -> str:
      """
      Process the text image based on OCR settings and equation presence.

      Args:
          text_element (TextElement): The text element containing the coordinates for extraction.
          text_image (Image): The cropped text image.
          pdf_page (Page): The PDF page containing the text element.

      Returns:
          str: The extracted text content.
      """
      # Use OCR or not based on the use_ocr flag and whether the text element contains equations
      if self.use_ocr:
          if text_element.has_equation_inside:
              return self._get_text_and_equations_with_ocr(text_image)
          else:
              return self._get_text_with_ocr(text_image)
      else:
          if text_element.has_equation_inside:
              return self._get_text_and_equations_without_ocr(text_element, text_image, pdf_page)
          else:
              return self._get_text_without_ocr(text_element, text_image, pdf_page)

    def extract(self, page: PDFPage, text_element: TextElement) -> TextElement:
        """
        Extracts text from a given page image based on the coordinates in the text element.

        Args:
            pdf_page (Page): The PDF page from which to extract the text.
            page_image (Image): The page image from which to extract the text.
            text_element (TextElement): The text element containing the coordinates for extraction.

        Returns:
            TextElement: The updated text element with the extracted text content.

        Raises:
            TypeError: If the types of the arguments are not as expected.
        """
        # Separate PIL and pdfplumber elements from the page
        pdf_page = page.get_pdf()
        page_image = page.get_image()
      
        # Verify the input variable types
        if not isinstance(page_image, Image.Image):
            raise TypeError("page_image must be a PIL.Image object")
        if not isinstance(text_element, TextElement):
            raise TypeError("text_element must be a TextElement object")

        # Extract the coordinates from the text element
        left = int(text_element.x_min * page_image.width)
        upper = int(text_element.y_min * page_image.height)
        right = int(text_element.x_max * page_image.width)
        lower = int(text_element.y_max * page_image.height)

        # Crop the image based on the coordinates
        text_image = page_image.crop((left, upper, right, lower))

        # Process the cropped text image and extract the text content
        text_content = self._process_text_image(text_element, text_image, pdf_page)

        # Update the text element with the extracted text content
        text_element.text_content = text_content

        return text_element

    def __str__(self) -> str:
        """
        Returns a string representation of the TextExtractor object.

        Returns:
            str: A string representation of the object.
        """
        return f"TextExtractor(use_ocr={self.use_ocr}, device={self.device}, lang={self.lang}, tolerance={self.tolerance}, text_ocr={self.text_ocr})"

    @property
    def tolerance(self) -> float:
        """
        Getter for the 'tolerance' attribute.

        Returns:
            float: The current tolerance level for text extraction.
        """
        return self._tolerance

    @tolerance.setter
    def tolerance(self, value: float):
        """
        Setter for the 'tolerance' attribute.

        Args:
            value (float): The new tolerance level to set for text extraction.

        Raises:
            TypeError: If the provided value is not a float.
        """
        if not isinstance(value, float):
            raise TypeError("Tolerance must be a float.")
        self._tolerance = value

    def _get_text_and_equations_with_ocr(self, equation_text_image: Image.Image) -> str:
        """
        Extracts text and equations from a given image using Optical Character Recognition (OCR).

        Args:
            equation_text_image (Image): The cropped image containing the text and equations to be extracted.

        Returns:
            str: The extracted text and equations content.

        Raises:
            TypeError: If the type of the equation_text_image is not PIL.Image.

        """

        # Verify the input variable type
        if not isinstance(equation_text_image, Image.Image):
            raise TypeError("equation_text_image must be a PIL.Image object")

        # Get OCR results for the text image
        ocr_results = self.text_equations_ocr(equation_text_image)

        # Loop through each OCR result box to process text and equations
        for box in ocr_results:
            if box['type'] == 'text':
                # Extract and normalize coordinates from the OCR box
                x_min, y_min, x_max, y_max = self._extract_coordinates(box)

                # Crop the image based on the coordinates to get the phrase
                cropped_phrase_image = equation_text_image.crop((x_min, y_min, x_max, y_max))

                # Use OCR to extract text from the cropped phrase image
                extracted_text = self.text_ocr(cropped_phrase_image)

                # Remove newline characters from the extracted text
                extracted_text = extracted_text.replace('\n', '')

                # Update the OCR box with the extracted text
                box['text'] = extracted_text

        # Merge the extracted texts into a single string
        final_extracted_text = self._merge_line_texts(ocr_results)

        return final_extracted_text

    def _get_text_with_ocr(self, cropped_text_image: Image.Image) -> str:
        """
        Extracts text from a given image using Optical Character Recognition (OCR).

        Args:
            cropped_text_image (Image): The cropped image containing the text to be extracted.

        Returns:
            str: The extracted text content.

        Raises:
            TypeError: If the type of the cropped_text_image is not PIL.Image.

        """

        # Verify the input variable type
        if not isinstance(cropped_text_image, Image.Image):
            raise TypeError("cropped_text_image must be a PIL.Image object")

        # Use OCR to extract text from the image
        extracted_text = self.text_ocr(cropped_text_image)

        return extracted_text

    def _get_text_without_ocr(self, text_element: TextElement, cropped_image: Image.Image, pdf_page: Page) -> str:
        """
        Extracts text from a given area in a PDF page without using OCR.

        Args:
            text_element (TextElement): The text element containing the coordinates for text extraction.
            cropped_image (Image): The cropped image containing the text.
            pdf_page (Page): The PDF page from which the text is to be extracted.

        Returns:
            str: The extracted text content.

        Raises:
            TypeError: If the types of the arguments are not as expected.
        """

        # Verify the input variable types
        if not isinstance(text_element, TextElement):
            raise TypeError("text_element must be an instance of TextElement")
        if not isinstance(cropped_image, Image.Image):
            raise TypeError("cropped_image must be a PIL.Image object")
        if not isinstance(pdf_page, Page):
            raise TypeError("pdf_page must be an instance of pdfplumber.Page")

        # Get the dimensions of the PDF page
        pdf_width, pdf_height = pdf_page.width, pdf_page.height

        # Calculate the coordinates for text extraction based on the TextElement
        x_min_coord = int(text_element.x_min * pdf_width)
        y_min_coord = int(text_element.y_min * pdf_height)
        x_max_coord = int(text_element.x_max * pdf_width)
        y_max_coord = int(text_element.y_max * pdf_height)

        # Define the bounding box for cropping the PDF page
        bounding_box = (x_min_coord, y_min_coord, x_max_coord, y_max_coord)

        # Crop the PDF page based on the bounding box
        cropped_pdf_page = pdf_page.crop(bounding_box)

        # Extract text from the cropped PDF page
        extracted_text = cropped_pdf_page.extract_text(x_tolerance=self.tolerance)

        # Remove newline characters from the extracted text
        cleaned_text = extracted_text.replace('\n', '')

        return cleaned_text

    def _extract_coordinates(self, box: Dict):
        """
        Extract coordinates from the OCR box.

        Args:
            box (Dict): The OCR result box.

        Returns:
            Tuple[float, float, float, float]: Coordinates (x_min, y_min, x_max, y_max).
        """
        x_min = min(box['position'][:, 0])
        y_min = min(box['position'][:, 1])
        x_max = max(box['position'][:, 0])
        y_max = max(box['position'][:, 1])

        return x_min, y_min, x_max, y_max

    def _extract_and_normalize_coordinates(self, box: Dict, image_size: Tuple[int, int]) -> Tuple[float, float, float, float]:
        """
        Extract and normalize coordinates from the OCR box.

        Args:
            box (Dict): The OCR result box.
            image_size (Tuple[int, int]): The size of the image.

        Returns:
            Tuple[float, float, float, float]: Normalized coordinates (x_min, y_min, x_max, y_max).
        """
        x_min = min(box['position'][:, 0])
        y_min = min(box['position'][:, 1])
        x_max = max(box['position'][:, 0])
        y_max = max(box['position'][:, 1])

        # Normalizing Coordinates
        x_min /= image_size[0]
        y_min /= image_size[1]
        x_max /= image_size[0]
        y_max /= image_size[1]

        return x_min, y_min, x_max, y_max

    def _convert_to_cropped_page_referential(self, x_min: float, y_min: float, x_max: float, y_max: float, cropped_page: Page) -> Tuple[int, int, int, int]:
        """
        Convert normalized coordinates to cropped page's referential.

        Args:
            x_min, y_min, x_max, y_max (float): Normalized coordinates.
            cropped_page (Page): The cropped PDF page.

        Returns:
            Tuple[int, int, int, int]: Coordinates in the cropped page's referential.
        """
        x_min = int(x_min * cropped_page.width)
        y_min = int(y_min * cropped_page.height)
        x_max = int(x_max * cropped_page.width)
        y_max = int(y_max * cropped_page.height)

        x0, y0, _, _ = cropped_page.bbox

        return x_min + int(x0), y_min + int(y0), x_max + int(x0), y_max + int(y0)

    def _merge_line_texts(self, ocr_results: List[Dict]) -> str:
        """
        Merges OCR results into a single string.

        This method takes a list of dictionaries containing OCR results and merges them into a single string.
        Each dictionary in the list is expected to contain text information, possibly among other details.

        Args:
            ocr_results (List[Dict]): A list of dictionaries containing OCR results.

        Returns:
            str: A single string that is the result of merging all the OCR text results.

        Note:
            This method currently uses the `pix2text.merge_line_texts` function for merging.
        """
        text = pix2text.merge_line_texts(ocr_results)
        return text

    def _get_text_and_equations_without_ocr(self, text_element: TextElement, text_image: Image.Image, pdf_page: Page) -> str:
        """
        Extracts text and equations from a given text image without using OCR to extract text.

        Args:
            text_element (TextElement): The text element containing the coordinates for extraction.
            text_image (Image): The cropped text image.
            pdf_page (Page): The PDF page containing the text element.

        Returns:
            str: The extracted text content.

        Raises:
            TypeError: If the types of the arguments are not as expected.
        """
        # Verify the input variable types
        if not isinstance(text_element, TextElement):
            raise TypeError("text_element must be a TextElement object")
        if not isinstance(text_image, Image.Image):
            raise TypeError("text_image must be a PIL.Image object")
        if not isinstance(pdf_page, Page):
            raise TypeError("pdf_page must be a pdfplumber.Page object")

        # Get PDF and image dimensions
        pdf_width, pdf_height = pdf_page.width, pdf_page.height
        image_width, image_height = text_image.size

        # Calculate PDF coordinates based on text element's normalized coordinates
        x_min_pdf = int(text_element.x_min * pdf_width)
        y_min_pdf = int(text_element.y_min * pdf_height)
        x_max_pdf = int(text_element.x_max * pdf_width)
        y_max_pdf = int(text_element.y_max * pdf_height)

        # Define the PDF crop box
        pdf_box = (x_min_pdf, y_min_pdf, x_max_pdf, y_max_pdf)

        # Crop the PDF page based on the calculated coordinates
        cropped_page = pdf_page.crop(pdf_box)

        # Get OCR results for the text image
        ocr_results = self.text_equations_ocr(text_image)

        # Process each OCR result box
        for box in ocr_results:
            if box['type'] == 'text':
                # Extract and normalize coordinates from the OCR box
                x_min, y_min, x_max, y_max = self._extract_and_normalize_coordinates(box, text_image.size)

                # Convert normalized coordinates to cropped page's referential
                x_min, y_min, x_max, y_max = self._convert_to_cropped_page_referential(x_min, y_min, x_max, y_max, cropped_page)

                # Extract text from the cropped PDF page based on the calculated coordinates
                segmentation = cropped_page.crop((x_min, y_min, x_max, y_max))
                extracted_text = segmentation.extract_text(x_tolerance=self.tolerance)
                extracted_text = extracted_text.replace('\n', '')

                # Update the OCR box with the extracted text
                box['text'] = extracted_text

        # Merge the extracted texts into a single string
        final_text = self._merge_line_texts(ocr_results)

        return final_text


scanipy/extractors/text.py
SVJLucas/Scanipy
import torch
import pix2text
from PIL import Image
from pdfplumber.page import Page
from .extractor import Extractor
from scanipy.pdfhandler import PDFPage
from scanipy.elements import TitleElement 
from scanipy.deeplearning.models import TextOCR
from typing import Union, Tuple, List, Dict


# Define the TitleExtractor class
class TitleExtractor(Extractor):
    """
    Represents a title extractor for extracting title from a document.
    """

    def __init__(self, use_ocr: bool, lang: str = 'en', tolerance: float = 1.5):
      """
      Initialize a TitleExtractor object.

      Args:
          use_ocr (bool): Whether to use OCR for title extraction or not.
          lang (str): The language of the title. Defaults to english.
          tolerance (float): The tolerance level for title extraction. Default is 1.5.

      Raises:
          TypeError: If the types of the arguments are not as expected.
      """

      # Verify the input variable types
      if not isinstance(use_ocr, bool):
          raise TypeError("use_ocr must be a boolean")

      # Initialize OCR settings and OCR model
      self.use_ocr = use_ocr

      # Check if CUDA is available and set the device accordingly
      self.device = 'cuda' if torch.cuda.is_available() else 'cpu'

      # Initialize OCR model for title and equations
      self.title_equations_ocr = pix2text.Pix2Text(device=self.device)

      # Set the language for OCR
      self.lang = lang

      # Initialize the TextOCR object
      self.title_ocr = TextOCR(self.lang, self.device)

      # Set the tolerance level for title extraction
      self.tolerance = tolerance

    def _process_title_image(self, title_element: TitleElement, title_image: Image.Image, pdf_page: Page) -> str:
      """
      Process the title image based on OCR settings and equation presence.

      Args:
          title_element (TitleElement): The title element containing the coordinates for extraction.
          title_image (Image): The cropped title image.
          pdf_page (Page): The PDF page containing the title element.

      Returns:
          str: The extracted title content.
      """
      # Use OCR or not based on the use_ocr flag and whether the title element contains equations
      if self.use_ocr:
          if title_element.has_equation_inside:
              return self._get_title_and_equations_with_ocr(title_image)
          else:
              return self._get_title_with_ocr(title_image)
      else:
          if title_element.has_equation_inside:
              return self._get_title_and_equations_without_ocr(title_element, title_image, pdf_page)
          else:
              return self._get_title_without_ocr(title_element, title_image, pdf_page)

    def extract(self, page: PDFPage, title_element: TitleElement) -> TitleElement:
        """
        Extracts title from a given page image based on the coordinates in the title element.

        Args:
            pdf_page (Page): The PDF page from which to extract the title.
            page_image (Image): The page image from which to extract the title.
            title_element (TitleElement): The title element containing the coordinates for extraction.

        Returns:
            TitleElement: The updated title element with the extracted title content.

        Raises:
            TypeError: If the types of the arguments are not as expected.
        """
        # Separate PIL and pdfplumber elements from the 
        pdf_page = page.get_pdf()
        page_image = page.get_image()

        # Verify the input variable types
        if not isinstance(page_image, Image.Image):
            raise TypeError("page_image must be a PIL.Image object")
        if not isinstance(title_element, TitleElement):
            raise TypeError("title_element must be a TitleElement object")

        # Extract the coordinates from the equation element
        left = int(title_element.x_min * page_image.width)
        upper = int(title_element.y_min * page_image.height)
        right = int(title_element.x_max * page_image.width)
        lower = int(title_element.y_max * page_image.height)

        # Crop the image based on the coordinates
        title_image = page_image.crop((left, upper, right, lower))

        # Process the cropped title image and extract the title content
        title_content = self._process_title_image(title_element, title_image, pdf_page)

        # Update the title element with the extracted title content
        title_element.title_content = title_content

        return title_element

    def __str__(self) -> str:
        """
        Returns a string representation of the TitleExtractor object.

        Returns:
            str: A string representation of the object.
        """
        return f"TitleExtractor(use_ocr={self.use_ocr}, device={self.device}, lang={self.lang}, tolerance={self.tolerance}, title_ocr={self.title_ocr})"

    @property
    def tolerance(self) -> float:
        """
        Getter for the 'tolerance' attribute.

        Returns:
            float: The current tolerance level for title extraction.
        """
        return self._tolerance

    @tolerance.setter
    def tolerance(self, value: float):
        """
        Setter for the 'tolerance' attribute.

        Args:
            value (float): The new tolerance level to set for title extraction.

        Raises:
            TypeError: If the provided value is not a float.
        """
        if not isinstance(value, float):
            raise TypeError("Tolerance must be a float.")
        self._tolerance = value

    def _get_title_and_equations_with_ocr(self, equation_title_image: Image.Image) -> str:
        """
        Extracts title and equations from a given image using Optical Character Recognition (OCR).

        Args:
            equation_title_image (Image): The cropped image containing the title and equations to be extracted.

        Returns:
            str: The extracted title and equations content.

        Raises:
            TypeError: If the type of the equation_title_image is not PIL.Image.

        """

        # Verify the input variable type
        if not isinstance(equation_title_image, Image.Image):
            raise TypeError("equation_title_image must be a PIL.Image object")

        # Get OCR results for the title image
        ocr_results = self.title_equations_ocr(equation_title_image)

        # Loop through each OCR result box to process title and equations
        for box in ocr_results:
            if box['type'] == 'text':
                # Extract and normalize coordinates from the OCR box
                x_min, y_min, x_max, y_max = self._extract_coordinates(box)

                # Crop the image based on the coordinates to get the phrase
                cropped_phrase_image = equation_title_image.crop((x_min, y_min, x_max, y_max))

                # Use OCR to extract title from the cropped phrase image
                extracted_title = self.title_ocr(cropped_phrase_image)

                # Remove newline characters from the extracted title
                extracted_title = extracted_title.replace('\n', '')

                # Update the OCR box with the extracted title
                box['text'] = extracted_title

        # Merge the extracted titles into a single string
        final_extracted_title = self._merge_line_titles(ocr_results)

        return final_extracted_title

    def _get_title_with_ocr(self, cropped_title_image: Image.Image) -> str:
        """
        Extracts title from a given image using Optical Character Recognition (OCR).

        Args:
            cropped_title_image (Image): The cropped image containing the title to be extracted.

        Returns:
            str: The extracted title content.

        Raises:
            TypeError: If the type of the cropped_title_image is not PIL.Image.

        """

        # Verify the input variable type
        if not isinstance(cropped_title_image, Image.Image):
            raise TypeError("cropped_title_image must be a PIL.Image object")

        # Use OCR to extract title from the image
        extracted_title = self.title_ocr(cropped_title_image)

        return extracted_title

    def _get_title_without_ocr(self, title_element: TitleElement, cropped_image: Image.Image, pdf_page: Page) -> str:
        """
        Extracts title from a given area in a PDF page without using OCR.

        Args:
            title_element (TitleElement): The title element containing the coordinates for title extraction.
            cropped_image (Image): The cropped image containing the title.
            pdf_page (Page): The PDF page from which the title is to be extracted.

        Returns:
            str: The extracted title content.

        Raises:
            TypeError: If the types of the arguments are not as expected.
        """

        # Verify the input variable types
        if not isinstance(title_element, TitleElement):
            raise TypeError("title_element must be an instance of TitleElement")
        if not isinstance(cropped_image, Image.Image):
            raise TypeError("cropped_image must be a PIL.Image object")
        if not isinstance(pdf_page, Page):
            raise TypeError("pdf_page must be an instance of pdfplumber.Page")

        # Get the dimensions of the PDF page
        pdf_width, pdf_height = pdf_page.width, pdf_page.height

        # Calculate the coordinates for title extraction based on the TitleElement
        x_min_coord = int(title_element.x_min * pdf_width)
        y_min_coord = int(title_element.y_min * pdf_height)
        x_max_coord = int(title_element.x_max * pdf_width)
        y_max_coord = int(title_element.y_max * pdf_height)

        # Define the bounding box for cropping the PDF page
        bounding_box = (x_min_coord, y_min_coord, x_max_coord, y_max_coord)

        # Crop the PDF page based on the bounding box
        cropped_pdf_page = pdf_page.crop(bounding_box)

        # Extract title from the cropped PDF page
        extracted_title = cropped_pdf_page.extract_text(x_tolerance=self.tolerance)

        # Remove newline characters from the extracted title
        cleaned_title = extracted_title.replace('\n', '')

        return cleaned_title

    def _extract_coordinates(self, box: Dict):
        """
        Extract coordinates from the OCR box.

        Args:
            box (Dict): The OCR result box.

        Returns:
            Tuple[float, float, float, float]: Coordinates (x_min, y_min, x_max, y_max).
        """
        x_min = min(box['position'][:, 0])
        y_min = min(box['position'][:, 1])
        x_max = max(box['position'][:, 0])
        y_max = max(box['position'][:, 1])

        return x_min, y_min, x_max, y_max

    def _extract_and_normalize_coordinates(self, box: Dict, image_size: Tuple[int, int]) -> Tuple[float, float, float, float]:
        """
        Extract and normalize coordinates from the OCR box.

        Args:
            box (Dict): The OCR result box.
            image_size (Tuple[int, int]): The size of the image.

        Returns:
            Tuple[float, float, float, float]: Normalized coordinates (x_min, y_min, x_max, y_max).
        """
        x_min = min(box['position'][:, 0])
        y_min = min(box['position'][:, 1])
        x_max = max(box['position'][:, 0])
        y_max = max(box['position'][:, 1])

        # Normalizing Coordinates
        x_min /= image_size[0]
        y_min /= image_size[1]
        x_max /= image_size[0]
        y_max /= image_size[1]

        return x_min, y_min, x_max, y_max

    def _convert_to_cropped_page_referential(self, x_min: float, y_min: float, x_max: float, y_max: float, cropped_page: Page) -> Tuple[int, int, int, int]:
        """
        Convert normalized coordinates to cropped page's referential.

        Args:
            x_min, y_min, x_max, y_max (float): Normalized coordinates.
            cropped_page (Page): The cropped PDF page.

        Returns:
            Tuple[int, int, int, int]: Coordinates in the cropped page's referential.
        """
        x_min = int(x_min * cropped_page.width)
        y_min = int(y_min * cropped_page.height)
        x_max = int(x_max * cropped_page.width)
        y_max = int(y_max * cropped_page.height)

        x0, y0, _, _ = cropped_page.bbox

        return x_min + int(x0), y_min + int(y0), x_max + int(x0), y_max + int(y0)

    def _merge_line_titles(self, ocr_results: List[Dict]) -> str:
        """
        Merges OCR results into a single string.

        This method takes a list of dictionaries containing OCR results and merges them into a single string.
        Each dictionary in the list is expected to contain title information, possibly among other details.

        Args:
            ocr_results (List[Dict]): A list of dictionaries containing OCR results.

        Returns:
            str: A single string that is the result of merging all the OCR title results.

        Note:
            This method currently uses the `pix2text.merge_line_texts` function for merging.
        """
        text = pix2text.merge_line_texts(ocr_results)
        return text

    def _get_title_and_equations_without_ocr(self, title_element: TitleElement, title_image: Image.Image, pdf_page: Page) -> str:
        """
        Extracts title and equations from a given title image without using OCR to extract title.

        Args:
            title_element (TitleElement): The title element containing the coordinates for extraction.
            title_image (Image): The cropped title image.
            pdf_page (Page): The PDF page containing the title element.

        Returns:
            str: The extracted title content.

        Raises:
            TypeError: If the types of the arguments are not as expected.
        """
        # Verify the input variable types
        if not isinstance(title_element, TitleElement):
            raise TypeError("title_element must be a TitleElement object")
        if not isinstance(title_image, Image.Image):
            raise TypeError("title_image must be a PIL.Image object")
        if not isinstance(pdf_page, Page):
            raise TypeError("pdf_page must be a pdfplumber.Page object")

        # Get PDF and image dimensions
        pdf_width, pdf_height = pdf_page.width, pdf_page.height
        image_width, image_height = title_image.size

        # Calculate PDF coordinates based on title element's normalized coordinates
        x_min_pdf = int(title_element.x_min * pdf_width)
        y_min_pdf = int(title_element.y_min * pdf_height)
        x_max_pdf = int(title_element.x_max * pdf_width)
        y_max_pdf = int(title_element.y_max * pdf_height)

        # Define the PDF crop box
        pdf_box = (x_min_pdf, y_min_pdf, x_max_pdf, y_max_pdf)

        # Crop the PDF page based on the calculated coordinates
        cropped_page = pdf_page.crop(pdf_box)

        # Get OCR results for the title image
        ocr_results = self.title_equations_ocr(title_image)

        # Process each OCR result box
        for box in ocr_results:
            if box['type'] == 'text':
                # Extract and normalize coordinates from the OCR box
                x_min, y_min, x_max, y_max = self._extract_and_normalize_coordinates(box, title_image.size)

                # Convert normalized coordinates to cropped page's referential
                x_min, y_min, x_max, y_max = self._convert_to_cropped_page_referential(x_min, y_min, x_max, y_max, cropped_page)

                # Extract title from the cropped PDF page based on the calculated coordinates
                segmentation = cropped_page.crop((x_min, y_min, x_max, y_max))
                extracted_title = segmentation.extract_text(x_tolerance=self.tolerance)
                extracted_title = extracted_title.replace('\n', '')

                # Update the OCR box with the extracted title
                box['text'] = extracted_title

        # Merge the extracted titles into a single string
        final_title = self._merge_line_titles(ocr_results)

        return final_title


scanipy/extractors/title.py
SVJLucas/Scanipy
import fitz
import logging

from .pdfhandler import PDFDocument
from .deeplearning.models import LayoutDetector, EquationFinder
from .elements import TitleElement, TextElement, TableElement, EquationElement, TitleElement, ImageElement
from .extractors import TitleExtractor, TextExtractor, TableDataExtractor, EquationExtractor, ImageExtractor
from .document import Document
from collections import defaultdict
import os


class Parser:
    """
    Parses a PDF file using PyMuPDF (fitz) and extracts its content into a Markdown file.

    Attributes:
        pdf_file (PyMuPDF.Document): The PyMuPDF Document object representing the PDF file.
    """

    def __init__(self): #TODO define device here
        """
        Initialize a new Parser instance.

        :param path: The path to the PDF file to parse.
        """
        self.layout_detector = LayoutDetector(device='cuda')
        self.table_extractor = TableDataExtractor()
        self.text_extractor = TextExtractor(use_ocr=False)
        self.title_extractor = TitleExtractor(use_ocr=False)
        self.image_extractor = ImageExtractor()
        self.equation_finder = EquationFinder(device='gpu')
        self.equation_extractor = EquationExtractor()
        # self.pipeline = [self.text_extractor, self.table_extractor, self.equation_extractor]
    
    def extract(self, path): #TODO add min and max pages
        pdfdoc = PDFDocument(path)
        elements_h = {}
        document = Document()
        image_number = 0

        for page in pdfdoc:
            elements = self.layout_detector(page.get_image())
            equations = self.equation_finder(page.get_image())

            logging.info(f'Detected {len(elements)} elements and {len(equations)} equations')

            for equation in equations:
                if equation.is_inside_text:
                    for element in elements:
                        if equation.is_in(element):
                            element.has_equation_inside = True
                            element.equation_inside = equation
                            break #TODO: what if two elements have the same equation? #BUG
            
            for element in elements:
                if isinstance(element, TextElement):
                    element = self.text_extractor.extract(page, element)
                elif isinstance(element, TableElement):
                    element = self.table_extractor.extract(page, element)
                elif isinstance(element, TitleElement):
                    element = self.title_extractor.extract(page, element)
                elif isinstance(element, ImageElement):
                    element = self.image_extractor.extract(page, element, unique_key=str(image_number))
                    image_number += 1

                document.add_element(page.page_number, element)
            for equation in equations:
                equation = self.equation_extractor.extract(page, equation)
                document.add_element(page.page_number, equation)
            elements_h[page.page_number] = [*elements,*equations]

        return document #elements_h
        

scanipy/parser.py
SVJLucas/Scanipy
from PIL import Image, ImageDraw
import pdfplumber
import pdfplumber.page
from typing import Tuple, Iterator, List

def draw_rectangle(image: Image.Image, coordinates: Tuple[float, float, float, float]) -> Image.Image:
    """
    Draw a rectangle on the page image and return the modified image.

    Args:
        image (PIL.Image.Image): The image on which to draw the rectangle.
        coordinates (Tuple[float, float, float, float]): A tuple containing the coordinates
            (xmin, ymin, xmax, ymax) of the rectangle, where each coordinate is in the range [0, 1].

    Returns:
        PIL.Image.Image: A modified PIL.Image.Image object with the rectangle drawn.
    """
    xmin, ymin, xmax, ymax = coordinates  # Unpack the coordinates tuple

    # Get the dimensions (width and height) of the image
    img_width, img_height = image.size

    # Scale the coordinates to match the image size
    if 0 <= xmin <= 1:
        xmin = xmin * img_width
        ymin = ymin * img_height
        xmax = xmax * img_width
        ymax = ymax * img_height

    # Create a copy of the image to avoid modifying the original
    modified_image = image.copy()

    # Create a drawing context on the copy of the image
    draw = ImageDraw.Draw(modified_image)

    # Draw the rectangle on the copy of the image
    draw.rectangle((xmin, ymin, xmax, ymax), outline="red", width=2)

    # Display or save the modified image as needed
    # You can add code here to display or save the image if necessary

    return modified_image

class PDFPage:
    def __init__(self, image: Image.Image, pdf_page: pdfplumber.page.Page, page_number: int) -> None:
        self.image: Image.Image = image
        self.pdf_page: pdfplumber.page.Page = pdf_page
        self.page_number: int = page_number

    def get_image(self) -> Image.Image:
        """Get an image of the page.

        Returns:
            PIL.Image.Image: A PIL.Image.Image object.
        """
        return self.image

    def get_pdf(self) -> pdfplumber.page.Page:
        """Get the pdfplumber Page object representing a page of the PDF.

        Returns:
            pdfplumber.page.Page: A pdfplumber Page object.
        """
        return self.pdf_page
    
    def draw_rectangle(self, coordinates: Tuple[float, float, float, float]) -> Image.Image:
        """
        Draw a rectangle on the page image and return the modified image.

        Args:
            coordinates (Tuple[float, float, float, float]): A tuple containing the coordinates
                (xmin, ymin, xmax, ymax) of the rectangle, where each coordinate is in the range [0, 1].

        Returns:
            PIL.Image.Image: A modified PIL.Image.Image object with the rectangle drawn.
        """
        modified_image = draw_rectangle(self.image, coordinates)
        return modified_image


class PDFDocument:
    """Represents a PDF document.

    Args:
        filepath (str): The path to the PDF file.

    Attributes:
        pdf_file (pdfplumber.pdf.PDF): A pdfplumber PDF object representing the PDF file.
        pages (List[PDFPage]): A list of PDFPage objects representing pages in the PDF.
    """
    def __init__(self, filepath: str):
        self.pdf_file = pdfplumber.open(filepath)
        self.pages = self._initialize_pages()

    def _initialize_pages(self) -> List[PDFPage]:
        """Initialize and return a list of PDFPage objects for each page in the PDF.

        Returns:
            List[PDFPage]: A list of PDFPage objects.
        """
        pages = []
        for page_number, pdf_page in enumerate(self.pdf_file.pages):
            image = pdf_page.to_image(resolution=200).original.copy()
            pages.append(PDFPage(image, pdf_page, page_number + 1))
        return pages

    def __iter__(self) -> Iterator[PDFPage]:
        """Iterator method to iterate over pages in the PDF.

        Returns:
            Iterator[PDFPage]: An iterator over PDFPage objects.
        """
        self.current_page_index = 0
        return self

    def __next__(self) -> PDFPage:
        """Get the next page in the PDF.

        Returns:
            PDFPage: The next page in the PDF.

        Raises:
            StopIteration: If there are no more pages to iterate.
        """
        if self.current_page_index < len(self.pages):
            current_page = self.pages[self.current_page_index]
            self.current_page_index += 1
            return current_page
        else:
            raise StopIteration
scanipy/pdfhandler.py
SVJLucas/Scanipy