こしあん
2023-10-13

Streamlitで動的に作ったコンテンツをダウンロードする方法


4k{icon} {views}


Streamlitでメソッドなどで作った動的なデータについて、ボタンをクリックするとダウンロードするような処理を実装します。これは2023年10月現在のStreamlitでは不可能ですが、ダウンロードしたいコンテンツをbase64でエンコードすることでこれを解決します。

はじめに

※この記事は2023年10月時点のもので、今後のStreamlitのアップデートではもっと便利にできる可能性があります

現在Streamlitではダウンロードリンクのボタンのコンポーネントが用意されています。しかし、以下のような使い方で、メソッドで作ったデータのように動的なコンテンツをダウンロードするということができません

以下は公式コードからですが、

import streamlit as st

@st.cache
def convert_df(df):
    # IMPORTANT: Cache the conversion to prevent computation on every rerun
    return df.to_csv().encode('utf-8')

csv = convert_df(my_large_df)

st.download_button(
    label="Download data as CSV",
    data=csv,
    file_name='large_df.csv',
    mime='text/csv',
)

ダウンロードしたいコンテンツを1回変数に落としてから実行する必要があり、ボタンを押す→コンテンツを生成する→ダウンロードするというプロセスが現状できません。

これに対する現状の解決策として、Base64でエンコードしたコンテンツを含んだリンクタグをインジェクションさせてダウンロードするというのを試します。

Streamlitのロードマップ

これには改善の方針がロードマップで示されており、現在は「Future」の中になります。今後はdataの部分をfunctionに対応される予定ですが、いつになるかわかりません。

これに関してはこちらのissueで議論されています。

Base64でエンコードしてHTMLタグ化

例えば、現在の日時とテキトーなメッセージを含むJSONをダウンロードさせるリンクを生成したいときは以下のようなコードになります。

import streamlit as st
import json
import base64
import datetime
from streamlit.components.v1 import html

def create_json_link():
    data = {
        "current": str(datetime.datetime.now()),
        "message": "Hello, world"
    }
    json_str = json.dumps(data, ensure_ascii=False, indent=4, separators=(',', ': '))
    json_str = base64.b64encode(json_str.encode()).decode()
    js_code = f"""<a href="data:text/csv;base64,{json_str}" download="sample.json">download</a>"""
    return js_code

def main():
    # Sidebarの場合
    # download_button = st.sidebar.button("Click to download")
    # container = st.sidebar.empty()

    download_button = st.button("Click to download")
    container = st.empty()
    if download_button:
        with container:
            html(create_json_link(), height=50)

if __name__ == "__main__":
    main()

これを「click_link.py」として保存し、Streamlitのアプリとして実行します。

streamlit run click_link.py

リンクをクリックすると以下のようなJSONがダウンロードされます。ボタンとリンクでクリックが二重になってしまうのがUI的にアレですが、とりあえず問題は解決できます。

{
    "current": "2023-10-13 14:41:26.077262",
    "message": "Hello, world"
}

画像をダウンロードする

Base64としてエンコードしてHTMLに落とし込めればなんでもいいので、画像もいけます。

import streamlit as st
import base64
from streamlit.components.v1 import html
import numpy as np
from PIL import Image
import io

def create_image_link():
    img_array = np.zeros((256, 256, 3), np.uint8)
    img_array[:, :, 0] = 255
    img_array[:, :, 1] = np.arange(256, dtype=np.uint8)[None, :]
    img_array[:, :, 2] = np.arange(256, dtype=np.uint8)[:, None]
    with io.BytesIO() as buf:
        with Image.fromarray(img_array) as img:
            img.save(buf, format="png")
        image_str = base64.b64encode(buf.getvalue()).decode()
        js_code = f"""<a href="data:png;base64,{image_str}" download="sample.png">download</a>"""
    return js_code

def main():
    download_button = st.button("Click to download")
    container = st.empty()
    if download_button:
        with container:
            html(create_image_link(), height=50)

if __name__ == "__main__":
    main()

ダウンロードすると次のような画像がでてきます。

base64に変換してダウンロードできるものならなんでもいけるのは便利ですね。最初「一時ストレージのようなバックエンドを立てないといけないのかな」と思ったのですが、これなしでいけるのは強いですね。



Shikoan's ML Blogの中の人が運営しているサークル「じゅ~しぃ~すくりぷと」の本のご案内

技術書コーナー

北海道の駅巡りコーナー


Add a Comment

メールアドレスが公開されることはありません。 が付いている欄は必須項目です