created dev branch and refactored code (#1)

- created dev branch and refactored code (31e0d1a8efa41e56aa53d9e147159a71066ee1bc)

.gitignore

@@ -0,0 +1 @@
+data/

README.md

@@ -8,6 +8,7 @@ The purpose of this project is to perform very basic intelligent document proces
 
 ### If you want to run the full pipeline and train the model from scratch
 1. You will need to install all of the necessary packages to run the setup.py script beforehand
+2. You will need to download pytesseract and add it to your Path if you are using Windows OS
 3. You will then need to run setup.py to create the data pipeline and train the model
 4. You will then need to run the frontend to use the model
 ```bash
@@ -17,7 +18,7 @@ streamlit run main.py
 ```
 
 ### If you want to just run the frontend
-1. You will need to install all of the necessary packages to run the setup.py script beforehand
+1. You will need to install all of the necessary packages to run the setup.py script beforehand and install pytesseract
 2. You will then need to run the frontend to use the model
 ```bash
 pip install -r requirements.txt
@@ -34,7 +35,8 @@ streamlit run main.py
 >   - build_features.py: script to prepare the dataset for training  
 >   - model.py: script to train model and predict  
 > - models: directory for trained models
->   - recommendation.pt: pytorch trained model for album recommendations 
+>   - trained_yolov8.pt: pytorch trained model for album recommendations 
+>   - gpt_model: directory to store the gpt model 
 > - data:  directory for project data
 >   - raw: directory for raw data 
 >   - processed: directory to store the processed data

config.yaml

@@ -1,6 +1,6 @@
 path: C:/Users/keese/term_project
-train: training/images
-val: validation/images
+train: data/processed/training/images
+val: data/processed/validation/images
 
 names:
   0: text

main.py

@@ -22,31 +22,199 @@ 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.
     
-if __name__ == '__main__':  
+    Inputs:
+    uploaded_image (UploadedFile): The image file uploaded through Streamlit.
     
-    st.header('Spotify Playlists')
+    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)
     
-    img1, img2 = st.columns(2)
+    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.
 
-    music_notes = Image.open('assets/music_notes.png')
-    img1.image(music_notes, use_column_width=True)
+    Inputs:
+        img (PIL.Image): The input image to be processed.
 
-    trumpet = Image.open('assets/trumpet.png')
-    img2.image(trumpet, use_column_width=True)
+    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:
-        playlist_name = st.selectbox(
-            "Playlist Selection",
-            (   list(set([1,2]))            )
-        )
-    
-    col1, col2 = st.columns(2)
-    with col1:
-        st.write(f'Artist')
-    with col2:
-        st.write(f'Album')
+        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())
 

setup.py

@@ -1,14 +1,14 @@
 import subprocess
 import sys
 
-script = 'make_dataset.py'
-command = f'{sys.executable} scripts/{script}'
-subprocess.run(command, shell=True)
+# script = 'make_dataset.py'
+# command = f'{sys.executable} scripts/{script}'
+# subprocess.run(command, shell=True)
 
 script = 'build_features.py'
-command = f'{sys.executable} python scripts/{script}'
+command = f'{sys.executable} scripts/{script}'
 subprocess.run(command, shell=True)
 
 script = 'model.py'
-command = f'{sys.executable} python scripts/{script}'
+command = f'{sys.executable} scripts/{script}'
 subprocess.run(command, shell=True)