Update app.py

app.py

@@ -1,17 +1,31 @@
 from PIL import Image
 import streamlit as st
-import pytesseract
+import torch
 from transformers import LayoutLMForTokenClassification, LayoutLMTokenizer
 
 # Load Hugging Face models
 layoutlm_model = LayoutLMForTokenClassification.from_pretrained("microsoft/layoutlm-base-uncased")
 layoutlm_tokenizer = LayoutLMTokenizer.from_pretrained("microsoft/layoutlm-base-uncased")
 
+
 # Function to extract text from image
 def extract_text_from_image(image):
-    # Convert image to text using Tesseract OCR
-    text = pytesseract.image_to_string(image)
-    return text
+    inputs = layoutlm_tokenizer(image, return_tensors="pt", padding=True, truncation=True)
+    with torch.no_grad():
+        outputs = layoutlm_model(**inputs)
+    
+    # Get the predicted labels (classifications) from the model output
+    predicted_labels = torch.argmax(outputs.logits, dim=2)
+    
+    # Convert labels to tokens using tokenizer
+    tokens = layoutlm_tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])
+    
+    # Filter out special tokens and format the extracted text
+    extracted_text = ""
+    for token, label in zip(tokens, predicted_labels[0]):
+        if label != layoutlm_tokenizer.label_pad_token_id:
+            extracted_text += token.replace("##", "") + " " # Removing '##' subwords
+    return extracted_text.strip()
 
 # Function to generate response
 def generate_response(prompt, text):
@@ -40,3 +54,4 @@ if uploaded_file is not None:
         else:
             st.warning("Please upload an image.")
 
+