354 lines
13 KiB
Python
354 lines
13 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
Subtitle Preview Contact Sheet · 字幕预览检查图生成器
|
||
==============================================
|
||
每次烧字幕后自动抽帧出检查图。
|
||
|
||
功能：
|
||
1. 从视频中按字幕时间点抽帧
|
||
2. 生成检查图（contact sheet）
|
||
3. 可选：在检查图上标注字幕框、问题区域
|
||
4. 生成 HTML 预览页面
|
||
|
||
依赖：
|
||
pip install opencv-python pillow numpy
|
||
|
||
用法：
|
||
# 从视频抽帧生成检查图
|
||
python subtitle-preview-contact-sheet.py --video input.mp4 --subtitle-ass subtitles.ass --output preview.png
|
||
|
||
# 生成 HTML 预览页面
|
||
python subtitle-preview-contact-sheet.py --video input.mp4 --subtitle-ass subtitles.ass --generate-html --output preview.html
|
||
|
||
# 批量处理多集
|
||
python subtitle-preview-contact-sheet.py --batch ./videos/ --subtitle-dir ./ass/ --output-dir ./previews/
|
||
|
||
# 自定义抽帧参数
|
||
python subtitle-preview-contact-sheet.py --video input.mp4 --subtitle-ass subtitles.ass --max-frames 20 --cols 4
|
||
|
||
输出：
|
||
- preview.png: 检查图（多帧拼接）
|
||
- preview.html: HTML 预览页面（可选）
|
||
- frames/: 抽出的关键帧（可选）
|
||
|
||
路径：
|
||
video-ai-system/engines/subtitle-pipeline/preview-tools/subtitle-preview-contact-sheet.py
|
||
"""
|
||
|
||
import argparse
|
||
import json
|
||
import os
|
||
import sys
|
||
from pathlib import Path
|
||
|
||
try:
|
||
import cv2
|
||
import numpy as np
|
||
from PIL import Image, ImageDraw, ImageFont
|
||
except ImportError as e:
|
||
print(f"[ERROR] 缺少依赖：{e}")
|
||
print("请先安装：pip install opencv-python pillow numpy")
|
||
sys.exit(1)
|
||
|
||
|
||
def extract_frames_by_subtitle(video_path: str, ass_path: str, output_dir: str, max_frames: int = 20) -> list:
|
||
"""
|
||
按字幕时间点从视频中抽帧
|
||
|
||
:param video_path: 视频文件路径
|
||
:param ass_path: ASS 字幕文件路径
|
||
:param output_dir: 输出目录
|
||
:param max_frames: 最大抽帧数量
|
||
:return: 抽出的帧文件路径列表
|
||
"""
|
||
# 打开视频
|
||
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))
|
||
|
||
print(f"[INFO] 视频信息：{total_frames} 帧，{fps} FPS")
|
||
|
||
# 读取 ASS 字幕文件，获取字幕时间点
|
||
subtitle_times = []
|
||
try:
|
||
with open(ass_path, "r", encoding="utf-8-sig") as f:
|
||
lines = f.readlines()
|
||
|
||
for line in lines:
|
||
if line.startswith("Dialogue:"):
|
||
# 解析 ASS 对话行
|
||
parts = line.split(",")
|
||
if len(parts) >= 10:
|
||
start_time = parts[1].strip() # 开始时间
|
||
end_time = parts[2].strip() # 结束时间
|
||
|
||
# 转换为秒
|
||
start_seconds = ass_time_to_seconds(start_time)
|
||
end_seconds = ass_time_to_seconds(end_time)
|
||
|
||
# 取字幕中间时间点作为抽帧时间
|
||
middle_time = (start_seconds + end_seconds) / 2
|
||
subtitle_times.append(middle_time)
|
||
|
||
except Exception as e:
|
||
print(f"[ERROR] 读取 ASS 文件失败：{e}")
|
||
return []
|
||
|
||
if not subtitle_times:
|
||
print("[WARN] 未找到字幕时间点，将按固定间隔抽帧")
|
||
# 按固定间隔抽帧
|
||
interval = total_frames / max_frames
|
||
subtitle_times = [i * interval / fps for i in range(max_frames)]
|
||
|
||
print(f"[INFO] 找到 {len(subtitle_times)} 个字幕时间点")
|
||
|
||
# 创建输出目录
|
||
frames_dir = os.path.join(output_dir, "frames")
|
||
os.makedirs(frames_dir, exist_ok=True)
|
||
|
||
# 抽帧
|
||
frame_paths = []
|
||
for i, timestamp in enumerate(subtitle_times[:max_frames]):
|
||
# 跳转到指定时间
|
||
cap.set(cv2.CAP_PROP_POS_MSEC, timestamp * 1000)
|
||
|
||
ret, frame = cap.read()
|
||
if not ret:
|
||
print(f"[WARN] 无法读取帧：{timestamp:.2f}s")
|
||
continue
|
||
|
||
# 保存帧
|
||
frame_path = os.path.join(frames_dir, f"frame_{i:04d}_{timestamp:.2f}s.png")
|
||
cv2.imwrite(frame_path, frame)
|
||
|
||
frame_paths.append(frame_path)
|
||
print(f"[INFO] 已抽帧：{timestamp:.2f}s -> {frame_path}")
|
||
|
||
cap.release()
|
||
|
||
print(f"[OK] 共抽出 {len(frame_paths)} 帧")
|
||
return frame_paths
|
||
|
||
|
||
def ass_time_to_seconds(ass_time: str) -> float:
|
||
"""
|
||
将 ASS 时间戳转换为秒
|
||
|
||
:param ass_time: ASS 时间戳（H:MM:SS.cc）
|
||
:return: 秒数
|
||
"""
|
||
parts = ass_time.split(":")
|
||
|
||
hours = int(parts[0])
|
||
minutes = int(parts[1])
|
||
seconds_parts = parts[2].split(".")
|
||
seconds = int(seconds_parts[0])
|
||
centiseconds = int(seconds_parts[1]) if len(seconds_parts) > 1 else 0
|
||
|
||
total_seconds = hours * 3600 + minutes * 60 + seconds + centiseconds / 100
|
||
return total_seconds
|
||
|
||
|
||
def generate_contact_sheet(frame_paths: list, output_path: str, cols: int = 4, padding: int = 10):
|
||
"""
|
||
生成检查图（多帧拼接）
|
||
|
||
:param frame_paths: 帧文件路径列表
|
||
:param output_path: 输出图片路径
|
||
:param cols: 列数
|
||
:param padding: 内边距（像素）
|
||
"""
|
||
if not frame_paths:
|
||
print("[ERROR] 没有帧文件")
|
||
return False
|
||
|
||
# 读取第一帧，获取尺寸
|
||
first_frame = Image.open(frame_paths[0])
|
||
frame_width, frame_height = first_frame.size
|
||
|
||
# 计算拼接图尺寸
|
||
rows = (len(frame_paths) + cols - 1) // cols
|
||
sheet_width = cols * (frame_width + padding) + padding
|
||
sheet_height = rows * (frame_height + padding) + padding
|
||
|
||
# 创建拼接图
|
||
sheet = Image.new("RGB", (sheet_width, sheet_height), (0, 0, 0))
|
||
|
||
# 拼接帧
|
||
for i, frame_path in enumerate(frame_paths):
|
||
frame = Image.open(frame_path)
|
||
|
||
# 计算位置
|
||
row = i // cols
|
||
col = i % cols
|
||
x = col * (frame_width + padding) + padding
|
||
y = row * (frame_height + padding) + padding
|
||
|
||
# 粘贴帧
|
||
sheet.paste(frame, (x, y))
|
||
|
||
# 添加时间戳标注
|
||
draw = ImageDraw.Draw(sheet)
|
||
timestamp = os.path.basename(frame_path).split("_")[1].replace(".png", "")
|
||
draw.text((x + 5, y + 5), f"{timestamp}", 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 generate_html_preview(frame_paths: list, output_path: str, video_path: str = None):
|
||
"""
|
||
生成 HTML 预览页面
|
||
|
||
:param frame_paths: 帧文件路径列表
|
||
:param output_path: 输出 HTML 文件路径
|
||
:param video_path: 视频文件路径（可选）
|
||
"""
|
||
html_lines = []
|
||
html_lines.append("<!DOCTYPE html>")
|
||
html_lines.append("<html>")
|
||
html_lines.append("<head>")
|
||
html_lines.append(" <meta charset='utf-8'>")
|
||
html_lines.append(" <title>字幕预览检查图</title>")
|
||
html_lines.append(" <style>")
|
||
html_lines.append(" body { font-family: Arial, sans-serif; margin: 20px; }")
|
||
html_lines.append(" .frame { margin: 10px; border: 1px solid #ccc; display: inline-block; }")
|
||
html_lines.append(" .timestamp { font-size: 12px; color: #666; text-align: center; }")
|
||
html_lines.append(" </style>")
|
||
html_lines.append("</head>")
|
||
html_lines.append("<body>")
|
||
html_lines.append(" <h1>字幕预览检查图</h1>")
|
||
|
||
if video_path:
|
||
html_lines.append(f" <p>视频：{os.path.basename(video_path)}</p>")
|
||
|
||
html_lines.append(f" <p>共 {len(frame_paths)} 帧</p>")
|
||
html_lines.append(" <div>")
|
||
|
||
for frame_path in frame_paths:
|
||
timestamp = os.path.basename(frame_path).split("_")[1].replace(".png", "")
|
||
rel_path = os.path.relpath(frame_path, os.path.dirname(output_path))
|
||
|
||
html_lines.append(" <div class='frame'>")
|
||
html_lines.append(f" <img src='{rel_path}' width='320'>")
|
||
html_lines.append(f" <div class='timestamp'>{timestamp}</div>")
|
||
html_lines.append(" </div>")
|
||
|
||
html_lines.append(" </div>")
|
||
html_lines.append("</body>")
|
||
html_lines.append("</html>")
|
||
|
||
# 写入文件
|
||
os.makedirs(os.path.dirname(os.path.abspath(output_path)), exist_ok=True)
|
||
with open(output_path, "w", encoding="utf-8") as f:
|
||
f.write("\n".join(html_lines))
|
||
|
||
print(f"[OK] HTML 预览页面已生成：{output_path}")
|
||
|
||
|
||
def main():
|
||
parser = argparse.ArgumentParser(
|
||
description="Subtitle Preview Contact Sheet · 字幕预览检查图生成器",
|
||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||
epilog="""
|
||
示例：
|
||
python subtitle-preview-contact-sheet.py --video input.mp4 --subtitle-ass subtitles.ass --output preview.png
|
||
python subtitle-preview-contact-sheet.py --video input.mp4 --subtitle-ass subtitles.ass --generate-html --output preview.html
|
||
python subtitle-preview-contact-sheet.py --batch ./videos/ --subtitle-dir ./ass/ --output-dir ./previews/
|
||
"""
|
||
)
|
||
parser.add_argument("--video", help="视频文件路径")
|
||
parser.add_argument("--subtitle-ass", help="ASS 字幕文件路径")
|
||
parser.add_argument("--output", help="输出文件路径（PNG 或 HTML）")
|
||
parser.add_argument("--generate-html", action="store_true", help="生成 HTML 预览页面")
|
||
parser.add_argument("--batch", help="批量处理视频目录")
|
||
parser.add_argument("--subtitle-dir", help="字幕文件目录（用于批量处理）")
|
||
parser.add_argument("--output-dir", help="输出目录（用于批量处理）")
|
||
parser.add_argument("--max-frames", type=int, default=20, help="最大抽帧数量")
|
||
parser.add_argument("--cols", type=int, default=4, help="检查图列数")
|
||
parser.add_argument("--padding", type=int, default=10, help="内边距（像素）")
|
||
|
||
args = parser.parse_args()
|
||
|
||
if args.video:
|
||
# 单文件处理
|
||
if not args.subtitle_ass:
|
||
print("[ERROR] 请指定 --subtitle-ass")
|
||
sys.exit(1)
|
||
|
||
if not args.output:
|
||
print("[ERROR] 请指定 --output")
|
||
sys.exit(1)
|
||
|
||
# 抽帧
|
||
output_dir = os.path.dirname(args.output)
|
||
if not output_dir:
|
||
output_dir = "."
|
||
|
||
frame_paths = extract_frames_by_subtitle(args.video, args.subtitle_ass, output_dir, args.max_frames)
|
||
|
||
if not frame_paths:
|
||
sys.exit(1)
|
||
|
||
# 生成检查图或 HTML
|
||
if args.generate_html:
|
||
generate_html_preview(frame_paths, args.output, args.video)
|
||
else:
|
||
generate_contact_sheet(frame_paths, args.output, args.cols, args.padding)
|
||
|
||
elif args.batch:
|
||
# 批量处理
|
||
if not args.subtitle_dir:
|
||
print("[ERROR] 批量处理需要指定 --subtitle-dir")
|
||
sys.exit(1)
|
||
|
||
if not args.output_dir:
|
||
print("[ERROR] 批量处理需要指定 --output-dir")
|
||
sys.exit(1)
|
||
|
||
os.makedirs(args.output_dir, exist_ok=True)
|
||
|
||
for video_file in os.listdir(args.batch):
|
||
if video_file.lower().endswith((".mp4", ".mov", ".avi")):
|
||
video_path = os.path.join(args.batch, video_file)
|
||
|
||
# 查找对应的字幕文件
|
||
subtitle_file = video_file.replace(".mp4", ".ass").replace(".mov", ".ass").replace(".avi", ".ass")
|
||
subtitle_path = os.path.join(args.subtitle_dir, subtitle_file)
|
||
|
||
if not os.path.isfile(subtitle_path):
|
||
print(f"[WARN] 未找到对应的字幕文件：{subtitle_file}")
|
||
continue
|
||
|
||
print(f"\n[INFO] 处理视频：{video_file}")
|
||
|
||
# 抽帧
|
||
frame_paths = extract_frames_by_subtitle(video_path, subtitle_path, args.output_dir, args.max_frames)
|
||
|
||
if frame_paths:
|
||
# 生成检查图
|
||
output_path = os.path.join(args.output_dir, video_file.replace(".mp4", "_preview.png"))
|
||
generate_contact_sheet(frame_paths, output_path, args.cols, args.padding)
|
||
|
||
# 生成 HTML
|
||
html_path = os.path.join(args.output_dir, video_file.replace(".mp4", "_preview.html"))
|
||
generate_html_preview(frame_paths, html_path, video_path)
|
||
|
||
else:
|
||
print("[ERROR] 请指定 --video 或 --batch")
|
||
sys.exit(1)
|
||
|
||
print("\n[OK] 处理完成")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|