xtlyxt commited on
Commit
8a4145e
1 Parent(s): b4bff53

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +126 -0
app.py ADDED
@@ -0,0 +1,126 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import streamlit as st
2
+ from PIL import Image
3
+ from transformers import pipeline
4
+ import pandas as pd
5
+ import matplotlib.pyplot as plt
6
+ import seaborn as sns
7
+ from pandas.plotting import parallel_coordinates
8
+
9
+ # Initialize session state for results, image names, and image sizes if not already present
10
+ if 'results' not in st.session_state:
11
+ st.session_state['results'] = []
12
+ if 'image_names' not in st.session_state:
13
+ st.session_state['image_names'] = []
14
+ if 'image_sizes' not in st.session_state:
15
+ st.session_state['image_sizes'] = []
16
+
17
+ # Disable PyplotGlobalUseWarning
18
+ st.set_option('deprecation.showPyplotGlobalUse', False)
19
+
20
+ # Create an image classification pipeline with scores
21
+ pipe = pipeline("image-classification", model="trpakov/vit-face-expression", top_k=None)
22
+
23
+ # Streamlit app
24
+ st.title("Emotion Recognition with vit-face-expression")
25
+
26
+ # Upload images
27
+ uploaded_images = st.file_uploader("Upload images", type=["jpg", "png"], accept_multiple_files=True)
28
+
29
+ # Display thumbnail images alongside file names and sizes in the sidebar
30
+ selected_images = []
31
+ if uploaded_images:
32
+ # Reset the image names and sizes lists each time new images are uploaded
33
+ st.session_state['image_names'] = [img.name for img in uploaded_images]
34
+ st.session_state['image_sizes'] = [round(img.size / 1024.0, 1) for img in uploaded_images]
35
+
36
+ # Add a "Select All" checkbox in the sidebar
37
+ select_all = st.sidebar.checkbox("Select All", False)
38
+
39
+ for idx, img in enumerate(uploaded_images):
40
+ image = Image.open(img)
41
+ checkbox_key = f"{img.name}_checkbox_{idx}" # Unique key for each checkbox
42
+ # Display thumbnail image and checkbox in sidebar
43
+ st.sidebar.image(image, caption=f"{img.name} {img.size / 1024.0:.1f} KB", width=40)
44
+ selected = st.sidebar.checkbox(f"Select {img.name}", value=select_all, key=checkbox_key)
45
+
46
+ if selected:
47
+ selected_images.append(image)
48
+
49
+ if st.button("Predict Emotions") and selected_images:
50
+ # Predict emotion for each selected image using the pipeline
51
+ st.session_state['results'] = [pipe(image) for image in selected_images]
52
+
53
+ # Generate DataFrame from results
54
+ if st.button("Generate HeatMap & DataFrame"):
55
+ # Access the results, image names, and sizes from the session state
56
+ results = st.session_state['results']
57
+ image_names = st.session_state['image_names']
58
+ image_sizes = st.session_state['image_sizes']
59
+ if results:
60
+ # Initialize an empty list to store all the data
61
+ data = []
62
+
63
+ # Iterate over the results and populate the list with dictionaries
64
+ for i, result_set in enumerate(results):
65
+ # Initialize a dictionary for the current set with zeros
66
+ current_data = {
67
+
68
+ 'Happy': 0,
69
+ 'Surprise': 0,
70
+ 'Neutral': 0,
71
+ 'Sad': 0,
72
+ 'Disgust': 0,
73
+ 'Angry': 0,
74
+ 'Fear': 0,
75
+
76
+
77
+
78
+ # Add other emotions if necessary
79
+ 'Image Name': image_names[i],
80
+ #'Image Size (KB)': image_sizes[i]
81
+ 'Image Size (KB)': f"{image_sizes[i]:.1f}" # Format the size to one decimal place
82
+ }
83
+
84
+ for result in result_set:
85
+ # Capitalize the label and update the score in the current set
86
+ emotion = result['label'].capitalize()
87
+ score = round(result['score'], 4) # Round the score to 4 decimal places
88
+ current_data[emotion] = score
89
+
90
+ # Append the current data to the data list
91
+ data.append(current_data)
92
+
93
+ # Convert the list of dictionaries into a pandas DataFrame
94
+ df_emotions = pd.DataFrame(data)
95
+
96
+ # Display the DataFrame
97
+ st.write(df_emotions)
98
+
99
+ # Plotting the heatmap for the first seven columns
100
+ plt.figure(figsize=(10, 10))
101
+ sns.heatmap(df_emotions.iloc[:, :7], annot=True, fmt=".1f", cmap='viridis')
102
+ plt.title('Heatmap of Emotion Scores')
103
+ plt.xlabel('Emotion Categories')
104
+ plt.ylabel('Data Points')
105
+ st.pyplot(plt)
106
+
107
+
108
+
109
+ # Optional: Save the DataFrame to a CSV file
110
+ df_emotions.to_csv('emotion_scores.csv', index=False)
111
+ st.success('DataFrame generated and saved as emotion_scores.csv')
112
+
113
+ with open('emotion_scores.csv', 'r') as f:
114
+ csv_file = f.read()
115
+
116
+ st.download_button(
117
+ label='Download Emotion Scores as CSV',
118
+ data=csv_file,
119
+ file_name='emotion_scores.csv',
120
+ mime='text/csv',
121
+ )
122
+
123
+ st.success('DataFrame generated and available for download.')
124
+
125
+ else:
126
+ st.error("No results to generate DataFrame. Please predict emotions first.")