Streamlit表单生成图片后，如何实现图片展示及下载按钮功能？

我有一个Streamlit表单应用，允许用户上传图片并据此生成多张图片，生成的图片保存到tmp目录，希望在页面展示图片并提供下载按钮。目前遇到两个问题：

• 在form内的generate方法后调用show_image()，会抛出StreamlitAPIException，提示st.download_button()不能在st.form()中使用；
• 在main()之后调用show_image()，该函数会在表单加载时立即执行，早于图片生成。

请问如何在generate方法执行完成后，实现图片展示及下载功能？

错误信息 #

StreamlitAPIException: st.download_button() can't be used in an st.form().

For more information, refer to the documentation for forms.

Traceback:
File "/Users/anup.rawka/Documents/GitHub/vdil-ai-framework/creative_automation/app/scripts/creative_app.py", line 65, in <module>
    main()
File "/Users/anup.rawka/Documents/GitHub/vdil-ai-framework/creative_automation/app/scripts/creative_app.py", line 62, in main
    show_image()
File "/Users/anup.rawka/Documents/GitHub/vdil-ai-framework/creative_automation/app/scripts/creative_app.py", line 35, in show_image
    st.download_button(

相关代码 #

def show_image():
    for i in range(1, 7):
        img = Image.open(r"{}{}.png".format(base_path, i))
        buf = BytesIO()
        img.save(buf, format="PNG")
        byte_im = buf.getvalue()
        st.image(img)
        st.download_button(
            label="Download image", data=byte_im,
            file_name="{}.png".format(i), mime="image/png")


def main():
    st.image('../images/1st_screen_masthead_3up_1824_412.png', caption='3up 1st Screen Masthead Template 1824x412')
    # side_bar, work_area = st.columns([0.2, 0.8])
    setup_sidebar()
    with st.form("US STVP 2Program 1stScreenMasthead 1824x412"):
        img_tpfm_shw_1 = st.file_uploader("Choose show 1 image")
        img_tpfm_shw_1_logo = st.file_uploader("Choose show 1 logo")
        img_tpfm_shw_2 = st.file_uploader("Choose show 2 image")
        img_tpfm_shw_2_logo = st.file_uploader("Choose show 2 logo")
        img_tpfm_shw_3 = st.file_uploader("Choose show 3 image")
        img_tpfm_shw_3_logo = st.file_uploader("Choose show 3 logo")
        submit_button = st.form_submit_button('Generate Creative')
        if submit_button:
            img_shw1 = iio.imread(img_tpfm_shw_1)
            img_shw2 = iio.imread(img_tpfm_shw_2)
            img_shw1_logo = iio.imread(img_tpfm_shw_1_logo)
            img_shw2_logo = iio.imread(img_tpfm_shw_2_logo)
            img_shw3 = iio.imread(img_tpfm_shw_3)
            img_shw3_logo = iio.imread(img_tpfm_shw_3_logo)
            generate(img_shw1, img_shw1_logo, img_shw2,
                     img_shw2_logo, img_shw3,
                     img_shw3_logo)


main()
show_image()

核心思路是利用Streamlit的st.session_state状态管理，标记图片是否已生成，只有在生成完成后才调用show_image()，同时避免在表单内部调用下载按钮。

修改步骤：

1. 在main()函数中初始化状态变量，标记图片未生成；
2. 当generate()执行完成后，将状态变量设为已生成；
3. 在全局代码中判断状态，仅当图片生成后才调用show_image()。

修改后的完整代码：

def show_image():
    for i in range(1, 7):
        img = Image.open(r"{}{}.png".format(base_path, i))
        buf = BytesIO()
        img.save(buf, format="PNG")
        byte_im = buf.getvalue()
        st.image(img)
        st.download_button(
            label="Download image", data=byte_im,
            file_name="{}.png".format(i), mime="image/png")


def main():
    st.image('../images/1st_screen_masthead_3up_1824_412.png', caption='3up 1st Screen Masthead Template 1824x412')
    setup_sidebar()
    
    # 初始化状态，标记图片未生成
    if 'images_generated' not in st.session_state:
        st.session_state['images_generated'] = False
        
    with st.form("US STVP 2Program 1stScreenMasthead 1824x412"):
        img_tpfm_shw_1 = st.file_uploader("Choose show 1 image")
        img_tpfm_shw_1_logo = st.file_uploader("Choose show 1 logo")
        img_tpfm_shw_2 = st.file_uploader("Choose show 2 image")
        img_tpfm_shw_2_logo = st.file_uploader("Choose show 2 logo")
        img_tpfm_shw_3 = st.file_uploader("Choose show 3 image")
        img_tpfm_shw_3_logo = st.file_uploader("Choose show 3 logo")
        submit_button = st.form_submit_button('Generate Creative')
        
        if submit_button:
            img_shw1 = iio.imread(img_tpfm_shw_1)
            img_shw2 = iio.imread(img_tpfm_shw_2)
            img_shw1_logo = iio.imread(img_tpfm_shw_1_logo)
            img_shw2_logo = iio.imread(img_tpfm_shw_2_logo)
            img_shw3 = iio.imread(img_tpfm_shw_3)
            img_shw3_logo = iio.imread(img_tpfm_shw_3_logo)
            generate(img_shw1, img_shw1_logo, img_shw2,
                     img_shw2_logo, img_shw3,
                     img_shw3_logo)
            # 标记图片已生成
            st.session_state['images_generated'] = True


main()

# 仅当图片生成后才展示图片和下载按钮
if st.session_state.get('images_generated', False):
    show_image()

关键说明 #

• Streamlit每次交互都会重跑整个脚本，st.session_state可以跨重跑保存状态，确保只有生成操作完成后才触发展示逻辑；
• st.download_button不支持在表单内部使用，将show_image()放在表单外部执行，完美规避这个限制。

内容的提问来源于stack exchange，提问作者Explorer