Update app.py

app.py

@@ -1,30 +1,175 @@
 import streamlit as st
-from transformers import pipeline
+from PIL import Image
+import io
+import torch
+from transformers import BlipProcessor, BlipForConditionalGeneration, pipeline
+from gtts import gTTS
+import tempfile
+import os
+import base64
 
-# Load the text classification model pipeline
-classifier = pipeline("text-classification",model='isom5240ust/bert-base-uncased-emotion', return_all_scores=True)
+# Set page config
+st.set_page_config(
+    page_title="StoryTime: Kids' Storyteller",
+    page_icon="📚",
+    layout="centered"
+)
 
-# Streamlit application title
-st.title("Text Classification for you")
-st.write("Classification for 6 emotions: sadness, joy, love, anger, fear, surprise")
+# Load and apply CSS
+def load_css(file_name):
+    with open(file_name) as f:
+        st.markdown(f'<style>{f.read()}</style>', unsafe_allow_html=True)
 
-# Text input for user to enter the text to classify
-text = st.text_area("Enter the text to classify", "")
+try:
+    load_css("style.css")
+except:
+    st.warning("Style file not found. Using default styling.")
 
-# Perform text classification when the user clicks the "Classify" button
-if st.button("Classify"):
-    # Perform text classification on the input text
-    results = classifier(text)[0]
+# Function to load image captioning model
+@st.cache_resource
+def load_captioning_model():
+    processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
+    model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")
+    return processor, model
 
-    # Display the classification result
-    max_score = float('-inf')
-    max_label = ''
+# Function to load story generation model
+@st.cache_resource
+def load_story_generator():
+    return pipeline('text-generation', model='gpt2')
 
-    for result in results:
-        if result['score'] > max_score:
-            max_score = result['score']
-            max_label = result['label']
+# Function to generate caption from image
+def generate_caption(image, processor, model):
+    inputs = processor(image, return_tensors="pt")
+    out = model.generate(**inputs, max_length=30)
+    caption = processor.decode(out[0], skip_special_tokens=True)
+    return caption
 
