assignment1 / app.py
WRX020510's picture
Update app.py
37a58fb verified
#Import part
from transformers import pipeline
import streamlit as st
import torch
# Use function for the implementation
# function part
# img2text
def img2text(img):
image_to_text_model = pipeline("image-to-text",
model="Salesforce/blip-image-captioning-base")
text = image_to_text_model(img)[0]["generated_text"]
return text
# text2story
def text2story(text):
generator = pipeline("text-generation", model="distilbert/distilgpt2")
story_text = generator(
text,
min_length=150, # min_length, # of tokens at least larger than100
max_length=300,
num_return_sequences=1,
no_repeat_ngram_size=2, # prevent repetition
early_stopping=False # prohibit early stopping
)[0]["generated_text"]
return story_text
# text2audio
def text2audio(story_text):
tts_pipeline = pipeline("text-to-speech", model="facebook/mms-tts-eng")
audio_data = tts_pipeline(story_text) # 直接返回字典
return audio_data
# tts_pipeline = pipeline("text-to-speech", model="suno/bark-small")
# audio_data = tts_pipeline(story_text)
# audio_buffer = io.BytesIO()
# wavfile.write(audio_buffer, rate=audio_data["sampling_rate"], data=audio_data["audio"])
# audio_buffer.seek(0)
# return {
# 'audio': audio_buffer.getvalue(),
# 'sampling_rate': audio_data["sampling_rate"]
# }
# processor = SpeechT5Processor.from_pretrained("microsoft/speecht5_tts")
# model = SpeechT5ForTextToSpeech.from_pretrained("microsoft/speecht5_tts")
# inputs = processor(text=story_text, return_tensors="pt")
# with torch.no_grad():
# speech = model.generate(**inputs)
# audio_data = speech.cpu().numpy().squeeze()
# audio_buffer = io.BytesIO()
# wavfile.write(audio_buffer, rate=16000, data=audio_data) # 16kHz 采样率
# audio_buffer.seek(0)
# return {'audio': audio_buffer.getvalue(), 'sampling_rate': 16000}
# program main part
st.set_page_config(page_title="Your Image to Audio Story",
page_icon="🦜")
st.header("Turn Your Image to Audio Story")
uploaded_file = st.file_uploader("Select an Image...")
if uploaded_file is not None:
print(uploaded_file)
bytes_data = uploaded_file.getvalue()
with open(uploaded_file.name, "wb") as file:
file.write(bytes_data)
st.image(uploaded_file, caption="Uploaded Image",
use_column_width=True)
#Stage 1: Image to Text
st.text('Processing img2text...')
scenario = img2text(uploaded_file.name)
st.write(scenario)
#Stage 2: Text to Story
st.text('Generating a story...')
story = text2story(scenario)
st.write(story)
#Stage 3: Story to Audio data
st.text('Generating audio data...')
audio_data =text2audio(story)
# Play button
if st.button("Play Audio"):
st.audio(audio_data['audio'],
format="audio/wav",
start_time=0,
sample_rate = audio_data['sampling_rate'])
#st.audio("kids_playing_audio.wav")