Update app.py

app.py

@@ -31,6 +31,12 @@ def process_and_show_completion(video_input_path, anomaly_threshold_input, fps,
 
         face_samples_frequent = [Image.open(path) for path in face_samples_frequent] if face_samples_frequent is not None else []
 
+        # Generate the correlation heatmap
+        correlation_heatmap = plot_correlation_heatmap(mse_embeddings, mse_posture, mse_voice)
+        
+        # Generate the 3D scatter plot
+        scatter_plot_3d = plot_3d_scatter(mse_embeddings, mse_posture, mse_voice)
+
         output = [
             exec_time, results_summary,
             df, mse_embeddings, mse_posture, mse_voice,
@@ -41,7 +47,8 @@ def process_and_show_completion(video_input_path, anomaly_threshold_input, fps,
             face_samples_frequent,
             aligned_faces_folder, frames_folder,
             mse_embeddings, mse_posture, mse_voice,
-            heatmap_video_path
+            heatmap_video_path,
+            correlation_heatmap, scatter_plot_3d
         ]
 
         return output
@@ -51,7 +58,7 @@ def process_and_show_completion(video_input_path, anomaly_threshold_input, fps,
         print(error_message)
         import traceback
         traceback.print_exc()
-        return [error_message] + [None] * 23
+        return [error_message] + [None] * 25
 
 def show_results(outputs):
     return gr.Group(visible=True)
@@ -97,6 +104,8 @@ with gr.Blocks() as iface:
 
             with gr.TabItem("Video with Heatmap"):
                 heatmap_video = gr.Video(label="Video with Anomaly Heatmap")
+                correlation_heatmap_plot = gr.Plot(label="Correlation Heatmap")
+                scatter_plot_3d_plot = gr.Plot(label="3D Scatter Plot of MSEs")
                 
     df_store = gr.State()
     mse_features_store = gr.State()
@@ -107,7 +116,7 @@ with gr.Blocks() as iface:
     mse_heatmap_embeddings_store = gr.State()
     mse_heatmap_posture_store = gr.State()
     mse_heatmap_voice_store = gr.State()
-
+    
     process_btn.click(
         process_and_show_completion,
         inputs=[video_input, anomaly_threshold, fps_slider],
@@ -121,7 +130,7 @@ with gr.Blocks() as iface:
             face_samples_most_frequent,
             aligned_faces_folder_store, frames_folder_store,
             mse_heatmap_embeddings_store, mse_heatmap_posture_store, mse_heatmap_voice_store,
-            heatmap_video
+            heatmap_video, correlation_heatmap_plot, scatter_plot_3d_plot  # Add these plots to the outputs
         ]
     ).then(
         show_results,