Firs Push

Model_loading.py

@@ -0,0 +1,24 @@
+from transformers import TrOCRProcessor, VisionEncoderDecoderModel
+
+# Singleton class for TrOCR model and processor
+class OCRModelSingleton:
+    _instance = None
+
+    def __init__(self):
+        if OCRModelSingleton._instance is not None:
+            raise Exception("This class is a singleton!")
+        else:
+            OCRModelSingleton._instance = self
+            self.processor = TrOCRProcessor.from_pretrained('microsoft/trocr-large-printed')
+            self.model = VisionEncoderDecoderModel.from_pretrained('microsoft/trocr-large-printed')
+
+    @staticmethod
+    def get_instance():
+        if OCRModelSingleton._instance is None:
+            OCRModelSingleton()
+        return OCRModelSingleton._instance
+
+# Automatically initialize the shared processor and model
+ocr_model_instance = OCRModelSingleton.get_instance()
+processor_tr_ocr = ocr_model_instance.processor
+trocr_model = ocr_model_instance.model

Models/TD_3k_yolo_150_wandb.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:eb91d6f1cdcaa3bf0799b22c2948438ba43526693dddca794a103fb5040506ff
+size 52040315

Models/invoice_yolo_100_7classes.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b31bbfff063a656ec5c7856c22cf88feae555ddd18b0984f8c85a83dc6832732
+size 52065770

TD.py

@@ -0,0 +1,66 @@
+import warnings
+warnings.filterwarnings("ignore")
+
+from ultralytics import YOLO
+import cv2
+import numpy as np
+from PIL import Image
+
+import os
+
+# Define a path relative to the script's location
+BASE_DIR = os.path.dirname(os.path.abspath(__file__))
+model_path = os.path.join(BASE_DIR, "Models", "TD_3k_yolo_150_wandb.pt")
+
+
+## Load the YOLO model
+model1 = YOLO(model_path)
+
+def TD_model1(image_arr, padding=5, confidence_threshold=0.5):
+    # Convert to RGB (for PIL cropping)
+    image_rgb = cv2.cvtColor(image_arr, cv2.COLOR_BGR2RGB)
+
+    # Perform object detection
+    results = model1.predict(image_arr)
+
+    # Check if any tables were detected
+    if len(results[0].boxes) == 0:
+        return []  # No tables detected
+
+    # Extract bounding boxes and confidence scores
+    boxes = results[0].boxes.xyxy.cpu().numpy()  # Bounding boxes
+    confidences = results[0].boxes.conf.cpu().numpy()  # Confidence scores
+
+    # Filter predictions by confidence threshold
+    valid_boxes = []
+    for i, box in enumerate(boxes):
+        if confidences[i] >= confidence_threshold:
+            valid_boxes.append((box, confidences[i]))
+
+    # If no valid boxes, return an empty list
+    if not valid_boxes:
+        return []
+
+    # Sort the boxes by confidence in descending order
+    valid_boxes = sorted(valid_boxes, key=lambda x: x[1], reverse=True)
+
+    # Crop the image based on the valid bounding boxes
+    cropped_images = []
+    image_pil = Image.fromarray(image_rgb)  # Convert NumPy array to PIL image
+
+    for i, (box, conf) in enumerate(valid_boxes):
+        x1, y1, x2, y2 = box
+        # Add padding to the bounding box coordinates (with boundary checks)
+        x1 = max(0, x1 - padding)
+        y1 = max(0, y1 - padding)
+        x2 = min(image_pil.width, x2 + padding)
+        y2 = min(image_pil.height, y2 + padding)
+
+        cropped_image = image_pil.crop((x1, y1, x2, y2))  # Crop using the bounding box coordinates
+        cropped_images.append(cropped_image)
+    
+    return cropped_images
+
+
+
+

Tr_ocr.py

