Update app.py

app.py

@@ -1,88 +1,51 @@
-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()
+# app.py  
+import streamlit as st  
+import cv2  
+import numpy as np  
+import pytesseract  
+from PIL import Image  
+import re  
+
+# Set the title of the webpage  
+st.title("OCR Text Extraction Tool")  
+
+# Uploading an image  
+uploaded_file = st.file_uploader("Upload an Image", type=["jpg", "jpeg", "png"])  
+
+if uploaded_file is not None:  
+    # Convert the uploaded file content to an image  
+    image = Image.open(uploaded_file)  
     
-    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
+    # Convert PIL Image to OpenCV format  
+    opencv_image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)  
 
-def search_text(text, keyword):
-    if not keyword or not text:
-        return []
+    # Display the image  
+    st.image(image, caption='Uploaded Image', use_column_width=True)  
     
-    # Case-insensitive search
-    pattern = re.compile(re.escape(keyword), re.IGNORECASE)
-    matches = list(pattern.finditer(text))
+    # Perform OCR  
+    text = pytesseract.image_to_string(opencv_image)  
     
-    # 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)
+    st.subheader("Extracted Text:")  
+    st.write(text)  
+
+    # Search functionality  
+    search_keyword = st.text_input("Enter a keyword to search in the extracted text:")  
+    if search_keyword:  
+        pattern = re.compile(re.escape(search_keyword), re.IGNORECASE)  
+        matches = list(pattern.finditer(text))  
         
-        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.")
+        if matches:  
+            st.markdown("### Keyword Found:")  
+            for match in matches:  
+                start, end = match.span()  
+                context_start = max(0, start - 50)  
+                context_end = min(len(text), end + 50)  
+                context = text[context_start:context_end]  
+                highlighted_text = (  
+                    context[:start-context_start] +  
+                    f"<span style='background-color: yellow;'>{context[start-context_start:end-context_start]}</span>" +  
+                    context[end-context_start:]  
+                )  
+                st.markdown(f"...{highlighted_text}...")  
+        else:  
+            st.write(f"Keyword '{search_keyword}' not found in the extracted text.")