batch_make_compare.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
批量调用 make_compare.py 为多个ID生成对比视频
"""
import subprocess
import sys
import os
import argparse
import tempfile
import shutil
from pathlib import Path

def read_id_list(id_file: str) -> list:
    """读取ID列表文件，返回ID数字列表（用于make_compare.py）"""
    ids = []
    with open(id_file, 'r', encoding='utf-8') as f:
        for line in f:
            line = line.strip()
            if line and not line.startswith('#'):
                # 处理 id_xxx 格式，提取数字部分
                if line.startswith('id_'):
                    id_num = line.replace('id_', '')
                    try:
                        ids.append(int(id_num))
                    except ValueError:
                        print(f"[WARN] 无法解析ID: {line}，跳过")
                else:
                    # 如果只是数字，直接使用
                    try:
                        ids.append(int(line))
                    except ValueError:
                        print(f"[WARN] 无法解析ID: {line}，跳过")
    return ids

def find_max_frame_in_dir(base_dir: str, id_val: int) -> int:
    """查找目录中最大的帧索引"""
    id_dir = os.path.join(base_dir, f"id_{id_val}")
    if not os.path.isdir(id_dir):
        return -1
    max_frame = -1
    for fname in os.listdir(id_dir):
        if fname.endswith('.png'):
            try:
                frame_num = int(fname.replace('.png', ''))
                max_frame = max(max_frame, frame_num)
            except ValueError:
                pass
    return max_frame

def run_make_compare_combined(id_val: int, args_dict: dict, dry_run: bool = False):
    """运行 make_compare.py 为单个ID生成合并视频（包含4fps和1fps）"""
    # 只使用 4fps 的 GT，合并 4fps 和 1fps 的 pred
    gt_dir = args_dict['gt']  # 只用 4fps 的 GT
    pred_dirs = args_dict['pred'].copy()  # 4fps 的 pred
    
    # 如果有 1fps 数据，只添加 pred（不添加 GT）
    if args_dict.get('pred_1fps'):
        pred_dirs.extend(args_dict['pred_1fps'])
    
    # 合并所有列：GT(4fps) + Pred(4fps) + Pred(1fps)
    all_cols = [gt_dir] + pred_dirs
    
    # 为 1fps 列创建缺失帧的符号链接（映射到对应的 1fps 帧）
    # 1fps 列每4帧更新一次，所以帧 0,1,2,3 -> 显示帧0，帧 4,5,6,7 -> 显示帧4，等等
    end_frame = args_dict['end']
    if args_dict.get('pred_1fps'):
        import tempfile
        import shutil
        
        # 为每个 1fps 列创建临时目录，包含所有需要的帧（通过符号链接）
        temp_dirs = []
        for pred_1fps_dir in args_dict['pred_1fps']:
            # 创建临时目录
            temp_dir = tempfile.mkdtemp(prefix=f'temp_1fps_{id_val}_')
            temp_dirs.append(temp_dir)
            
            # 复制 id 目录结构
            src_id_dir = os.path.join(pred_1fps_dir, f"id_{id_val}")
            dst_id_dir = os.path.join(temp_dir, f"id_{id_val}")
            os.makedirs(dst_id_dir, exist_ok=True)
            
            # 先找到 1fps 列中实际存在的最大帧
            max_available_frame = find_max_frame_in_dir(pred_1fps_dir, id_val)
            
            # 为每一帧创建符号链接（映射到对应的 1fps 帧）
            # 映射规则：4fps 的帧 fid -> 1fps 的帧 (fid // 4)
            # 例如：4fps 帧 0,1,2,3 -> 1fps 帧 0
            #       4fps 帧 4,5,6,7 -> 1fps 帧 1
            #       4fps 帧 8,9,10,11 -> 1fps 帧 2
            start_frame = args_dict.get('start') if args_dict.get('start') is not None else 0
            for fid in range(start_frame, end_frame + 1):
                # 映射规则：fid -> fid // 4（每4帧更新一次，1fps的1帧对应4fps的4帧）
                mapped_fid = fid // 4
                
                # 如果映射的帧超过了实际存在的最大帧，使用最大帧
                if mapped_fid > max_available_frame:
                    mapped_fid = max_available_frame
                
                src_frame = os.path.join(src_id_dir, f"{mapped_fid}.png")
                dst_frame = os.path.join(dst_id_dir, f"{fid}.png")
                
                if os.path.isfile(src_frame):
                    # 创建符号链接
                    if os.path.exists(dst_frame):
                        os.remove(dst_frame)
                    os.symlink(src_frame, dst_frame)
                else:
                    # 如果源帧不存在，尝试找最近的帧（向下查找）
                    found = False
                    for check_fid in range(mapped_fid, -1, -1):
                        check_src = os.path.join(src_id_dir, f"{check_fid}.png")
                        if os.path.isfile(check_src):
                            if os.path.exists(dst_frame):
                                os.remove(dst_frame)
                            os.symlink(check_src, dst_frame)
                            found = True
                            break
                    if not found:
                        print(f"[WARN] id_{id_val} 1fps 列找不到任何帧")
            
            # 更新 pred_dirs 中的路径，使用临时目录
            pred_dirs[pred_dirs.index(pred_1fps_dir)] = temp_dir
        
        # 保存临时目录列表，以便后续清理
        args_dict['_temp_dirs'] = temp_dirs
    
    # 使用用户提供的标签（如果提供了完整的标签列表）
    if args_dict.get('labels_combined'):
        labels = args_dict['labels_combined']
    else:
        # 否则自动生成标签
        labels = args_dict['labels'].copy()  # GT + 4fps pred
        if args_dict.get('pred_1fps'):
            # 为 1fps 的 pred 添加后缀
            labels_1fps = [f"{label}_1fps" for label in args_dict['labels'][1:]]  # 跳过 GT
            labels = labels + labels_1fps
    
    # 使用 4fps 的 fps 值（因为合并后视频统一使用一个 fps）
    fps = args_dict.get('fps', 1)
    output_name = f'compare_id{id_val}.mp4'
    
    cmd = [
        sys.executable,
        'make_compare.py',
        '--gt', gt_dir,  # 只用 4fps 的 GT
        '--pred'] + pred_dirs + [  # 4fps pred + 1fps pred
        '--id', str(id_val),
        '--out_dir', args_dict['out_dir'],
        '--end', str(end_frame),
        '--labels'] + labels + [
        '--fps', str(int(fps)) if fps == int(fps) else str(fps),
        '--out', output_name,
    ]
    
    if args_dict.get('safe_even', False):
        cmd.append('--safe_even')
    
    if args_dict.get('start', None):
        cmd.extend(['--start', str(args_dict['start'])])
    
    print(f"\n{'='*80}")
    print(f"处理 ID: id_{id_val} (合并 4fps + 1fps, fps={fps})")
    print(f"列数: {len(all_cols)} (GT_4fps + {len(args_dict['pred'])} pred_4fps + {len(args_dict.get('pred_1fps', []))} pred_1fps)")
    print(f"命令: {' '.join(cmd)}")
    print(f"{'='*80}")
    
    if dry_run:
        print("[DRY RUN] 不会实际执行")
        return True
    
    try:
        result = subprocess.run(cmd, check=True, cwd=args_dict.get('work_dir', '.'))
        print(f"[OK] id_{id_val} (合并视频) 完成")
        
        # 清理临时目录
        if args_dict.get('_temp_dirs'):
            for temp_dir in args_dict['_temp_dirs']:
                try:
                    shutil.rmtree(temp_dir)
                except Exception as e:
                    print(f"[WARN] 清理临时目录失败 {temp_dir}: {e}")
        
        return True
    except subprocess.CalledProcessError as e:
        print(f"[ERROR] id_{id_val} (合并视频) 失败: {e}", file=sys.stderr)
        
        # 清理临时目录
        if args_dict.get('_temp_dirs'):
            for temp_dir in args_dict['_temp_dirs']:
                try:
                    shutil.rmtree(temp_dir)
                except Exception:
                    pass
        
        return False
    except Exception as e:
        print(f"[ERROR] id_{id_val} (合并视频) 出错: {e}", file=sys.stderr)
        
        # 清理临时目录
        if args_dict.get('_temp_dirs'):
            for temp_dir in args_dict['_temp_dirs']:
                try:
                    shutil.rmtree(temp_dir)
                except Exception:
                    pass
        
        return False

def main():
    parser = argparse.ArgumentParser(description='批量生成对比视频')
    parser.add_argument('--id_file', type=str, required=True,
                       help='ID列表文件路径')
    parser.add_argument('--gt', type=str, required=True,
                       help='GT rollout目录')
    parser.add_argument('--pred', nargs='+', required=True,
                       help='预测 rollout 目录（一个或多个）')
    parser.add_argument('--out_dir', type=str, required=True,
                       help='输出目录')
    parser.add_argument('--labels', nargs='+', required=True,
                       help='标签列表（与GT+pred数量对应）')
    parser.add_argument('--end', type=int, default=51,
                       help='结束帧索引')
    parser.add_argument('--start', type=int, default=None,
                       help='开始帧索引（可选）')
    parser.add_argument('--fps', type=float, default=1,
                       help='视频FPS（用于rollout_4fps，默认1）')
    parser.add_argument('--gt_1fps', type=str, default=None,
                       help='GT rollout_1fps目录（已废弃，不再使用）')
    parser.add_argument('--pred_1fps', nargs='+', default=None,
                       help='预测 rollout_1fps 目录（一个或多个，将合并到视频中，但不包含GT）')
    parser.add_argument('--fps_1fps', type=float, default=0.25,
                       help='视频FPS（用于rollout_1fps，默认0.25）')
    parser.add_argument('--safe_even', action='store_true',
                       help='使用安全偶数尺寸')
    parser.add_argument('--work_dir', type=str, default=None,
                       help='工作目录（默认当前目录）')
    parser.add_argument('--dry_run', action='store_true',
                       help='仅显示命令，不实际执行')
    parser.add_argument('--skip_existing', action='store_true',
                       help='跳过已存在的视频文件')
    parser.add_argument('--max_workers', type=int, default=1,
                       help='最大并行数（默认1，串行执行）')
    
    args = parser.parse_args()
    
    # 读取ID列表
    ids = read_id_list(args.id_file)
    print(f"[INFO] 从 {args.id_file} 读取到 {len(ids)} 个ID")
    
    if not ids:
        print("[ERROR] 没有找到有效的ID", file=sys.stderr)
        sys.exit(1)
    
    # 确定处理模式：合并模式（如果指定了 1fps pred 数据）
    combine_mode = args.pred_1fps is not None
    if combine_mode:
        if len(args.pred_1fps) != len(args.pred):
            print(f"[ERROR] rollout_1fps 的预测目录数量 ({len(args.pred_1fps)}) 与 rollout_4fps ({len(args.pred)}) 不匹配", file=sys.stderr)
            sys.exit(1)
    
    # 准备参数
    args_dict = {
        'gt': args.gt,
        'pred': args.pred,
        'out_dir': args.out_dir,
        'labels': args.labels,
        'labels_combined': args.labels if combine_mode and len(args.labels) == (1 + len(args.pred) + len(args.pred_1fps)) else None,
        'end': args.end,
        'start': args.start,
        'fps': args.fps,
        'fps_1fps': args.fps_1fps,
        'gt_1fps': args.gt_1fps,
        'pred_1fps': args.pred_1fps,
        'safe_even': args.safe_even,
        'work_dir': args.work_dir or os.getcwd(),
        'dry_run': args.dry_run,
    }
    
    # 验证标签数量
    if combine_mode:
        expected_labels = 1 + len(args.pred) + len(args.pred_1fps)  # GT + 4fps pred + 1fps pred
        if len(args.labels) != expected_labels:
            print(f"[WARN] 标签数量 ({len(args.labels)}) 与预期 ({expected_labels}) 不匹配")
            print(f"      预期: GT + {len(args.pred)} 个 4fps 预测方法 + {len(args.pred_1fps)} 个 1fps 预测方法")
    else:
        expected_labels = 1 + len(args.pred)  # GT + pred数量
        if len(args.labels) != expected_labels:
            print(f"[WARN] 标签数量 ({len(args.labels)}) 与预期 ({expected_labels}) 不匹配")
            print(f"      预期: GT + {len(args.pred)} 个预测方法")
    
    # 检查输出目录
    os.makedirs(args.out_dir, exist_ok=True)
    
    # 执行批量处理
    success_count = 0
    skip_count = 0
    fail_count = 0
    total_tasks = len(ids)
    current_task = 0
    
    if combine_mode:
        print(f"\n开始批量处理 {len(ids)} 个ID（合并模式：4fps + 1fps）...")
    else:
        print(f"\n开始批量处理 {len(ids)} 个ID（仅 4fps）...")
    print(f"输出目录: {args.out_dir}")
    print(f"并行数: {args.max_workers}")
    print(f"跳过已存在: {args.skip_existing}")
    print(f"总任务数: {total_tasks}\n")
    
    for i, id_val in enumerate(ids, 1):
        current_task += 1
        
        # 检查是否已存在
        if args.skip_existing:
            video_path = os.path.join(args.out_dir, f"id_{id_val}", f"compare_id{id_val}.mp4")
            if os.path.exists(video_path):
                print(f"[SKIP] id_{id_val} 已存在，跳过 ({current_task}/{total_tasks})")
                skip_count += 1
                continue
        
        print(f"\n[{current_task}/{total_tasks}] 处理 id_{id_val}...")
        if combine_mode:
            success = run_make_compare_combined(id_val, args_dict, args.dry_run)
        else:
            # 如果没有 1fps 数据，只生成 4fps 视频（简化版本）
            cmd = [
                sys.executable,
                'make_compare.py',
                '--gt', args_dict['gt'],
                '--pred'] + args_dict['pred'] + [
                '--id', str(id_val),
                '--out_dir', args_dict['out_dir'],
                '--end', str(args_dict['end']),
                '--labels'] + args_dict['labels'] + [
                '--fps', str(int(args_dict['fps'])) if args_dict['fps'] == int(args_dict['fps']) else str(args_dict['fps']),
            ]
            if args_dict.get('safe_even', False):
                cmd.append('--safe_even')
            if args_dict.get('start', None):
                cmd.extend(['--start', str(args_dict['start'])])
            
            if args_dict.get('dry_run', False):
                print(f"[DRY RUN] 命令: {' '.join(cmd)}")
                success = True
            else:
                try:
                    result = subprocess.run(cmd, check=True, cwd=args_dict.get('work_dir', '.'))
                    success = True
                except Exception as e:
                    print(f"[ERROR] id_{id_val} 失败: {e}", file=sys.stderr)
                    success = False
        
        if success:
            success_count += 1
        else:
            fail_count += 1
            if not args.dry_run:
                # 询问是否继续
                response = input(f"\nid_{id_val} 失败，是否继续？(y/n): ").strip().lower()
                if response != 'y':
                    print("用户取消")
                    return
    
    # 总结
    print(f"\n{'='*80}")
    print("批量处理完成！")
    if combine_mode:
        print(f"总计: {len(ids)} 个ID（每个包含 4fps + 1fps 合并视频）")
    else:
        print(f"总计: {len(ids)} 个ID（仅 4fps）")
    print(f"成功: {success_count} 个")
    print(f"跳过: {skip_count} 个")
    print(f"失败: {fail_count} 个")
    print(f"{'='*80}")

if __name__ == '__main__':
    main()