Update app.py

app.py

@@ -1,21 +1,284 @@
-from streamlit_webrtc import webrtc_streamer
-#import os
-#from twilio.rest import Client
+import streamlit as st
+import cv2
+import time
+from streamlit_webrtc import VideoTransformerBase, webrtc_streamer
+from PIL import Image
+from transformers import pipeline
+import os
+from twilio.rest import Client
+from collections import Counter
+import base64
 
+
+# ======================
+# 模型加载函数（缓存）
+# ======================
+
+@st.cache_resource
+def load_smoke_pipeline():
+    """初始化并缓存吸烟图片分类 pipeline。"""
+    return pipeline("image-classification", model="ccclllwww/smoker_cls_base_V9", use_fast=True)
+
+@st.cache_resource
+def load_gender_pipeline():
+    """初始化并缓存性别图片分类 pipeline。"""
+    return pipeline("image-classification", model="rizvandwiki/gender-classification-2", use_fast=True)
+
+@st.cache_resource
+def load_age_pipeline():
+    """初始化并缓存年龄图片分类 pipeline。"""
+    return pipeline("image-classification", model="akashmaggon/vit-base-age-classification", use_fast=True)
+
+# 预先加载所有模型
+load_smoke_pipeline()
+load_gender_pipeline()
+load_age_pipeline()
+
+# ======================
+# remote settings
+# ======================
 # Find your Account SID and Auth Token at twilio.com/console
 # and set the environment variables. See http://twil.io/secure
