Update app.py

app.py

@@ -1,8 +1,8 @@
 # Import necessary libraries
-import streamlit as st
-from transformers import pipeline
-from gtts import gTTS
-import os
+import streamlit as st  # For building the web application
+from transformers import pipeline  # For using pre-trained models (image-to-text and text-generation)
+from gtts import gTTS  # For converting text to speech
+import os  # For file handling (saving and deleting temporary files)
 
 # Function to convert image to text using Hugging Face's BLIP model
 def img2text(url):
@@ -13,17 +13,24 @@ def img2text(url):
         url (str): Path to the image file.
     
     Returns:
-        str: Generated text caption from the image, without words like "illustration".
+        str: Generated text caption from the image, without unwanted words like "illustration".
     """
-    image_to_text_model = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")
-    text = image_to_text_model(url)[0]["generated_text"]
-    
-    # Remove unwanted words like "illustration"
-    unwanted_words = ["illustration", "drawing", "sketch", "picture", "dream", "imagination"]
-    for word in unwanted_words:
-        text = text.replace(word, "")
-    
-    return text.strip()
+    try:
+        # Load the image-to-text model
+        image_to_text_model = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")
+        
+        # Generate text caption from the image
+        text = image_to_text_model(url)[0]["generated_text"]
+        
+        # Remove unwanted words like "illustration" to make the caption cleaner
+        unwanted_words = ["illustration", "drawing", "sketch", "picture", "dream", "imagination"]
+        for word in unwanted_words:
+            text = text.replace(word, "")
+        
+        return text.strip()  # Return the cleaned caption
+    except Exception as e:
+        st.error(f"Error processing image: {e}")  # Display error message if something goes wrong
+        return None
 
 # Function to generate a kid-friendly superhero story from the text caption
 def text2story(text):
@@ -36,21 +43,28 @@ def text2story(text):
     Returns:
         str: Generated superhero story suitable for kids aged 3-10, within 100 words.
     """
-    # Load the text generation model
-    story_generator = pipeline("text-generation", model="pranavpsv/gpt2-genre-story-generator")
-    
-    # Generate the story with the superhero genre
-    prompt = f"<BOS> <superhero> {text}"
-    story = story_generator(prompt, max_length=100, num_return_sequences=1)[0]['generated_text']
-    
-    # Remove <BOS> and <superhero> tags from the generated story
-    story = story.replace("<BOS>", "").replace("<superhero>", "").strip()
-    
-    # Remove the input text (scenario) from the generated story
-    if text in story:
-        story = story.replace(text, "").strip()
-    
-    return story
+    try:
+        # Load the text-generation model
+        story_generator = pipeline("text-generation", model="pranavpsv/gpt2-genre-story-generator")
+        
+        # Generate the story with the superhero genre
+        prompt = f"<BOS> <superhero> {text}"  # Add genre tags to the prompt
+        story = story_generator(prompt, max_length=100, num_return_sequences=1)[0]['generated_text']
+        
+        # Remove <BOS> and <superhero> tags from the generated story
+        story = story.replace("<BOS>", "").replace("<superhero>", "").strip()
+        
+        # Remove the input text (scenario) from the generated story to avoid redundancy
+        if text in story:
+            story = story.replace(text, "").strip()
+        
+        # Ensure the story is within 100 words by truncating if necessary
+        story = " ".join(story.split()[:100])
+        
+        return story
+    except Exception as e:
+        st.error(f"Error generating story: {e}")  # Display error message if something goes wrong
+        return None
 
 # Function to convert text to audio using gTTS
 def text2audio(story_text):
