Update app.py

app.py

@@ -1,282 +1,268 @@
+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 collections import Counter
+import base64
 
 # ======================
-# Import Section
+# 模型加载函数（缓存）
 # ======================
 
-# Core Libraries
-import io # Input/output operations for byte streams
-
-# AI/ML Frameworks
-from transformers import pipeline  # Hugging Face transformers pipeline
-import torch  # PyTorch tensor operations
-
-# Audio Processing
-import numpy as np
-from scipy.io.wavfile import write as write_wav  # Audio file I/O operations
-
-# Image Processing
-from PIL import Image  # Image manipulation library
+@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)
 
-# Web Interface
-import streamlit as st  # Web app framework
+@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()
 
 # ======================
-# Model Loading Functions
+# 音频加载函数（缓存）
 # ======================
 
 @st.cache_resource
-def load_caption_pipeline():
-    """Initialize and cache the image captioning pipeline.
-    
-    Returns:
-        Pipeline: BLIP model for image-to-text generation
-    """
-    return pipeline("image-to-text", model="Salesforce/blip-image-captioning-base",use_fast=True)
-
-@st.cache_resource
-def load_story_pipeline():
-    """Initialize and cache the story generation pipeline.
-    
-    Returns:
-        Pipeline: Fine-tuned LLaMA model for children's story generation
-    """
-    return pipeline("text-generation", model="aspis/gpt2-genre-story-generation",use_fast=True)
-
-@st.cache_resource
-def load_tts_pipeline():
-    """Initialize and cache the text-to-speech pipeline.
-    
-    Returns:
-        Pipeline: Microsoft's SpeechT5 for high-quality speech synthesis
-    """
-    return pipeline("text-to-speech", model="facebook/mms-tts-eng",use_fast=True)
-
+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()
 
 # ======================
-# Core Processing Functions
+# 核心处理函数
 # ======================
 
 @st.cache_data(show_spinner=False, max_entries=3)
-def generate_image_caption(image: Image.Image) -> str:
-    """Generate descriptive caption for uploaded image.
-    
-    Args:
-        image (PIL.Image): RGB formatted input image
-        
-    Returns:
-        str: Generated image caption
-        
-    Raises:
-        StreamlitError: If caption generation fails
-    """
+def smoking_classification(image: Image.Image) -> str:
+    """接受 PIL 图片并利用吸烟分类 pipeline 进行判定，返回标签（如 "smoking"）。"""
     try:
-        img2caption = load_caption_pipeline()
-        # Generate caption
-        caption = img2caption(image)[0]['generated_text']
-        return caption
+        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"🔍 The caption fairy is confused about the picture!  says: {str(e)}")
+        st.error(f"🔍 图像处理错误: {str(e)}")
         st.stop()
 
 @st.cache_data(show_spinner=False, max_entries=3)
-def generate_story(caption: str) -> str:
-    """Generate child-friendly story from image caption.
-    
-    Args:
-        caption (str): Image description from previous step
-        
-    Returns:
-        str: Generated story (60-80 words) with happy ending
-        
-    Raises:
-        StreamlitError: If story generation fails
-    """
+def gender_classification(image: Image.Image) -> str:
+    """进行性别分类，返回模型输出的性别（依模型输出）。"""
     try:
-        cap2story = load_story_pipeline()
-        output = cap2story(f"<BOS> <adventure>{caption}",truncation=True,max_length=80,do_sample=True,repetition_penalty=1.5, temperature=1.2,top_p=0.95,top_k=50,no_repeat_ngram_size=2)
-        story = output[0]['generated_text'].replace("<BOS> <adventure>","")
-        return story
+        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"🧚 The writing fairy is sleeping! says: {str(e)}")
+        st.error(f"🔍 图像处理错误: {str(e)}")
         st.stop()
 
 @st.cache_data(show_spinner=False, max_entries=3)
-def read_story(story):
-    """Convert generated story to speech audio.
-    
-    Args:
-        story (str): Generated story text
-        
-    Returns:
-        io.BytesIO: Audio buffer in WAV format
-        
-    Raises:
-        StreamlitError: If audio generation fails
-    """
+def age_classification(image: Image.Image) -> str:
+    """进行年龄分类，返回年龄范围，例如 "10-19" 等。"""
     try:
-        text2speech = load_tts_pipeline()
-        audio_data = text2speech(story)
-        audio_buffer = io.BytesIO()
-        audio_array = np.array(audio_data["audio"])
-        if len(audio_array.shape) == 2:
-            audio_array = audio_array[:, 0]
-        if audio_array.dtype == np.float32:
-            audio_array = (audio_array * 32767).astype(np.int16)
-        elif audio_array.dtype != np.int16:
-            audio_array = audio_array.astype(np.int16)
-        channels = 1
-        bit_depth = 16
-        sample_rate = int(audio_data["sampling_rate"])
-        block_align = channels * (bit_depth // 8)
-        bytes_per_second = sample_rate * block_align
-        write_wav(audio_buffer, sample_rate, audio_array)
-        audio_buffer.seek(0)
-        return audio_buffer
+        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"🔊 The reading fairy is sneezing! says: {str(e)}")
-        st.stop()  
+        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  # 返回原始帧以供前端显示
 
 # ======================