+account_sid = os.environ['TWILIO_ACCOUNT_SID']
+auth_token = os.environ['TWILIO_AUTH_TOKEN']
+client = Client(account_sid, auth_token)
+
+token = client.tokens.create()
+
+
+# ======================
+# 音频加载函数（缓存）
+# ======================
+
+@st.cache_resource
+def load_all_audios():
+    """加载 audio 目录中的所有 .wav 文件，并返回一个字典，
+    键为文件名（不带扩展名），值为音频字节数据。"""
+    audio_dir = "audio"
+    audio_files = [f for f in os.listdir(audio_dir) if f.endswith(".wav")]
+    audio_dict = {}
+    for audio_file in audio_files:
+        file_path = os.path.join(audio_dir, audio_file)
+        with open(file_path, "rb") as af:
+            audio_bytes = af.read()
+        # 去掉扩展名作为键
+        key = os.path.splitext(audio_file)[0]
+        audio_dict[key] = audio_bytes
+    return audio_dict
+
+# 应用启动时加载所有音频
+audio_data = load_all_audios()
+
+# ======================
+# 核心处理函数
+# ======================
+
+@st.cache_data(show_spinner=False, max_entries=3)
+def smoking_classification(image: Image.Image) -> str:
+    """接受 PIL 图片并利用吸烟分类 pipeline 进行判定，返回标签（如 "smoking"）。"""
+    try:
+        smoke_pipeline = load_smoke_pipeline()
+        output = smoke_pipeline(image)
+        status = max(output, key=lambda x: x["score"])['label']
+        return status
+    except Exception as e:
+        st.error(f"🔍 图像处理错误: {str(e)}")
+        st.stop()
+
+@st.cache_data(show_spinner=False, max_entries=3)
+def gender_classification(image: Image.Image) -> str:
+    """进行性别分类，返回模型输出的性别（依模型输出）。"""
+    try:
+        gender_pipeline = load_gender_pipeline()
+        output = gender_pipeline(image)
+        status = max(output, key=lambda x: x["score"])['label']
+        return status
+    except Exception as e:
+        st.error(f"🔍 图像处理错误: {str(e)}")
+        st.stop()
+
+@st.cache_data(show_spinner=False, max_entries=3)
+def age_classification(image: Image.Image) -> str:
+    """进行年龄分类，返回年龄范围，例如 "10-19" 等。"""
+    try:
+        age_pipeline = load_age_pipeline()
+        output = age_pipeline(image)
+        age_range = max(output, key=lambda x: x["score"])['label']
+        return age_range
+    except Exception as e:
+        st.error(f"🔍 图像处理错误: {str(e)}")
+        st.stop()
+
+# ======================
+# 自定义JS播放音频函数
+# ======================
+
+@st.cache_resource
+def play_audio_via_js(audio_bytes):
+    """
+    利用自定义 HTML 和 JavaScript 播放音频。
+    将二进制音频数据转换为 Base64 后嵌入 audio 标签，
+    并用 JS 在页面加载后模拟点击进行播放。
+    """
+    audio_base64 = base64.b64encode(audio_bytes).decode("utf-8")
+    html_content = f"""
+    <audio id="audio_player" controls style="width: 100%;">
+        <source src="data:audio/wav;base64,{audio_base64}" type="audio/wav">
+        Your browser does not support the audio element.
+    </audio>
+    <script type="text/javascript">
+        // 等待 DOMContentLoaded 事件，并在1秒后自动调用 play() 方法
+        window.addEventListener('DOMContentLoaded', function() {{
+            setTimeout(function() {{
+                var audioElement = document.getElementById("audio_player");
+                if (audioElement) {{
+                    audioElement.play().catch(function(e) {{
+                        console.log("播放被浏览器阻止:", e);
+                    }});
+                }}
+            }}, 1000);
+        }});
+    </script>
+    """
+    st.components.v1.html(html_content, height=150)
+
+# ======================
+# VideoTransformer 定义：处理摄像头帧与快照捕获
+# ======================
+
+class VideoTransformer(VideoTransformerBase):
+    def __init__(self):
+        self.snapshots = []  # 存储捕获的快照
+        self.last_capture_time = time.time()  # 上次捕获时间
+        self.capture_interval = 0.5  # 每0.5秒捕获一张快照
+
+    def transform(self, frame):
+        """从摄像头流捕获单帧图像，并转换为 PIL Image。"""
+        img = frame.to_ndarray(format="bgr24")
+        current_time = time.time()
+        # 每隔 capture_interval 秒捕获一张快照，直到捕获20张
+        if current_time - self.last_capture_time >= self.capture_interval and len(self.snapshots) < 20:
+            img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+            self.snapshots.append(Image.fromarray(img_rgb))
+            self.last_capture_time = current_time
+            st.write(f"已捕获快照 {len(self.snapshots)}/20")
+        return img  # 返回原始帧以供前端显示
+
+# ======================
+# 主函数：整合视频流、自动图片分类并展示结果
+# ======================
+
+def main():
+    st.title("Streamlit-WebRTC 自动图片分类示例")
+    st.write("程序在一分钟内捕获20张快照进行图片分类，首先判定是否吸烟。若检测到吸烟的快照超过2次，则展示年龄与性别分类结果。")
+
+    # 创建用于显示进度文字和进度条的占位容器
+    capture_text_placeholder = st.empty()
+    capture_progress_placeholder = st.empty()
+    classification_text_placeholder = st.empty()
+    classification_progress_placeholder = st.empty()
+    detection_info_placeholder = st.empty()  # 用于显示“开始侦测”
+
+    # 启动实时视频流
+    # Then, pass the ICE server information to webrtc_streamer().
+
+    ctx = webrtc_streamer(key="unique_example", video_transformer_factory=VideoTransformer,rtc_configuration={"iceServers": token.ice_servers})
+    image_placeholder = st.empty()
+    audio_placeholder = st.empty()
+
+    capture_target = 10  # 本轮捕获目标张数
+
+    if ctx.video_transformer is not None:
+        classification_result_placeholder = st.empty()  # 用于显示分类结果
+        detection_info_placeholder.info("开始侦测")
+
+        while True:
+            snapshots = ctx.video_transformer.snapshots
+
+            # 更新捕获阶段进度：同时显示文字和进度条
+            if len(snapshots) < capture_target:
+                capture_text_placeholder.text(f"捕获进度: {len(snapshots)}/{capture_target} 张快照")
+                progress_value = int(len(snapshots) / capture_target * 100)
+                capture_progress_placeholder.progress(progress_value)
+            else:
+                # 捕获完成，清空捕获进度条，并显示完成提示
+                capture_text_placeholder.text("捕获进度: 捕获完成！")
+                capture_progress_placeholder.empty()
+                detection_info_placeholder.empty()  # 清除“开始侦测”提示
+
+                # ---------- 分类阶段进度 ----------
+                total = len(snapshots)
+                classification_text_placeholder.text("分类进度: 正在分类...")
+                classification_progress = classification_progress_placeholder.progress(0)
+
+                # 1. 吸烟分类 (0 ~ 33%)
+                smoke_results = []
+                for idx, img in enumerate(snapshots):
+                    smoke_results.append(smoking_classification(img))
+                smoking_count = sum(1 for result in smoke_results if result.lower() == "smoking")
+                classification_progress.progress(33)
+
+                # 2. 若吸烟次数超过2，再进行性别和年龄分类 (33% ~ 100%)
+                if smoking_count > 2:
+                    gender_results = []
+                    for idx, img in enumerate(snapshots):
+                        gender_results.append(gender_classification(img))
+                    classification_progress.progress(66)
+
+                    age_results = []
+                    for idx, img in enumerate(snapshots):
+                        age_results.append(age_classification(img))
+                    classification_progress.progress(100)
+                    classification_text_placeholder.text("分类进度: 分类完成！")
+
+                    most_common_gender = Counter(gender_results).most_common(1)[0][0]
+                    most_common_age = Counter(age_results).most_common(1)[0][0]
+
+                    result_text = (
+                        f"**吸烟状态：** Smoking (检测到 {smoking_count} 次)\n\n"
+                        f"**性别：** {most_common_gender}\n\n"
+                        f"**年龄范围：** {most_common_age}"
+                    )
+                    classification_result_placeholder.markdown(result_text)
+
+                    # 选择第一张分类结果为 "smoking" ��快照，如未检测到，则显示第一张
+                    smoking_image = None
+                    for idx, label in enumerate(smoke_results):
+                        if label.lower() == "smoking":
+                            smoking_image = snapshots[idx]
+                            break
+                    if smoking_image is None:
+                        smoking_image = snapshots[0]
+                    image_placeholder.image(smoking_image, caption="捕获的快照示例", use_container_width=True)
+
+                    # 清空播放区域后再播放对应音频
+                    audio_placeholder.empty()
+                    audio_key = f"{most_common_age} {most_common_gender.lower()}"
+                    if audio_key in audio_data:
+                        audio_bytes = audio_data[audio_key]
+                        play_audio_via_js(audio_bytes)
+                    else:
+                        st.error(f"音频文件不存在: {audio_key}.wav")
+                else:
+                    result_text = "**吸烟状态：** Not Smoking"
+                    classification_result_placeholder.markdown(result_text)
+                    image_placeholder.empty()
+                    audio_placeholder.empty()
+                    classification_text_placeholder.text("分类进度: 分类完成！")
+                    classification_progress.progress(100)
+
+                # 分类阶段结束后清空分类进度占位区
+                time.sleep(1)
+                classification_progress_placeholder.empty()
+                classification_text_placeholder.empty()
+                capture_text_placeholder.empty()
+
+
+                # 重置快照列表，准备下一轮捕获
+                detection_info_placeholder.info("开始侦测")
+                ctx.video_transformer.snapshots = []
+                ctx.video_transformer.last_capture_time = time.time()
+            time.sleep(0.1)
 
-#account_sid = os.environ['AC9b487ef91db53af2b4afdaff07982831']
-#auth_token = os.environ['06578b544f0b268cfce25330d3b81eb3']
-#client = Client(account_sid, auth_token)
-#token = client.tokens.create()
-
-webrtc_streamer(
-    key="sample",
-    # ...
-    rtc_configuration=
-    {
-        "iceServers": [{"urls": ["stun:stun.l.google.com:19302"]}]
-    }
-    # ...
-)
+if __name__ == "__main__":
+    main()