如何通过st.download_button的on_click回调按需预处理下载数据？

我运行的一款应用中，模型输出结果为np.ndarray，目前通过st.image(result_matrix)展示该结果。我希望添加下载功能，让用户能下载此图片，但需要将其转换为PIL.Image并传入buffer.getvalue()给下载按钮。由于用户下载频率很低，为节省算力和负载，我不想每次都将结果转换为PIL.Image。

是否支持按需处理数据后再下载的功能？

我尝试了以下代码，但出现了不接受numpy数组的明显错误：

import streamlit as st
from PIL import Image
import numpy as np
from io import BytesIO

st.session_state['result'] = some_numpy_RGB_array

def process_image():
    img = Image.fromarray(st.session_state['result'])
    buffer = BytesIO()
    img.save(buffer, format="jpeg")
    st.session_state['result'] = buffer.getvalue()

_ = st.download_button(label="Download",data=st.session_state['result'],file_name="image.jpeg",mime="image/jpeg",on_click=process_image)

你当前代码的问题是download_button初始化时就会读取st.session_state['result']里的numpy数组，而不是等点击后再处理。要实现按需转换，调整逻辑即可：

• 把原始numpy数组单独存在session_state里，别覆盖它
• 点击下载时，在回调函数里生成二进制下载数据，存入另一个session_state键
• 让download_button的data参数指向这个仅在点击后才生成的二进制数据

修改后的代码如下：

import streamlit as st
from PIL import Image
import numpy as np
from io import BytesIO

# 模拟模型输出的RGB numpy数组（替换成你的实际数据）
if 'result_np' not in st.session_state:
    st.session_state['result_np'] = np.random.randint(0, 255, (200, 200, 3), dtype=np.uint8)

# 展示图片
st.image(st.session_state['result_np'], caption="模型输出结果")

def process_for_download():
    # 基于原始numpy数组生成下载用的二进制数据
    img = Image.fromarray(st.session_state['result_np'])
    buffer = BytesIO()
    img.save(buffer, format="JPEG")
    # 将二进制数据存入单独的session_state键，不覆盖原始数组
    st.session_state['download_data'] = buffer.getvalue()

# 下载按钮：初始化时data设为None，点击后通过回调生成数据再触发下载
st.download_button(
    label="下载图片",
    data=st.session_state.get('download_data'),
    file_name="result.jpeg",
    mime="image/jpeg",
    on_click=process_for_download
)

关键说明 #

• 原始numpy数组始终存在result_np中，不会被修改，避免无意义的重复转换
• 只有用户点击下载按钮时，才会执行process_for_download生成二进制数据
• Streamlit会在回调执行后重新渲染页面，此时download_data已有值，下载按钮就能正常输出文件

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