Update app.py

app.py

@@ -1,115 +1,119 @@
 import streamlit as st
+from transformers import pipeline
 from PIL import Image
-import os
-import tempfile
-import sys
+import io
+from gtts import gTTS
+import time
 
-# function part
-# img2text with a model that doesn't require sentencepiece
-def img2text(image_path):
+# Set page title
+st.set_page_config(page_title="Kids Story Generator")
+
+# Title and introduction
+st.title("Kids Story Generator")
+st.write("Upload a picture and let's create a magical story!")
+
+# Initialize models
+@st.cache_resource
+def load_models():
+    image_to_text = pipeline("image-to-text", model="microsoft/git-base-coco")
+    story_generator = pipeline("text-generation", model="gpt2")
+    return image_to_text, story_generator
+
+image_to_text, story_generator = load_models()
+
+# Function to generate caption from image
+def generate_caption(image):
     try:
-        from transformers import pipeline
-        
-        # Use the Salesforce model instead of Donut to avoid sentencepiece issues
-        st.info("Using Salesforce/blip-image-captioning-base model for image-to-text")
-        image_to_text_model = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")
-        
-        # Open the image file
-        image = Image.open(image_path)
-        
-        # Extract text from the image
-        result = image_to_text_model(image)
-        
-        # Get the generated text
-        text = result[0]["generated_text"] if result else "No text detected"
-        return text
+        caption = image_to_text(image)[0]['generated_text']
+        return caption
     except Exception as e:
-        st.error(f"Error processing image: {str(e)}")
-        return f"Error: {str(e)}"
-
-# text2story
-def text2story(text):
-    # For now, just return the extracted text as the story
-    story_text = f"Here's a story based on the text: {text}"
-    return story_text
+        st.error(f"Error generating caption: {str(e)}")
+        return "children playing in a colorful park"
 
-# text2audio using Google Text-to-Speech
-def text2audio(story_text):
+# Function to generate story from caption (less than 100 words)
+def generate_story(caption):
     try:
-        from gtts import gTTS
+        prompt = f"Once upon a time, {caption} "
         
-        # Create a temporary file
-        temp_audio = tempfile.NamedTemporaryFile(delete=False, suffix='.mp3')
-        temp_audio_path = temp_audio.name
-        temp_audio.close()
+        # Set parameters for faster generation while keeping quality
+        story = story_generator(
+            prompt, 
+            max_length=100, 
+            do_sample=True,
+            temperature=0.7,  # Lower temperature for faster generation
+            top_p=0.9,
+            num_return_sequences=1
+        )[0]['generated_text']
         
-        # Initialize gTTS and generate audio
-        tts = gTTS(text=story_text, lang='en', slow=False)
+        # Ensure story doesn't exceed 100 words
+        words = story.split()
+        if len(words) > 100:
+            words = words[:100]
+            story = " ".join(words)
+            # Add period to the end if needed
+            if not story.endswith(('.', '!', '?')):
+                story += '.'
         
-        # Save to the temporary file
-        tts.save(temp_audio_path)
-        
-        return temp_audio_path
+        return story
+    except Exception as e:
+        st.error(f"Error generating story: {str(e)}")
+        return f"Once upon a time, {caption}. Something magical happened and everyone lived happily ever after."
+
+# Function to convert text to speech
+def text_to_speech(text):
+    try:
+        tts = gTTS(text=text, lang='en', slow=False)
+        audio_file = "story_audio.mp3"
+        tts.save(audio_file)
+        return audio_file
     except Exception as e:
         st.error(f"Error generating audio: {str(e)}")
         return None
 
-# main part
-st.set_page_config(page_title="Your Image to Audio Story",
-                   page_icon="🦜")
-st.header("Turn Your Image to Audio Story")
-st.subheader("Image to Text to Audio Conversion")
-
-uploaded_file = st.file_uploader("Select an Image...", type=['png', 'jpg', 'jpeg', 'gif', 'bmp', 'webp'])
+# File uploader
+uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
 
 if uploaded_file is not None:
-    # Save the uploaded file temporarily
-    bytes_data = uploaded_file.getvalue()
-    image_temp_path = os.path.join(tempfile.gettempdir(), uploaded_file.name)
-    with open(image_temp_path, "wb") as file:
-        file.write(bytes_data)
-
-    # Display the uploaded image
-    st.image(uploaded_file, caption="Uploaded Image",
-             use_column_width=True)
-
-    # Stage 1: Image to Text
-    with st.spinner('Processing img2text...'):
-        extracted_text = img2text(image_temp_path)
-        st.subheader("Extracted Text:")
-        st.write(extracted_text)
-
-    # Stage 2: Text to Story
-    with st.spinner('Generating a story...'):
-        story = text2story(extracted_text)
-        st.subheader("Generated Story:")
-        st.write(story)
-
-    # Stage 3: Story to Audio data
-    audio_file_path = None
-    with st.spinner('Generating audio data...'):
-        audio_file_path = text2audio(story)
-
-    # Remove the temporary image file
-    if os.path.exists(image_temp_path):
-        os.remove(image_temp_path)
-
-    # Play button
-    if st.button("Play Audio"):
-        if audio_file_path and os.path.exists(audio_file_path):
-            # Play the generated audio
-            with open(audio_file_path, "rb") as audio_file:
-                audio_bytes = audio_file.read()
-            st.audio(audio_bytes, format="audio/mp3")
+    try:
+        # Display the uploaded image
+        image = Image.open(uploaded_file)
+        st.image(image, caption='Uploaded Image', use_container_width=True)
+        
+        # Generate button
+        if st.button("Generate Story"):
+            # Use progress bar for better UX
+            progress_bar = st.progress(0)
+            
+            # Generate caption
+            progress_bar.progress(25)
+            st.text('Analyzing image...')
+            caption = generate_caption(image)
+            st.write("Image caption:", caption)
+            
+            # Generate story
+            progress_bar.progress(50)
+            st.text('Creating story...')
+            story = generate_story(caption)
+            word_count = len(story.split())
+            st.write(f"### Your Story ({word_count} words)")
+            st.write(story)
+            
+            # Generate audio
+            progress_bar.progress(75)
+            st.text('Generating audio...')
+            audio_file = text_to_speech(story)
             
-            # Clean up the audio file after playing
-            try:
-                os.remove(audio_file_path)
-            except:
-                pass
-        else:
-            st.warning("Audio generation failed. Playing a placeholder audio.")
-            try:
-                st.audio("kids_playing_audio.wav")
-            except FileNotFoundError:
-                st.error("Placeholder audio file not found. Audio playback is unavailable.")
+            # Display audio
+            if audio_file:
+                progress_bar.progress(100)
+                st.write("### Listen to your story")
+                st.audio(audio_file)
+            
+            # Clear progress when done
+            progress_bar.empty()
+    
+    except Exception as e:
+        st.error(f"An error occurred: {str(e)}")
+
+st.markdown("---")
+st.write("Created for ISOM5240 Assignment")