Update visualization.py

visualization.py

@@ -241,28 +241,26 @@ def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, mse_v
     mse_voice = np.interp(np.linspace(0, len(mse_voice) - 1, total_frames), 
                           np.arange(len(mse_voice)), mse_voice)
     
+    # Normalize the MSE values
     mse_embeddings_norm = (mse_embeddings - np.min(mse_embeddings)) / (np.max(mse_embeddings) - np.min(mse_embeddings))
     mse_posture_norm = (mse_posture - np.min(mse_posture)) / (np.max(mse_posture) - np.min(mse_posture))
     mse_voice_norm = (mse_voice - np.min(mse_voice)) / (np.max(mse_voice) - np.min(mse_voice))
     
-    combined_mse = np.zeros((3, total_frames))
+    combined_mse = np.full((3, total_frames), np.nan)
     combined_mse[0] = mse_embeddings_norm
     combined_mse[1] = mse_posture_norm
     combined_mse[2] = mse_voice_norm
 
+    # Create custom colormap
     cdict = {
-        '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)]
+        'red':   [(0.0,  0.5, 0.5), (1.0,  1.0, 1.0)],
+        'green': [(0.0,  0.5, 0.5), (1.0,  0.0, 0.0)],
+        'blue':  [(0.0,  0.5, 0.5), (1.0,  0.0, 0.0)],
     }
-
     custom_cmap = LinearSegmentedColormap('custom_cmap', segmentdata=cdict, N=256)
     
     fig, ax = plt.subplots(figsize=(width/100, 2))
-    im = ax.imshow(combined_mse, aspect='auto', cmap=custom_cmap, extent=[0, total_frames, 0, 3])
+    im = ax.imshow(combined_mse, aspect='auto', cmap=custom_cmap, extent=[0, total_frames, 0, 3], vmin=0, vmax=1)
     ax.set_yticks([0.5, 1.5, 2.5])
     ax.set_yticklabels(['Face', 'Posture', 'Voice'])
     ax.set_xticks([])
@@ -273,7 +271,7 @@ def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, mse_v
         
         if hasattr(ax, 'line'):
             ax.lines.pop(0)
-        ax.axvline(x=frame_count, color='r', linewidth=2)
+        ax.axvline(x=frame_count, color='blue', linewidth=2)
         
         canvas = FigureCanvasAgg(fig)
         canvas.draw()
@@ -304,8 +302,8 @@ def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, mse_v
     final_clip = VideoClip(combine_video_and_heatmap, duration=video.duration)
     final_clip = final_clip.set_audio(video.audio)
 
-    # Write the final video
-    final_clip.write_videofile(heatmap_video_path, codec='libx264', audio_codec='aac', fps=video.fps)
+    # Write the final video with progress bar
+    final_clip.write_videofile(heatmap_video_path, codec='libx264', audio_codec='aac', fps=video.fps, progress_bar=progress)
     
     # Close the video clips
     video.close()