如何用Matplotlib Animation绘制GroupBy DataFrame并叠加OpenCV视频

问题背景 #

你已经能单独绘制分组后的轨迹（用df.groupby('TrackID')遍历绘图），但想把这个逻辑整合到Matplotlib Animation里，让轨迹随着视频帧更新动态显示，目前视频能正常播放，但轨迹始终画不出来，希望用df.plot()替代空的ax.plot()来实现需求。

现有代码的核心问题 #

你的updatefig函数里用track.set_data(drawPlot())是行不通的：

1. drawPlot()返回的是GroupBy对象，不是plot()需要的x/y数据数组
2. 单个track对象只能对应一条线，但分组后是多个TrackID的轨迹，没法用一个对象承载所有线条
3. 动画更新时没有清除旧的轨迹，也没有正确把新绘制的线条加入动画渲染队列

解决方案思路 #

要在动画里实现分组轨迹绘制，需要每次更新帧时：

• 先清除轴上之前绘制的轨迹线条（保留视频背景）
• 筛选当前帧范围内的数据集，分组后遍历每个TrackID，用group.plot()或直接ax.plot()绘制轨迹
• 收集所有需要渲染的艺术家对象（视频图像+所有轨迹线条），返回给动画函数

修改后的完整代码 #

import cv2
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
import matplotlib.ticker as ticker

def getImage():
    ret, frame = cap.read()
    if not ret:
        # 处理视频读取完毕的情况，返回当前图像避免报错
        return im.get_array()
    image = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
    return image

def updatefig(*fargs):
    global frame_id
    # 1. 更新视频背景
    im.set_array(getImage())
    
    # 2. 清除之前的轨迹线条（保留imshow的图像）
    for artist in ax.lines:
        artist.remove()
    
    # 3. 筛选当前帧范围内的数据（这里保留最近45帧的轨迹，可按需调整）
    frame_list = [x for x in range(frame_id - 44, frame_id + 1) if x > 0]
    df = dataset_sorted[dataset_sorted['FrameID'].isin(frame_list)]
    
    # 4. 分组绘制每个TrackID的轨迹
    for ID, group in df.groupby('TrackID'):
        # 用df.plot()绘制，指定当前ax，设置样式和标签
        group.plot(kind='line', x='HeadX', y='HeadY', ax=ax, 
                   label=f'ID {ID}', rot=60, linewidth=2)
    
    # 5. 保持坐标范围和视频匹配（y轴翻转，因为Matplotlib和OpenCV的y轴方向相反）
    ax.set_xlim(0, 1920)
    ax.set_ylim(1080, 0)
    
    frame_id += 1
    # 返回所有需要渲染的艺术家：视频图像 + 所有轨迹线条
    return [im] + ax.lines

# 初始化数据和视频（替换成你的文件路径）
csv_file = "your_data.csv"
video_file = "your_video.mp4"

dataset = pd.read_csv(csv_file, sep=',', usecols=range(1, 9))
dataset_sorted = dataset.sort_values('FrameID')
cap = cv2.VideoCapture(video_file)

# 设置画布和轴参数
w = cap.get(cv2.CAP_PROP_FRAME_WIDTH) / 2
h = cap.get(cv2.CAP_PROP_FRAME_HEIGHT) / 2
fig, ax = plt.subplots(1, 1, figsize=(w, h))
ax.xaxis.tick_top()
x_axis = list(np.linspace(0, 1920, num=16))
y_axis = list(np.linspace(1080, 0, num=9))
ax.xaxis.set_major_locator(ticker.FixedLocator(x_axis))
ax.yaxis.set_major_locator(ticker.FixedLocator(y_axis))

frame_id = 0
# 初始化视频背景
im = plt.imshow(getImage(), animated=True)

# 创建动画：blit设为False，因为每次返回的艺术家数量可能变化（TrackID数量变化）
ani = animation.FuncAnimation(fig, updatefig, interval=33, blit=False, repeat=False)
plt.legend()
plt.show()

关键修改点说明 #

1. 清除旧轨迹：每次更新时遍历ax.lines移除旧线条，避免轨迹重叠堆积
2. 分组绘图整合：在updatefig里直接处理分组逻辑，用group.plot()绘制每个轨迹，指定ax确保画在视频所在的轴上
3. 动画参数调整：把blit设为False，因为每次返回的艺术家数量可能变化（不同帧的TrackID数量可能不同），blit=True要求返回的艺术家列表长度固定，这里不适用
4. 坐标范围锁定：每次更新都重置x/y轴范围，避免轨迹绘制导致轴范围偏移
5. 视频读取容错：添加了视频读取完毕的处理，防止报错

内容的提问来源于stack exchange，提问作者tcastrog10