Update visualization.py

visualization.py

@@ -227,13 +227,17 @@ def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, outpu
     frame_count = 0
     total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
     
+    line = None  # Initialize line object
+    
     while True:
         ret, frame = cap.read()
         if not ret:
             break
         
         # Update the heatmap with the current frame position
-        ax.axvline(x=frame_count, color='r', linewidth=2)
+        if line:
+            line.remove()  # Remove the previous line
+        line = ax.axvline(x=frame_count, color='r', linewidth=2)
         
         # Convert the matplotlib figure to an image
         canvas = FigureCanvasAgg(fig)
@@ -251,9 +255,6 @@ def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, outpu
         
         out.write(combined_frame)
         frame_count += 1
-        
-        # Remove the vertical line for the next iteration
-        ax.lines.pop()
     
     cap.release()
     out.release()