Update app.py

app.py

@@ -1,34 +1,53 @@
 import streamlit as st
-import io
 from PIL import Image
+import io
 import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, AutoProcessor, AutoModelForSeq2SeqLM
 from easyocr import Reader
 
-
+# Load the OCR model and text generation model
 ocr_reader = Reader(['en'])
 text_generator = AutoModelForCausalLM.from_pretrained("gpt2")
 text_tokenizer = AutoTokenizer.from_pretrained("gpt2")
+
+# Load the image captioning model
+processor = AutoProcessor.from_pretrained("Salesforce/blip-image-captioning-large")
+caption_model = AutoModelForSeq2SeqLM.from_pretrained("Salesforce/blip-image-captioning-large")
+
+# Define a function to extract text from an image using OCR
 def extract_text(image):
     return ocr_reader.readtext(image)
-def explain_text(text):
-    input_ids = text_tokenizer.encode(text, return_tensors="pt")
-    explanation_ids = text_generator.generate(input_ids, max_length=100, num_return_sequences=1)
-    explanation = text_tokenizer.decode(explanation_ids[0], skip_special_tokens=True)
-    return explanation
 
-st.title("Text Classification Model")
+# Define a function to explain the extracted text using text generation
+def explain_text(text, caption_model, processor):
+    # Extracted text
+    extracted_text = " ".join([res[1] for res in text])
+
+    # Generate an image caption using the image captioning model
+    inputs = processor(extracted_text, return_tensors="pt", padding="max_length", max_length=100, truncation=True)
+    input_ids = inputs["input_ids"]
+    caption = caption_model.generate(input_ids, max_length=50, num_return_sequences=1, no_repeat_ngram_size=2)
+
+    # Decode and return the generated caption
+    generated_caption = processor.decode(caption[0], skip_special_tokens=True)
+    return generated_caption
+
+# Create a Streamlit layout
+st.title("Text Extraction and Explanation")
+
+# Allow users to upload an image
 uploaded_file = st.file_uploader("Upload an image:")
 
+# Extract text from the uploaded image and explain it
 if uploaded_file is not None:
     image = Image.open(uploaded_file)
     ocr_results = extract_text(image)
-    extracted_text = " ".join([res[1] for res in ocr_results])
-    explanation = explain_text(extracted_text)
+    explanation = explain_text(ocr_results, caption_model, processor)
+
     st.markdown("**Extracted text:**")
-    st.markdown(extracted_text)
+    st.markdown(" ".join([res[1] for res in ocr_results]))
 
-    st.markdown("**Explanation:**")
+    st.markdown("**Explanation (Image Caption):**")
     st.markdown(explanation)
 
 else: