Update app.py

app.py

@@ -4,30 +4,20 @@ from qwen_vl_utils import process_vision_info
 import torch
 from PIL import Image
 
+@st.cache_resource
+def load_model():
+    # Load model on CPU
+    model = Qwen2VLForConditionalGeneration.from_pretrained(
+        "Qwen/Qwen2-VL-2B-Instruct", torch_dtype=torch.float32, device_map=None
+    ).to("cpu")  # type:ignore # Ensure the model is on CPU
 
+    min_pixels = 256*28*28
+    max_pixels = 1280*28*28
+    processor = AutoProcessor.from_pretrained("Qwen/Qwen2-VL-2B-Instruct", min_pixels=min_pixels, max_pixels=max_pixels)
+    return model, processor
+    
 
-# Load model on CPU
-model = Qwen2VLForConditionalGeneration.from_pretrained(
-    "Qwen/Qwen2-VL-2B-Instruct", torch_dtype=torch.float32, device_map=None
-).to("cpu")  # Ensure the model is on CPU
-
-min_pixels = 256*28*28
-max_pixels = 1280*28*28
-processor = AutoProcessor.from_pretrained("Qwen/Qwen2-VL-2B-Instruct", min_pixels=min_pixels, max_pixels=max_pixels)
-
-# Streamlit app
-st.title("OCR Application with Keyword Search")
-
-# Upload image
-uploaded_file = st.file_uploader("Choose an image...", type=["png", "jpg", "jpeg"])
-
-if uploaded_file is not None:
-    # Convert the uploaded file to an image
-    img = Image.open(uploaded_file)
-
-    # Display the uploaded image
-    st.image(img, caption="Uploaded Image", use_column_width=True)
-
+def process_file(img, model, processor):
     # Prepare the image for the model
     messages = [
         {
@@ -72,17 +62,52 @@ if uploaded_file is not None:
         generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
     )
     
+    return output_text[0]
+
+
+# Streamlit app
+st.title("OCR Application with Keyword Search")
+
+
+# Initialize session state variables
+if 'current_image' not in st.session_state:
+    st.session_state.current_image = None
+if 'extracted_text' not in st.session_state:
+    st.session_state.extracted_text = None
+
+
+model, processor = load_model()
+
+# Upload image
+uploaded_file = st.file_uploader("Choose an image...", type=["png", "jpg", "jpeg"])
+
+if uploaded_file is not None:
+    # Convert the uploaded file to an image
+    img = Image.open(uploaded_file)
+
+    if st.session_state.current_image != uploaded_file:
+        st.session_state.current_image = uploaded_file
+        st.session_state.extracted_text = process_file(img, model, processor)
+    
+    # Display the uploaded image
+    st.image(img, caption="Uploaded Image", use_column_width=True)
+
+    # if 'extracted_text' not in st.session_state:
+    #     st.session_state.extracted_text = process_file(img, model, processor)
+    
     # Display the extracted text
-    extracted_text = output_text[0]
     st.subheader("Extracted Text")
-    st.write(extracted_text)
-
-    # Keyword Search
-    keyword = st.text_input("Enter keyword to search in the extracted text")
-    if keyword:
-        if keyword.lower() in extracted_text.lower():
-            highlighted_text = extracted_text.replace(keyword, f"**{keyword}**")
-            st.subheader("Keyword Found")
-            st.write(highlighted_text, unsafe_allow_html=True)
-        else:
-            st.write("Keyword not found in the extracted text.")
+    st.write(st.session_state.extracted_text)
+
+# Keyword Search
+keyword = st.text_input("Enter keyword to search in the extracted text")
+if keyword and st.session_state.extracted_text:
+    if keyword.lower() in st.session_state.extracted_text.lower():
+        highlighted_text = st.session_state.extracted_text.replace(keyword, f"**{keyword}**")
+        st.subheader("Keyword Found")
+        st.markdown(highlighted_text, unsafe_allow_html=True)
+    else:
+        st.write("Keyword not found in the extracted text.")
+elif keyword:
+    st.write("Please upload an image first to extract text.")
+