-    st.write("Text:", text)
-    st.write("Label:", max_label)
-    st.write("Score:", max_score)
+# Function to generate story from caption
+def generate_story(caption, generator):
+    prompt = f"{caption} "
+    story = generator(
+        prompt, 
+        max_length=150, 
+        num_return_sequences=1,
+        temperature=0.8,
+        top_k=50
+    )[0]['generated_text']
+    
+    # Clean up the story
+    story = story.replace('\n', ' ')
+    sentences = story.split('.')
+    if len(sentences) > 5:
+        story = '.'.join(sentences[:5]) + '.'
+    
+    # Strictly control word count between 50-100 words
+    words = story.split()
+    word_count = len(words)
+    
+    if word_count < 50:
+        # If story is too short, generate more content
+        additional_content = generator(
+            story + " Then, ", 
+            max_length=100, 
+            num_return_sequences=1,
+            temperature=0.8,
+            top_k=50
+        )[0]['generated_text']
+        
+        # Add only what's needed to reach 50 words
+        additional_words = additional_content.split()[word_count:]
+        words_needed = 50 - word_count
+        story = ' '.join(words + additional_words[:words_needed])
+        
+    if word_count > 100:
+        # If story is too long, truncate to exactly 100 words
+        story = ' '.join(words[:100])
+    
+    # Ensure the story ends with a period
+    if not story.endswith('.'):
+        story += '.'
+    
+    return story
+
+# Function to convert text to speech
+def text_to_speech(text):
+    tts = gTTS(text=text, lang='en', slow=False)
+    fp = tempfile.NamedTemporaryFile(delete=False, suffix='.mp3')
+    tts.save(fp.name)
+    return fp.name
+
+# Add background decorations
+def add_background_decorations():
+    st.markdown(
+        """
+        <div style="position: fixed; top: 0; right: 0; z-index: -1; opacity: 0.3;">
+            <img src="data:image/svg+xml;base64,PHN2ZyB4bWxucz0iaHR0cDovL3d3dy53My5vcmcvMjAwMC9zdmciIHdpZHRoPSIyMDAiIGhlaWdodD0iMjAwIiB2aWV3Qm94PSIwIDAgMjAwIDIwMCI+PHBhdGggZD0iTTEwMCAxOTBjNTEuNCAwIDkwLTQwLjYgOTAtOTBTMTUxLjQgMTAgMTAwIDEwIDEwIDUwLjYgMTAgMTAwczM4LjYgOTAgOTAgOTB6IiBmaWxsPSIjNzZiNWM1IiBvcGFjaXR5PSIwLjIiLz48cGF0aCBkPSJNMTgwIDEwMGMwIDQ0LjEtMzUuOSA4MC04MCA4MHMtODAtMzUuOS04MC04MCAzNS45LTgwIDgwLTgwIDgwIDM1LjkgODAgODB6IiBmaWxsPSIjM2Q4NWM2IiBvcGFjaXR5PSIwLjIiLz48L3N2Zz4=" width="200"/>
+        </div>
+        <div style="position: fixed; bottom: 0; left: 0; z-index: -1; opacity: 0.3;">
+            <img src="data:image/svg+xml;base64,PHN2ZyB4bWxucz0iaHR0cDovL3d3dy53My5vcmcvMjAwMC9zdmciIHdpZHRoPSIyMDAiIGhlaWdodD0iMjAwIiB2aWV3Qm94PSIwIDAgMjAwIDIwMCI+PHBhdGggZD0iTTEwMCAxOTBjNTEuNCAwIDkwLTQwLjYgOTAtOTBTMTUxLjQgMTAgMTAwIDEwIDEwIDUwLjYgMTAgMTAwczM4LjYgOTAgOTAgOTB6IiBmaWxsPSIjNzZiNWM1IiBvcGFjaXR5PSIwLjIiLz48cGF0aCBkPSJNMTgwIDEwMGMwIDQ0LjEtMzUuOSA4MC04MCA4MHMtODAtMzUuOS04MC04MCAzNS45LTgwIDgwLTgwIDgwIDM1LjkgODAgODB6IiBmaWxsPSIjM2Q4NWM2IiBvcGFjaXR5PSIwLjIiLz48L3N2Zz4=" width="200"/>
+        </div>
+        """,
+        unsafe_allow_html=True
+    )
+
+# Main UI
+def main():
+    add_background_decorations()
+    
+    st.title("📚 StoryTime: Kids' Storyteller")
+    st.markdown("### Upload a picture and listen to a magical story!")
+    
+    # Load models
+    with st.spinner("Loading models... This might take a moment!"):
+        processor, caption_model = load_captioning_model()
+        story_generator = load_story_generator()
+    
+    # Create columns for better layout
+    col1, col2 = st.columns([1, 1])
+    
+    # Image upload
+    with col1:
+        uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
+    
+    if uploaded_file is not None:
+        # Display the uploaded image
+        image_bytes = uploaded_file.getvalue()
+        original_image = Image.open(io.BytesIO(image_bytes))
+        
+        # Display the image
+        with col1:
+            st.image(original_image, caption='Your Magical Picture', use_column_width=True)
+        
+        # Generate caption and story
+        with st.spinner("Looking at your picture and thinking of a story..."):
+            caption = generate_caption(original_image, processor, caption_model)
+            story = generate_story(caption, story_generator)
+        
+        # Display the story and audio
+        with col2:
+            st.markdown("### Here's your story:")
+            st.write(story)
+            
+            # Word count display
+            word_count = len(story.split())
+            st.caption(f"Story length: {word_count} words")
+            
+            # Convert to speech and play
+            with st.spinner("Creating the storytelling voice..."):
+                audio_file = text_to_speech(story)
+                
+            st.audio(audio_file, format='audio/mp3')
+            st.success("Story created! Click the play button to listen!")
+            
+            # Add download button for audio
+            with open(audio_file, "rb") as f:
+                audio_bytes = f.read()
+            
+            audio_b64 = base64.b64encode(audio_bytes).decode()
+            href = f'<a href="data:audio/mp3;base64,{audio_b64}" download="story.mp3">Download the story audio</a>'
+            st.markdown(href, unsafe_allow_html=True)
+            
+            # Clean up the temp file
+            os.unlink(audio_file)
+
+if __name__ == "__main__":
+    main()