add video process

app.py

@@ -4,6 +4,7 @@ import streamlit as st
 from PIL import Image
 import pandas as pd
 import numpy as np
+import cv2
 
 
 
@@ -15,108 +16,154 @@ pipe_emotions_refined = pipeline("image-classification", model="felixwf/fine_tun
 
 st.title("Online Teaching Effect Monitor")
 
-file_name = st.file_uploader("Upload a image or a video")
+file_name = st.file_uploader("Upload an image or a video")
 
 if file_name is not None:
-    face_image = Image.open(file_name)
-    st.image(face_image)
-    output = pipe_yolos(face_image)
-
-    data = output
-    # 过滤出所有标签为 "person" 的项
-    persons = [item for item in data if item['label'] == 'person']
-    
-    # 打印结果
-    print(persons)
-    st.text(persons)
-    
-    # 假设有一张原始图片，加载图片并截取出每个 "person" 的部分
-    # original_image_path = 'input_image.jpg'  # 替换为实际图片路径
-    # original_image = Image.open(original_image_path)
-    original_image = face_image
-    persons_image_list = []
-    
-    # 截取每个 "person" 的部分并保存
-    for idx, person in enumerate(persons):
-        box = person['box']
-        cropped_image = original_image.crop((box['xmin'], box['ymin'], box['xmax'], box['ymax']))
-        cropped_image.save(f'person_{idx}.jpg')
-        cropped_image.show()
-        persons_image_list.append(cropped_image)
-    
-    
-    # Calculate the number of rows needed for 3 columns
-    num_images = len(persons)
-    num_cols = 3
-    num_rows = (num_images + num_cols - 1) // num_cols  # Ceiling division
-
-    # Create a new canvas to stitch all person images in a grid with 3 columns
-    fig, axes = plt.subplots(num_rows, num_cols, figsize=(15, 5 * num_rows))
-
-    # Flatten the axes array for easy iteration
-    axes = axes.flatten()
-
-    # Crop each "person" part and plot it on the grid
-    for idx, person in enumerate(persons):
-        box = person['box']
-        cropped_image = original_image.crop((box['xmin'], box['ymin'], box['xmax'], box['ymax']))
-        axes[idx].imshow(cropped_image)
-        axes[idx].axis('off')
-        axes[idx].set_title(f'Person {idx}')
-
-    # Turn off any unused subplots
-    for ax in axes[num_images:]:
-        ax.axis('off')
-
-    # 识别每个人的表情
-    output_list_emotions = []
-    output_list_emotions_refined = []
-    
-    for idx, face in enumerate(persons_image_list):
-      print(f"processing {idx}")
-      output = pipe_emotions(face)
-      output_list_emotions.append(output[0])
-      output = pipe_emotions_refined(face)
-      output_list_emotions_refined.append(output[0])
-    
-    print(output_list_emotions)
-    st.subheader("Emotions by model: dima806/facial_emotions_image_detection")
-    st.text(output_list_emotions)
-    print(output_list_emotions_refined)
-    st.subheader("Emotions by model: felixwf/fine_tuned_face_emotion_model")
-    st.text(output_list_emotions_refined)
-
-    
-    # 统计各种标签的数量
-    label_counts = {}
-    for item in output_list_emotions:
-        label = item['label']
-        if label in label_counts:
-            label_counts[label] += 1
-        else:
-            label_counts[label] = 1
-    for item in output_list_emotions_refined:
-        label = item['label']
-        if label in label_counts:
-            label_counts[label] += 1
-        else:
-            label_counts[label] = 1
-
-    # 绘制饼状图
-    labels = list(label_counts.keys())
-    sizes = list(label_counts.values())
-    
-    pie_fig, ax = plt.subplots()
-    ax.pie(sizes, labels=labels, autopct='%1.1f%%', startangle=140)
-    ax.set_title('Distribution of Emotions')
-    ax.axis('equal')  # 确保饼状图为圆形
-    # plt.show()
-    # Use Streamlit columns to display the images and pie chart side by side
-    col1, col2 = st.columns(2)
-
-    with col1:
-        st.pyplot(fig)  # Display the stitched person images
-
-    with col2:
-        st.pyplot(pie_fig)  # Display the pie chart
-
+    if file_name.type.startswith('image'):
+        # Process image
+        face_image = Image.open(file_name)
+        st.image(face_image)
+        output = pipe_yolos(face_image)
+
+        data = output
+        # 过滤出所有标签为 "person" 的项
+        persons = [item for item in data if item['label'] == 'person']
+        
+        # 打印结果
+        print(persons)
+        st.text(persons)
+        
+        # 假设有一张原始图片，加载图片并截取出每个 "person" 的部分
+        original_image = face_image
+        persons_image_list = []
+        
+        # 截取每个 "person" 的部分并保存
+        for idx, person in enumerate(persons):
+            box = person['box']
+            cropped_image = original_image.crop((box['xmin'], box['ymin'], box['xmax'], box['ymax']))
+            cropped_image.save(f'person_{idx}.jpg')
+            cropped_image.show()
+            persons_image_list.append(cropped_image)
+        
+        # Calculate the number of rows needed for 3 columns
+        num_images = len(persons)
+        num_cols = 3
+        num_rows = (num_images + num_cols - 1) // num_cols  # Ceiling division
+
+        # Create a new canvas to stitch all person images in a grid with 3 columns
+        fig, axes = plt.subplots(num_rows, num_cols, figsize=(15, 5 * num_rows))
+
+        # Flatten the axes array for easy iteration
+        axes = axes.flatten()
+
+        # Crop each "person" part and plot it on the grid
+        for idx, person in enumerate(persons):
+            box = person['box']
+            cropped_image = original_image.crop((box['xmin'], box['ymin'], box['xmax'], box['ymax']))
+            axes[idx].imshow(cropped_image)
+            axes[idx].axis('off')
+            axes[idx].set_title(f'Person {idx}')
+
+        # Turn off any unused subplots
+        for ax in axes[num_images:]:
+            ax.axis('off')
+
+        # 识别每个人的表情
+        output_list_emotions = []
+        output_list_emotions_refined = []
+        
+        for idx, face in enumerate(persons_image_list):
+          print(f"processing {idx}")
+          output = pipe_emotions(face)
+          output_list_emotions.append(output[0])
+          output = pipe_emotions_refined(face)
+          output_list_emotions_refined.append(output[0])
+        
+        print(output_list_emotions)
+        st.subheader("Emotions by model: dima806/facial_emotions_image_detection")
+        st.text(output_list_emotions)
+        print(output_list_emotions_refined)
+        st.subheader("Emotions by model: felixwf/fine_tuned_face_emotion_model")
+        st.text(output_list_emotions_refined)
+
+        
+        # 统计各种标签的数量
+        label_counts = {}
+        for item in output_list_emotions:
+            label = item['label']
+            if label in label_counts:
+                label_counts[label] += 1
+            else:
+                label_counts[label] = 1
+        for item in output_list_emotions_refined:
+            label = item['label']
+            if label in label_counts:
+                label_counts[label] += 1
+            else:
+                label_counts[label] = 1
+
+        # 绘制饼状图
+        labels = list(label_counts.keys())
+        sizes = list(label_counts.values())
+        
+        pie_fig, ax = plt.subplots()
+        ax.pie(sizes, labels=labels, autopct='%1.1f%%', startangle=140)
+        ax.set_title('Distribution of Emotions')
+        ax.axis('equal')  # 确保饼状图为圆形
+        # plt.show()
+        # Use Streamlit columns to display the images and pie chart side by side
+        col1, col2 = st.columns(2)
+
+        with col1:
+            st.pyplot(fig)  # Display the stitched person images
+
+        with col2:
+            st.pyplot(pie_fig)  # Display the pie chart
+
+    elif file_name.type.startswith('video'):
+        # Process video
+        video = cv2.VideoCapture(file_name)
+        frame_count = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
+        frame_rate = int(video.get(cv2.CAP_PROP_FPS))
+        frame_interval = frame_rate  # Process one frame per second
+
+        frame_emotions = []
+
+        for frame_idx in range(0, frame_count, frame_interval):
+            video.set(cv2.CAP_PROP_POS_FRAMES, frame_idx)
+            ret, frame = video.read()
+            if not ret:
+                break
+
+            # Convert frame to PIL Image
+            frame_image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
+            output = pipe_yolos(frame_image)
+
+            data = output
+            persons = [item for item in data if item['label'] == 'person']
+            persons_image_list = []
+
+            for person in persons:
+                box = person['box']
+                cropped_image = frame_image.crop((box['xmin'], box['ymin'], box['xmax'], box['ymax']))
+                persons_image_list.append(cropped_image)
+
+            # Recognize emotions for each person in the frame
+            frame_emotion = []
+            for face in persons_image_list:
+                output = pipe_emotions(face)
+                frame_emotion.append(output[0]['label'])
+            frame_emotions.append(frame_emotion)
+
+        # Plot emotions over frames
+        fig, ax = plt.subplots(figsize=(10, 5))
+        ax.plot(range(len(frame_emotions)), [len(emotions) for emotions in frame_emotions], label='Number of Persons Detected')
+        ax.set_xlabel('Frame')
+        ax.set_ylabel('Number of Persons')
+        ax.set_title('Number of Persons Detected Over Frames')
+        ax.legend()
+
+        st.pyplot(fig)
+
+    else:
+        st.error("Unsupported file type. Please upload an image or a video.")