Update app.py

app.py

@@ -1,159 +1,54 @@
-import streamlit as st
-from PIL import Image
-from transformers import pipeline
-import pandas as pd
-import matplotlib.pyplot as plt
-import seaborn as sns
-import xlsxwriter
-import io
-
-# 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]
-
-# Initialize an empty DataFrame outside of the button press condition
-df_emotions = pd.DataFrame()
+# ... [rest of your code] ...
 
 # Generate DataFrame from results
 if st.button("Generate HeatMap & 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 = {
-                'Happy': 0,
-                'Surprise': 0,
-                'Neutral': 0,
-                'Sad': 0,
-                'Disgust': 0,
-                'Angry': 0,
-                'Fear': 0,
-                'Image Name': image_names[i],
-                'Image Size (KB)': f"{image_sizes[i]:.1f}"  # Format the size to one decimal place
-            }
-            
-            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)
-
-        # Add a placeholder for the 'Image View' column
-        df_emotions['Image View'] = [''] * len(df_emotions)
-
-        # Display the DataFrame
-        st.write(df_emotions)
-
-        # Plotting the heatmap for the first seven columns
-        plt.figure(figsize=(10, 10))
-        sns.heatmap(df_emotions.iloc[:, :7], annot=True, fmt=".1f", cmap='viridis')
-        plt.title('Heatmap of Emotion Scores')
-        plt.xlabel('Emotion Categories')
-        plt.ylabel('Data Points')
-        st.pyplot(plt)
-
-        # Save the DataFrame to a CSV file without the 'Image View' column
-        df_emotions.drop(columns=['Image View']).to_csv('emotion_scores.csv', index=False)
-        st.success('DataFrame generated and saved as emotion_scores.csv')
-
-        with open('emotion_scores.csv', 'r') as f:
-            csv_file = f.read()
-
-        st.download_button(
-            label='Download Emotion Scores as CSV',
-            data=csv_file,
-            file_name='emotion_scores.csv',
-            mime='text/csv',
-        )
-
-        # Create a BytesIO buffer for the Excel file
-        output = io.BytesIO()
-
-        # Create a new Excel writer object using the buffer as the file
-        writer = pd.ExcelWriter(output, engine='xlsxwriter')
-        df_emotions.to_excel(writer, index=False, header=True)
-
-        # Access the xlsxwriter workbook and worksheet objects
-        workbook  = writer.book
-        worksheet = writer.sheets['Sheet1']
-
-        # Set the column width and row height
-        worksheet.set_column('A:I', 20)  # Set width for columns A-I
-        worksheet.set_column('J:J', 30)  # Set width for column J (Image View)
-        for row_num in range(len(df_emotions) + 1):  # +1 to include the header row
-            worksheet.set_row(row_num, 38)  # Set the row height to 38
-
-        # Iterate over the images and insert them into the 'Image View' column
-        for idx, image in enumerate(selected_images):
-            # Convert the image to a format that can be inserted into Excel
-            image_stream = io.BytesIO()
-            image.save(image_stream, format='PNG')
-            image_stream.seek(0)
-            worksheet.insert_image(f'J{idx + 2}', 'image.png', {'image_data': image_stream})
-
-        # Close the writer object
-        writer.close()
-
-        # Rewind the buffer
-        output.seek(0)
-
-        # Use Streamlit's download button to offer the Excel file for download
-        st.download_button(
-            label='Download Emotion Scores as Excel',
-            data=output,
-            file_name='emotion_scores.xlsx',
-            mime='application/vnd.openxmlformats-officedocument.spreadsheetml.sheet',
-        )
+    # ... [rest of your code for generating df_emotions] ...
+
+    # Create a BytesIO buffer for the Excel file
+    output = io.BytesIO()
+
+    # Create a new Excel writer object using the buffer as the file
+    writer = pd.ExcelWriter(output, engine='xlsxwriter')
+    df_emotions.to_excel(writer, index=False, header=True)
+
+    # Access the xlsxwriter workbook and worksheet objects
+    workbook  = writer.book
+    worksheet = writer.sheets['Sheet1']
+
+    # Set the column width and row height
+    worksheet.set_column('A:G', 8)  # Set width for columns A-G
+    worksheet.set_column('H:H', 22)  # Set width for column H (Image Name)   
+    worksheet.set_column('I:I', 14)  # Set width for column I (Image Size)   
+    worksheet.set_column('J:J', 12)  # Set width for column J (Image View)   
+    worksheet.set_column('K:K', 12)  # Set width for column K (Image Num)
+    for row_num in range(len(df_emotions) + 1):  # +1 to include the header row
+        worksheet.set_row(row_num, 38)  # Set the row height to 38
+
+    # Add a new column for the images and image numbers
+    df_emotions['Image View'] = ''  # This will create a placeholder column for the images
+    df_emotions['Image Num'] = range(len(df_emotions))  # This will create a sequence of numbers
+
+    # Write the DataFrame to the Excel writer object
+    df_emotions.to_excel(writer, sheet_name='Sheet1', startrow=0, startcol=0, index=False, header=True)
+
+    # Iterate over the images and insert them into the 'Image View' column
+    for idx, image in enumerate(selected_images):
+        # Convert the image to a format that can be inserted into Excel
+        image_stream = io.BytesIO()
+        image.save(image_stream, format='PNG')
+        image_stream.seek(0)
+        worksheet.insert_image(f'J{idx + 2}', 'image.png', {'image_data': image_stream})
+
+    # Close the writer object
+    writer.close()
+
+    # Rewind the buffer
+    output.seek(0)
+
+    # Use Streamlit's download button to offer the Excel file for download
+    st.download_button(
+        label='Download Emotion Scores as Excel',
+        data=output,
+        file_name='emotion_scores.xlsx',
+        mime='application/vnd.openxmlformats-officedocument.spreadsheetml.sheet',
+    )