kash shi

app.py

@@ -38,9 +38,11 @@ st.write("Upload an image for OCR processing. Using GOT-OCR for English translat
 st.sidebar.header("Configuration")
 model_choice = st.sidebar.selectbox("Select OCR Model:", ("For English + Hindi", "For English (GOT-OCR)"))
 
-# Store the uploaded image in session state
+# Store the uploaded image and extracted text in session state
 if 'uploaded_image' not in st.session_state:
     st.session_state.uploaded_image = None
+if 'extracted_text' not in st.session_state:
+    st.session_state.extracted_text = ""
 
 uploaded_file = st.sidebar.file_uploader("Choose an image...", type=["png", "jpg", "jpeg"])
 
@@ -60,37 +62,59 @@ if st.session_state.uploaded_image:
         # Display a smaller preview of the uploaded image (set width to 300px)
         col1.image(image, caption='Uploaded Image', use_column_width=False, width=300)
 
+# Handle predictions
 if predict_button and st.session_state.uploaded_image:
-    with col2:
-        with st.spinner("Processing..."):
-            # Save the uploaded file temporarily
-            with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as temp_file:
-                temp_file.write(st.session_state.uploaded_image.getvalue())
-                temp_file_path = temp_file.name
-
-            image = Image.open(temp_file_path)
-            image = image.convert("RGB")  
-
-            if model_choice == "For English + Hindi":
-                langs = ["en", "hi"]  
-                predictions = run_ocr([image], [langs], det_model, det_processor, rec_model, rec_processor)
-                text_list = re.findall(r"text='(.*?)'", str(predictions[0]))
-                extracted_text = ' '.join(text_list)
-
-                with col2:
-                    st.subheader("Extracted Text (Surya):")
-                    st.write(extracted_text)
-
-            elif model_choice == "For English (GOT-OCR)":
-                image_file = temp_file_path
-                res = got_model.chat(tokenizer, image_file, ocr_type='ocr')
-
-                with col2:
-                    st.subheader("Extracted Text (GOT-OCR):")
-                    st.write(res)
-
-            # Delete the temporary file after processing
-            if os.path.exists(temp_file_path):
-                os.remove(temp_file_path)
-# else:
-#     st.sidebar.warning("Please upload an image before predicting.")
+    # with col2:
+    with st.spinner("Processing..."):
+        # Save the uploaded file temporarily
+        with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as temp_file:
+            temp_file.write(st.session_state.uploaded_image.getvalue())
+            temp_file_path = temp_file.name
+
+        image = Image.open(temp_file_path)
+        image = image.convert("RGB")  
+
+        if model_choice == "For English + Hindi":
+            langs = ["en", "hi"]  
+            predictions = run_ocr([image], [langs], det_model, det_processor, rec_model, rec_processor)
+            text_list = re.findall(r"text='(.*?)'", str(predictions[0]))
+            extracted_text = ' '.join(text_list)
+
+            st.session_state.extracted_text = extracted_text  # Save extracted text in session state
+
+            # with col2:
+            #     st.subheader("Extracted Text (Surya):")
+            #     st.write(extracted_text)
+
+        elif model_choice == "For English (GOT-OCR)":
+            image_file = temp_file_path
+            res = got_model.chat(tokenizer, image_file, ocr_type='ocr')
+
+            st.session_state.extracted_text = res  # Save extracted text in session state
+
+            # with col2:
+            #     st.subheader("Extracted Text (GOT-OCR):")
+            #     st.write(res)
+
+        # Delete the temporary file after processing
+        if os.path.exists(temp_file_path):
+            os.remove(temp_file_path)
+
+# Search functionality
+if st.session_state.extracted_text:
+    search_query = st.text_input("Search in extracted text:", key="search_query", placeholder="Type to search...")
+
+    # Create a pattern to find the search query in a case-insensitive way
+    if search_query:
+        pattern = re.compile(re.escape(search_query), re.IGNORECASE)
+        highlighted_text = st.session_state.extracted_text
+        
+        # Replace matching text with highlighted version (bright green)
+        highlighted_text = pattern.sub(lambda m: f"<span style='background-color: limegreen;'>{m.group(0)}</span>", highlighted_text)
+
+        st.markdown("### Highlighted Search Results:")
+        st.markdown(highlighted_text, unsafe_allow_html=True)
+    else:
+        # If no search query, show the original extracted text
+        st.markdown("### Extracted Text:")
+        st.markdown(st.session_state.extracted_text, unsafe_allow_html=True)