Update app.py

app.py

@@ -10,7 +10,11 @@ layoutlm_tokenizer = LayoutLMTokenizer.from_pretrained(model_name)
 
 # Function to extract text from image using LayoutLM
 def extract_text_from_image(image):
-    inputs = layoutlm_tokenizer(image, return_tensors="pt")
+    # Convert image to text using OCR (Optical Character Recognition) before tokenizing
+    # For simplicity, we'll assume the image is converted to text directly without OCR
+    # You may need to use OCR libraries like pytesseract for real-world scenarios
+    text = image_to_text(image)
+    inputs = layoutlm_tokenizer(text, return_tensors="pt")
     with torch.no_grad():
         outputs = layoutlm_model(**inputs)
     predicted_tokens = layoutlm_tokenizer.convert_ids_to_tokens(outputs.logits.argmax(2).squeeze().tolist())
@@ -18,31 +22,31 @@ def extract_text_from_image(image):
     extracted_text = " ".join([token for token in predicted_tokens if token not in ['[CLS]', '[SEP]', '[PAD]']])
     return extracted_text
 
+# Function to convert image to text (replace with OCR library if needed)
+def image_to_text(image):
+    # For simplicity, return a placeholder text
+    return "Invoice text extracted from the image."
+
 # Streamlit App
 st.set_page_config(page_title="Invoice Extractor")
-st.header("Gemini Application")
+st.header("Invoice Extractor")
 
-prompt = st.text_input("Input Prompt: ", key="input")
-uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
-image = ""
+prompt = st.text_input("Input Prompt: ", "Please provide details about this invoice.")
+uploaded_file = st.file_uploader("Upload an invoice image...", type=["jpg", "jpeg", "png"])
+image = None
 
 if uploaded_file is not None:
     image = Image.open(uploaded_file)
     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.
-"""
+submit = st.button("Extract and Generate")
 
 if submit:
-    if image:
+    if image is None:
+        st.warning("Please upload an image.")
+    else:
         extracted_text = extract_text_from_image(image)
         # Implement your response generation logic here based on the extracted text and prompt
         response = "Response generated based on the extracted text and prompt."
         st.subheader("Generated Response:")
         st.write(response)
-    else:
-        st.warning("Please upload an image.")