invoice_extractor.py

import os
import logging
from PIL import Image, ImageDraw
import traceback
import torch
from docquery import pipeline
from docquery.document import load_bytes, load_document, ImageDocument
from docquery.ocr_reader import get_ocr_reader
from pdf2image import convert_from_path

os.environ["TOKENIZERS_PARALLELISM"] = "false"

# Initialize the logger
logging.basicConfig(filename="invoice_extraction.log", level=logging.DEBUG)  # Create a log file

# Checkpoint for different models
CHECKPOINTS = {
    "LayoutLMv1 for Invoices 🧾": "impira/layoutlm-invoices",
}
PIPELINES = {}


class InvoiceKeyValuePair():

    """
    This class provides a utility to extract key-value pairs from invoices using LayoutLM.
    """

    def __init__(self):

      self.fields = {
                "Vendor Name": ["Vendor Name - Logo?", "Vendor Name - Address?"],
                "Vendor Address": ["Vendor Address?"],
                "Customer Name": ["Customer Name?"],
                "Customer Address": ["Customer Address?"],
                "Invoice Number": ["Invoice Number?"],
                "Invoice Date": ["Invoice Date?"],
                "Due Date": ["Due Date?"],
                "Subtotal": ["Subtotal?"],
                "Total Tax": ["Total Tax?"],
                "Invoice Total": ["Invoice Total?"],
                "Amount Due": ["Amount Due?"],
                "Payment Terms": ["Payment Terms?"],
                "Remit To Name": ["Remit To Name?"],
                "Remit To Address": ["Remit To Address?"],
            }
      self.model = list(CHECKPOINTS.keys())[0]

    def ensure_list(self, x):
        try:
            # Log the function entry
            logging.info(f'Entering ensure_list with x={x}')

            # Check if 'x' is already a list
            if isinstance(x, list):
                return x
            else:
                # If 'x' is not a list, wrap it in a list and return
                return [x]
        except Exception as e:
            # Log exceptions
            logging.error("An error occurred:", exc_info=True)
            return []

    def construct_pipeline(self, task, model):
        try:
            # Log the function entry
            logging.info(f'Entering construct_pipeline with task={task} and model={model}')

            # Global dictionary to cache pipelines based on model checkpoint names
            global PIPELINES

            # Check if a pipeline for the specified model already exists in the cache
            if model in PIPELINES:
                # If it exists, return the cached pipeline
                return PIPELINES[model]
            try:
                # Determine the device to use for inference (GPU if available, else CPU)
                device = "cuda" if torch.cuda.is_available() else "cpu"

                # Create the pipeline using the specified task and model checkpoint
                ret = pipeline(task=task, model=CHECKPOINTS[model], device=device)

                # Cache the created pipeline for future use
                PIPELINES[model] = ret

                # Return the constructed pipeline
                return ret
            except Exception as e:
                # Handle exceptions and log the error message
                logging.error("An error occurred:", exc_info=True)
                return None
        except Exception as e:
            # Log exceptions
            logging.error("An error occurred:", exc_info=True)
            return None

    def run_pipeline(self, model, question, document, top_k):
        try:
            # Log the function entry
            logging.info(f'Entering run_pipeline with model={model}, question={question}, and document={document}')

            # Use the construct_pipeline method to get or create a pipeline for the specified model
            pipeline = self.construct_pipeline("document-question-answering", model)

            # Use the constructed pipeline to perform question-answering on the document
            # Pass the question, document context, and top_k as arguments to the pipeline
            return pipeline(question=question, **document.context, top_k=top_k)
        except Exception as e:
            # Log exceptions
            logging.error("An error occurred:", exc_info=True)
            return None

    def lift_word_boxes(self, document, page):
        try:
            # Log the function entry
            logging.info(f'Entering lift_word_boxes with document={document} and page={page}')

            # Extract the word boxes for the specified page from the document's context
            return document.context["image"][page][1]
        except Exception as e:
            # Log exceptions
            logging.error("An error occurred:", exc_info=True)
            return []

    def expand_bbox(self, word_boxes):
        try:
            # Log the function entry
            logging.info(f'Entering expand_bbox with word_boxes={word_boxes}')

            # Check if the input list of word boxes is empty
            if len(word_boxes) == 0:
                return None

            # Extract the minimum and maximum coordinates of the word boxes
            min_x, min_y, max_x, max_y = zip(*[x[1] for x in word_boxes])

            # Calculate the overall minimum and maximum coordinates
            min_x, min_y, max_x, max_y = [min(min_x), min(min_y), max(max_x), max(max_y)]

            # Return the expanded bounding box as [min_x, min_y, max_x, max_y]
            return [min_x, min_y, max_x, max_y]
        except Exception as e:
            # Log exceptions
            logging.error("An error occurred:", exc_info=True)
            return None

    def normalize_bbox(self, box, width, height, padding=0.005):
        try:
            # Log the function entry
            logging.info(f'Entering normalize_bbox with box={box}, width={width}, height={height}, and padding={padding}')

            # Extract the bounding box coordinates and convert them from millimeters to fractions
            min_x, min_y, max_x, max_y = [c / 1000 for c in box]

            # Apply padding if specified (as a fraction of image dimensions)
            if padding != 0:
                min_x = max(0, min_x - padding)
                min_y = max(0, min_y - padding)
                max_x = min(max_x + padding, 1)
                max_y = min(max_y + padding, 1)

            # Scale the normalized coordinates to match the image dimensions
            return [min_x * width, min_y * height, max_x * width, max_y * height]
        except Exception as e:
            # Log exceptions
            logging.error("An error occurred:", exc_info=True)
            return None

    def annotate_page(self, prediction, pages, document):
        try:
            # Log the function entry
            logging.info(f'Entering annotate_page with prediction={prediction}, pages={pages}, and document={document}')

            # Check if a prediction exists and contains word_ids
            if prediction is not None and "word_ids" in prediction:

                # Get the image of the page where the prediction was made
                image = pages[prediction["page"]]

                # Create a drawing object for the image
                draw = ImageDraw.Draw(image, "RGBA")

                # Extract word boxes for the page
                word_boxes = self.lift_word_boxes(document, prediction["page"])

                # Expand and normalize the bounding box of the predicted words
                x1, y1, x2, y2 = self.normalize_bbox(
                    self.expand_bbox([word_boxes[i] for i in prediction["word_ids"]]),
                    image.width,
                    image.height,
                )

                # Draw a semi-transparent green rectangle around the predicted words
                draw.rectangle(((x1, y1), (x2, y2)), fill=(0, 255, 0, int(0.4 * 255)))
        except Exception as e:
            # Log exceptions
            logging.error("An error occurred:", exc_info=True)

    def process_fields(self, document, fields, model=list(CHECKPOINTS.keys())[0]):
        try:
            # Log the function entry
            logging.info(f'Entering process_fields with document={document}, fields={fields}, and model={model}')

            # Convert preview pages of the document to RGB format
            pages = [x.copy().convert("RGB") for x in document.preview]

            # Initialize dictionaries to store results
            ret = {}
            table = []

            # Iterate through the fields and associated questions
            for (field_name, questions) in fields.items():

                # Extract answers for each question and filter based on score
                answers = [
                    a
                    for q in questions
                    for a in self.ensure_list(self.run_pipeline(model, q, document, top_k=1))
                    if a.get("score", 1) > 0.5
                ]

                # Sort answers by score (higher score first)
                answers.sort(key=lambda x: -x.get("score", 0) if x else 0)

                # Get the top answer (if any)
                top = answers[0] if len(answers) > 0 else None

                # Annotate the page with the top answer's bounding box
                self.annotate_page(top, pages, document)

                # Store the top answer for the field and add it to the table
                ret[field_name] = top
                table.append([field_name, top.get("answer") if top is not None else None])

            # Return the table of key-value pairs
            return table
        except Exception as e:
            # Log exceptions
            logging.error("An error occurred:", exc_info=True)
            return []

    def process_document(self, document, fields, model, error=None):
        try:
            # Log the function entry
            logging.info(f'Entering process_document with document={document}, fields={fields}, model={model}, and error={error}')

            # Check if the document is not None and no error occurred during processing
            if document is not None and error is None:

                # Process the fields in the document using the specified model
                table = self.process_fields(document, fields, model)
                return table
        except Exception as e:
            # Log exceptions
            logging.error("An error occurred:", exc_info=True)
            return []

    def process_path(self, path, fields, model):
        try:
            # Log the function entry
            logging.info(f'Entering process_path with path={path}, fields={fields}, and model={model}')

            # Initialize error and document variables
            error = None
            document = None

            # Check if a file path is provided
            if path:
                try:
                    # Load the document from the specified file path
                    document = load_document(path)
                except Exception as e:
                    # Handle exceptions and store the error message
                    logging.error("An error occurred:", exc_info=True)
                    error = str(e)

            # Process the loaded document and extract key-value pairs
            return self.process_document(document, fields, model, error)
        except Exception as e:
            # Log exceptions
            logging.error("An error occurred:", exc_info=True)
            return []

    def pdf_to_image(self, file_path):
        try:
            # Log the function entry
            logging.info(f'Entering pdf_to_image with file_path={file_path}')

            # Convert PDF to a list of image objects (one for each page)
            images = convert_from_path(file_path)

            # Loop through each image and save it
            for i, image in enumerate(images):
                image_path = f'page_{i + 1}.png'
            
            return image_path
        except Exception as e:
            # Log exceptions
            logging.error("An error occurred:", exc_info=True)
            return []

    def process_upload(self, file):
        try:
            # Log the function entry
            logging.info(f'Entering process_upload with file={file}')

            # Get the model and fields from the instance
            model = self.model
            fields = self.fields

            # Convert the uploaded PDF file to a list of image files
            image = self.pdf_to_image(file)

            # Use the first generated image file as the file path for processing
            file = image

            # Process the document (image) and extract key-value pairs
            return self.process_path(file if file else None, fields, model)
        except Exception as e:
            # Log exceptions
            logging.error("An error occurred:", exc_info=True)
            return []

    def extract_key_value_pair(self, invoice_file):
        try:
            # Log the function entry
            logging.info(f'Entering extract_key_value_pair with invoice_file={invoice_file}')

            # Process the uploaded invoice PDF file and extract key-value pairs
            data = self.process_upload(invoice_file.name)

            # Iterate through the extracted key-value pairs and print them
            for item in data:
                key, value = item
                return f'{key}: {value}'

        except Exception as e:
            # Log exceptions
            logging.error("An error occurred:", exc_info=True)