@@ -63,63 +77,78 @@ def text2audio(story_text):
     Returns:
         str: Path to the generated audio file.
     """
-    # Convert text to speech
-    tts = gTTS(text=story_text, lang='en')
-    audio_file = "story_audio.mp3"
-    tts.save(audio_file)
-    
-    return audio_file
+    try:
+        # Convert text to speech using gTTS
+        tts = gTTS(text=story_text, lang='en')
+        audio_file = "story_audio.mp3"  # Define the output audio file name
+        tts.save(audio_file)  # Save the audio file
+        return audio_file
+    except Exception as e:
+        st.error(f"Error generating audio: {e}")  # Display error message if something goes wrong
+        return None
+
+# Main application function
+def main():
+    """
+    Main function to run the Streamlit application.
+    """
+    # Configure the Streamlit app page
+    st.set_page_config(page_title="Picture Stories 🎨📖", page_icon="🦄")
+    st.title("Picture Stories 🎨📖")
+    st.markdown("### Turn your pictures into fun superhero stories and listen to them! 🎉")
 
-# Main application
-st.set_page_config(page_title="Picture Stories 🎨📖", page_icon="🦄")
-st.title("Picture Stories 🎨📖")
-st.markdown("### Turn your pictures into fun superhero stories and listen to them! 🎉")
+    # Instructions for kids
+    st.markdown("""
+    **How to use this app:**
+    1. **Upload a picture** of something fun, like your favorite toy, a park, or your pet.
+    2. Wait for the app to **create a superhero story** from your picture.
+    3. **Listen to the story** by clicking the "Play Audio" button.
+    4. Enjoy your fun superhero story! 🎧
+    """)
 
-# Instructions for kids
-st.markdown("""
-**How to use this app:**
-1. **Upload a picture** of something fun, like your favorite toy, a park, or your pet.
-2. Wait for the app to **create a superhero story** from your picture.
-3. **Listen to the story** by clicking the "Play Audio" button.
-4. Enjoy your fun superhero story! 🎧
-""")
+    # Upload image
+    uploaded_file = st.file_uploader("📷 **Upload your picture here!**", type=["jpg", "jpeg", "png"])
 
-# Upload image
-uploaded_file = st.file_uploader("📷 **Upload your picture here!**", type=["jpg", "jpeg", "png"])
+    if uploaded_file is not None:
+        # Save the uploaded file to disk
+        image_bytes = uploaded_file.getvalue()
+        with open(uploaded_file.name, "wb") as file:
+            file.write(image_bytes)
 
-if uploaded_file is not None:
-    # Save the uploaded file
-    bytes_data = uploaded_file.getvalue()
-    with open(uploaded_file.name, "wb") as file:
-        file.write(bytes_data)
+        # Display the uploaded image in the app
+        st.image(uploaded_file, caption="Your awesome picture!", use_container_width=True)
 
-    # Display the uploaded image
-    st.image(uploaded_file, caption="Your awesome picture!", use_container_width=True)
+        # Stage 1: Image to Text
+        with st.spinner('✨ Turning your picture into words...'):
+            scenario = img2text(uploaded_file.name)  # Generate text caption from the image
+            if scenario:
+                st.write("**What we see:**", scenario)  # Display the generated caption
 
-    # Stage 1: Image to Text
-    st.text('✨ Turning your picture into words...')
-    scenario = img2text(uploaded_file.name)
-    st.write("**What we see:**", scenario)
+                # Stage 2: Text to Story
+                with st.spinner('📖 Creating a fun superhero story for you...'):
+                    story = text2story(scenario)  # Generate a superhero story from the caption
+                    if story:
+                        st.write("**Your superhero story:**", story)  # Display the generated story
 
-    # Stage 2: Text to Story
-    st.text('📖 Creating a fun superhero story for you...')
-    story = text2story(scenario)
-    st.write("**Your superhero story:**", story)
+                        # Stage 3: Story to Audio
+                        with st.spinner('🎧 Turning your story into audio...'):
+                            # Generate audio file if it doesn't already exist in the session state
+                            if 'audio_file' not in st.session_state:
+                                st.session_state.audio_file = text2audio(story)
+                            
+                            # Play button for the generated audio
+                            if st.button("🎵 **Play Audio**"):
+                                if os.path.exists(st.session_state.audio_file):
+                                    st.audio(st.session_state.audio_file, format="audio/mp3")  # Play the audio
+                                else:
+                                    st.error("Audio file not found. Please try again.")  # Display error if audio file is missing
 
-    # Stage 3: Story to Audio
-    st.text('🎧 Turning your story into audio...')
-    
-    # Use session state to avoid regenerating audio on button click
-    if 'audio_file' not in st.session_state:
-        st.session_state.audio_file = text2audio(story)
-    
-    # Play button for the generated audio
-    if st.button("🎵 **Play Audio**"):
-        if os.path.exists(st.session_state.audio_file):
-            st.audio(st.session_state.audio_file, format="audio/mp3")
-        else:
-            st.error("Audio file not found. Please try again.")
+        # Clean up temporary files (uploaded image and generated audio)
+        if os.path.exists(uploaded_file.name):
+            os.remove(uploaded_file.name)  # Delete the uploaded image file
+        if 'audio_file' in st.session_state and os.path.exists(st.session_state.audio_file):
+            os.remove(st.session_state.audio_file)  # Delete the generated audio file
 
-    # Clean up the generated audio file and uploaded image
-    if os.path.exists(uploaded_file.name):
-        os.remove(uploaded_file.name)
+# Run the application
+if __name__ == "__main__":
+    main()