Update app.py

app.py

@@ -1,3 +1,4 @@
+
 from transformers import AutoModel, AutoTokenizer
 import streamlit as st
 from PIL import Image
@@ -5,61 +6,44 @@ import re
 import os
 import uuid
 
+# Set the page layout to wide
+st.set_page_config(layout="wide")
+
 # Load the model and tokenizer only once
-if "model" not in st.session_state or "tokenizer" not in st.session_state:
-    @st.cache_resource
-    def load_model(model_name):
-        if model_name == "OCR for English or Hindi (CPU)":
-            tokenizer = AutoTokenizer.from_pretrained('srimanth-d/GOT_CPU', trust_remote_code=True)
-            model = AutoModel.from_pretrained('srimanth-d/GOT_CPU', trust_remote_code=True, use_safetensors=True, pad_token_id=tokenizer.eos_token_id)
-            model = model.eval()
-        elif model_name == "OCR for English (GPU)":
-            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, device_map='cuda', use_safetensors=True, pad_token_id=tokenizer.eos_token_id)
-            model = model.eval().to('cuda')
-        return model, tokenizer
+@st.cache_resource
+def load_model(model_name):
+    if model_name == "OCR on CPU":
+        tokenizer = AutoTokenizer.from_pretrained('srimanth-d/GOT_CPU', trust_remote_code=True)
+        model = AutoModel.from_pretrained('srimanth-d/GOT_CPU', trust_remote_code=True, use_safetensors=True, pad_token_id=tokenizer.eos_token_id).eval()
+    else:
+        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, device_map='cuda', use_safetensors=True, pad_token_id=tokenizer.eos_token_id).eval().to('cuda')
+    return model, tokenizer
 
-    # Load and store in session state
-    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
-else:
-    model = st.session_state["model"]
-    tokenizer = st.session_state["tokenizer"]
+if "model" not in st.session_state or "tokenizer" not in st.session_state:
+    model, tokenizer = load_model("OCR for English or Hindi (CPU)")
+    st.session_state.update({"model": model, "tokenizer": tokenizer})
 
 # Function to run the GOT model for multilingual OCR
 def run_ocr(image, model, tokenizer):
-    unique_id = str(uuid.uuid4())
-    image_path = f"{unique_id}.png"
-    
-    # Save image to disk
+    image_path = f"{uuid.uuid4()}.png"
     image.save(image_path)
 
     try:
-        # Use the model to extract text from the image
         res = model.chat(tokenizer, image_path, ocr_type='ocr')
-        if isinstance(res, str):
-            return res
-        else:
-            return str(res)
+        return res if isinstance(res, str) else str(res)
     except Exception as e:
         return f"Error: {str(e)}"
     finally:
-        # Clean up the saved image
-        if os.path.exists(image_path):
-            os.remove(image_path)
+        os.remove(image_path)
 
 # Function to highlight keyword in text
 def highlight_text(text, search_term):
-    if not search_term:
-        return text
-    pattern = re.compile(re.escape(search_term), re.IGNORECASE)
-    return pattern.sub(lambda m: f'<span style="background-color: red;">{m.group()}</span>', text)
+    return re.sub(re.escape(search_term), lambda m: f'<span style="background-color: red;">{m.group()}</span>', text, flags=re.IGNORECASE) if search_term else text
 
 # Streamlit App
-st.title("GOT-OCR Multilingual Demo")
-st.write("Upload an image for OCR")
+st.title(":blue[Object character recognition Application]")
+st.write("Give your Image")
 
 # Create two columns
 col1, col2 = st.columns(2)
@@ -73,30 +57,22 @@ with col1:
 
 # 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)"])
+    model_option = st.selectbox("Select Model", ["OCR on CPU", "OCR on GPU"])
     
-    if st.button("Run OCR"):
-        with st.spinner("Processing..."):
-            # Run OCR and store the result in session state
-            if uploaded_image:
+    if st.button("DO OCR "):
+        if uploaded_image:
+            with st.spinner("Processing..."):
+                model, tokenizer = load_model(model_option)
                 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
+                    st.session_state["extracted_text"] = result_text
                 else:
                     st.error(result_text)
-            else:
-                st.error("Please upload an image before running OCR.")
+        else:
+            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"]
-
-        # 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.subheader("Extracted Text:")
-        st.markdown(f'<div style="white-space: pre-wrap;">{highlighted_text}</div>', unsafe_allow_html=True)
+        st.markdown(f'<div style="white-space: pre-wrap;">{highlight_text(st.session_state["extracted_text"], search_term)}</div>', unsafe_allow_html=True)