Streamlit 教程02 Iris Flower类型预测

338 阅读1分钟

本例展示如何将sklearn跟Streamlit相结合,使用sklearn中默认提供的Iris数据集训练一个分类模型,然后预测用户提供的新数据为哪类花

注:此处只是演示streamlit可视化,真实的机器学习模型不该如此简单粗暴的训练

最终效果

截屏2022-11-18 下午5.08.06.png

代码

requirements.txt中

streamlit
scikit-learn
pandas

import pandas as pd
import streamlit as st
from sklearn import datasets
from sklearn.ensemble import RandomForestClassifier


st.write("""
# Simple Iris Flower Prediction App
This app predicts the **Iris flower** type!
"""
)

st.sidebar.header('User Input Parameters')

def user_input_features():
    sepal_length = st.sidebar.slider('sepal_length', 4.3, 7.9, 5.4)
    sepal_width = st.sidebar.slider('sepal_widht', 2.0, 4.4, 3.4)
    petal_length = st.sidebar.slider('petal_length', 1.0, 6.9, 1.3)
    petal_width = st.sidebar.slider('petal_width', 0.1, 2.5, 0.2)
    
    data = {
        'sepal_length': sepal_length,
        'sepal_width': sepal_width,
        'petal_length': petal_length,
        'petal_width': petal_width,
    }
    
    return pd.DataFrame(data, index=[0])
    

df = user_input_features()

st.subheader('User Input parameters')
st.write(df)

# load Iris Dataset
iris = datasets.load_iris()
X = iris.data
y = iris.target

clf = RandomForestClassifier()
clf.fit(X, y)

pred = clf.predict(df)
pred_prob = clf.predict_proba(df)

st.subheader('class labels and their corresponding index number')
st.write(iris.target_names)

st.subheader('Prediction')
st.write(iris.target_names[pred])

st.subheader('Prediction Probability')
st.write(pred_prob)