main.py

"""
Attribution: https://github.com/AIPI540/AIPI540-Deep-Learning-Applications/

Jon Reifschneider
Brinnae Bent 

"""

import streamlit as st
from PIL import Image
import numpy as np
import os
import numpy as np
import pandas as pd
import pandas as pd
import os
import json
import pandas as pd
import torch
import numpy as np
import pandas as pd
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt
from ultralytics import YOLO
from PIL import Image, ImageDraw, ImageFont
import numpy as np
import cv2
import pytesseract
from PIL import ImageEnhance
import numpy as np
import os
import json
from transformers import GPT2LMHeadModel, GPT2Tokenizer, Trainer, TrainingArguments
from datasets import load_dataset
from transformers import DataCollatorForLanguageModeling
from PIL import Image, ImageEnhance
from io import StringIO


def crop_image(model, original_image):
    """
    Crop the region of interest (table) from an image using a YOLO model.

    Inputs:
        model (YOLO): The YOLO model used for object detection.
        image_file (str): Path to the image file to be processed.

    Returns:
        PIL.Image: The cropped image containing the detected table.
    """
    image_array = np.array(image)
    results = model(image_array)

    for r in results:
        boxes = r.boxes  
        
        for box in boxes:
            if box.cls == 3:  
                x1, y1, x2, y2 = box.xyxy[0]  
                x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
                
                table_image = original_image.crop((x1, y1, x2, y2))
                
                return table_image
        return

def process_image(model, image):
    """
    Process the uploaded image with YOLO model and draw bounding boxes with class-specific colors.
    
    Inputs:
    uploaded_image (UploadedFile): The image file uploaded through Streamlit.
    
    Returns:
    PIL.Image: The processed image with bounding boxes and labels.
    """
    colors = {'title': (255, 0, 0),
            'text': (0, 255, 0),
            'figure': (0, 0, 255),
            'table': (255, 255, 0),
            'list': (0, 255, 255)}

    image_array = np.array(image)
    results = model(image_array)
    
    for result in results:
        boxes = result.boxes.cpu().numpy()
        for box in boxes:
            r = box.xyxy[0].astype(int)
            label = result.names[int(box.cls)]
            color = colors.get(label.lower(), (255, 255, 255))
            
            cv2.rectangle(image_array, r[:2], r[2:], color, 2)
            
            label_size, baseline = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
            top_left = (r[0], r[1] - label_size[1] - baseline)
            bottom_right = (r[0] + label_size[0], r[1])
            cv2.rectangle(image_array, top_left, bottom_right, color, cv2.FILLED)
            cv2.putText(image_array, label, (r[0], r[1] - baseline), 
                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1)
    
    return Image.fromarray(image_array)

def improve_ocr_accuracy(img):
    """
    Preprocess the image to improve OCR accuracy.

    This function resizes the image, increases contrast, and applies thresholding
    to enhance the image for better OCR results.

    Inputs:
        img (PIL.Image): The input image to be processed.

    Returns:
        numpy.ndarray: A binary thresholded image as a numpy array.
    """
    img = img.resize((img.width * 4, img.height * 4))
    
    enhancer = ImageEnhance.Contrast(img)
    img = enhancer.enhance(2)

    _, thresh = cv2.threshold(np.array(img), 127, 255, cv2.THRESH_BINARY_INV)
    
    return thresh

def ocr_core(image):
    """
    Perform OCR on the given image and process the extracted text.

    This function uses pytesseract to extract text from the image and then
    processes the extracted data to format it with appropriate line breaks
    and spacing.

    Inputs:
        image (numpy.ndarray): The preprocessed image as a numpy array.

    Returns:
        str: The extracted and formatted text from the image.
    """
    data = pytesseract.image_to_data(image, output_type=pytesseract.Output.DICT)
    df = pd.DataFrame(data)
    df = df[df['conf'] != -1]
    df['left_diff'] = df.groupby('block_num')['left'].diff().fillna(0).astype(int)
    df['prev_width'] = df['width'].shift(1).fillna(0).astype(int)
    df['spacing'] = (df['left_diff'] - df['prev_width']).fillna(0).astype(int)
    df['text'] = df.apply(lambda x: '\n' + x['text'] if (x['word_num'] == 1) & (x['block_num'] != 1) else x['text'], axis=1)
    df['text'] = df.apply(lambda x: ',' + x['text'] if x['spacing'] > 80 else x['text'], axis=1)
    ocr_text = ""
    for text in df['text']:
        ocr_text += text + ' '
    return ocr_text

def generate_csv_from_text(tokenizer, model, ocr_text):
    """
    Generate CSV text from OCR extracted text using the gpt model

    This function takes the OCR extracted text, processes it through a language model,
    and generates CSV formatted text.

    Inputs:
        tokenizer: The tokenizer for the gpt model
        model: The gpt model used for csv
        ocr_text (str): The text extracted from OCR

    Returns:
        str: The generated CSV formatted text.
    """
    inputs = tokenizer.encode(ocr_text, return_tensors='pt')
    outputs = model.generate(inputs, max_length=1000, num_return_sequences=1)
    csv_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    
    return csv_text

if __name__ == '__main__':
    pytesseract.pytesseract.tesseract_cmd = r'C:/Program Files/Tesseract-OCR/tesseract.exe'  # Update this path for your system

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
  
    model = YOLO(os.getcwd() + '/models/trained_yolov8.pt')
    gpt_model = GPT2LMHeadModel.from_pretrained(os.getcwd() + '/models/gpt_model')
    tokenizer = GPT2Tokenizer.from_pretrained(os.getcwd() + '/models/gpt_model')    
    
    st.header('''
        Intelligent Document Processing: Table Extraction
    ''')
    
    header_img = Image.open('assets/header_img.png')
    st.image(header_img, use_column_width=True)
    
    with st.sidebar:
        user_image = st.file_uploader("Upload an image", type=["png", "jpg", "jpeg"])

    if user_image is not None:
        st.divider() 
        image = Image.open(user_image)
        st.image(image, caption='Uploaded Image', use_column_width=True)
        
        st.divider() 
        st.subheader("Document Classes:")
        processed_image = process_image(model, image)
        st.image(processed_image, caption='Processed Image', use_column_width=True)
        
        st.divider() 
        st.subheader("Table Cropped Image:")
        cropped_table = crop_image(model, image)
        st.image(cropped_table, caption='Cropped Table', use_column_width=True)

        st.divider() 
        st.subheader("OCR Text:")
        improved_image = improve_ocr_accuracy(cropped_table)
        ocr_text = ocr_core(improved_image)
        st.write(ocr_text)
        
        st.divider() 
        st.subheader("CSV Output:")
        csv_output = generate_csv_from_text(tokenizer,gpt_model,ocr_text)
        data = StringIO(csv_output)
        st.dataframe(pd.read_csv(data, sep=",").head())