|
""" |
|
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. |
|
original_image (PIL.image): The image 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: |
|
model: The trained YOLO model |
|
image: 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__': |
|
|
|
|
|
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) |
|
|
|
st.subheader("Please upload an image of a scanned document with a table using the sidebar") |
|
|
|
with st.sidebar: |
|
user_image = st.file_uploader("Upload an image of a scanned document", 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) |
|
|
|
try: |
|
cropped_table = crop_image(model, image) |
|
st.divider() |
|
st.subheader("Table Cropped Image:") |
|
st.image(cropped_table, caption='Cropped Table', use_column_width=True) |
|
|
|
improved_image = improve_ocr_accuracy(cropped_table) |
|
st.divider() |
|
st.subheader("Improved Table Image:") |
|
st.image(improved_image, caption='Improved Table Image', use_column_width=True) |
|
|
|
ocr_text = ocr_core(improved_image) |
|
st.divider() |
|
st.subheader("OCR Text:") |
|
st.write(ocr_text) |
|
|
|
csv_output = generate_csv_from_text(tokenizer,gpt_model,ocr_text) |
|
st.divider() |
|
st.subheader("CSV Output:") |
|
st.write(csv_output.encode('utf-8')) |
|
except: |
|
st.divider() |
|
st.subheader("Error:") |
|
st.write("Please upload a scanned document with a table") |
|
|
|
|
|
|
|
|
