streamlit：算法工程師快速編寫demo的利器

本文首發於：行者AI

在工做當中，算法工程師常常須要快速編寫一些演示demo，例如快速演示一些算法，或者須要編寫數據標註的工具等。常見的實現方式是算法工程師試用flask等框架編寫api，再由前端工程師編寫相關的網頁調用api。這樣毫無疑問是很是耗時，效率低下的。

然而，使用streamlit框架就能夠快速搭建demo頁面，算法工程師只使用python語言就能夠編寫比較精美的demo頁面。streamlit框架內部負責處理網頁的渲染工做，算法工程師將重點放在算法的流程方面就好。

框架安裝

streamlit框架的安裝很是簡單，使用pip就能夠安裝：

pip install streamlit

安裝完成以後，可使用命令進行版本驗證：

streamlit --version

在這篇文章當中，我會用一個實際工做中的例子簡單介紹streamlit框架的使用流程。

項目準備

Collaborative-Distillation是一個支持高分辨率的圖像風格化方案，該模型的輸入是風格圖片以及待處理的圖片，輸出是風格化以後的圖片。

在這個代碼倉庫當中，咱們須要使用比較複雜的命令行命令來進行風格化操做：

# use original VGG-19, normal images
CUDA_VISIBLE_DEVICES=0 python WCT.py --debug --mode original

# use original VGG-19, ultra-res images
CUDA_VISIBLE_DEVICES=0 python WCT.py --debug --mode original --UHD

# use our pruned VGG-19, normal images
CUDA_VISIBLE_DEVICES=0 python WCT.py --debug --mode 16x

# use our pruned VGG-19, ultra-res images
CUDA_VISIBLE_DEVICES=0 python WCT.py --debug --mode 16x --UHD

# If your RAM cannot afford some large images, you can change the content and style size via '--content_size' and '--style_size'
CUDA_VISIBLE_DEVICES=0 python WCT.py --debug --mode 16x --UHD --content_size 3000 --style_size 2000

可是這樣的操做對於用戶來講至關複雜，因此咱們可使用streamlit編寫一個demo頁面，方便用戶使用。

在這個demo頁面當中，會用到streamlit的如下幾種組件：

• streamlit.file_uploader：文件上傳組件，具體見下圖

  該組件支持拖拽上傳文件和文件管理器選擇文件，相對來講比較方便，使用方法以下列代碼所示：

  style_file = st.file_uploader("請上傳風格化圖片")
  
  if style_file:
      stringio = style_file.getvalue()
      style_file_path = 'style_file/'+ style_file.name
      with open(style_file_path,'wb') as f:
          f.write(stringio)

使用文件上傳組件上傳文件以後，可使用上面的代碼將文件保存到特定路徑等待使用。

• streamlit.image：圖片顯示組件，具體見下圖：

  該組件能夠在demo頁面中根據圖片路徑顯示圖片。

  style_file_path = 'style_file/'+ style_file.name
  st.image(style_file_path)

• streamlit.write：文字顯示組件，該組件能夠在網頁上顯示一些提示信息。

  st.write('高分辨率風格化demo')

• streamlit.button：按鈕組件，點擊以後能夠進行一些任務。

if st.button('開始進行風格化處理'):
    style_func()

• streamlit.progress：進度顯示組件，能夠用來顯示任務的進度。

for i in range(0,100,10):
    st.progress(i + 1)

streamlit中還有一些重要的組件，例如：

• streamlit.cache：數據緩存組件，該組件能夠做爲裝飾器使用，用處是緩存數據，加快數據載入速度。能夠用在須要反覆加載數據或者進行計算的函數當中。

  @st.cache
  def load_dataset(data_link):
      dataset = pd.read_csv(data_link)
      return dataset

• streamlit.audio：音頻展現組件，能夠根據音頻地址播放音頻。

with open('audio.mp3','rb') as f:
    st.audio(f,format="audio/mp3")

• streamlit.audio：選擇組件，該組件可讓用戶從多個選項中選擇一項。

model_choose = st.radio('請選擇分離模型：',['人聲+伴奏','人聲+鋼琴+吉他+鼓+其餘'],0)

其中參數0表示默認選擇第一項。

streamlit支持的組件仍是不少的，若是感興趣，請參考官方文檔。

項目編寫

這個demo頁面的主要功能是讓用戶分別上傳style圖片和content圖片，而後後臺進行風格化操做，風格化操做完成以後顯示結果圖片。這樣用戶就能夠快速的進行風格化操做並知道結果。

streamlit應用是用python語言編寫的。在python文件開頭，須要導入streamlit包。

import streamlit as st

接着進行文件的上傳與預處理：

style_file = st.file_uploader("請上傳風格化圖片")
content_file = st.file_uploader("請上傳待處理圖片")
image_slot = st.empty()

if style_file:
    stringio = style_file.getvalue()
    style_file_path = 'style_file/'+ style_file.name
    with open(style_file_path,'wb') as f:
        f.write(stringio)
    image_slot.image(style_file_path)

if content_file:
    stringio = content_file.getvalue()
    content_file_path = 'content_file/'+ content_file.name
    with open(content_file_path,'wb') as f:
        f.write(stringio)

if content_file and style_file:
    img1 = Image.open( style_file_path)
    img1 = img1.resize((640, 640))

    img2 = Image.open( content_file_path)
    img2 = img2.resize((640, 640))
    new_img = Image.new('RGB', (1280, 640), 255)
    new_img.paste(img1,(0,0))
    new_img.paste(img2,(640,0))
    new_img.save('concrate_file/' + os.path.basename(style_file_path))
    image_slot.image('concrate_file/' + os.path.basename(style_file_path))

最後寫一個按鈕，執行風格化操做，並顯示最終結果，同時添加一個進度條：

if st.button('開始進行風格化處理'):
    my_bar = st.progress(10)

    UHD_content_folder_path = 'PytorchWCT/content/UHD_content'
    output_path = WCT_func.process(content_file_path,style_file_path)
    for i in range(0,100,10):
        my_bar.progress(i + 1)
    my_bar.progress(100)
    st.write('風格化以後的圖片')
    st.image(output_path)

項目的運行和部署

streamlit框架的運行方式很是簡單，直接在命令行執行：

$ streamlit run streamlit_demo.py

就能夠在瀏覽器中進行訪問了。

總結

streamlit框架很是適合快速編寫流程不太複雜且須要可視化操做的demo，做者從開始編寫到編寫完成這個demo用時不到半個小時，編寫代碼不到50行，並且運行部署起來很是方便，頁面看起來要比使用flask之類的框架渲染出的網頁美觀許多，實乃算法工程師的利器。