cot_generate/batch_add_timestamp.py

"""
批量给视频加时间戳水印 — 从 Video-R1 JSON 提取唯一视频路径,
保持原始目录结构输出到新目录, 支持多进程并行和断点续跑。

用法:
    python batch_add_timestamp.py \
        --input_json /path/to/Video-R1-260k.json \
        --video_root /path/to/video_root \
        --output_root /path/to/output_root \
        --workers 8
"""

import argparse
import json
import os
from pathlib import Path
from multiprocessing import Pool
from functools import partial

from add_timestamp_to_video import add_timestamp_to_video


def extract_unique_videos(json_path: str) -> list:
    """从 Video-R1 JSON 中提取所有唯一的视频相对路径"""
    with open(json_path, encoding="utf-8") as f:
        data = json.load(f)
    videos = set()
    for item in data:
        rel = item.get("path", "").lstrip("./")
        if rel:
            videos.add(rel)
    return sorted(videos)


def process_one(rel_path: str, video_root: str, output_root: str) -> dict:
    """处理单个视频: 加时间戳水印"""
    src = os.path.join(video_root, rel_path)
    dst = os.path.join(output_root, rel_path)

    if not os.path.exists(src):
        return {"path": rel_path, "status": "missing", "msg": f"Source not found: {src}"}

    if os.path.exists(dst):
        return {"path": rel_path, "status": "skipped", "msg": "Already exists"}

    try:
        ok = add_timestamp_to_video(src, dst)
        if ok:
            return {"path": rel_path, "status": "done"}
        else:
            return {"path": rel_path, "status": "error", "msg": "add_timestamp returned False"}
    except Exception as e:
        return {"path": rel_path, "status": "error", "msg": str(e)}


def main():
    parser = argparse.ArgumentParser(description="批量加时间戳水印 (保持目录结构)")
    parser.add_argument("--input_json", required=True,
                        help="Video-R1 格式的 JSON 文件路径")
    parser.add_argument("--video_root", required=True,
                        help="原始视频的根目录")
    parser.add_argument("--output_root", required=True,
                        help="输出水印视频的根目录 (保持原始子目录结构)")
    parser.add_argument("--workers", type=int, default=4,
                        help="并行处理进程数 (默认 4)")
    args = parser.parse_args()

    videos = extract_unique_videos(args.input_json)
    print(f"[INFO] Found {len(videos)} unique videos in {args.input_json}")

    # 统计已完成的
    already_done = sum(1 for v in videos if os.path.exists(os.path.join(args.output_root, v)))
    if already_done > 0:
        print(f"[INFO] {already_done} already timestamped, {len(videos) - already_done} remaining")

    os.makedirs(args.output_root, exist_ok=True)

    worker_fn = partial(process_one,
                        video_root=args.video_root,
                        output_root=args.output_root)

    done = 0
    errors = 0
    skipped = 0

    if args.workers <= 1:
        for i, v in enumerate(videos, 1):
            result = worker_fn(v)
            if result["status"] == "done":
                done += 1
            elif result["status"] == "skipped":
                skipped += 1
            else:
                errors += 1
                print(f"  [ERROR] {result['path']}: {result.get('msg', '')}")
            if i % 50 == 0:
                print(f"  Progress: {i}/{len(videos)} (done={done}, skipped={skipped}, errors={errors})")
    else:
        with Pool(processes=args.workers) as pool:
            for i, result in enumerate(pool.imap_unordered(worker_fn, videos), 1):
                if result["status"] == "done":
                    done += 1
                elif result["status"] == "skipped":
                    skipped += 1
                else:
                    errors += 1
                    print(f"  [ERROR] {result['path']}: {result.get('msg', '')}")
                if i % 50 == 0:
                    print(f"  Progress: {i}/{len(videos)} (done={done}, skipped={skipped}, errors={errors})")

    total = done + skipped
    print(f"\n[DONE] Processed: {done} new, {skipped} skipped, {errors} errors")
    print(f"[DONE] Total timestamped: {total}/{len(videos)}")
    print(f"[DONE] Output: {args.output_root}")


if __name__ == "__main__":
    main()