tapvid-点跟踪数据集学习-处理

wahahoo

于 2025-06-23 10:23:18 发布

阅读量723

点赞数 10

CC 4.0 BY-SA版权

文章标签：学习人工智能目标跟踪数据集深度学习

本文链接：https://blog.csdn.net/qq_60199131/article/details/148840602

TAP-Vid 是一个用于评估视频中任意点跟踪性能的数据集和基准，由 Google DeepMind 和牛津大学的研究人员创建。

地址：https://github.com/google-deepmind/tapnet/tree/main/tapnet/tapvid

数据集构成

TAP-Vid 数据集包含真实世界视频和合成视频，具体分为以下几个子集：

TAP-Vid-DAVIS：包含 30 个来自 DAVIS 验证集的视频，这些视频经过精心标注和多次迭代优化，以确保标注质量。
TAP-Vid-Kinetics：包含 1000 个来自 Kinetics 验证集的视频，这些视频从 YouTube 获取，标注经过严格监控、筛选和优化。
TAP-Vid-Kubric：包含大量合成视频，这些视频来自 Kubric 数据集，提供了完美的点轨迹标注，主要用于模型训练。
TAP-Vid-RGB-Stacking：包含 50 个合成视频，用于评估。

数据集特点

点级精度：与以往的边界框或分割跟踪不同，TAP-Vid 要求算法能够精确跟踪视频中任意可跟踪点。
长期跟踪：与光流方法不同，TAP-Vid 要求算法能够在长视频片段中跟踪点的运动。
类别无关性：与人类关键点跟踪等特定类别任务不同，TAP-Vid 要求算法能够跟踪任何物体表面的点。
标注质量高：真实世界视频的标注经过人工多次审核和优化，确保标注的准确性和可靠性。
TAP-Vid-DAVIS：包含 30 个来自 DAVIS 验证集的视频，这些视频经过精心标注和多次迭代优化，以确保标注质量。
TAP-Vid-Kinetics：包含 1000 个来自 Kinetics 验证集的视频，这些视频从 YouTube 获取，标注经过严格监控、筛选和优化。
TAP-Vid-Kubric：包含大量合成视频，这些视频来自 Kubric 数据集，提供了完美的点轨迹标注，主要用于模型训练。
TAP-Vid-RGB-Stacking：包含 50 个合成视频，用于评估。

1、以tapvid_davis.pkl为例，查看并处理数据集
这段代码的目的是探tapvid_davis.pkl 文件的内部数据结构，特别是针对 'car-roundabout' 这个视频片段。

数据结构信息:
数据类型: <class 'dict'>

字典键 (视频ID): ['goat', 'car-roundabout', 'motocross-jump', 'breakdance', 'drift-chicane', 'drift-straight', 'judo', 'soapbox', 'dogs-jump', 'parkour', 'india', 'pigs', 'cows', 'gold-fish', 'paragliding-launch', 'camel', 'blackswan', 'dog', 'bike-packing', 'shooting', 'lab-coat', 'kite-surf', 'bmx-trees', 'dance-twirl', 'car-shadow', 'libby', 'scooter-black', 'mbike-trick', 'loading', 'horsejump-high']

正在检查视频ID: car-roundabout
视频数据键: ['points', 'occluded', 'video']

键 'points':
  类型: <class 'numpy.ndarray'>
  形状: (20, 75, 2)

键 'occluded':
  类型: <class 'numpy.ndarray'>
  形状: (20, 75)

键 'video':
  类型: <class 'numpy.ndarray'>
  形状: (75, 480, 854, 3)

数据结构是一个 Python 字典 (dict)。包含'goat', 'car-roundabout'等视频片段，例如ID: car-roundabout中

视频数据键有三个类型 ['points', 'occluded', 'video']，

（1）'points' 键对应的数据是一个 NumPy 数组。形状是 (20, 75, 2)

20：表示这个视频跟踪了 20 个不同的点。

75：表示这个视频总共有 75 帧。

2：表示每个点的坐标由 2 个值（x，y）组成。

