using gradio instead of streamlit

app.py

@@ -1,4 +1,4 @@
-import streamlit as st
+import gradio as gr
 from PIL import Image
 import torch
 from transformers import Qwen2VLForConditionalGeneration, AutoProcessor, pipeline
@@ -19,44 +19,28 @@ login(token=hf_token)
 try:
     image_to_text_pipeline = pipeline("image-to-text", model="google/paligemma-3b-mix-448", device=0 if torch.cuda.is_available() else -1)
 except Exception as e:
-    st.error(f"Error loading image-to-text model: {e}")
-    st.stop()
+    raise Exception(f"Error loading image-to-text model: {e}")
 
 # Load ColPali model with Hugging Face token
 try:
     model_colpali = ColPali.from_pretrained("vidore/colpali-v1.2", torch_dtype=torch.bfloat16).to(device)
     processor_colpali = ColPaliProcessor.from_pretrained("google/paligemma-3b-mix-448")
 except Exception as e:
-    st.error(f"Error loading ColPali model or processor: {e}")
-    st.stop()
+    raise Exception(f"Error loading ColPali model or processor: {e}")
 
 # Load Qwen model
 try:
     model_qwen = Qwen2VLForConditionalGeneration.from_pretrained("Qwen/Qwen2-VL-7B-Instruct").to(device)
     processor_qwen = AutoProcessor.from_pretrained("Qwen/Qwen2-VL-7B-Instruct")
 except Exception as e:
-    st.error(f"Error loading Qwen model or processor: {e}")
-    st.stop()
+    raise Exception(f"Error loading Qwen model or processor: {e}")
 
-# Streamlit UI
-st.title("OCR and Document Search Web Application")
-st.write("Upload an image containing text in both Hindi and English for OCR processing and keyword search.")
-
-# File uploader for the image
-uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
-
-if uploaded_file is not None:
+# Function to process the image and extract text
+def process_image(image, keyword):
     try:
-        image = Image.open(uploaded_file)
-        st.image(image, caption='Uploaded Image.', use_column_width=True)
-        st.write("")
-
         # Use the image-to-text pipeline to extract text from the image
         output_text_img_to_text = image_to_text_pipeline(image)
 
-        st.write("Extracted Text from Image:")
-        st.write(output_text_img_to_text)
-
         # Prepare input for Qwen model for image description
         conversation = [{"role": "user", "content": [{"type": "image"}, {"type": "text", "text": "Describe this image."}]}]
         text_prompt = processor_qwen.apply_chat_template(conversation, add_generation_prompt=True)
@@ -68,18 +52,40 @@ if uploaded_file is not None:
             generated_ids_qwen = [output_ids_qwen[len(input_ids):] for input_ids, output_ids_qwen in zip(inputs_qwen.input_ids, output_ids_qwen)]
             output_text_qwen = processor_qwen.batch_decode(generated_ids_qwen, skip_special_tokens=True, clean_up_tokenization_spaces=True)
 
-        st.write("Qwen Model Description:")
-        st.write(output_text_qwen)
+        extracted_text = output_text_img_to_text[0]['generated_text']
 
         # Keyword search in the extracted text
-        keyword = st.text_input("Enter a keyword to search in the extracted text:")
+        keyword_found = ""
         if keyword:
-            if keyword.lower() in output_text_img_to_text[0]['generated_text'].lower():
-                st.write(f"Keyword '{keyword}' found in the text.")
+            if keyword.lower() in extracted_text.lower():
+                keyword_found = f"Keyword '{keyword}' found in the text."
             else:
-                st.write(f"Keyword '{keyword}' not found in the text.")
-    except Exception as e:
-        st.error(f"An error occurred: {e}")
+                keyword_found = f"Keyword '{keyword}' not found in the text."
 
+        return extracted_text, output_text_qwen[0], keyword_found
+    except Exception as e:
+        return str(e), "", ""
+
+# Define Gradio Interface
+title = "OCR and Document Search Web Application"
+description = "Upload an image containing text in both Hindi and English for OCR processing and keyword search."
+
+# Gradio interface for input and output
+image_input = gr.inputs.Image(type="pil")
+keyword_input = gr.inputs.Textbox(label="Enter a keyword to search in the extracted text (Optional)")
+output_textbox = gr.outputs.Textbox(label="Extracted Text")
+output_description = gr.outputs.Textbox(label="Qwen Model Description")
+output_keyword_search = gr.outputs.Textbox(label="Keyword Search Result")
+
+# Set up Gradio interface layout
+interface = gr.Interface(
+    fn=process_image,  # Function to call when button is pressed
+    inputs=[image_input, keyword_input],  # Input types (image and keyword)
+    outputs=[output_textbox, output_description, output_keyword_search],  # Outputs (text boxes for results)
+    title=title,
+    description=description
+)
+
+# Launch the Gradio app
 if __name__ == "__main__":
-    st.write("Deploying the web application...")
+    interface.launch(share=True)

requirements.txt

@@ -1,4 +1,4 @@
-streamlit
+gradio
 Pillow
 torch
 transformers