Update app.py

app.py

@@ -1,46 +1,53 @@
+
+
+#Invoice Extractor
+import os
 import streamlit as st
 from PIL import Image
-import torch
-from transformers import LayoutLMv2Processor, LayoutLMv2ForQuestionAnswering
+from transformers import AutoProcessor, AutoModelForTokenClassification
 
-@st.cache(allow_output_mutation=True)
-def load_model():
-    processor = LayoutLMv2Processor.from_pretrained("microsoft/layoutlmv2-base-uncased")
-    model = LayoutLMv2ForQuestionAnswering.from_pretrained("microsoft/layoutlmv2-base-uncased")
-    return processor, model
+#function to load gemini pro vision model and get response
+def get_model_response(input,image,prompt):  #input: how LLM model behave like, #image: To extract info, prompt: ask something
+    #load a model
+    processor = AutoProcessor.from_pretrained("monuirctc/invoice-extraction")
+    model = AutoModelForTokenClassification.from_pretrained("monuirctc/invoice-extraction")
+    response=model.generate_content([input,image[0],prompt]) #get response from model
+    return response.text
 
-# Function to process the image and return the extracted information
-def process_image(image, question, processor, model):
-    # Prepare the image for the model
-    encoding = processor(image, question, return_tensors="pt", truncation=True, padding="max_length", max_length=512)
-    # Make sure we have the pixel values (image), input IDs, attention mask, and token type IDs
-    inputs = {k: v.to(torch.device('cuda' if torch.cuda.is_available() else 'cpu')) for k, v in encoding.items()}
-    # Move model to the appropriate device
-    model.to(torch.device('cuda' if torch.cuda.is_available() else 'cpu'))
-    # Perform the forward pass
-    outputs = model(**inputs)
-    start_logits = outputs.start_logits
-    end_logits = outputs.end_logits
-    # Get the most likely beginning and end of answer with the argmax of the logits
-    all_tokens = processor.tokenizer.convert_ids_to_tokens(inputs['input_ids'].squeeze().tolist())
-    answer = ' '.join(all_tokens[torch.argmax(start_logits):torch.argmax(end_logits)+1])
-    return answer
+#function to provide image
+def input_image_setup(uploaded_file):
+    if uploaded_file is not None:
+        #read te file into byte
+        bytes_data = uploaded_file.getvalue()        
+        image_parts=[
+            {
+                "mime_type": uploaded_file.type,
+                "data": bytes_data
+            }
+        ]
+        return image_parts
+    else:
+        raise FileNotFoundError("No file uploaded")
 
+#Streamlit App
 st.set_page_config(page_title="Invoice Extractor")
-st.header("Invoice Extractor using LayoutLMv2")
-
-question = st.text_input("Enter your question about the invoice:")
-uploaded_file = st.file_uploader("Choose an invoice image...", type=["jpg", "jpeg", "png"])
-image = None
+st.header("Gemini Application")
+input=st.text_input("Input Prompt: ",key="input")
+uploaded_file = st.file_uploader("choose an image...", type=["jpg","jpeg","png"])
+image=""
 if uploaded_file is not None:
     image = Image.open(uploaded_file)
-    st.image(image, caption="Uploaded Image.", use_column_width=True)
+    st.image(image, caption="Uploaded Image.", use_column_width=True) 
+submit = st.button("Tell me about the invoice")
+input_prompt="""
+You are an expert in understanding invoices. you will receive input images as invoices and 
+you will have to answer questions based on the input image.
+"""
+#if submit button is clicked,
+if submit:
+    image_data=input_image_setup(uploaded_file)
+    response=get_model_response(input_prompt,image_data,input)
+    st.subheader("The response is..")
+    st.write(response)
 
-submit = st.button("Get Answer")
 
-if submit and image is not None and question:
-    st.write("Processing...")
-    processor, model = load_model()
-    answer = process_image(image, question, processor, model)
-    st.subheader("The response is:")
-    st.write(answer)