（2）'occluded' 数组的形状是 (20, 75)，20个点，75帧。每个点在每一帧是否被遮挡的布尔值（True/False）或 0/1 值。

（3）video 数组的形状是 (75, 480, 854, 3)，

75：表示这个视频总共有 75 帧。

480：表示每帧图像的高度是 480 像素。

854：表示每帧图像的宽度是 854 像素。

3：表示图像是3个颜色通道（通常是RGB或BGR）。

（突然发现apvid_rgb_stacking.pkl数据结构信息:数据类型: <class 'list'>是个列表）

import pickle
def check_pkl_structure(pkl_path, video_id_to_check=None):
    """
    检查pickle文件的数据结构
    Args:
        pkl_path: pickle文件路径
        video_id_to_check: 要检查的特定视频ID，None表示打印所有视频的概览
    """
    print(f"正在读取文件: {pkl_path}")
    with open(pkl_path, 'rb') as f:
        data = pickle.load(f)
    
    print("\n数据结构信息:")
    print(f"数据类型: {type(data)}")
    
    if isinstance(data, dict):
        print(f"\n字典键 (视频ID): {list(data.keys())}")
        
        if video_id_to_check:
            if video_id_to_check in data:
                video_data = data[video_id_to_check]
                print(f"\n正在检查视频ID: {video_id_to_check}")
                print(f"视频数据键: {list(video_data.keys())}")
                
                for key, value in video_data.items():
                    print(f"\n键 '{key}':")
                    print(f"  类型: {type(value)}")
                    if hasattr(value, 'shape'):
                        print(f"  形状: {value.shape}")
                    elif isinstance(value, list):
                        print(f"  长度: {len(value)}")
                        if len(value) > 0:
                            print(f"  第一个元素类型: {type(value[0])}")
                            if hasattr(value[0], 'shape'):
                                print(f"  第一个元素形状: {value[0].shape}")
                    elif isinstance(value, dict):
                        print(f"  字典键: {list(value.keys())}")
                        for sub_key, sub_value in value.items():
                            print(f"    子键 '{sub_key}':")
                            print(f"      类型: {type(sub_value)}")
                            if hasattr(sub_value, 'shape'):
                                print(f"      形状: {sub_value.shape}")
            else:
                print(f"视频ID '{video_id_to_check}' 未找到。")
        else:
            # 如果没有指定视频ID，则打印每个视频的概览
            print("\n所有视频概览:")
            for video_id, video_data in data.items():
                print(f"视频ID: {video_id}, 键: {list(video_data.keys())}")
                if 'points' in video_data and hasattr(video_data['points'], 'shape'):
                    print(f"  Points shape: {video_data['points'].shape}")
if __name__ == "__main__":
    pkl_path = "datasets/tapvid_davis/tapvid_davis.pkl"
    # 检查特定视频，例如 'car-roundabout'
    check_pkl_structure(pkl_path, video_id_to_check='car-roundabout')

2、可视化部分视频片段

读取 TAP-Vid 数据集中的视频和跟踪点信息，然后将每一帧可视化，并在图像上绘制跟踪点及其相关信息，最后保存为图片。

正在读取文件: datasets\tapvid_davis\tapvid_davis_val.pkl
将处理前 2 个视频: ['goat', 'car-roundabout']
处理视频序列:   0%|            | 0/2 [00:00<?, ?it/s] 已保存视频 goat 的所有帧到 output/tapvid_visualization\goat
处理视频序列:  50%|██  | 1/2 [00:00<00:00,  4.34it/s]
DEBUG: car-roundabout - Frame 3 - Image Shape: (480, 854, 3)
DEBUG: car-roundabout - Frame 3 - Non-Occluded Point Coordinates:
  Point 0: (0.7731770873069763, 0.3967592716217041)
  Point 1: (0.7579905986785889, 0.41443660855293274)
  Point 2: (0.7809486389160156, 0.41652432084083557)
  Point 3: (0.8156262636184692, 0.40799999237060547)
  Point 4: (0.6892624497413635, 0.47359153628349304)
  Point 10: (0.5361979007720947, 0.45787036418914795)
  Point 12: (0.5257812738418579, 0.44398146867752075)
  Point 14: (0.5315104126930237, 0.4393518567085266)
  Point 15: (0.5283854007720947, 0.4810185134410858)
  Point 16: (0.4611979126930237, 0.5217592716217041)
  Point 17: (0.35494792461395264, 0.5134259462356567)
  Point 18: (0.3304687440395355, 0.6393518447875977)
