File size: 3,846 Bytes
68ebbad
 
b2639c4
1423812
 
cfa8b24
a6c4560
030c58e
 
a6c4560
 
 
 
030c58e
48d2197
 
 
3dbcc1d
da3ed73
68ebbad
 
 
 
b0b6588
 
 
d766871
d7dd53f
d6bb7f6
a6c4560
 
 
 
3c9434b
 
 
5426655
d6bb7f6
5426655
0c020f3
 
3c9434b
 
d6bb7f6
 
a24e27f
 
48d2197
 
a1ed311
384bc42
e99de43
a6c4560
030c58e
a6c4560
 
030c58e
a6c4560
 
030c58e
a6c4560
 
030c58e
a6c4560
 
 
 
 
 
 
 
 
 
030c58e
 
 
 
 
a6c4560
030c58e
a6c4560
 
030c58e
a6c4560
 
030c58e
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
import streamlit as st
from PIL import Image
from transformers import pipeline
import pandas as pd
import matplotlib.pyplot as plt

# Initialize session state for results, image names, and image sizes if not already present
if 'results' not in st.session_state:
    st.session_state['results'] = []
if 'image_names' not in st.session_state:
    st.session_state['image_names'] = []
if 'image_sizes' not in st.session_state:
    st.session_state['image_sizes'] = []

# Disable PyplotGlobalUseWarning
st.set_option('deprecation.showPyplotGlobalUse', False)

# Create an image classification pipeline with scores
pipe = pipeline("image-classification", model="trpakov/vit-face-expression", top_k=None)

# Streamlit app
st.title("Emotion Recognition with vit-face-expression")

# Upload images
uploaded_images = st.file_uploader("Upload images", type=["jpg", "png"], accept_multiple_files=True)

# Display thumbnail images alongside file names and sizes in the sidebar
selected_images = []
if uploaded_images:
    # Reset the image names and sizes lists each time new images are uploaded
    st.session_state['image_names'] = [img.name for img in uploaded_images]
    st.session_state['image_sizes'] = [round(img.size / 1024.0, 1) for img in uploaded_images]

    # Add a "Select All" checkbox in the sidebar
    select_all = st.sidebar.checkbox("Select All", False)
    
    for idx, img in enumerate(uploaded_images):
        image = Image.open(img)
        checkbox_key = f"{img.name}_checkbox_{idx}"  # Unique key for each checkbox
        # Display thumbnail image and checkbox in sidebar
        st.sidebar.image(image, caption=f"{img.name} {img.size / 1024.0:.1f} KB", width=40)
        selected = st.sidebar.checkbox(f"Select {img.name}", value=select_all, key=checkbox_key)         
        
        if selected:
            selected_images.append(image)

if st.button("Predict Emotions") and selected_images:
    # Predict emotion for each selected image using the pipeline
    st.session_state['results'] = [pipe(image) for image in selected_images]

# Generate DataFrame from results
if st.button("Generate DataFrame"):
    # Access the results, image names, and sizes from the session state
    results = st.session_state['results']
    image_names = st.session_state['image_names']
    image_sizes = st.session_state['image_sizes']
    if results:
        # Initialize an empty list to store all the data
        data = []

        # Iterate over the results and populate the list with dictionaries
        for i, result_set in enumerate(results):
            # Initialize a dictionary for the current set with zeros
            current_data = {
                'Neutral': 0,
                'Happy': 0,
                'Surprise': 0,
                'Disgust': 0,
                'Angry': 0,
                # Add other emotions if necessary
                'Image Name': image_names[i],
                'Image Size (KB)': image_sizes[i]
            }
            
            for result in result_set:
                # Capitalize the label and update the score in the current set
                emotion = result['label'].capitalize()
                score = round(result['score'], 4)  # Round the score to 4 decimal places
                current_data[emotion] = score
            
            # Append the current data to the data list
            data.append(current_data)

        # Convert the list of dictionaries into a pandas DataFrame
        df_emotions = pd.DataFrame(data)

        # Display the DataFrame
        st.write(df_emotions)

        # Optional: Save the DataFrame to a CSV file
        df_emotions.to_csv('emotion_scores.csv', index=False)
        st.success('DataFrame generated and saved as emotion_scores.csv')
    else:
        st.error("No results to generate DataFrame. Please predict emotions first.")