#!/usr/bin/env python3 """ Subtitle Style A/B Test · 字幕样式 A/B 测试器 ==========================​================= 同一底片一次生成多个字幕风格版本，对比选，不要一版一版盲调。 功能： 1. 读取多个字幕样式配置 2. 用同一个底片和 SRT，生成多个字幕版本 3. 生成对比预览图（contact sheet） 4. 生成 HTML 对比页面 依赖： pip install pysrt 用法： # 单视频 A/B 测试 python subtitle-style-ab-test.py --video input.mp4 --srt input.srt --styles style1.json style2.json --output-dir ./ab-test/ # 批量 A/B 测试 python subtitle-style-ab-test.py --batch ./videos/ --srt-dir ./srt/ --styles style1.json style2.json --output-dir ./ab-tests/ # 使用预设样式 python subtitle-style-ab-test.py --video input.mp4 --srt input.srt --preset-styles reference-drama short-drama-bold --output-dir ./ab-test/ 输出目录结构： ./ab-test/ ├── style1/ │ ├── subtitles.ass │ ├── output.mp4 │ └── preview.png ├── style2/ │ ├── subtitles.ass │ ├── output.mp4 │ └── preview.png ├── comparison.png # 对比预览图 └── comparison.html # 对比 HTML 页面 路径： video-ai-system/engines/subtitle-pipeline/ab-test-tools/subtitle-style-ab-test.py """ import argparse import json import os import sys import subprocess from pathlib import Path try: import pysrt except ImportError: print("[ERROR] 缺少依赖：pysrt") print("请先安装：pip install pysrt") sys.exit(1) # 预设样式列表 PRESET_STYLES = { "reference-drama": { "font_family": "PingFang SC", "font_size": 38, "font_color": "&HFFFFFF&", "bold": True, "stroke_enabled": True, "stroke_width": 2, "stroke_color": "&H000000&", "alignment": 2, "margin_left": 100, "margin_right": 100, "margin_vertical": 50, }, "short-drama-bold": { "font_family": "PingFang SC", "font_size": 42, "font_color": "&HFFFFFF&", "bold": True, "stroke_enabled": True, "stroke_width": 1, "stroke_color": "&H121212&", "alignment": 2, "margin_left": 80, "margin_right": 80, "margin_vertical": 60, }, "clean-white": { "font_family": "PingFang SC", "font_size": 36, "font_color": "&HFFFFFF&", "bold": False, "stroke_enabled": False, "stroke_width": 0, "stroke_color": "&H000000&", "alignment": 2, "margin_left": 100, "margin_right": 100, "margin_vertical": 50, }, "black-box": { "font_family": "PingFang SC", "font_size": 36, "font_color": "&HFFFFFF&", "bold": False, "stroke_enabled": False, "stroke_width": 0, "stroke_color": "&H000000&", "background_enabled": True, "background_color": "&H000000&", "background_opacity": 0.6, "alignment": 2, "margin_left": 100, "margin_right": 100, "margin_vertical": 50, }, } def load_style_config(style_path: str) -> dict: """ 加载样式配置 :param style_path: 样式文件路径或预设样式名称 :return: 样式字典 """ # 检查是否是预设样式 if style_path in PRESET_STYLES: print(f"[INFO] 使用预设样式：{style_path}") return PRESET_STYLES[style_path] # 加载自定义样式文件 if not os.path.isfile(style_path): print(f"[ERROR] 样式文件不存在：{style_path}") return None try: with open(style_path, "r", encoding="utf-8") as f: config = json.load(f) print(f"[INFO] 已加载样式配置：{style_path}") return config except Exception as e: print(f"[ERROR] 加载样式配置失败：{e}") return None def generate_ass_file_for_style(srt_path: str, output_path: str, style: dict) -> bool: """ 为指定样式生成 ASS 字幕文件 :param srt_path: SRT 文件路径 :param output_path: 输出 ASS 文件路径 :param style: 样式字典 :return: 是否成功 """ if not os.path.isfile(srt_path): print(f"[ERROR] SRT 文件不存在：{srt_path}") return False # 加载 SRT try: subs = pysrt.open(srt_path, encoding="utf-8") except Exception as e: print(f"[ERROR] 加载 SRT 失败：{e}") return False print(f"[INFO] 为样式生成 ASS 文件：{os.path.basename(output_path)}") # 构建 ASS 文件内容（简化版，只生成 Dialogues） ass_lines = [] # 1. [Script Info] 节 ass_lines.append("[Script Info]") ass_lines.append("; Script generated by Subtitle Style A/B Test") ass_lines.append("PlayResX: 1080") ass_lines.append("PlayResY: 1920") ass_lines.append("") # 2. [V4+ Styles] 节 ass_lines.append("[V4+ Styles]") ass_lines.append( "Format: Name, Fontname, Fontsize, PrimaryColour, SecondaryColour, OutlineColour, BackColour, " "Bold, Italic, Underline, StrikeOut, ScaleX, ScaleY, Spacing, Angle, " "BorderStyle, Outline, Shadow, Alignment, MarginL, MarginR, MarginV, Encoding" ) # 构建样式行 font_name = style.get("font_family", "PingFang SC") font_size = style.get("font_size", 38) primary_colour = style.get("font_color", "&HFFFFFF&") outline_colour = style.get("stroke_color", "&H000000&") if style.get("background_enabled", False): bg_opacity = int(style.get("background_opacity", 0.6) * 255) back_colour = f"&H{bg_opacity:02X}000000&" else: back_colour = "&H00000000&" bold = -1 if style.get("bold", False) else 0 outline_width = style.get("stroke_width", 2) if style.get("stroke_enabled", True) else 0 alignment = style.get("alignment", 2) margin_l = style.get("margin_left", 100) margin_r = style.get("margin_right", 100) margin_v = style.get("margin_vertical", 50) style_line = ( f"Style: Default,{font_name},{font_size},{primary_colour},&H000000FF&,{outline_colour},{back_colour}," f"{bold},0,0,0,100,100,0,0," f"1,{outline_width},0,{alignment},{margin_l},{margin_r},{margin_v},1" ) ass_lines.append(style_line) ass_lines.append("") # 3. [Events] 节 ass_lines.append("[Events]") ass_lines.append("Format: Layer, Start, End, Style, Name, MarginL, MarginR, MarginV, Effect, Text") for sub in subs: start_time = srt_time_to_ass(sub.start) end_time = srt_time_to_ass(sub.end) text = sub.text.strip() # 转义特殊字符 text = text.replace("\\", "\\\\") text = text.replace("{", "\\{") text = text.replace("}", "\\}") text = text.replace("\n", "\\N") dialogue_line = f"Dialogue: 0,{start_time},{end_time},Default,,0,0,0,,{text}" ass_lines.append(dialogue_line) # 写入 ASS 文件 try: os.makedirs(os.path.dirname(os.path.abspath(output_path)), exist_ok=True) with open(output_path, "w", encoding="utf-8-sig") as f: f.write("\n".join(ass_lines)) print(f"[OK] ASS 字幕文件已生成：{output_path}") return True except Exception as e: print(f"[ERROR] 写入 ASS 文件失败：{e}") return False def srt_time_to_ass(srt_time) -> str: """ 将 pysrt 时间对象转换为 ASS 时间戳格式 :param srt_time: pysrt 时间对象 :return: ASS 时间戳字符串（H:MM:SS.cc） """ hours = srt_time.hours minutes = srt_time.minutes seconds = srt_time.seconds milliseconds = srt_time.milliseconds centiseconds = milliseconds // 10 return f"{hours}:{minutes:02d}:{seconds:02d}.{centiseconds:02d}" def burn_subtitles_with_ffmpeg(video_path: str, ass_path: str, output_path: str) -> bool: """ 用 FFmpeg + libass 烧录字幕到视频 :param video_path: 输入视频路径 :param ass_path: ASS 字幕文件路径 :param output_path: 输出视频路径 :return: 是否成功 """ if not os.path.isfile(video_path): print(f"[ERROR] 视频文件不存在：{video_path}") return False if not os.path.isfile(ass_path): print(f"[ERROR] ASS 字幕文件不存在：{ass_path}") return False # FFmpeg 命令 # 正确语法：subtitles=filename='path' （需要 filename= 前缀） ass_path_escaped = ass_path.replace(":", "\\:").replace("'", "'\\''") cmd = [ "ffmpeg", "-y", "-i", video_path, "-vf", f"subtitles=filename='{ass_path_escaped}'", "-c:v", "libx264", "-pix_fmt", "yuv420p", "-c:a", "copy", "-shortest", output_path ] print(f"[INFO] 烧录字幕：{os.path.basename(output_path)}") try: result = subprocess.run(cmd, check=False, capture_output=True, text=True) except FileNotFoundError: print("[ERROR] 未找到 ffmpeg") return False if result.returncode != 0: print("[ERROR] FFmpeg 烧录字幕失败") print(result.stderr[-2000:]) return False if os.path.isfile(output_path) and os.path.getsize(output_path) > 0: print(f"[OK] 字幕已烧录：{output_path}") return True return False def extract_sample_frames(video_path: str, output_dir: str, num_frames: int = 4) -> list: """ 从视频中抽取样本帧（用于预览） :param video_path: 视频文件路径 :param output_dir: 输出目录 :param num_frames: 抽取帧数 :return: 帧文件路径列表 """ import cv2 cap = cv2.VideoCapture(video_path) if not cap.isOpened(): print(f"[ERROR] 无法打开视频：{video_path}") return [] fps = cap.get(cv2.CAP_PROP_FPS) total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) # 计算采样间隔 interval = total_frames // (num_frames + 1) frame_paths = [] for i in range(1, num_frames + 1): frame_num = i * interval # 跳转到指定帧 cap.set(cv2.CAP_PROP_POS_FRAMES, frame_num) ret, frame = cap.read() if not ret: continue # 保存帧 frame_path = os.path.join(output_dir, f"sample_{i:02d}.png") cv2.imwrite(frame_path, frame) frame_paths.append(frame_path) cap.release() print(f"[INFO] 已抽取 {len(frame_paths)} 个样本帧") return frame_paths def generate_comparison_preview(style_dirs: list, output_path: str): """ 生成对比预览图 :param style_dirs: 样式目录列表 :param output_path: 输出图片路径 """ try: from PIL import Image except ImportError: print("[ERROR] 缺少依赖：Pillow") print("请先安装：pip install Pillow") return False # 收集所有样本帧 all_frames = [] style_names = [] for style_dir in style_dirs: style_name = os.path.basename(style_dir) style_names.append(style_name) # 查找样本帧 frame_files = [f for f in os.listdir(style_dir) if f.startswith("sample_") and f.endswith(".png")] frame_files.sort() for frame_file in frame_files: frame_path = os.path.join(style_dir, frame_file) all_frames.append((style_name, frame_path)) if not all_frames: print("[WARN] 未找到样本帧，跳过对比预览图生成") return False # 读取第一帧，获取尺寸 first_frame = Image.open(all_frames[0][1]) frame_width, frame_height = first_frame.size # 计算拼接图尺寸 num_styles = len(style_dirs) num_samples = len(all_frames) // num_styles # 横向拼接：每个样式一行，每行包含多个样本帧 sheet_width = num_samples * (frame_width + 10) + 10 sheet_height = num_styles * (frame_height + 30) + 10 # 30px 用于样式名称标注 # 创建拼接图 sheet = Image.new("RGB", (sheet_width, sheet_height), (0, 0, 0)) # 拼接帧 for i, (style_name, frame_path) in enumerate(all_frames): frame = Image.open(frame_path) # 计算位置 style_idx = style_names.index(style_name) sample_idx = i % num_samples x = sample_idx * (frame_width + 10) + 10 y = style_idx * (frame_height + 30) + 10 # 粘贴帧 sheet.paste(frame, (x, y + 20)) # +20px 用于样式名称标注 # 添加样式名称标注 draw = ImageDraw.Draw(sheet) draw.text((x + 5, y + 5), style_name, fill=(255, 255, 255)) # 保存 os.makedirs(os.path.dirname(os.path.abspath(output_path)), exist_ok=True) sheet.save(output_path) print(f"[OK] 对比预览图已生成：{output_path}") return True def main(): parser = argparse.ArgumentParser( description="Subtitle Style A/B Test · 字幕样式 A/B 测试器", formatter_class=argparse.RawDescriptionHelpFormatter, epilog=""" 示例： python subtitle-style-ab-test.py --video input.mp4 --srt input.srt --styles style1.json style2.json --output-dir ./ab-test/ python subtitle-style-ab-test.py --video input.mp4 --srt input.srt --preset-styles reference-drama short-drama-bold --output-dir ./ab-test/ """ ) parser.add_argument("--video", help="视频文件路径") parser.add_argument("--srt", help="SRT 字幕文件路径") parser.add_argument("--styles", nargs="+", help="样式配置文件路径列表") parser.add_argument("--preset-styles", nargs="+", help="预设样式名称列表") parser.add_argument("--batch", help="批量处理视频目录") parser.add_argument("--srt-dir", help="SRT 文件目录（用于批量处理）") parser.add_argument("--output-dir", required=True, help="输出目录") parser.add_argument("--num-samples", type=int, default=4, help="每个样式抽取的样本帧数") args = parser.parse_args() # 收集样式配置 styles = [] style_names = [] if args.styles: for style_path in args.styles: style = load_style_config(style_path) if style: styles.append(style) style_names.append(os.path.splitext(os.path.basename(style_path))[0]) if args.preset_styles: for preset_name in args.preset_styles: style = load_style_config(preset_name) if style: styles.append(style) style_names.append(preset_name) if not styles: print("[ERROR] 请指定至少一种样式（--styles 或 --preset-styles）") sys.exit(1) print(f"[INFO] 将测试 {len(styles)} 种样式：{style_names}") if args.video: # 单视频处理 if not args.srt: print("[ERROR] 请指定 --srt") sys.exit(1) # 为每个样式生成字幕和视频 style_dirs = [] for i, (style, style_name) in enumerate(zip(styles, style_names)): print(f"\n[INFO] 处理样式 {i+1}/{len(styles)}：{style_name}") # 创建样式目录 style_dir = os.path.join(args.output_dir, style_name) os.makedirs(style_dir, exist_ok=True) style_dirs.append(style_dir) # 生成 ASS 文件 ass_path = os.path.join(style_dir, "subtitles.ass") ok = generate_ass_file_for_style(args.srt, ass_path, style) if ok: # 烧录字幕到视频 output_video = os.path.join(style_dir, os.path.basename(args.video)) burn_subtitles_with_ffmpeg(args.video, ass_path, output_video) # 抽取样本帧 extract_sample_frames(output_video, style_dir, args.num_samples) # 生成对比预览图 comparison_preview = os.path.join(args.output_dir, "comparison.png") generate_comparison_preview(style_dirs, comparison_preview) elif args.batch: # 批量处理 if not args.srt_dir: print("[ERROR] 批量处理需要指定 --srt-dir") sys.exit(1) for video_file in os.listdir(args.batch): if not video_file.lower().endswith((".mp4", ".mov", ".avi")): continue video_path = os.path.join(args.batch, video_file) # 查找对应的 SRT 文件 srt_file = video_file.replace(".mp4", ".srt").replace(".mov", ".srt").replace(".avi", ".srt") srt_path = os.path.join(args.srt_dir, srt_file) if not os.path.isfile(srt_path): print(f"[WARN] 未找到对应的 SRT 文件：{srt_file}") continue print(f"\n[INFO] 批量处理：{video_file}") # 为每个样式生成字幕和视频 video_output_dir = os.path.join(args.output_dir, os.path.splitext(video_file)[0]) for i, (style, style_name) in enumerate(zip(styles, style_names)): print(f"\n[INFO] 处理样式 {i+1}/{len(styles)}：{style_name}") style_dir = os.path.join(video_output_dir, style_name) os.makedirs(style_dir, exist_ok=True) # 生成 ASS 文件 ass_path = os.path.join(style_dir, "subtitles.ass") ok = generate_ass_file_for_style(srt_path, ass_path, style) if ok: # 烧录字幕到视频 output_video = os.path.join(style_dir, video_file) burn_subtitles_with_ffmpeg(video_path, ass_path, output_video) # 抽取样本帧 extract_sample_frames(output_video, style_dir, args.num_samples) print("\n[OK] A/B 测试完成") if __name__ == "__main__": main()