LCNada commited on
Commit
e8bbfee
·
verified ·
1 Parent(s): f65c21b

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +101 -0
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()