ASMT_1 / app.py
dustinfeng's picture
Update app.py
1715179 verified
import streamlit as st
from transformers import pipeline
from gtts import gTTS
import os
def img2text(url):
image_to_text_model = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")
text = image_to_text_model(url)[0]["generated_text"]
return text
# text2story
def text2story(text):
# 添加 trust_remote_code=True 参数
story_generator = pipeline("text-generation", model="facebook/opt-125m", trust_remote_code=True)
# 增大 max_length 并设置 min_length
story = story_generator(text, num_return_sequences=1, max_length=300, min_length=150,
temperature=0.7, top_k=50, top_p=0.9, no_repeat_ngram_size=2)[0]["generated_text"]
# 截取前 100 词左右的内容,如果想保留完整生成内容可注释掉下面两行
# words = story.split()
# story = " ".join(words[:100])
return story
# text2audio using gTTS
def text2audio(story_text):
# 创建 gTTS 对象
tts = gTTS(text=story_text, lang='en')
# 保存音频文件
audio_file_path = "story_audio.mp3"
tts.save(audio_file_path)
return audio_file_path
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:
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_file_path = text2audio(story)
# Play button
if st.button("Play Audio"):
audio_file = open(audio_file_path, "rb")
audio_bytes = audio_file.read()
st.audio(audio_bytes, format="audio/mp3")
audio_file.close()
# 删除临时音频文件
os.remove(audio_file_path)