Update app.py

app.py

@@ -1,4 +1,4 @@
-from transformers import AutoModel, AutoTokenizer
+ffrom transformers import AutoModel, AutoTokenizer
 import streamlit as st
 from PIL import Image
 import re
@@ -20,7 +20,7 @@ if "model" not in st.session_state or "tokenizer" not in st.session_state:
         return model, tokenizer
 
     # Load and store in session state
-    model_option = st.selectbox("Select Model", ["OCR for English or Hindi (CPU)", "OCR for English (GPU)"])
+    model_option = "OCR for English or Hindi (CPU)"  # Default value for loading purposes
     model, tokenizer = load_model(model_option)
     st.session_state["model"] = model
     st.session_state["tokenizer"] = tokenizer
@@ -61,34 +61,42 @@ def highlight_text(text, search_term):
 st.title("GOT-OCR Multilingual Demo")
 st.write("Upload an image for OCR")
 
-# Upload image
-uploaded_image = st.file_uploader("Upload Image", type=["png", "jpg", "jpeg"])
+# Create two columns
+col1, col2 = st.columns(2)
+
+# Left column - Display the uploaded image
+with col1:
+    uploaded_image = st.file_uploader("Upload Image", type=["png", "jpg", "jpeg"])
+    if uploaded_image:
+        image = Image.open(uploaded_image)
+        st.image(image, caption='Uploaded Image', use_column_width=True)
 
-if uploaded_image:
-    image = Image.open(uploaded_image)
-    st.image(image, caption='Uploaded Image', use_column_width=True)
+# Right column - Model selection, options, and displaying extracted text
+with col2:
+    model_option = st.selectbox("Select Model", ["OCR for English or Hindi (CPU)", "OCR for English (GPU)"])
     
     if st.button("Run OCR"):
         with st.spinner("Processing..."):
             # Run OCR and store the result in session state
-            result_text = run_ocr(image, model, tokenizer)
-            if "Error" not in result_text:
-                st.session_state["extracted_text"] = result_text  # Store the result in session state
+            if uploaded_image:
+                result_text = run_ocr(image, model, tokenizer)
+                if "Error" not in result_text:
+                    st.session_state["extracted_text"] = result_text  # Store the result in session state
+                else:
+                    st.error(result_text)
             else:
-                st.error(result_text)
+                st.error("Please upload an image before running OCR.")
 
-# Display the extracted text if it exists in session state
-if "extracted_text" in st.session_state:
-    extracted_text = st.session_state["extracted_text"]
+    # Display the extracted text if it exists in session state
+    if "extracted_text" in st.session_state:
+        extracted_text = st.session_state["extracted_text"]
 
-    st.subheader("Extracted Text:")
-    st.text(extracted_text)  # Display the raw extracted text
-
-    # Keyword input for search
-    search_term = st.text_input("Enter a word or phrase to highlight:")
-    
-    # Highlight keyword in the extracted text
-    if search_term:
+        # Keyword input for search
+        search_term = st.text_input("Enter a word or phrase to highlight:")
+        
+        # Highlight keyword in the extracted text
         highlighted_text = highlight_text(extracted_text, search_term)
+        
         # Display the highlighted text using markdown
-        st.markdown(highlighted_text, unsafe_allow_html=True)
+        st.subheader("Extracted Text:")
+        st.markdown(f'<div style="white-space: pre-wrap;">{highlighted_text}</div>', unsafe_allow_html=True)