@@ -0,0 +1,45 @@
+## Tr OCR
+import numpy as np
+from tqdm.auto import tqdm
+from Model_loading import processor_tr_ocr, trocr_model
+import torch
+from PIL import Image
+
+
+
+def apply_TRocr(cropped_image,cell_coordinates):
+    """
+    Apply TrOCR to the given cell coordinates.
+    
+    Parameters:
+    - cell_coordinates: List of coordinates for cropping cells from an image.
+    
+    Returns:
+    - data: A list of extracted text for each cell.
+    """
+    data = []  # Initialize a list to store data for each row
+
+    for idx, row in enumerate(tqdm(cell_coordinates)):
+        row_text = []  # List to store text for the current row
+        for cell in row["cells"]:
+            # Crop cell out of the image
+            cell_image = np.array(cropped_image.crop(cell["cell"]))
+            
+            
+            # Convert the cell image to RGB format for TrOCR processing
+            pil_image = Image.fromarray(cell_image).convert("RGB")
+            pixel_values = processor_tr_ocr(images=pil_image, return_tensors="pt").pixel_values
+
+            # Generate text predictions using TrOCR model
+            generated_ids = trocr_model.generate(pixel_values)
+            text = processor_tr_ocr.batch_decode(generated_ids, skip_special_tokens=True)[0]
+            
+            if not text:
+                text = "NAN"  # Append "NAN" if no text is detected
+            
+            row_text.append(text)
+
+        # Append the row's text list to the data list
+        data.append(row_text)
+
+    return data

app.py

@@ -0,0 +1,148 @@
+import streamlit as st
+from PIL import Image
+import numpy as np
+import io
+import pandas as pd
+
+import os
+import subprocess
+
+try:
+    import ultralytics
+except ImportError:
+    subprocess.run(["pip", "install", "ultralytics"])
+
+# Import your models and utility functions
+from TD import TD_model1
+from tsr import TSR, get_cell_coordinates_by_row, apply_ocr
+from info_det_ocr import info_det_and_ocr
+from Tr_ocr import apply_TRocr
+
+# Set page config
+st.set_page_config(
+    page_title="INDRA OCR",
+    page_icon="🧾",
+    layout="wide",
+)
+
+# Customizing background color and font color
+st.markdown("""
+    <style>
+    body {
+        background-color: blue; /* Change to your desired background color */
+        color: #333333; /* Change to your desired font color */
+    }
+    </style>
+""", unsafe_allow_html=True)
+
+# App Header
+st.markdown(
+    """
+    <style>
+    .main-title {
+        font-size: 3.5rem; /* Increased font size */
+        color: white;
+        text-align: center;
+        font-weight: bold;
+        margin-bottom: 10px;
+    }
+    .subtitle {
+        font-size: 1.2rem;
+        color: #555555;
+        text-align: center;
+        margin-bottom: 20px;
+    }
+    </style>
+    <div class="main-title">INDRA OCR: Intelligent Invoice Data Recognition and Automation</div>
+    """,
+    unsafe_allow_html=True,
+)
+
+# Create two columns for layout: one for the file uploader and another for the image and extracted data
+col1, col2 = st.columns([1, 3])  # Adjust the ratio to control the width of the columns
+
+with col1:
+    # File uploader for image (placed in the left column, smaller button)
+    uploaded_file = st.file_uploader(
+        "Upload Your Invoice (JPG, JPEG, PNG):", type=["jpg", "jpeg", "png"], label_visibility="collapsed"
+    )
+
+with col2:
+    # Display the uploaded image (right column)
+    if uploaded_file is not None:
+        # If a new image is uploaded, reset the session state
+        if 'uploaded_image' not in st.session_state or st.session_state.uploaded_image != uploaded_file:
+            st.session_state.uploaded_image = uploaded_file
+            st.session_state.processed_data = None
+            st.session_state.cropped_image = None
+            st.session_state.cell_coordinates = None
+            st.session_state.df = None
+
+        # Process the image only if it hasn't been processed yet
+        if st.session_state.processed_data is None:
+            # Read the image file as a PIL image
+            image = Image.open(uploaded_file)
+
+            # Resize the image to make it smaller for display
+            image.thumbnail((800, 800))  # Resize to fit within the 800px limit
+
+            # Convert the image to OpenCV format (NumPy array)
+            image_arr = np.array(image)
+
+            # Display the uploaded image
+            st.image(image, caption="Uploaded Image", width= 500)
+
+            # Run info detection and OCR
+            detected_data = info_det_and_ocr(image_arr)
+            st.session_state.processed_data = detected_data
+
+            # Run the Table detection model and crop images
+            cropped_image = TD_model1(image_arr)
+
+            if isinstance(cropped_image, list):
+                cropped_image = cropped_image[0]  # Extract the image from the list
+
+            st.session_state.cropped_image = cropped_image
+
+            # Run Table Structure Recognition (TSR)
+            output_image, cells = TSR(cropped_image)
+
+            # Get cell coordinates and perform OCR on table cells
+            cell_coordinates = get_cell_coordinates_by_row(cells)
+            st.session_state.cell_coordinates = cell_coordinates
+
+            # Use TR OCR for extracting table data
+            data = apply_TRocr(cropped_image, cell_coordinates)
+
+            # Store the DataFrame in session state for CSV download
+            st.session_state.df = pd.DataFrame(data)
+
+        # Show extracted data if the dataframe is available
+        if st.session_state.df is not None:
+            st.markdown("### Extracted Data Table:")
+            st.dataframe(st.session_state.df)  # Display the dataframe as a table
+
+            # Show download button if the data has been processed
+            st.markdown("### Download Extracted Data")
+            # Convert DataFrame to CSV and store in-memory buffer
+            csv_buffer = io.StringIO()
+            st.session_state.df.to_csv(csv_buffer, index=False)
+
+            # Get the CSV data as bytes
+            csv_data = csv_buffer.getvalue().encode('utf-8')
+
+            # Create a download button in Streamlit to download the CSV file
+            st.download_button(
+                label="Download CSV",
+                data=csv_data,
+                file_name="extracted_table_data.csv",
+                mime="text/csv",
+            )
+
+# Footer
+st.markdown(
+    """
+    <div class="footer">© 2024 INDRA OCR | Designed for efficient invoice processing</div>
+    """,
+    unsafe_allow_html=True,
+)

info_det_ocr.py

@@ -0,0 +1,75 @@
+import torch
+import cv2
+import numpy as np
+from PIL import Image
+import matplotlib.pyplot as plt
+from ultralytics import YOLO
+from Model_loading import processor_tr_ocr, trocr_model
+
+import os
+
+# Define a path relative to the script's location
+BASE_DIR = os.path.dirname(os.path.abspath(__file__))
+model_path = os.path.join(BASE_DIR, "Models", "invoice_yolo_100_7classes.pt")
+
+
+# Load the YOLO model
+
+info_model = YOLO(model_path)
+
+
+
+# Function for performing OCR with TrOCR
+def ocr_with_transformer(image):
+    # Convert the image to RGB and process it for TrOCR
+    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+    pixel_values = processor_tr_ocr(images=image_rgb, return_tensors="pt").pixel_values
+    generated_ids = trocr_model.generate(pixel_values)
+    
+    # Decode the generated text from the model
+    generated_text = processor_tr_ocr.batch_decode(generated_ids, skip_special_tokens=True)[0]
+    return generated_text
+
+# Main function for detection, cropping, and OCR
+def info_det_and_ocr(image_arr, conf_threshold=0.5):
+    # Initialize the dictionary to store detected class names and text
+    detected_data = {}
+
+    # Perform inference using YOLO
+    results = info_model.predict(image_arr)
+    
+    boxes = results[0].boxes.xyxy.cpu().numpy()  # Bounding boxes
+    scores = results[0].boxes.conf.cpu().numpy()  # Confidence scores
+    class_ids = results[0].boxes.cls.cpu().numpy().astype(int)  # Class IDs
+    class_names = results[0].names  # Class names dictionary
+    
+    # Iterate over the detected objects
+    for i, box in enumerate(boxes):
+        score = scores[i]
+        if score < conf_threshold:
+            continue  # Skip low confidence detections
+
+        x1, y1, x2, y2 = map(int, box)  # Convert bounding box to integers
+        class_id = class_ids[i]
+        class_name = class_names[class_id]
+
+        # Crop the detected region
+        cropped_region = image_arr[y1:y2, x1:x2]
+
+        # Perform OCR using TrOCR on the cropped region
+        detected_text = ocr_with_transformer(cropped_region)
+
+        # Save the detected text into the dictionary with class name as key
+        detected_data[class_name] = detected_text
+
+        # Optional: Display cropped image for debugging
+        #plt.imshow(cv2.cvtColor(cropped_region, cv2.COLOR_BGR2RGB))
+        #plt.title(f"Detected: {class_name}")
+        #plt.show()
+
+    return detected_data
+
+# You can now call this function by passing the image as a numpy array.
+# Example:
+# extracted_data = info_det_and_ocr(image_arr)
+# print(extracted_data)