-# Main Application
+# 主函数：整合视频流、自动图片分类并展示结果
 # ======================
 
 def main():
-    """Main application flow and UI configuration."""
-    
-    # Configure page settings
-    st.set_page_config(
-    page_title="Magic Story Time",
-    page_icon="🧚",
-    layout="centered",
-    initial_sidebar_state="expanded"
-    )
-
-    # Custom CSS styling
-    st.markdown("""
-    <style>
-    .story-box {
-    background: linear-gradient(145deg, #fff1eb 0%, #ace0f9 100%);
-    border-radius: 15px;
-    padding: 25px;
-    font-size: 1.1em;
-    line-height: 1.8;
-    color: #2c3e50;
-    box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
-    margin: 20px 0;
-    }
-    .upload-section {
-    border: 2px dashed #4CAF50;
-    border-radius: 10px;
-    padding: 20px;
-    background: rgba(76, 175, 80, 0.05);
-    }
-    </style>
-    """, unsafe_allow_html=True)
-
-    
-    # Sidebar - Image Upload
-    with st.sidebar:
-        st.header("🖼️ Upload Your Magic Drawing Paper")
-        uploaded_image = st.file_uploader(
-            label="Upload an image",
-            type=["jpg", "jpeg", "png"],
-            help="Format in JPEG/JPG/PNG, max 1MB",
-            key="image_uploader",
-            accept_multiple_files=False
-        )
-
-        if uploaded_image:
-          st.success(f"🔍 The caption fairy received your image: {uploaded_image.name}")
-  
-    # Main Content Area
-    # App title
-    st.title("🧚 Magic Story Camp")
-    st.markdown("---")
-
-
-    # input validation
-    if uploaded_image:
-        
-        # Validate file specifications
-        if uploaded_image.size > 1024* 1024:
-            st.error("🔍 The caption fairy says the image is too big! please give me image under 1MB")
-            st.stop()
-        if uploaded_image.type not in ["image/jpeg", "image/png"]:
-            st.error("🔍 The caption fairy says only JPG/PNG allowed!")
-            st.stop()
-
-        # Processing pipeline
-        with st.spinner("🧙 The fairies are casting magic spells, it may take some time⏳..."):
-            try:
-                # Convert to RGB format for model compatibility
-                image = Image.open(uploaded_image).convert("RGB") 
-
-                # Display processing UI elements
-                status_display = st.empty()
-                progress_bar = st.progress(0)
-
-                # Image preview expander
-                with st.expander("view the image", expanded=True):
-                    st.image(image, use_container_width=True)
-        
-                
-                # Processing stages
-                # Stage 1: Image Captioning
-                status_display.markdown("🔍 **The caption fairy is viewing the image...**")
-                progress_bar.progress(25)
-                caption = generate_image_caption(image)
-                
-        
-                # Stage 2: Story Generation
-                status_display.markdown("🧚 **The writing fairy is writing the story...**")
-                progress_bar.progress(50)
-                story = generate_story(caption)
-                
-        
-                # Stage 3: Audio Synthesis
-                status_display.markdown("🔊 **The reading fairy is preparing audio magic...**")
-                progress_bar.progress(75)
-                speech = read_story(story)
-                
-        
-                # Finish
-                progress_bar.progress(100)
-                status_display.markdown("🧚 **The Story is ready!**")
-
-                # Display formatted story
-                st.markdown("### 📖 Your Magic story")
-                st.markdown(f'<div class="story-box">{story}</div>', unsafe_allow_html=True)
-        
-                # Audio playback and download
-                st.audio(speech, format="audio/wav")
-                st.download_button(
-                      "🎵 Download Story",
-                      data=speech,
-                      file_name="magic_story.wav",
-                      mime="audio/wav",
-                      help="click to download your story"
-                      )
-            except Exception as e:
-                st.error(f"💥 The magic spell broke! Please try again. {str(e)}")
-                st.stop()
-    else:
-        # Page instructions
-        st.markdown("""
-        <div class="upload-section">
-          <h3 style="color:#4CAF50; text-align:center;">❓ guidance</h3>
-          1. 🖼️ Upload Your Picture in the sidebar<br>
-          2. Wait for the magic sparkles ✨）<br>
-          3. Read/listen to your story and download with 🎵 button!<br>
-          <br>
-          Note: First-time model loading may take longer.<br> 
-          Please have a glass of juice and be patient for a few moments<br>
-        </div>
-        """, unsafe_allow_html=True)
+    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()  # 用于显示“开始侦测”
+
+    # 启动实时视频流
+    ctx = webrtc_streamer(key="unique_example", video_transformer_factory=VideoTransformer)
+    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)
 
 if __name__ == "__main__":
     main()