Update app.py

app.py

@@ -1,39 +1,49 @@
 import streamlit as st
 from PIL import Image
-import torch
-from transformers import AutoFeatureExtractor, AutoModelForCausalLM, AutoTokenizer
+import io
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+from easyocr import Reader
 
-# Load the image captioning model
-model_name = "Salesforce/blip-image-captioning-large"
-feature_extractor = AutoFeatureExtractor.from_pretrained(model_name)
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name)
+# Load the OCR model and text explanation model
+ocr_reader = Reader(['en'])
 
-# Streamlit app title and description
-st.title("Image Information Extractor")
-st.write("Upload an image, and this app will extract information from it.")
+# Load the text explanation model
+text_generator = AutoModelForCausalLM.from_pretrained("HuggingFaceH4/zephyr-7b-alpha")
+text_tokenizer = AutoTokenizer.from_pretrained("HuggingFaceH4/zephyr-7b-alpha")
 
-# Upload image
-image = st.file_uploader("Upload an image", type=["jpg", "png", "jpeg"])
+# Define a function to extract text from an image using OCR
+def extract_text(image):
+    return ocr_reader.readtext(image)
 
-if image:
-    # Display the uploaded image
-    st.image(image, use_column_width=True, caption="Uploaded Image")
+# Define a function to explain the extracted text using text generation
+def explain_text(text, text_generator, text_tokenizer):
+    # Extracted text
+    extracted_text = " ".join([res[1] for res in text])
 
-    # Process the image and extract information
-    image = Image.open(image)
-    st.write("Extracting information from the image...")
+    # Generate an explanation using the text explanation model
+    input_ids = text_tokenizer.encode(extracted_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
 
-    # Use the image captioning model to generate a description
-    inputs = tokenizer(image, return_tensors="pt")
-    with st.spinner("Generating caption..."):
-        captions = model.generate(**inputs)
-    caption = tokenizer.decode(captions[0], skip_special_tokens=True)
+# Create a Streamlit layout
+st.title("Text Extraction and Explanation")
 
-    # Display the image description
-    st.write("Image Description:")
-    st.write(caption)
+# Allow users to upload an image
+uploaded_file = st.file_uploader("Upload an image:")
 
-    # You can add more processing or explanations here if needed
+# 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)
+    explanation = explain_text(ocr_results, text_generator, text_tokenizer)
 
-st.sidebar.text("Built with ❤️ by Your Name")
+    st.markdown("**Extracted text:**")
+    st.markdown(" ".join([res[1] for res in ocr_results]))
+
+    st.markdown("**Explanation:**")
+    st.markdown(explanation)
+
+else:
+    st.markdown("Please upload an image to extract text and get an explanation.")