req.txt

@@ -0,0 +1,12 @@
+
+streamlit
+Pillow
+numpy
+pandas
+transformers
+ultralytics
+opencv-python-headless
+tqdm
+torch
+torchvision
+easyocr

tsr.py

@@ -0,0 +1,212 @@
+import torch
+import cv2
+from PIL import Image
+import numpy as np
+from io import BytesIO
+import torch
+from ultralytics import YOLO
+from torchvision import transforms
+from transformers import TableTransformerForObjectDetection
+from PIL import ImageDraw
+import numpy as np
+import csv
+import easyocr
+from tqdm.auto import tqdm
+import csv
+
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+# new v1.1 checkpoints require no timm anymore
+structure_model = TableTransformerForObjectDetection.from_pretrained("microsoft/table-structure-recognition-v1.1-all")
+structure_model.to(device)
+print("")
+
+from torchvision import transforms
+
+class MaxResize(object):
+    def __init__(self, max_size=800):
+        self.max_size = max_size
+
+    def __call__(self, image):
+        width, height = image.size
+        current_max_size = max(width, height)
+        scale = self.max_size / current_max_size
+        resized_image = image.resize((int(round(scale*width)), int(round(scale*height))))
+
+        return resized_image
+
+detection_transform = transforms.Compose([
+    MaxResize(800),
+    transforms.ToTensor(),
+    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+])
+
+# for output bounding box post-processing
+def box_cxcywh_to_xyxy(x):
+    x_c, y_c, w, h = x.unbind(-1)
+    b = [(x_c - 0.5 * w), (y_c - 0.5 * h), (x_c + 0.5 * w), (y_c + 0.5 * h)]
+    return torch.stack(b, dim=1)
+
+
+def rescale_bboxes(out_bbox, size):
+    img_w, img_h = size
+    b = box_cxcywh_to_xyxy(out_bbox)
+    b = b * torch.tensor([img_w, img_h, img_w, img_h], dtype=torch.float32)
+    return b
+
+
+# update id2label to include "no object"
+id2label = structure_model.config.id2label
+id2label[len(structure_model.config.id2label)] = "no object"
+
+
+def outputs_to_objects(outputs, img_size, id2label):
+    m = outputs.logits.softmax(-1).max(-1)
+    pred_labels = list(m.indices.detach().cpu().numpy())[0]
+    pred_scores = list(m.values.detach().cpu().numpy())[0]
+    pred_bboxes = outputs['pred_boxes'].detach().cpu()[0]
+    pred_bboxes = [elem.tolist() for elem in rescale_bboxes(pred_bboxes, img_size)]
+
+    objects = []
+    for label, score, bbox in zip(pred_labels, pred_scores, pred_bboxes):
+        class_label = id2label[int(label)]
+        if not class_label == 'no object':
+            objects.append({'label': class_label, 'score': float(score),
+                            'bbox': [float(elem) for elem in bbox]})
+
+    return objects
+
+structure_transform = transforms.Compose([
+    MaxResize(1000),
+    transforms.ToTensor(),
+    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+])
+
+def TSR(cropped_image):
+    pixel_values = structure_transform(cropped_image).unsqueeze(0)
+    pixel_values = pixel_values.to(device)
+    print(pixel_values.shape)
+
+
+    # forward pass
+    with torch.no_grad():
+        outputs = structure_model(pixel_values)
+
+    # update id2label to include "no object"
+    structure_id2label = structure_model.config.id2label
+    structure_id2label[len(structure_id2label)] = "no object"
+
+    cells = outputs_to_objects(outputs, cropped_image.size, structure_id2label)
+    #print(cells)
+
+    cropped_table_visualized = cropped_image.copy()
+    draw = ImageDraw.Draw(cropped_table_visualized)
+
+    for cell in cells:
+        draw.rectangle(cell["bbox"], outline="red")
+
+    return cropped_table_visualized , cells
+############# Visualizing rows and columns on cropped image
+
+
+## Modified
+
+def get_cell_coordinates_by_row(table_data):
+    # Extract rows and columns
+    rows = [entry for entry in table_data if entry['label'] == 'table row']
+    columns = [entry for entry in table_data if entry['label'] == 'table column']
+
+    # Sort rows and columns by their Y and X coordinates, respectively
+    rows.sort(key=lambda x: x['bbox'][1])
+    columns.sort(key=lambda x: x['bbox'][0])
+
+    # Function to find cell coordinates
+    def find_cell_coordinates(row, column):
+        # Use the row's Y coordinates for the cell's top and bottom
+        cell_ymin = row['bbox'][1]
+        cell_ymax = row['bbox'][3]  # Adjust as needed for better height
+
+        # Use the column's X coordinates for the cell's left and right
+        cell_xmin = column['bbox'][0]
+        cell_xmax = column['bbox'][2]
+
+        return [cell_xmin, cell_ymin, cell_xmax, cell_ymax]
+
+    # Generate cell coordinates and count cells in each row
+    cell_coordinates = []
+
+    for row in rows:
+        row_cells = []
+        for column in columns:
+            cell_bbox = find_cell_coordinates(row, column)
+            row_cells.append({'column': column['bbox'], 'cell': cell_bbox})
+
+        # Sort cells in the row by X coordinate
+        row_cells.sort(key=lambda x: x['column'][0])
+
+        # Append row information to cell_coordinates
+        cell_coordinates.append({'row': row['bbox'], 'cells': row_cells, 'cell_count': len(row_cells)})
+
+    # Sort rows from top to bottom
+    cell_coordinates.sort(key=lambda x: x['row'][1])
+
+    # Debugging: Print intermediate results
+    #for i, row_info in enumerate(cell_coordinates):
+     #   print(f"Row {i}: {row_info['row']}, Cell Count: {row_info['cell_count']}")
+      #  for cell in row_info['cells']:
+       #     print(f"  Cell Bounding Box: {cell['cell']}, Column Bounding Box: {cell['column']}")
+
+    return cell_coordinates
+
+
+# Initialize the EasyOCR reader
+reader = easyocr.Reader(['en'])  # this needs to run only once to load the model into memory
+
+def apply_ocr(cell_coordinates,cropped_image):
+    # Initialize a list to store data for each row
+    data = []
+
+    for idx, row in enumerate(tqdm(cell_coordinates)):
+        row_text = []  # List to store text for the current row
+        for cell in row["cells"]:
+            # Crop cell out of the image
+            cell_image = np.array(cropped_image.crop(cell["cell"]))
+            # Apply OCR
+            result = reader.readtext(np.array(cell_image))
+            if len(result) > 0:
+                # Extract and join the detected text
+                text = " ".join([x[1] for x in result])
+                row_text.append(text)
+            else:
+                row_text.append("NAN")  # Append empty string if no text is detected
+
+        # Append the row's text list to the data list
+        data.append(row_text)
+
+    return data
+
+
+# Print the extracted text for each row
+"""for idx, row_data in enumerate(data):
+    print(f"Row {idx + 1}: {row_data}")"""
+
+def op_csv(data):
+
+# Define the output CSV file path
+    output_csv_file = 'extract.csv'
+
+# Write the data to a CSV file
+    try:
+        with open(output_csv_file, mode='w', newline='', encoding='utf-8') as file:
+            writer = csv.writer(file)
+
+            for row in data:
+                writer.writerow(row)  # Write each row individually
+
+        print(f"Data successfully written to {output_csv_file}")
+        return output_csv_file
+    except Exception as e:
+        print(f"An error occurred: {e}")
+
+
+