559 lines
18 KiB
Python
559 lines
18 KiB
Python
#!/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()
|