Update app.py

app.py

@@ -9,11 +9,15 @@ import matplotlib.pyplot as plt
 import gradio as gr
 from fer import FER
 import cv2
+import os
 
 # Dictionaries to store emotion data over time
 text_dataDict = {"Time": [], "Emotion": [], "Confidence Score": []}
 face_dataDict = {"Time": [], "Emotion": [], "Confidence Score": []}
 
+# List of temporary files to clean up
+temp_files = []
+
 # Load model and tokenizer directly from HuggingFace
 emotionDetectModel = AutoModelForSequenceClassification.from_pretrained("borisn70/bert-43-multilabel-emotion-detection")
 tokenizer = AutoTokenizer.from_pretrained("borisn70/bert-43-multilabel-emotion-detection")  # Load tokenizer directly from model
@@ -37,8 +41,13 @@ def emotionAnalysis(message, face):
     Returns:
         tuple: (str, plt) Contains the emotion results text and the updated plot
     """
-    if (message.lower() == "quit"):
-        return "Quitting...", displayResults()
+    if (message.lower() == "finish"):
+        graph = displayResults()
+        filename = "Emotion_Timeline.png"
+        graph.savefig(filename)
+        temp_files.append(filename)
+        download = gr.DownloadButton(label="Download Emotion Timeline", value=filename, visible=True)
+        return "Quitting...", graph, download
     
     # Process text emotion
     result = pipe(message)
@@ -77,8 +86,12 @@ def emotionAnalysis(message, face):
     face_dataDict["Emotion"].append(face_emotion)  # Now face_emotion will always be a string
     face_dataDict["Confidence Score"].append(face_score)
 
+    data = displayResults()
+    file_name = "unfinishedPLT.png"
+    data.savefig(file_name)
+    dL = gr.DownloadButton(label="Download Emotion Timeline", value=file_name, visible=False)
     # Return both the text result and the updated plot
-    return f"Text: {text_emotion} | Face: {face_emotion}", displayResults()
+    return f"Text: {text_emotion} | Face: {face_emotion}", data, dL
 
 def displayResults():
     """
@@ -131,6 +144,18 @@ def process_webcam(img):
         print(f"Error processing image: {str(e)}")
         return img
 
+def cleanUp_Files():
+    """
+    Removes temporary plot files created during the application's runtime
+    """
+    for file in temp_files:
+        try:
+            if os.path.exists(file):
+                os.remove(file)
+                print(f"Cleaned up {file}")
+        except Exception as e:
+            print(f"Error cleaning up {file}: {str(e)}")
+
 '''
 2 rows, 2 columns
 column 1: inputs
@@ -158,13 +183,13 @@ with gr.Blocks(title="Emotion Reader", theme=gr.themes.Ocean()) as emotion_reade
         
     5. In the "Emotion Results" box, you will see something like "Text: (emotion) | Face: (emotion) " and the timeline will update
         
-    6. You can press "Stop" to turn off the camera or type "quit" as your message to stop the program
+    6. You can press "Stop" to turn off the camera or type "finish" as your message to be able to download your results
     """
     )
     with gr.Row():
         with gr.Column(): #user text input
             text_input = gr.Textbox(
-                label="Type your thoughts here. Type 'quit' to see final results.",
+                label="Type your thoughts here. Type 'finish' to see final results.",
                 placeholder="Enter text"
             )
             examples = gr.Examples(
@@ -177,6 +202,7 @@ with gr.Blocks(title="Emotion Reader", theme=gr.themes.Ocean()) as emotion_reade
             )
         with gr.Column(): #emotion results
             emotion_result = gr.Textbox(label="Emotion Results")
+            download_button = gr.DownloadButton(label="Download Emotion Timeline", visible=False)
     
     with gr.Row():
         with gr.Column(): #camera live feed
@@ -199,7 +225,12 @@ with gr.Blocks(title="Emotion Reader", theme=gr.themes.Ocean()) as emotion_reade
     text_input.submit(
         emotionAnalysis,
         inputs=[text_input, output_img],
-        outputs=[emotion_result, plot_output]
+        outputs=[emotion_result, plot_output, download_button]
+    )
+
+    #cleanup files
+    download_button.click(
+        cleanUp_Files
     )
 
 # Launch the interface