Update app.py

app.py

@@ -1,29 +1,33 @@
 import streamlit as st
-from transformers import pipeline as transformers_pipeline, AutoTokenizer, AutoModelForCausalLM, AutoTokenizer, AutoModelForSequenceClassification
+from transformers import pipeline as transformers_pipeline, AutoTokenizer, AutoModelForCausalLM, AutoModelForSequenceClassification
 import pandas as pd
 import torch
+import requests
+from PIL import Image
+import io
 
-# Load the tokenizer and models
+# Load tokenizer and models for similarity and story generation
 similarity_tokenizer = AutoTokenizer.from_pretrained("AidenYan/MiniLM_L6_v2_finetuned_ISOM5240_Group27")
 similarity_model = AutoModelForSequenceClassification.from_pretrained("AidenYan/MiniLM_L6_v2_finetuned_ISOM5240_Group27")
-
 story_tokenizer = AutoTokenizer.from_pretrained("pranavpsv/genre-story-generator-v2")
 story_model = AutoModelForCausalLM.from_pretrained("pranavpsv/genre-story-generator-v2")
 
-labels_df = pd.read_csv("labels_to_image_urls.csv")  # Update this path
+# Load the CSV file into a dataframe
+labels_df = pd.read_csv("path_to_your_csv_file/labels_to_image_urls.csv")  # Make sure to update this path
 
 def image_to_text_from_url(image_url):
     """
-    Generates a caption from the image at the given URL.
+    Generates a caption from the image at the given URL using an image-to-text pipeline.
     """
     image_to_text_pipeline = transformers_pipeline("image-to-text", model="nlpconnect/vit-gpt2-image-captioning")
     return image_to_text_pipeline(image_url)[0]['generated_text']
 
 def generate_mask_from_result(input_text):
     """
-    Simulate generating a mask from the result. Replace with actual logic.
+    Placeholder for generating a mask from the result. This should be replaced with your actual logic.
     """
-    return "This is a placeholder text based on the input: " + input_text
+    # Placeholder logic, replace with actual text processing if needed
+    return "Processed input: " + input_text
 
 def generate_story_from_text(input_text):
     """
@@ -35,7 +39,7 @@ def generate_story_from_text(input_text):
 
 def select_closest_sentence(generated_text):
     """
-    Predict the similarity label for the generated text.
+    Predicts the similarity label for the generated text using the similarity model.
     """
     inputs = similarity_tokenizer(generated_text, return_tensors="pt")
     outputs = similarity_model(**inputs)
@@ -45,7 +49,7 @@ def select_closest_sentence(generated_text):
 
 def get_image_url_for_label(label):
     """
-    Returns the image URL for a given label.
+    Returns the image URL for a given label from the labels dataframe.
     """
     row = labels_df[labels_df['Label'] == label]
     if not row.empty:
@@ -53,29 +57,47 @@ def get_image_url_for_label(label):
     else:
         return None
 
+def display_image_from_url(image_url):
+    """
+    Displays an image in the Streamlit app given its URL.
+    """
+    try:
+        response = requests.get(image_url)
+        image = Image.open(io.BytesIO(response.content))
+        st.image(image, use_column_width=True)
+    except Exception as e:
+        st.error(f"Failed to load image from URL: {e}")
+
 def main():
     st.title("SmartCart (Product Recommender)")
 
+    # User input for text or URL
     input_option = st.radio("Select input option:", ("Text", "URL"))
 
+    # Handling input via text
     if input_option == "Text":
         text_input = st.text_input("Enter the text:")
         if st.button("Generate Story and Image") and text_input:
-            generated_text = generate_mask_from_result(text_input)
-            story_text = generate_story_from_text(generated_text)
-            st.success(f'Generated Caption: {text_input}')
-            st.success(f'Generated Text: {generated_text}')
-            st.text_area('Generated Story:', story_text, height=200)
+            processed_text = generate_mask_from_result(text_input)
+            story_text = generate_story_from_text(processed_text)
+            st.text_area('Generated Story:', story_text, height=300)
+            closest_label = select_closest_sentence(processed_text)
+            image_url = get_image_url_for_label(closest_label)
+            if image_url:
+                display_image_from_url(image_url)
 
+    # Handling input via image URL
     elif input_option == "URL":
         image_url = st.text_input("Enter the image URL:")
         if st.button("Generate Story and Image") and image_url:
             image_text = image_to_text_from_url(image_url)
-            generated_text = generate_mask_from_result(image_text)
-            story_text = generate_story_from_text(generated_text)
-            st.success(f'Generated Caption: {image_text}')
-            st.success(f'Generated Text: {generated_text}')
-            st.text_area('Generated Story:', story_text, height=200)
+            processed_text = generate_mask_from_result(image_text)
+            story_text = generate_story_from_text(processed_text)
+            st.text_area('Generated Story:', story_text, height=300)
+            closest_label = select_closest_sentence(processed_text)
+            mapped_image_url = get_image_url_for_label(closest_label)
+            if mapped_image_url:
+                display_image_from_url(mapped_image_url)
 
 if __name__ == "__main__":
     main()