Update app.py

app.py

@@ -1,27 +1,38 @@
-# Import convention
-import streamlit as st
-from transformers import pipeline
-from PIL import Image
+# Import necessary libraries
+import streamlit as st  # Streamlit for web application
+from transformers import pipeline  # Hugging Face transformer pipeline
+from PIL import Image  # Python Imaging Library for image handling
 
-# Load the image-to-text pipeline
-image_read = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")
+# Set the title of the Streamlit app
+st.set_page_config(page_title="Storytelling Friend",
+                   page_icon="haha")
+ # Title of the application
 
-def read(image):
-    results = image_read(image)
-    return results
+# Create a file uploader for the image
+uploaded_file = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])  # User uploads an image
 
-# Streamlit UI
-st.title("Storytelling from images for 3-10 year-old kids")
+# Load the image captioning model
+caption_model = pipeline("image-captioning", model="facebook/blip-image-captioning-base")  # Load pre-trained model
 
-uploaded_file = st.file_uploader("Upload an image", type=["jpg", "png", "jpeg"])
+# Load the text generation model
+story_model = pipeline("text-generation", model="gpt2")  # Load a text generation model
 
-if uploaded_file is not None:
-    image = Image.open(uploaded_file).convert("RGB")
-    st.image(image, caption="Uploaded Image", use_column_width=True)
+# Function to generate a story from the caption
+def generate_story(caption):  # Define a function to generate a story
+    story_input = f"Once upon a time, {caption}"  # Create a story prompt
+    story = story_model(story_input, max_length=150, num_return_sequences=1)[0]['generated_text']  # Generate the story
+    return story  # Return the generated story
 
-    # Read image
-    image_content = read(image)
+# Process the uploaded image and generate story
+if uploaded_file is not None:  # Check if a file is uploaded
+    image = Image.open(uploaded_file)  # Open the uploaded image
+    st.image(image, caption="Uploaded Image", use_column_width=True)  # Display the uploaded image
 
-    # Display results
-    st.subheader("Image Content:")
-    st.write(f"Image Content: {image_content[0]}")
+    caption = caption_model(image)[0]['caption']  # Generate caption from the image
+    st.subheader("Image Caption:")  # Subtitle for the caption
+    st.write(caption)  # Display the caption
+
+    # Generate story based on the caption
+    story = generate_story(caption)  # Call the story generation function
+    st.subheader("Generated Story:")  # Subtitle for the generated story
+    st.write(story)  # Display the generated story