已保存视频 car-roundabout 的所有帧到 output/tapvid_visualization\car-roundabout

比如第三帧中，Point 0: (0.773..., 0.396...) 表示归一化后的坐标，第 0 个点在图像宽度方向的 77.3% 处，高度方向的 39.6% 处。

白色文字：视频ID和帧号（左上角）

绿色圆点：跟踪点位置

绿色文字：跟踪点标签

红色文字：遮挡信息（底部）

import pickle
import os
import cv2
import numpy as np
from tqdm import tqdm
import argparse
def visualize_tapvid_frames(pkl_path, output_dir, num_videos=None):
    """
    读取tapvid pickle文件并可视化每一帧
    
    Args:
        pkl_path: pickle文件路径
        output_dir: 输出目录路径
        num_videos: 要处理的视频数量，None表示处理所有视频
    """
    # 创建输出目录
    os.makedirs(output_dir, exist_ok=True)
    
    # 读取pickle文件
    print(f"正在读取文件: {pkl_path}")
    with open(pkl_path, 'rb') as f:
        data = pickle.load(f)
    
    # 获取要处理的视频列表
    video_ids = list(data.keys())
    if num_videos is not None:
        video_ids = video_ids[:num_videos]
        print(f"将处理前 {num_videos} 个视频: {video_ids}")
    
    # 遍历每个视频序列
    for video_id in tqdm(video_ids, desc="处理视频序列"):
        video_data = data[video_id]
        # 为每个视频创建子目录
        video_dir = os.path.join(output_dir, video_id)
        os.makedirs(video_dir, exist_ok=True)
        
        # 获取视频帧
        frames = video_data['video']  # 视频帧存储在'video'键下
        points = video_data['points']  # 跟踪点存储在'points'键下
        occluded = video_data['occluded']  # 遮挡信息存储在'occluded'键下
        
        # 保存每一帧
        for frame_idx, frame in enumerate(frames):
            # 确保帧是uint8类型
            if frame.dtype != np.uint8:
                frame = (frame * 255).astype(np.uint8)
            
            # 在帧上绘制跟踪点和标签信息
            frame_with_points = frame.copy()
            
            # 获取图像的宽度和高度
            height, width, _ = frame.shape
            # 添加视频ID和帧号信息
            cv2.putText(frame_with_points, f"Video: {video_id}", (10, 30), 
                       cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2)
            cv2.putText(frame_with_points, f"Frame: {frame_idx}", (10, 70), 
                       cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2)
            
            # --- DEBUG START ---
            if video_id == 'car-roundabout' and frame_idx == 3:
                print(f"\nDEBUG: {video_id} - Frame {frame_idx} - Image Shape: {frame.shape}")
                print(f"DEBUG: {video_id} - Frame {frame_idx} - Non-Occluded Point Coordinates:")
                for p_idx in range(points.shape[0]):
                    if not occluded[p_idx, frame_idx]:
                        x, y = points[p_idx, frame_idx]
                        print(f"  Point {p_idx}: ({x}, {y})")
            # --- DEBUG END ---
            # 绘制跟踪点和标签
            for point_idx in range(points.shape[0]):
                if not occluded[point_idx, frame_idx]:  # 只绘制未被遮挡的点
                    # 将归一化坐标转换为像素坐标
                    x_norm, y_norm = points[point_idx, frame_idx]
                    x_pixel = int(x_norm * width)
                    y_pixel = int(y_norm * height)
                    # 绘制跟踪点
                    cv2.circle(frame_with_points, (x_pixel, y_pixel), 3, (0, 255, 0), -1)
                    # 添加点标签
                    cv2.putText(frame_with_points, f"P{point_idx}", 
                              (x_pixel + 5, y_pixel - 5), 
                              cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1)
            
            # 添加遮挡信息
            occluded_text = "Occluded points: "
            occluded_points_list = [f"P{i}" for i in range(points.shape[0]) if occluded[i, frame_idx]]
            if occluded_points_list:
                occluded_text += ", ".join(occluded_points_list)
            else:
                occluded_text += "None"
            
            # 在图像底部添加遮挡信息
            cv2.putText(frame_with_points, occluded_text, (10, frame.shape[0] - 20), 
                       cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 1)
            
            # 保存帧
            frame_path = os.path.join(video_dir, f"frame_{frame_idx:04d}.jpg")
            cv2.imwrite(frame_path, frame_with_points)
            
        print(f"已保存视频 {video_id} 的所有帧到 {video_dir}")
