Create app.py
Browse files
app.py
ADDED
|
@@ -0,0 +1,101 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import streamlit as st
|
| 2 |
+
from PIL import Image
|
| 3 |
+
from transformers import pipeline
|
| 4 |
+
|
| 5 |
+
# ----------------------------
|
| 6 |
+
# 生成图像描述函数
|
| 7 |
+
# ----------------------------
|
| 8 |
+
def generate_caption(image_file):
|
| 9 |
+
"""
|
| 10 |
+
使用 Hugging Face pipeline 的 image-to-text 模型生成图片描述
|
| 11 |
+
参数:
|
| 12 |
+
image_file: 上传的图片文件(文件对象或文件路径)
|
| 13 |
+
返回:
|
| 14 |
+
caption: 生成的图片描述文本
|
| 15 |
+
"""
|
| 16 |
+
# 打开图片(如果上传的是文件流,可以直接传给 pipeline)
|
| 17 |
+
image = Image.open(image_file)
|
| 18 |
+
# 利用 image-to-text pipeline 加载 Salesforce/blip-image-captioning-base 模型
|
| 19 |
+
caption_generator = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")
|
| 20 |
+
# 直接将图片传入 pipeline,返回结果是一个列表,每个元素是一个字典
|
| 21 |
+
caption_results = caption_generator(image)
|
| 22 |
+
caption = caption_results[0]['generated_text'] # 取第一个结果
|
| 23 |
+
return caption
|
| 24 |
+
|
| 25 |
+
# ----------------------------
|
| 26 |
+
# 基于图片描述生成完整故事的函数
|
| 27 |
+
# ----------------------------
|
| 28 |
+
def generate_story(caption):
|
| 29 |
+
"""
|
| 30 |
+
基于图片描述生成完整故事,确保生成的故事至少包含100个单词。
|
| 31 |
+
参数:
|
| 32 |
+
caption: 图片描述文本
|
| 33 |
+
返回:
|
| 34 |
+
story: 生成的故事文本
|
| 35 |
+
"""
|
| 36 |
+
# 使用 text-generation pipeline 加载 GPT-2 模型
|
| 37 |
+
story_generator = pipeline("text-generation", model="gpt2")
|
| 38 |
+
# 构建生成故事的提示语
|
| 39 |
+
prompt = f"Based on the following image caption: '{caption}', generate a complete fairy tale story for children with at least 100 words. "
|
| 40 |
+
|
| 41 |
+
# 生成故事文本
|
| 42 |
+
result = story_generator(prompt, max_length=300, num_return_sequences=1)
|
| 43 |
+
story = result[0]['generated_text']
|
| 44 |
+
|
| 45 |
+
# 简单检查生成的故事单词数是否达到100,否则再生成部分文本补充
|
| 46 |
+
if len(story.split()) < 100:
|
| 47 |
+
additional = story_generator(prompt, max_length=350, num_return_sequences=1)[0]['generated_text']
|
| 48 |
+
story += " " + additional
|
| 49 |
+
return story
|
| 50 |
+
|
| 51 |
+
# ----------------------------
|
| 52 |
+
# 文字转语音 (TTS) 函数
|
| 53 |
+
# ----------------------------
|
| 54 |
+
def text_to_speech(text, output_file="output.mp3"):
|
| 55 |
+
"""
|
| 56 |
+
将文本转换为语音并保存为 mp3 文件
|
| 57 |
+
参数:
|
| 58 |
+
text: 要转换的文本
|
| 59 |
+
output_file: 保存的音频文件名
|
| 60 |
+
返回:
|
| 61 |
+
output_file: 转换后的音频文件路径
|
| 62 |
+
"""
|
| 63 |
+
from gtts import gTTS
|
| 64 |
+
# 这里语言参数设为英语 "en",
|
| 65 |
+
# 如需中文可修改 lang="zh-cn",但对应文本生成模型也需生成中文
|
| 66 |
+
tts = gTTS(text=text, lang="en")
|
| 67 |
+
tts.save(output_file)
|
| 68 |
+
return output_file
|
| 69 |
+
|
| 70 |
+
# ----------------------------
|
| 71 |
+
# 主函数:构建 Streamlit 界面
|
| 72 |
+
# ----------------------------
|
| 73 |
+
def main():
|
| 74 |
+
st.title("儿童故事生成应用")
|
| 75 |
+
st.write("上传一张图片,我们将根据图片生成有趣的故事,并转换成语音播放!")
|
| 76 |
+
|
| 77 |
+
uploaded_file = st.file_uploader("选择一张图片", type=["png", "jpg", "jpeg"])
|
| 78 |
+
|
| 79 |
+
if uploaded_file is not None:
|
| 80 |
+
# 显示上传的图片
|
| 81 |
+
image = Image.open(uploaded_file)
|
| 82 |
+
st.image(image, caption="上传的图片", use_column_width=True)
|
| 83 |
+
|
| 84 |
+
# 生成图片描述
|
| 85 |
+
with st.spinner("正在生成图片描述..."):
|
| 86 |
+
caption = generate_caption(uploaded_file)
|
| 87 |
+
st.write("图片描述:", caption)
|
| 88 |
+
|
| 89 |
+
# 根据图片描述生成完整故事
|
| 90 |
+
with st.spinner("正在生成故事..."):
|
| 91 |
+
story = generate_story(caption)
|
| 92 |
+
st.write("生成的故事:")
|
| 93 |
+
st.write(story)
|
| 94 |
+
|
| 95 |
+
# 文本转语音
|
| 96 |
+
with st.spinner("正在转换成语音..."):
|
| 97 |
+
audio_file = text_to_speech(story)
|
| 98 |
+
st.audio(audio_file, format="audio/mp3")
|
| 99 |
+
|
| 100 |
+
if __name__ == "__main__":
|
| 101 |
+
main()
|