Update app.py

app.py

@@ -9,13 +9,19 @@ import torch
 import tempfile
 import os
 import re
-import base64
+import json
 from groq import Groq
 
 # Page configuration
 st.set_page_config(page_title="DualTextOCRFusion", page_icon="🔍", layout="wide")
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
+# Directories for images and results
+IMAGES_DIR = "images"
+RESULTS_DIR = "results"
+os.makedirs(IMAGES_DIR, exist_ok=True)
+os.makedirs(RESULTS_DIR, exist_ok=True)
+
 # Load Surya OCR Models (English + Hindi)
 det_processor, det_model = load_det_processor(), load_det_model()
 det_model.to(device)
@@ -44,25 +50,18 @@ def init_qwen_model():
 
 # Text Cleaning AI - Clean spaces, handle dual languages
 def clean_extracted_text(text):
-    # Remove extra spaces
     cleaned_text = re.sub(r'\s+', ' ', text).strip()
     cleaned_text = re.sub(r'\s([?.!,])', r'\1', cleaned_text)
     return cleaned_text
 
 # Polish the text using a model
 def polish_text_with_ai(cleaned_text):
-    prompt = f"Remove unwanted spaces between and inside words to join incomplete words, creating a meaningful sentence in either Hindi, English, or Hinglish without altering any words from the given extracted text. Then, return the corrected text with adjusted spaces, keeping it as close to the original as possible, along with relevant details or insights that an AI can provide about the extracted text. Extracted Text : {cleaned_text}"
+    prompt = f"Remove unwanted spaces between and inside words to join incomplete words, creating a meaningful sentence in either Hindi, English, or Hinglish without altering any words from the given extracted text. Then, return the corrected text with adjusted spaces."
     client = Groq(api_key="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")
     chat_completion = client.chat.completions.create(
-    messages=[
-            {
-                "role": "system",
-                "content": "You are a pedantic sentence corrector. Remove extra spaces between and within words to make the sentence meaningful in English, Hindi, or Hinglish, according to the context of the sentence, without changing any words."
-            },
-            {
-                "role": "user",
-                "content": prompt,
-            }
+        messages=[
+            {"role": "system", "content": "You are a pedantic sentence corrector."},
+            {"role": "user", "content": prompt},
         ],
         model="gemma2-9b-it",
     )
@@ -88,11 +87,9 @@ def extract_text_qwen(image_file, model, processor):
 
 # Highlight keyword search
 def highlight_text(text, search_term):
-    if not search_term:  # If no search term is provided, return the original text
+    if not search_term:
         return text
-    # Use a regular expression to search for the term, case insensitive
     pattern = re.compile(re.escape(search_term), re.IGNORECASE)
-    # Highlight matched terms with yellow background
     return pattern.sub(lambda m: f'<span style="background-color: yellow;">{m.group()}</span>', text)
 
 # Title and UI
@@ -106,89 +103,73 @@ model_choice = st.sidebar.selectbox("Select OCR Model:", ("GOT_CPU", "GOT_GPU",
 
 # Upload Section
 uploaded_file = st.sidebar.file_uploader("Choose an image...", type=["png", "jpg", "jpeg"])
+clipboard_text = st.sidebar.text_area("Paste image path from clipboard:")
+
+if uploaded_file or clipboard_text:
+    image_path = None
+    if uploaded_file:
+        image_path = os.path.join(IMAGES_DIR, uploaded_file.name)
+        with open(image_path, "wb") as f:
+            f.write(uploaded_file.getvalue())
+    elif clipboard_text:
+        image_path = clipboard_text.strip()
+
+    # Predict button
+    predict_button = st.sidebar.button("Predict")
+
+    # Main columns
+    col1, col2 = st.columns([2, 1])
+
+    # Check if result JSON already exists
+    result_json_path = os.path.join(RESULTS_DIR, f"{os.path.basename(image_path)}_result.json") if image_path else None
+
+    if predict_button and image_path:
+        if os.path.exists(result_json_path):
+            with open(result_json_path, "r") as json_file:
+                result_data = json.load(json_file)
+            polished_text = result_data.get("polished_text", "")
+        else:
+            with st.spinner("Processing..."):
+                image = Image.open(image_path).convert("RGB")
+
+                if model_choice == "GOT_CPU":
+                    got_model, tokenizer = init_got_model()
+                    extracted_text = extract_text_got(image_path, got_model, tokenizer)
+
+                elif model_choice == "GOT_GPU":
+                    got_gpu_model, tokenizer = init_got_gpu_model()
+                    extracted_text = extract_text_got(image_path, got_gpu_model, tokenizer)
+
+                elif model_choice == "Qwen":
+                    qwen_model, qwen_processor = init_qwen_model()
+                    extracted_text = extract_text_qwen(image_path, qwen_model, qwen_processor)
+
+                elif model_choice == "Surya (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)
+
+                cleaned_text = clean_extracted_text(extracted_text)
+                polished_text = polish_text_with_ai(cleaned_text) if model_choice in ["GOT_CPU", "GOT_GPU"] else cleaned_text
+
+                # Save result to JSON
+                with open(result_json_path, "w") as json_file:
+                    json.dump({"polished_text": polished_text}, json_file)
+
+        # Display image preview and text
+        if image_path:
+            with col1:
+                col1.image(image_path, caption='Uploaded Image', use_column_width=False, width=300)
 
-# Input from clipboard
-if st.sidebar.button("Paste from Clipboard"):
-    try:
-        clipboard_data = st.experimental_get_clipboard()
-        if clipboard_data:
-            # Assuming clipboard data is base64 encoded image
-            image_data = base64.b64decode(clipboard_data)
-            uploaded_file = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
-            uploaded_file.write(image_data)
-            uploaded_file.seek(0)
-    except:
-        st.sidebar.warning("Clipboard data is not an image.")
-
-# Input from camera
-camera_file = st.sidebar.camera_input("Capture from Camera")
-if camera_file:
-    uploaded_file = camera_file
-
-# Predict button
-predict_button = st.sidebar.button("Predict")
-
-# Main columns
-col1, col2 = st.columns([2, 1])
-
-# Display image preview
-if uploaded_file:
-    image = Image.open(uploaded_file)
-    with col1:
-        col1.image(image, caption='Uploaded Image', use_column_width=False, width=300)
-
-# Handle predictions
-if predict_button and uploaded_file:
-    with st.spinner("Processing..."):
-        # Save uploaded image
-        with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as temp_file:
-            temp_file.write(uploaded_file.getvalue())
-            temp_file_path = temp_file.name
-
-        image = Image.open(temp_file_path)
-        image = image.convert("RGB")
-
-        if model_choice == "GOT_CPU":
-            got_model, tokenizer = init_got_model()
-            extracted_text = extract_text_got(temp_file_path, got_model, tokenizer)
-
-        elif model_choice == "GOT_GPU":
-            got_gpu_model, tokenizer = init_got_gpu_model()
-            extracted_text = extract_text_got(temp_file_path, got_gpu_model, tokenizer)
-
-        elif model_choice == "Qwen":
-            qwen_model, qwen_processor = init_qwen_model()
-            extracted_text = extract_text_qwen(temp_file_path, qwen_model, qwen_processor)
-
-        elif model_choice == "Surya (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)
-
-        # Clean extracted text
-        cleaned_text = clean_extracted_text(extracted_text)
-
-        # Optionally, polish text with AI model for better language flow
-        polished_text = polish_text_with_ai(cleaned_text) if model_choice in ["GOT_CPU", "GOT_GPU"] else cleaned_text
-
-        # Delete temp file
-        if os.path.exists(temp_file_path):
-            os.remove(temp_file_path)
-
-        # Display extracted text and search functionality
         st.subheader("Extracted Text (Cleaned & Polished)")
         st.markdown(polished_text, unsafe_allow_html=True)
 
         # Input box for real-time search
-        search_query = st.text_input("Search in extracted text:", key="search_query", placeholder="Type to search...")
-
-        # Update results dynamically based on the search term
+        search_query = st.text_input("Search in extracted text:", key="search_query", placeholder="Type to search...", on_change=lambda: st.session_state.update(search_query) disabled=not uploaded_file)
+        
+        # Highlight the search term in the text
         if search_query:
-            # Highlight the search term in the text
             highlighted_text = highlight_text(polished_text, search_query)
             st.markdown("### Highlighted Search Results:")
             st.markdown(highlighted_text, unsafe_allow_html=True)
-        else:
-            st.markdown("### Extracted Text:")
-            st.markdown(polished_text)