Create app.py

app.py

@@ -0,0 +1,88 @@
+import streamlit as st
+from PIL import Image
+import torch
+from transformers import AutoProcessor, AutoModelForVision2Seq
+import re
+
+# Load OCR model
+model_name = "microsoft/trocr-large-handwritten"  # Using a larger model for better performance
+processor = AutoProcessor.from_pretrained(model_name)
+model = AutoModelForVision2Seq.from_pretrained(model_name)
+
+@st.cache_resource
+def load_model():
+    return processor, model
+
+def perform_ocr(image):
+    processor, model = load_model()
+    
+    try:
+        # Preprocess the image
+        pixel_values = processor(image, return_tensors="pt").pixel_values
+        
+        # Generate text
+        generated_ids = model.generate(pixel_values)
+        generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+        
+        return generated_text
+    except Exception as e:
+        st.error(f"An error occurred during OCR processing: {str(e)}")
+        return None
+
+def search_text(text, keyword):
+    if not keyword or not text:
+        return []
+    
+    # Case-insensitive search
+    pattern = re.compile(re.escape(keyword), re.IGNORECASE)
+    matches = list(pattern.finditer(text))
+    
+    # Get context around matches
+    context_size = 50  # characters before and after the match
+    results = []
+    for match in matches:
+        start = max(0, match.start() - context_size)
+        end = min(len(text), match.end() + context_size)
+        context = text[start:end]
+        highlighted = re.sub(pattern, lambda m: f"**{m.group()}**", context)
+        results.append(highlighted)
+    
+    return results
+
+st.title("OCR and Text Search Application")
+
+uploaded_file = st.file_uploader("Choose an image file", type=["jpg", "jpeg", "png"])
+
+if uploaded_file is not None:
+    try:
+        image = Image.open(uploaded_file)
+        st.image(image, caption="Uploaded Image", use_column_width=True)
+        
+        if st.button("Perform OCR"):
+            with st.spinner("Processing image..."):
+                extracted_text = perform_ocr(image)
+            
+            if extracted_text:
+                st.success("OCR completed successfully!")
+                st.subheader("Extracted Text:")
+                st.write(extracted_text)
+                
+                # Save extracted text to session state
+                st.session_state.extracted_text = extracted_text
+            else:
+                st.error("Failed to extract text from the image. Please try another image.")
+    except Exception as e:
+        st.error(f"An error occurred while processing the image: {str(e)}")
+
+# Search functionality
+search_keyword = st.text_input("Enter a keyword to search:")
+if search_keyword and 'extracted_text' in st.session_state:
+    search_results = search_text(st.session_state.extracted_text, search_keyword)
+    if search_results:
+        st.success(f"Found {len(search_results)} matches for '{search_keyword}':")
+        for i, result in enumerate(search_results, 1):
+            st.markdown(f"{i}. ...{result}...")
+    else:
+        st.warning(f"No matches found for '{search_keyword}'.")
+elif search_keyword:
+    st.info("Please perform OCR on an image before searching.")