Update visualization.py

visualization.py

@@ -208,7 +208,7 @@ def plot_posture(df, posture_scores, color='blue', anomaly_threshold=3):
     return fig
 
 
-def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, mse_voice, output_path, desired_fps, largest_cluster):
+def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, mse_voice, output_path, desired_fps, largest_cluster, progress=None):
     # Filter the DataFrame to only include frames from the largest cluster
     df_largest_cluster = df[df['Cluster'] == largest_cluster]
     
@@ -247,32 +247,30 @@ def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, mse_v
     combined_mse[0] = mse_embeddings_norm
     combined_mse[1] = mse_posture_norm
     combined_mse[2] = mse_voice_norm
-    
+
     # Custom colormap definition
     cdict = {
-        'red':   [(0.0,  0.5, 0.5),  # Low MSE: 50% red (gray)
-                  (1.0,  1.0, 1.0)], # High MSE: Full red
-        'green': [(0.0,  0.5, 0.5),  # Low MSE: 50% green (gray)
-                  (1.0,  0.0, 0.0)], # High MSE: No green
-        'blue':  [(0.0,  0.5, 0.5),  # Low MSE: 50% blue (gray)
-                  (1.0,  0.0, 0.0)]  # High MSE: No blue
+        'red':   [(0.0,  0.0, 0.0),
+                  (1.0,  1.0, 1.0)],
+        'green': [(0.0,  1.0, 1.0),
+                  (1.0,  0.0, 0.0)],
+        'blue':  [(0.0,  1.0, 1.0),
+                  (1.0,  0.0, 0.0)]
     }
 
     custom_cmap = LinearSegmentedColormap('custom_cmap', segmentdata=cdict, N=256)
     
     fig, ax = plt.subplots(figsize=(width/100, 2))
-    # Use the custom colormap in the heatmap generation
     im = ax.imshow(combined_mse, aspect='auto', cmap=custom_cmap, extent=[0, total_frames, 0, 3])
     ax.set_yticks([0.5, 1.5, 2.5])
     ax.set_yticklabels(['Face', 'Posture', 'Voice'])
     ax.set_xticks([])
     plt.tight_layout()
     
+    if progress:
+        progress(0, desc="Generating video with heatmap")
+
     line = None
-    
-    # Add progress tracking
-    progress(0.9, desc="Generating video with heatmap")
-    
     for frame_count in range(total_frames):
         cap.set(cv2.CAP_PROP_POS_FRAMES, frame_count)
         ret, frame = cap.read()
@@ -281,7 +279,7 @@ def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, mse_v
         
         if line:
             line.remove()
-        line = ax.axvline(x=frame_count, color='blue', linewidth=3)
+        line = ax.axvline(x=frame_count, color='r', linewidth=2)
         
         canvas = FigureCanvasAgg(fig)
         canvas.draw()
@@ -289,9 +287,6 @@ def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, mse_v
         heatmap_img = heatmap_img.reshape(canvas.get_width_height()[::-1] + (3,))
         heatmap_img = cv2.resize(heatmap_img, (width, 200))
         
-        # Convert heatmap_img from RGB to BGR
-        heatmap_img = cv2.cvtColor(heatmap_img, cv2.COLOR_RGB2BGR)
-        
         combined_frame = np.vstack((frame, heatmap_img))
         
         seconds = frame_count / original_fps
@@ -299,9 +294,9 @@ def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, mse_v
         cv2.putText(combined_frame, f"Time: {timecode}", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2)
         
         out.write(combined_frame)
-        
-        # Update progress
-        progress(0.9 + (0.1 * (frame_count + 1) / total_frames), desc="Generating video with heatmap")
+
+        if progress:
+            progress((frame_count + 1) / total_frames, desc="Generating video with heatmap")
     
     cap.release()
     out.release()