【3DCV】Re10K数据集：抽帧处理、数据类构造

1. 视频下载

import os
import yt_dlp

def download_youtube_video(url, output_name="re10k_video"):
    # 配置参数：核心是outtmpl
    ydl_opts = {
        # outtmpl指定保存路径和文件名，%(ext)s会自动替换为视频格式（mp4/webm等）
        'outtmpl': f're10k/videos/train/{output_name}.%(ext)s',
        # 可选：指定下载格式（优先mp4），避免格式混乱
        'format': 'bestvideo[ext=mp4]+bestaudio[ext=m4a]/best[ext=mp4]/best',
    }
    
    with yt_dlp.YoutubeDL(ydl_opts) as ydl:
        # 下载并返回视频信息（可选）
        video_info = ydl.extract_info(url, download=True)
        print(f"下载完成！文件名为：{output_name}.{video_info['ext']}")


train_txt_dir = 're10k\\txt\\train'
file_ids = [f.split('.')[0] for f in os.listdir(train_txt_dir)]
file_paths = [f're10k\\txt\\train\\{f}'for f in os.listdir(train_txt_dir)]
for idx, file in enumerate(file_paths[:1]): # 测试，只下载一个视频
    with open(file, 'r') as f:
        lines = f.readlines()
        vedio_url = lines[0].split(' ')[0]
        print("正在下载：", vedio_url)
        download_youtube_video(vedio_url, file_ids[idx]) 

1.1 使用 yt-dlp 默认下载的代码无法播放

当你使用最简单的 ydl.download([url]) 时，yt-dlp 为了给你最高画质，通常会下载一个 .webm 格式的视频轨道和一个音频轨道。如果没有正确合并，或者你的播放器不支持该格式，就会出现无法播放、只有声音没画面或只有画面没声音的情况。

以下是解决方法：

1. 核心原因：缺少 FFmpeg
   yt-dlp 本身不具备合并音视频的能力，它依赖一个名为 FFmpeg 的外部工具。

• 现象：下载完成后，文件夹里可能出现了两个文件（一个视频，一个音频），或者一个损坏的 .webm 文件。

• 解决方法：

  1. 去 FFmpeg 官网 下载对应系统的二进制文件。

  2. 解压并将 bin 文件夹路径添加到系统的 环境变量 (Path) 中。

  3. 安装好后，在终端输入 ffmpeg -version 确认成功。

  4. 有了 FFmpeg 后，yt-dlp 会自动将音视频合并为一个完整的 .mp4 或 .mkv。

2. 视频抽帧

import cv2
import os
import glob

def extract_frames(txt_path, video_dir, output_dir):
    # 1. 获取对应的视频文件路径
    # 假设 txt 文件名与视频文件名一致（例如 0a0a998c176713fd.mp4）
    video_id = os.path.splitext(os.path.basename(txt_path))[0]
    video_path = os.path.join(video_dir, f"{video_id}.mp4")
    
    if not os.path.exists(video_path):
        print(f"跳过：未找到视频文件 {video_path}")
        return

    # 2. 创建该视频对应的图片保存目录
    video_output_dir = os.path.join(output_dir, video_id)
    os.makedirs(video_output_dir, exist_ok=True)

    # 3. 读取视频
    cap = cv2.VideoCapture(video_path)
    
    # 4. 解析 txt 文件并抽帧
    with open(txt_path, 'r') as f:
        lines = f.readlines()
        
    print(f"正在处理视频: {video_id}, 计划提取 {len(lines)} 帧")

    for line in lines[1:]: # 跳过第一行url
        parts = line.split()
        if not parts: continue
        
        # 第一列是时间戳（通常为微秒）
        timestamp_us = int(parts[0])
        # 转换为毫秒 (ms)，因为 OpenCV 设置位置使用毫秒
        timestamp_ms = timestamp_us / 1000.0
        
        # 定位到指定时间戳
        cap.set(cv2.CAP_PROP_POS_MSEC, timestamp_ms)
        ret, frame = cap.read()
        
        if ret:
            # 按时间戳命名保存图片
            img_name = f"{timestamp_us}.jpg"
            save_path = os.path.join(video_output_dir, img_name)
            cv2.imwrite(save_path, frame)
        else:
            print(f"无法读取时间戳 {timestamp_us} 处的帧")

    cap.release()

