2. 视频字幕对齐与数据集构建

任务说明

构建一个视频-字幕对齐数据集,将原始视频文件与对应的字幕文件(SRT格式)进行时间轴对齐,清洗无效条目,最终生成视频语言多模态模型训练所需的JSON标注文件。每个标注条目包含视频路径、字幕文本、起始时间与结束时间。

环境依赖

需要 Python 3.8+,无需第三方库(仅使用标准库 json、re、pathlib)。字幕文件需为UTF-8编码的SRT格式。

数据集目录结构

dataset/
├── videos/
│   ├── video_001.mp4
│   ├── video_002.mp4
│   └── ...
└── subtitles/
    ├── video_001.srt
    ├── video_002.srt
    └── ...

完整实现代码

import json
import re
from pathlib import Path

def parse_srt(srt_path):
    """解析SRT文件,返回[(start_ms, end_ms, text), ...]"""
    raw = Path(srt_path).read_text(encoding='utf-8')
    # 按空行分割字幕块
    blocks = re.split(r'\n\s*\n', raw.strip())
    entries = []
    for block in blocks:
        lines = [l.strip() for l in block.split('\n') if l.strip()]
        if len(lines) < 2:
            continue
        # 跳过序号行(纯数字)
        if lines[0].isdigit():
            lines = lines[1:]
        if not lines:
            continue
        # 解析时间行: 00:00:01,000 --> 00:00:04,000
        time_match = re.match(
            r'(\d{2}):(\d{2}):(\d{2}),(\d{3})\s*-->\s*'
            r'(\d{2}):(\d{2}):(\d{2}),(\d{3})',
            lines[0]
        )
        if not time_match:
            continue
        t = time_match.groups()
        start_ms = (int(t[0])*3600 + int(t[1])*60 + int(t[2]))*1000 + int(t[3])
        end_ms = (int(t[4])*3600 + int(t[5])*60 + int(t[6]))*1000 + int(t[7])
        text = ' '.join(lines[1:]).strip()
        if text and end_ms > start_ms:
            entries.append((start_ms, end_ms, text))
    return entries

def build_dataset(video_dir, subtitle_dir, output_json):
    video_dir = Path(video_dir)
    subtitle_dir = Path(subtitle_dir)
    dataset = []
    skipped = []

    for srt_file in sorted(subtitle_dir.glob('*.srt')):
        video_file = video_dir / (srt_file.stem + '.mp4')
        if not video_file.exists():
            # 尝试其他常见视频扩展名
            candidates = list(video_dir.glob(srt_file.stem + '.*'))
            video_file = candidates[0] if candidates else None
        if video_file is None:
            skipped.append(f'{srt_file.name}: 未找到对应视频文件')
            continue

        entries = parse_srt(srt_file)
        if not entries:
            skipped.append(f'{srt_file.name}: 无有效字幕条目')
            continue

        for start_ms, end_ms, text in entries:
            dataset.append({
                'video_path': str(video_file),
                'subtitle_text': text,
                'start_ms': start_ms,
                'end_ms': end_ms
            })

    with open(output_json, 'w', encoding='utf-8') as f:
        json.dump(dataset, f, ensure_ascii=False, indent=2)

    print(f'共生成 {len(dataset)} 条对齐样本')
    if skipped:
        print(f'跳过 {len(skipped)} 个文件:')
        for msg in skipped:
            print(f'  - {msg}')

if __name__ == '__main__':
    build_dataset('dataset/videos', 'dataset/subtitles', 'dataset/annotations.json')

运行方式

python build_dataset.py

将上述代码保存为 build_dataset.py,与 dataset 目录同级,直接运行即可。

验证方法与通过条件

  1. 运行无报错:脚本正常结束,控制台输出“共生成 N 条对齐样本”。
  2. 检查JSON结构:打开 dataset/annotations.json,确认每条记录包含 video_path、subtitle_text、start_ms、end_ms 四个字段。
  3. 抽样核对:随机选取一条记录,用播放器打开对应视频,跳转到 start_ms/1000 秒处,确认字幕文本与视频画面内容一致。
  4. 时间轴正确性:确认 end_ms 大于 start_ms,且时间值在视频总时长范围内。

若以上条件全部满足,则数据集构建成功,可直接用于后续模型训练的数据加载。