Update app.py

app.py

@@ -1,3 +1,9 @@
+# -*- coding: utf-8 -*-
+"""
+Created on Tue May 20 11:00:14 2025
+
+@author: ColinWang
+"""
 import streamlit as st
 import cv2
 import time
@@ -5,10 +11,10 @@ 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
-
+#from twilio.rest import Client
+from collections import Counter
+import uuid
 
 # ======================
 # 模型加载函数（缓存）
@@ -30,15 +36,17 @@ def load_age_pipeline():
     return pipeline("image-classification", model="akashmaggon/vit-base-age-classification", use_fast=True)
 
 # 预先加载所有模型
-load_smoke_pipeline()
-load_gender_pipeline()
-load_age_pipeline()
+smoke_pipeline = load_smoke_pipeline()
+gender_pipeline = load_gender_pipeline()
+age_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)
@@ -69,15 +77,46 @@ def load_all_audios():
 # 应用启动时加载所有音频
 audio_data = load_all_audios()
 
+
 # ======================
-# 核心处理函数
+# 照片檢測处理函数
+# ======================
+
+def smoking_detection(image: Image.Image) -> str:
+    try:
+        output = smoke_pipeline(image)
+        status = output[0]["label"]
+        return status
+    except Exception as e:
+        st.error(f"🔍 图像处理错误: {str(e)}")
+        st.stop()
+        
+def gender_detection(image: Image.Image) -> str:
+    try:
+        output = gender_pipeline(image)
+        status = output[0]["label"]
+        return status
+    except Exception as e:
+        st.error(f"🔍 图像处理错误: {str(e)}")
+        st.stop()
+        
+def age_detection(image: Image.Image) -> str:
+    try:
+        output = age_pipeline(image)
+        status = output[0]["label"]
+        return status
+    except Exception as e:
+        st.error(f"🔍 图像处理错误: {str(e)}")
+        st.stop()
+        
+# ======================
+# 實時檢測核心处理函数
 # ======================
 
 @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
@@ -89,7 +128,6 @@ def smoking_classification(image: Image.Image) -> str:
 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
@@ -101,7 +139,6 @@ def gender_classification(image: Image.Image) -> str:
 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
@@ -122,15 +159,14 @@ def play_audio_via_js(audio_bytes):
     """
     audio_base64 = base64.b64encode(audio_bytes).decode("utf-8")
     html_content = f"""
-    <audio id="audio_player" controls style="width: 100%;">
+    <audio id="audio_player_{uuid.uuid4()}" 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");
+                var audioElement = document.getElementById("audio_player_{uuid.uuid4()}");
                 if (audioElement) {{
                     audioElement.play().catch(function(e) {{
                         console.log("播放被浏览器阻止:", e);
@@ -165,11 +201,57 @@ class VideoTransformer(VideoTransformerBase):
         return img  # 返回原始帧以供前端显示
 
 # ======================
-# 主函数：整合视频流、自动图片分类并展示结果
+# 照片检测页面
 # ======================
 
-def main():
-    st.title("Streamlit-WebRTC 自动图片分类示例")
+def photo_detection_page():
+    st.title("照片检测")
+    st.write("上传一张图片或使用摄像头拍摄，检测是否吸烟，若检测到吸烟则进一步分析性别和年龄。")
+
+    # 提供上传和摄像头选项
+    option = st.radio("选择输入方式", ["上传图片", "使用摄像头拍摄"])
+
+    image = None
+    if option == "上传图片":
+        uploaded_file = st.file_uploader("选择一张图片", type=["jpg", "jpeg", "png"])
+        if uploaded_file is not None:
+            image = Image.open(uploaded_file)
+            st.image(image, caption="上传的图片", use_container_width=True)
+    else:
+        # 摄像头拍摄
+        enable = st.checkbox("启用摄像头")
+        camera_file = st.camera_input("拍摄照片", disabled=not enable)
+        if camera_file is not None:
+            image = Image.open(camera_file)
+            st.image(image, caption="拍攝的圖片", use_container_width=True)
+
+    if image is not None:
+
+        # 吸烟分类
+        with st.spinner("Wait for smoking detection"):
+            smoke_result = smoking_detection(image)
+            st.success("The smoke result is:")
+        st.write(smoke_result)
+        
+        if smoke_result.lower() == "smoking":
+            # 性别分类
+            with st.spinner("Wait for gender detection"):
+                gender_result = gender_detection(image)
+                st.success("The gender result is:")
+            st.write(gender_result)
+
+            # 年龄分类
+            with st.spinner("Wait for age detection"):
+                age_result = age_detection(image)
+                st.success("The age result is:")
+            st.write(age_result)
+
+# ======================
+# 实时检测页面（原主函数）
+# ======================
+
+def real_time_detection_page():
+    st.title("实时视频检测")
     st.write("程序在一分钟内捕获20张快照进行图片分类，首先判定是否吸烟。若检测到吸烟的快照超过2次，则展示年龄与性别分类结果。")
 
     # 创建用于显示进度文字和进度条的占位容器
@@ -177,48 +259,43 @@ def main():
     capture_progress_placeholder = st.empty()
     classification_text_placeholder = st.empty()
     classification_progress_placeholder = st.empty()
-    detection_info_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})
+    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  # 本轮捕获目标张数
+    capture_target = 20
 
     if ctx.video_transformer is not None:
-        classification_result_placeholder = st.empty()  # 用于显示分类结果
+        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()  # 清除“开始侦测”提示
+                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):
@@ -241,7 +318,6 @@ def main():
                     )
                     classification_result_placeholder.markdown(result_text)
 
-                    # 选择第一张分类结果为 "smoking" 的快照，如未检测到，则显示第一张
                     smoking_image = None
                     for idx, label in enumerate(smoke_results):
                         if label.lower() == "smoking":
@@ -251,7 +327,6 @@ def main():
                         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:
@@ -267,18 +342,28 @@ def main():
                     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)
 
+# ======================
+# 主函数：多页面导航
+# ======================
+
+def main():
+    st.sidebar.title("导航")
+    page = st.sidebar.selectbox("选择页面", ["照片检测", "实时视频检测"])
+
+    if page == "照片检测":
+        photo_detection_page()
+    else:
+        real_time_detection_page()
+
 if __name__ == "__main__":
-    main()
+    main()