txt_dir = r're10k/txt/train'  # txt 所在目录
video_dir = r're10k/videos/train'          # 视频所在目录
output_base_dir = r're10k/images/train'    # 图片保存目录

# 获取所有 txt 文件
txt_files = glob.glob(os.path.join(txt_dir, "*.txt"))

for txt_file in txt_files:
    extract_frames(txt_file, video_dir, output_base_dir)

关键细节说明：

• 时间戳转换 ：RE10K 的原始数据通常以微秒记录。而 cv2.CAP_PROP_POS_MSEC 接收的是毫秒，所以代码中进行了 timestamp_us / 1000.0 的转换。

3. 数据类的构造

import torch
from torch.utils.data import Dataset
import numpy as np
import os
from PIL import Image
from torchvision import transforms

class RE10KDataset(Dataset):
    def __init__(self, image_root, text_root, transform=None):
        """
        Args:
            image_root (str): 抽帧后图片存放的根目录 (例如 './image')
            text_root (str): 原始 txt 文件存放目录 (例如 './data/re10k/train')
            transform: 对单张图片的 torchvision transforms 处理
        """
        self.image_root = image_root
        self.text_root = text_root
        
        # 如果没有传入 transform，默认转换为 Tensor
        if transform is None:
            self.transform = transforms.Compose([
                transforms.ToTensor(), # 自动将 [0, 255] 转为 [0.0, 1.0] 且维度变为 [C, H, W]
            ])
        else:
            self.transform = transform
        
        # 获取所有场景 ID
        self.scene_ids = sorted([os.path.splitext(f)[0] for f in os.listdir(text_root) if f.endswith('.txt')])

    def __len__(self):
        return len(self.scene_ids)

    def __getitem__(self, scene_id):
        txt_path = os.path.join(self.text_root, f"{scene_id}.txt")
        scene_img_dir = os.path.join(self.image_root, scene_id)

        images = []
        intrinsics = []
        distortions = []
        poses = []
        timestamps = []

        if not os.path.exists(txt_path):
            return None

        with open(txt_path, 'r') as f:
            lines = f.readlines()[1:] # 跳过第一行url
            for line in lines:
                data = list(map(float, line.split()))
                if not data: continue
                
                ts_raw = data[0]
                timestamp_int = int(ts_raw)
                img_path = os.path.join(scene_img_dir, f"{timestamp_int}.jpg")
                
                # 检查图片是否存在
                if not os.path.exists(img_path):
                    continue

                # 1. 处理图像：打开 -> 转换 -> Tensor
                img = Image.open(img_path).convert('RGB')
                img_tensor = self.transform(img) # 这里执行了 ToTensor()
                
                # 2. 内参矩阵 K (3x3)
                # data[1:5] 分别是 fx, fy, cx, cy
                K = torch.tensor([
                    [data[1], 0,       data[3]],
                    [0,       data[2], data[4]],
                    [0,       0,       1]
                ], dtype=torch.float32)

                # 3. 畸变参数 (k1, k2)
                dist = torch.tensor([data[5], data[6]], dtype=torch.float32)

                # 4. 外参矩阵 (Camera-to-World, 3x4)
                # data[7:16] 是旋转矩阵 R, data[16:19] 是平移 T
                R = torch.tensor(data[7:16]).view(3, 3)
                T = torch.tensor(data[16:19]).view(3, 1)
                c2w = torch.cat([R, T], dim=1) 

                images.append(img_tensor)
                intrinsics.append(K)
                distortions.append(dist)
                poses.append(c2w)
                timestamps.append(ts_raw)

        # 将列表转换为堆叠的 Tensor
        return {
            "scene_id": scene_id,
            "images": torch.stack(images),         # [N, 3, H, W]
            "intrinsics": torch.stack(intrinsics),   # [N, 3, 3]
            "distortions": torch.stack(distortions), # [N, 2]
            "poses": torch.stack(poses),             # [N, 3, 4]
            "timestamps": torch.tensor(timestamps)   # [N]
        }

txt_dir = r're10k/txt/train'  # txt 所在目录
image_dir = r're10k/images/train'    # 图片保存目录
re10kdataset = RE10KDataset(image_dir, txt_dir)
scene_id = '0000cc6d8b108390'
re10kdataset[scene_id] # 查看某个场景的数据