5分钟极速部署:用Streamlit打造高交互安全帽检测Web应用

在工业安全领域,快速验证和展示AI模型的能力往往比模型本身的精度更为关键。想象一下这样的场景:你刚刚完成了一个基于YOLOv8的安全帽检测模型训练,准确率达到了93%,现在需要向非技术背景的项目经理或客户进行演示。传统的方式可能需要前后端开发团队数周的协作,而今天我将展示如何用Python的Streamlit框架,在5分钟内将这个模型转化为一个功能完整的Web应用。

1. 环境准备与模型加载

首先确保你的Python环境已经安装了PyTorch和Streamlit。如果你使用conda管理环境,可以这样创建:

conda create -n helmet_detection python=3.8
conda activate helmet_detection
pip install torch torchvision ultralytics streamlit opencv-python

接下来是模型加载部分。我们假设你已经训练好了YOLOv8模型,得到了一个best.pt权重文件。创建一个detection.py文件,写入以下代码:

from ultralytics import YOLO
import cv2

class HelmetDetector:
    def __init__(self, model_path):
        self.model = YOLO(model_path)
        self.class_names = {0: "安全帽", 1: "未佩戴"}
    
    def detect(self, image):
        results = self.model(image)
        return self._parse_results(results)
    
    def _parse_results(self, results):
        detections = []
        for result in results:
            for box in result.boxes:
                x1, y1, x2, y2 = map(int, box.xyxy[0].tolist())
                conf = float(box.conf[0])
                cls = int(box.cls[0])
                detections.append({
                    "bbox": [x1, y1, x2, y2],
                    "confidence": conf,
                    "class_name": self.class_names[cls]
                })
        return detections

这个类封装了YOLOv8模型的加载和推理过程,将原始输出转换为更易处理的字典格式。

2. 构建Streamlit交互界面

Streamlit的魅力在于它能让开发者用最少的代码创建丰富的Web界面。在同一个文件中继续添加:

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

st.set_page_config(layout="wide")
st.title("安全帽检测演示系统")

@st.cache_resource
def load_model():
    return HelmetDetector("best.pt")

detector = load_model()

upload_option = st.radio(
    "选择输入源:",
    ("上传图片", "使用摄像头")
)

