wangX1/Intelligent-Logistics-Defect-Detection

1import streamlit as st
2import cv2
3import numpy as np
4import os
5from ultralytics import YOLO
6
7# -------------------------- 页面配置 --------------------------
8st.set_page_config(page_title="包裹缺陷检测演示", page_icon="📦", layout="wide")
9st.title("📦 智能物流包裹缺陷批量检测演示")
10st.markdown("支持：孔洞、污渍、划痕、破损 四类缺陷识别")
11
12
13# -------------------------- 加载模型（缓存加速） --------------------------
14@st.cache_resource
15def load_model():
16    # 你的模型文件，把best.pt或者best.onnx和代码放在一起
17    return YOLO("best.pt")  # 改成你的模型文件名
18
19
20model = load_model()
21
22# -------------------------- 网页功能：上传并检测 --------------------------
23uploaded_files = st.file_uploader(
24    "上传需要检测的包裹图片（可多选）",
25    type=["jpg", "jpeg", "png"],
26    accept_multiple_files=True
27)
28
29if uploaded_files:
30    st.success(f"已上传 {len(uploaded_files)} 张图片，开始检测...")
31
32    # 分两列显示原图和检测结果
33    col1, col2 = st.columns(2)
34
35    for idx, file in enumerate(uploaded_files):
36        # 读取图片
37        file_bytes = np.asarray(bytearray(file.read()), dtype=np.uint8)
38        img = cv2.imdecode(file_bytes, cv2.IMREAD_COLOR)
39        img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
40
41        # 模型推理
42        results = model(img_rgb, conf=0.3)
43        annotated_img = results[0].plot()
44
45        # 显示原图和检测结果
46        with col1:
47            st.subheader(f"原图 {idx + 1}")
48            st.image(img_rgb, use_column_width=True)
49        with col2:
50            st.subheader(f"检测结果 {idx + 1}")
51            st.image(annotated_img, use_column_width=True)
52
53        # 统计信息
54        defect_count = len(results[0].boxes)
55        if defect_count > 0:
56            st.warning(f"图片 {idx + 1} 检测到 {defect_count} 个缺陷！")
57        else:
58            st.success(f"图片 {idx + 1} 未检测到缺陷")