Update visualization.py

visualization.py

@@ -221,6 +221,10 @@ def create_video_with_heatmap(video_path, mse_heatmap_embeddings_img, mse_heatma
     fourcc = cv2.VideoWriter_fourcc(*'mp4v')
     out = cv2.VideoWriter(output_path, fourcc, desired_fps, (width, height + 400))
     
+    # Resize heatmap images to match the width of the video frames
+    mse_heatmap_embeddings_img = cv2.resize(mse_heatmap_embeddings_img, (width, 200))
+    mse_heatmap_posture_img = cv2.resize(mse_heatmap_posture_img, (width, 200))
+    
     for frame_count in range(0, total_frames):
         cap.set(cv2.CAP_PROP_POS_FRAMES, frame_count)
         ret, frame = cap.read()
@@ -241,4 +245,4 @@ def create_video_with_heatmap(video_path, mse_heatmap_embeddings_img, mse_heatma
     out.release()
     plt.close('all')
     
-    return output_path
+    return output_path