使用streamlit快速创建网站应用

一只大鸽子

发布于 2023-09-02 17:36:38

9580

发布于 2023-09-02 17:36:38

文章被收录于专栏：Python基础、进阶与实战Python基础、进阶与实战

简介

streamlit是一个Python库，可以只用Python（无需前端）创建一个网页应用。只要几行代码就可以为我们的应用创建一个界面，很适合做一些演示，比如展示数据、演示模型等。

streamlit在几分钟内将数据脚本转变为Web 应用程序。仅需Python，无需前端经验。--https://streamlit.io/

下面是用streamlit实现图片分类、目标检测的例子，效果如下：

代码

为了使用streamlit，需要用pip安装： pip install streamlit

此外还需要torch torchvision，使用pip安装： pip install torch torchvision

图片分类

import streamlit as st
import torch
import torchvision.models as models
from PIL import Image

weights = models.VGG16_Weights.DEFAULT
img_preprocess = weights.transforms()
categories = weights.meta["categories"]


@st.cache_resource
def get_model():
    model = models.vgg16(weights='IMAGENET1K_V1')
    model.eval()
    return model


model = get_model()


def make_prediction(img):
    img_processed = img_preprocess(img)
    prediction = model(img_processed.unsqueeze(0))
    index = torch.argmax(prediction, 1)
    return  categories[index]


# Dashboard
st.title("图片分类")  # 设置标题
upload = st.file_uploader(label="上传图片", type=["png", "jpg", "jpeg"])  # 上传文件
if upload:
    img = Image.open(upload)
    st.image(img)
    prediction = make_prediction(img)
    st.header("预测概率")
    st.write(prediction)  # 写入各种东西，文字、表格、图片

运行： streamlit run classcification_app.py

运行后会自动打开浏览器（如果没有自动打开，就手动复制提示的地址访问。）

使用的函数如下：

• st.title：创建标题
• st.file_uploader：文件上传
• st.image：显示图片
• st.write：显示文字表格等内容
• st.header：标题

我们只用了很少的几个函数就创建了一个网页。并且同一个局域网的其他人也可以访问。（如果有公网ip就可以让所有人访问）

目标检测

目标检测的实现类似，内容差不多只是目标检测看起来复杂一点，但是streamlit的部分是差不多的。

import matplotlib.pyplot as plt
import numpy as np
import streamlit as st
import torch
from PIL import Image
from torchvision.models.detection import fasterrcnn_resnet50_fpn_v2, FasterRCNN_ResNet50_FPN_V2_Weights
from torchvision.utils import draw_bounding_boxes

# 使用fasterRCNN 进行目标检测，使用预训练的权重
weights = FasterRCNN_ResNet50_FPN_V2_Weights.DEFAULT
categories = weights.meta["categories"]
img_preprocess = weights.transforms()


@st.cache_resource
def load_model():
    model = fasterrcnn_resnet50_fpn_v2(weights=weights, box_score_thresh=0.8)
    model.eval()
    return model


model = load_model()


def make_prediction(img):
    img_processed = img_preprocess(img)
    prediction = model(img_processed.unsqueeze(0))
    prediction = prediction[0]
    prediction["labels"] = [categories[label] for label in prediction["labels"]]
    return prediction


def create_image_with_bboxes(img, prediction):
    img_tensor = torch.tensor(img)
    img_with_bboxes = draw_bounding_boxes(img_tensor, boxes=prediction["boxes"], labels=prediction["labels"],
                                          colors=["red" if label == "person" else "green" for label in
                                                  prediction["labels"]], width=2)
    img_with_bboxes_up = img_with_bboxes.detach().numpy().transpose(1, 2, 0)

    return img_with_bboxes_up


# Dashboard
st.title("目标检测")  # 设置标题
upload = st.file_uploader(label="上传图片", type=["png", "jpg", "jpeg"])  # 上传文件
if upload:
    img = Image.open(upload)
    prediction = make_prediction(img)
    img_with_bbox = create_image_with_bboxes(np.array(img).transpose(2, 0, 1), prediction)
    fig = plt.figure(figsize=(12, 12))
    ax = fig.add_subplot(111)
    plt.imshow(img_with_bbox)
    plt.xticks([], [])
    plt.yticks([], [])
    ax.spines[["top", "bottom", "right", "left"]].set_visible(False)
    st.pyplot(fig, use_container_width=True)  # 绘图，显示matplotlib创建的图

    del prediction["boxes"]
    st.header("预测概率")
    st.write(prediction) # 写入各种东西，文字、表格、图片