Upload 2 files

OCR.py

@@ -0,0 +1,50 @@
+import os
+import torch
+from transformers import AutoProcessor, PaliGemmaForConditionalGeneration
+from PIL import Image
+import io
+
+# Set environment variable
+os.environ['PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION'] = 'python'
+
+# Model and device setup
+device = "cuda:0" if torch.cuda.is_available() else "cpu"
+model_id = "google/paligemma-3b-mix-224"
+
+# Load model and processor
+model = PaliGemmaForConditionalGeneration.from_pretrained(model_id).to(device)
+processor = AutoProcessor.from_pretrained(model_id)
+
+def extract_text_from_image(image_content):
+    image = Image.open(io.BytesIO(image_content))
+
+    # Prompt for detecting text
+    prompt = "Extract all relevant details from this invoice."
+
+    # Prepare inputs for the model
+    inputs = processor(text=prompt, images=image, return_tensors="pt").to(device)
+    input_len = inputs["input_ids"].shape[-1]
+
+    with torch.inference_mode():
+        # Generate the output
+        generation = model.generate(**inputs, max_new_tokens=100, do_sample=False)
+        generation = generation[0][input_len:]
+        decoded = processor.decode(generation, skip_special_tokens=True)
+
+    return decoded
+
+def extract_text_from_pdf(pdf_content):
+    # For simplicity, let's assume you're converting the PDF to images first
+    # You may use libraries like pdf2image to convert PDF pages to images
+    # Then call extract_text_from_image for each image
+    pass
+
+def extract_invoice_details(text):
+    # Implement your logic to extract invoice details from the text
+    details = {}
+    # Example extraction logic
+    details['Invoice Number'] = re.search(r'Invoice Number: (\S+)', text).group(1) if re.search(r'Invoice Number: (\S+)', text) else 'N/A'
+    details['Amount'] = re.search(r'Total Amount Due: (\S+)', text).group(1) if re.search(r'Total Amount Due: (\S+)', text) else 'N/A'
+    details['Invoice Date'] = re.search(r'Invoice Date: (\S+)', text).group(1) if re.search(r'Invoice Date: (\S+)', text) else 'N/A'
+    details['Due Date'] = re.search(r'Due Date: (\S+)', text).group(1) if re.search(r'Due Date: (\S+)', text) else 'N/A'
+    return details

app.py

@@ -0,0 +1,64 @@
+import streamlit as st
+from dataclasses import dataclass
+import pytesseract
+from PIL import Image
+import io
+import re
+import cv2
+import numpy as np
+import OCR
+
+from  OCR import *
+
+# Initialize chat history
+if "messages" not in st.session_state:
+    st.session_state.messages = [{"role": "Invoice Reader", "content": "Submit an invoice and I will read it."}]
+
+# Display chat messages from history on app rerun
+for message in st.session_state.messages:
+    with st.chat_message(message["role"]):
+        st.markdown(message["content"])
+
+USER = "user"
+ASSISTANT = "Invoice Reader"
+
+# Accept file uploads
+uploaded_file = st.file_uploader("Upload an invoice", type=["pdf", "png", "jpg", "jpeg"])
+if uploaded_file is not None:
+    # Display uploaded file content
+    file_content = uploaded_file.getvalue()
+    st.session_state.messages.append({"role": USER, "content": f"Uploaded file: {uploaded_file.name}"})
+    with st.chat_message(USER):
+        st.markdown(f"Uploaded file: {uploaded_file.name}")
+
+    # Preprocess and extract text from image or PDF
+    try:
+        if uploaded_file.type == "application/pdf":
+            text = extract_text_from_pdf(file_content)
+        else:
+            text = extract_text_from_image(file_content)
+
+        # Extract specific details
+        details = extract_invoice_details(text)
+
+        # Create and display assistant's response to extracted text
+        assistant_response = (
+            f"Extracted text from the uploaded file:\n\n{text}\n\n"
+            f"**Extracted Details:**\n"
+            f"**Invoice Number:** {details['Invoice Number']}\n"
+            
+            f"**Amount:** {details['Amount']}\n"
+            f"**Invoice Date:** {details['Invoice Date']}\n"
+            f"**Due Date:** {details['Due Date']}"
+        )
+        st.session_state.messages.append({"role": ASSISTANT, "content": assistant_response})
+        with st.chat_message(ASSISTANT):
+            st.markdown(assistant_response)
+    except Exception as e:
+        error_message = f"An error occurred while processing the file: {e}"
+        st.session_state.messages.append({"role": ASSISTANT, "content": error_message})
+        with st.chat_message(ASSISTANT):
+            st.markdown(error_message)
+
+
+#streamlit run C:/Users/leahw/PycharmProjects/Int-to-Artificial-Intelligence-Final-Project/app.py