if upload_option == "上传图片":
    uploaded_file = st.file_uploader("选择图片...", type=["jpg", "png", "jpeg"])
    if uploaded_file is not None:
        image = Image.open(uploaded_file)
        image_np = np.array(image)
        
        with st.spinner("检测中..."):
            detections = detector.detect(image_np)
        
        draw_image = image_np.copy()
        for det in detections:
            x1, y1, x2, y2 = det["bbox"]
            color = (0, 255, 0) if det["class_name"] == "安全帽" else (0, 0, 255)
            cv2.rectangle(draw_image, (x1, y1), (x2, y2), color, 2)
            label = f"{det['class_name']} {det['confidence']:.2f}"
            cv2.putText(draw_image, label, (x1, y1-10), 
                       cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
        
        col1, col2 = st.columns(2)
        with col1:
            st.image(image, caption="原始图片", use_column_width=True)
        with col2:
            st.image(draw_image, caption="检测结果", use_column_width=True)
        
        st.json(detections)

这段代码创建了一个简单的界面,允许用户上传图片并查看检测结果。我们使用了两栏布局来并排显示原始图片和检测结果。

3. 添加实时摄像头支持

为了让演示更具冲击力,我们添加实时摄像头检测功能。在之前的代码后面继续添加:

elif upload_option == "使用摄像头":
    run_camera = st.checkbox("开启摄像头")
    FRAME_WINDOW = st.image([])
    camera = cv2.VideoCapture(0)
    
    while run_camera:
        ret, frame = camera.read()
        if not ret:
            st.error("无法获取摄像头画面")
            break
        
        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        detections = detector.detect(frame)
        
        for det in detections:
            x1, y1, x2, y2 = det["bbox"]
            color = (0, 255, 0) if det["class_name"] == "安全帽" else (255, 0, 0)
            cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)
            label = f"{det['class_name']} {det['confidence']:.2f}"
            cv2.putText(frame, label, (x1, y1-10), 
                       cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
        
        FRAME_WINDOW.image(frame)
    
    camera.release()

这段代码会打开用户的默认摄像头,实时显示检测结果。Streamlit的st.image组件可以动态更新,非常适合这种实时应用场景。

4. 增强用户体验的功能扩展

基础功能完成后,我们可以添加一些提升用户体验的功能:

# 在load_model()后添加
confidence_threshold = st.slider(
    "置信度阈值", 
    min_value=0.0, 
    max_value=1.0, 
    value=0.5,
    help="调整检测结果的严格程度"
)

# 修改detect方法
def detect(self, image):
    results = self.model(image, conf=confidence_threshold)
    return self._parse_results(results)

# 添加统计信息显示
if detections:
    stats = {
        "安全帽": sum(1 for d in detections if d["class_name"] == "安全帽"),
        "未佩戴": sum(1 for d in detections if d["class_name"] == "未佩戴")
    }
    st.metric(label="检测统计", value=f"{stats['安全帽']}安全帽 / {stats['未佩戴']}未佩戴")

# 添加模型切换功能
model_version = st.selectbox(
    "选择模型版本",
    ("YOLOv8n", "YOLOv7", "YOLOv6", "YOLOv5"),
    index=0
)

这些改进让用户可以调整检测灵敏度、查看统计信息,并在不同版本的YOLO模型间切换(假设你已经训练了多个版本的模型)。

5. 部署与分享你的应用

完成开发后,你可以通过以下命令本地运行应用:

streamlit run detection.py

要分享给他人,最简单的部署方式是使用Streamlit Community Cloud:

  1. 将代码推送到GitHub仓库
  2. 登录share.streamlit.io
  3. 选择仓库和主文件路径
  4. 点击Deploy

几分钟后,你的应用就会有一个公开URL,可以分享给任何人。对于企业内部分享,你也可以使用Docker容器化部署:

FROM python:3.8-slim
WORKDIR /app
COPY requirements.txt .
RUN pip install -r requirements.txt
COPY . .
CMD ["streamlit", "run", "detection.py", "--server.port=8501", "--server.address=0.0.0.0"]

构建并运行容器:

docker build -t helmet-detection .
docker run -p 8501:8501 helmet-detection

6. 性能优化技巧

当你的应用需要处理更高分辨率或更多并发请求时,可以考虑以下优化:

模型推理优化:

# 使用半精度浮点数加速推理
def __init__(self, model_path):
    self.model = YOLO(model_path)
    if torch.cuda.is_available():
        self.model = self.model.half().to("cuda")
    
# 修改detect方法
def detect(self, image):
    if torch.cuda.is_available():
        image = torch.from_numpy(image).half().to("cuda") / 255.0
    results = self.model(image, conf=confidence_threshold)
    return self._parse_results(results)

视频处理优化:

# 修改摄像头处理部分
frame_skip = 2  # 每3帧处理1帧
frame_count = 0

while run_camera:
    ret, frame = camera.read()
    frame_count += 1
    if frame_count % frame_skip != 0:
        FRAME_WINDOW.image(frame)
        continue
    
    # 剩余处理逻辑不变

缓存策略优化:

@st.cache_data(max_entries=3)
def process_image(image):
    return detector.detect(image)

这些优化可以显著提升应用的响应速度,特别是在资源有限的环境中。

7. 实际应用案例

在某建筑工地的实际部署中,这个应用展示了惊人的效果。安全主管通过平板电脑访问Web界面,可以:

  • 实时检查工地各个区域的监控画面
  • 拍照抽查特定工作点的安全合规情况
  • 生成每日安全报告,统计佩戴率变化趋势
  • 在不同型号的检测算法间切换,比较效果

项目经理则利用这个工具向客户演示工地的安全管理系统,无需安排现场参观就能展示安全措施的落实情况。相比传统方案,这个方案:

对比项传统方案Streamlit方案
开发周期2-3周1天
部署难度需要专业运维一键部署
维护成本
灵活性
硬件要求专用服务器普通PC/云端

这种快速原型开发能力让AI技术能够更快地落地到实际场景中,缩短了从实验室到工地的距离。

Logo

Agent 垂直技术社区,欢迎活跃、内容共建。

更多推荐