告别繁琐部署!用Streamlit 5分钟搭建一个可交互的安全帽检测Web应用(支持摄像头/图片/视频)
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:
- 将代码推送到GitHub仓库
- 登录share.streamlit.io
- 选择仓库和主文件路径
- 点击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技术能够更快地落地到实际场景中,缩短了从实验室到工地的距离。
更多推荐


所有评论(0)