Update app.py

app.py

@@ -1,7 +1,7 @@
 import streamlit as st
-import cv2
-import numpy as np
+from streamlit_webrtc import VideoTransformerBase, webrtc_streamer
 import tensorflow as tf
+import numpy as np
 from tensorflow.keras.applications.mobilenet import preprocess_input, decode_predictions
 
 st.title("物体识别应用")
@@ -10,43 +10,24 @@ st.write("通过摄像头识别物体，从左到右显示主要物体的名称"
 # 加载 MobileNet 预训练模型
 model = tf.keras.applications.MobileNet(weights="imagenet")
 
-# 打开摄像头函数
-def open_camera():
-    video_capture = cv2.VideoCapture(0)
-    return video_capture
+class ObjectDetectionTransformer(VideoTransformerBase):
+    def transform(self, frame):
+        img = frame.to_ndarray(format="bgr24")
 
-# 物体识别函数
-def detect_objects(frame):
-    image_resized = cv2.resize(frame, (224, 224))
-    image_array = np.expand_dims(image_resized, axis=0)
-    processed_image = preprocess_input(image_array)
-    
-    preds = model.predict(processed_image)
-    decoded_preds = decode_predictions(preds, top=3)[0]
-    objects = [f"{label}: {round(score * 100, 2)}%" for (_, label, score) in decoded_preds]
-    return objects
+        # 将图像调整大小并进行预处理
+        image_resized = cv2.resize(img, (224, 224))
+        image_array = np.expand_dims(image_resized, axis=0)
+        processed_image = preprocess_input(image_array)
 
-# 按钮控制摄像头启动
-if st.button("开始摄像头"):
-    video_capture = open_camera()
-    stframe = st.empty()
-    
-    while video_capture.isOpened():
-        ret, frame = video_capture.read()
-        if not ret:
-            st.write("无法读取摄像头数据。")
-            break
-        
-        objects = detect_objects(frame)
+        # 进行物体识别
+        preds = model.predict(processed_image)
+        decoded_preds = decode_predictions(preds, top=3)[0]
+        objects = [f"{label}: {round(score * 100, 2)}%" for (_, label, score) in decoded_preds]
         detected_text = " | ".join(objects)
-        
-        # 显示检测结果在帧上
-        cv2.putText(frame, detected_text, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2, cv2.LINE_AA)
-        
-        # 转换并显示在 Streamlit
-        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-        stframe.image(frame_rgb, caption="检测到的物体", channels="RGB")
 
-    video_capture.release()
-else:
-    st.write("点击上方按钮以启动摄像头")
+        # 将检测结果写在图像上
+        cv2.putText(img, detected_text, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2, cv2.LINE_AA)
+        return img
+
+# 使用 streamlit-webrtc 启动摄像头流
+webrtc_streamer(key="object-detection", video_transformer_factory=ObjectDetectionTransformer)