Update app.py

app.py

@@ -1,23 +1,23 @@
+
 import streamlit as st
 from PIL import Image
 import re
-from transformers import VisionEncoderDecoderModel, AutoProcessor
+from transformers import AutoModel, AutoTokenizer
 
 st.set_page_config(page_title="OCR Application", page_icon="🖼️", layout="wide")
-device = "cpu"
 
 @st.cache_resource
 def load_model():
-    # Load processor and model for OCR
-    processor = AutoProcessor.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True)
-    model = VisionEncoderDecoderModel.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True, device_map='cpu', low_cpu_mem_usage=True)
-    return processor, model
-
-def extract_text(image, processor, model):
-    # Preprocess the image and extract text
-    pixel_values = processor(images=image, return_tensors="pt").pixel_values
-    generated_ids = model.generate(pixel_values)
-    extracted_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+    # Load the tokenizer and model for processing images
+    tokenizer = AutoTokenizer.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True)
+    model = AutoModel.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True, low_cpu_mem_usage=True, use_safetensors=True, pad_token_id=tokenizer.eos_token_id)
+    return tokenizer, model
+
+def extract_text(image, tokenizer, model):
+    # Preprocess the image and extract text using the model
+    inputs = tokenizer(images=image, return_tensors="pt")  # Adjust based on how the model expects inputs
+    generated_ids = model.generate(**inputs)
+    extracted_text = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
     return extracted_text
 
 def highlight_matches(text, keywords):
@@ -27,10 +27,10 @@ def highlight_matches(text, keywords):
     return highlighted_text
 
 def main():
-    st.title("OCR Text Extractor using Hugging Face Model")
+    st.title("OCR Text Extractor using Qwen Model")
 
-    # Load model and processor
-    processor, model = load_model()
+    # Load model and tokenizer
+    tokenizer, model = load_model()
 
     # Upload Image
     uploaded_file = st.file_uploader("Upload an image for OCR", type=["png", "jpg", "jpeg"])
@@ -39,9 +39,9 @@ def main():
         image = Image.open(uploaded_file)
         st.image(image, caption="Uploaded Image", use_column_width=True)
 
-        # Extract text from the image
+        # Extract text from the image using the model
         with st.spinner("Extracting text from the image..."):
-            extracted_text = extract_text(image, processor, model)
+            extracted_text = extract_text(image, tokenizer, model)
         
         st.subheader("Extracted Text")
         st.text_area("Text from Image", extracted_text, height=300)
@@ -57,3 +57,4 @@ def main():
 
 if __name__ == "__main__":
     main()
+