if __name__ == "__main__":
    # 创建命令行参数解析器
    parser = argparse.ArgumentParser(description='可视化TAP-Vid数据集')
    parser.add_argument('--num_videos', type=int, default=2,
                      help='要处理的视频数量（默认：2）')
    parser.add_argument('--output_dir', type=str, default='output/tapvid_visualization',
                      help='输出目录路径（默认：output/tapvid_visualization）')
    args = parser.parse_args()
    
    # 设置输入输出路径
    pkl_path = r"datasets\tapvid_davis\tapvid_davis_val.pkl"
    
    # 执行可视化
    visualize_tapvid_frames(pkl_path, args.output_dir, args.num_videos)

3、在原tapvid_davis_val.pkl数据集中提取几个片段作为小数据使用

从一个大型的 TAP-Vid DAVIS 数据集中提取一个较小的子集，并将其保存为一个新的 .pkl 文件。这个子集通常用于快速测试或开发，避免处理整个庞大的数据集。

正在加载原始数据集...
正在提取视频: ['goat', 'car-roundabout']
视频 goat 信息:
- 帧数: 90
- 分辨率: (480, 854)
- 跟踪点数: 5
-------------------
视频 car-roundabout 信息:
- 帧数: 75
- 分辨率: (480, 854)
- 跟踪点数: 20
-------------------
正在保存子集到 datasets/tapvid_davis/tapvid_davis_val.pkl
完成!


import pickle
import os
import numpy as np
from PIL import Image
import io
def decode_frame(frame):
    """解码JPEG字节流为numpy数组"""
    if isinstance(frame, bytes):
        byteio = io.BytesIO(frame)
        img = Image.open(byteio)
        return np.array(img)
    return frame
def main():
    # 输入和输出路径
    input_path = 'datasets/tapvid_davis/tapvid_davis.pkl'
    output_path = 'datasets/tapvid_davis/tapvid_davis_val.pkl'
    
    # 确保输出目录存在
    os.makedirs(os.path.dirname(output_path), exist_ok=True)
    
    print("正在加载原始数据集...")
    with open(input_path, 'rb') as f:
        data = pickle.load(f)
    
    # 获取前两个视频的键
    video_keys = list(data.keys())[:2]
    subset_data = {}
    
    print(f"正在提取视频: {video_keys}")
    for key in video_keys:
        video_data = data[key]
        
        # 提取视频帧
        frames = video_data['video']
        if isinstance(frames[0], bytes):
            frames = np.array([decode_frame(frame) for frame in frames])
        
        # 提取对应的点和遮挡信息
        points = video_data['points']
        occluded = video_data['occluded']
        
        # 存储到新的字典中
        subset_data[key] = {
            'video': frames,
            'points': points,
            'occluded': occluded
        }
        
        print(f"视频 {key} 信息:")
        print(f"- 帧数: {len(frames)}")
        print(f"- 分辨率: {frames.shape[1:3]}")
        print(f"- 跟踪点数: {points.shape[0]}")
        print("-------------------")
    
    print(f"正在保存子集到 {output_path}")
    with open(output_path, 'wb') as f:
        pickle.dump(subset_data, f)
    
    print("完成!")
if __name__ == '__main__':
    main()