100 lines
3.4 KiB
Python
100 lines
3.4 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
从已有 OCR 数据重新生成 SRT，不重新调用 OCR API。
|
||
用法: python rerun_srt.py
|
||
"""
|
||
import sys, os, json
|
||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||
|
||
from difflib import SequenceMatcher
|
||
from tool.log.log_manager import UnifiedLogger
|
||
from tool.getsrt import (
|
||
analyze_subtitle_blocks, analyze_h,
|
||
extract_highest_subtitle, check_subtitle_quality_enhanced,
|
||
evaluate_text_quality_rules, generate_srt_from_results
|
||
)
|
||
|
||
COMPOSITE_DIR = sys.argv[1] if len(sys.argv) > 1 else "./composite/BV1aL4y1p7Bi_p11"
|
||
JSONL_PATH = os.path.join(COMPOSITE_DIR, "ocr_result.jsonl")
|
||
META_PATH = os.path.join(COMPOSITE_DIR, "metadata.json")
|
||
OUTPUT_SRT = os.path.join(COMPOSITE_DIR, "BV1aL4y1p7Bi_p11.new.srt")
|
||
INTERVAL = 0.33
|
||
|
||
logger = UnifiedLogger()
|
||
logger.reset()
|
||
logger.init_logging("rerun_log.json")
|
||
logger.start_program("rerun")
|
||
|
||
# 读取已有 OCR 数据，将旧路径修正为当前实际路径
|
||
all_pic_ocr_pair = []
|
||
with open(JSONL_PATH, 'r', encoding='utf-8') as f:
|
||
lines = f.readlines()
|
||
|
||
i = 0
|
||
while i < len(lines):
|
||
if i + 1 < len(lines) and lines[i].startswith("pic_path: "):
|
||
raw_path = lines[i].strip()[len("pic_path: "):]
|
||
filename = os.path.basename(raw_path)
|
||
actual_path = os.path.join(COMPOSITE_DIR, filename)
|
||
try:
|
||
ocr_json = json.loads(lines[i + 1])
|
||
all_pic_ocr_pair.append([actual_path, ocr_json])
|
||
i += 2
|
||
continue
|
||
except json.JSONDecodeError:
|
||
pass
|
||
i += 1
|
||
|
||
print(f"读取 {len(all_pic_ocr_pair)} 帧的OCR数据")
|
||
|
||
# 分析字幕网格位置
|
||
best_grids = analyze_subtitle_blocks(all_pic_ocr_pair, META_PATH, logger)
|
||
h = analyze_h(all_pic_ocr_pair, META_PATH)
|
||
|
||
if not best_grids.get("composite"):
|
||
print("ERROR: 未找到合适的字幕区域")
|
||
sys.exit(1)
|
||
|
||
# 逐帧提取字幕
|
||
results = []
|
||
DEBUG_KEYWORD = "曹小月从角落里拿出雨伞"
|
||
for idx in range(len(all_pic_ocr_pair)):
|
||
print(f"\r处理 {idx + 1}/{len(all_pic_ocr_pair)}", end="", flush=True)
|
||
text, conf, flag, sim = extract_highest_subtitle(
|
||
idx, all_pic_ocr_pair, best_grids, META_PATH, h, logger
|
||
)
|
||
if text and not check_subtitle_quality_enhanced(text) and conf <= 90:
|
||
text = ""
|
||
# 打印包含关键词的帧的原始OCR和提取结果
|
||
if DEBUG_KEYWORD in (text or ""):
|
||
pic_path, ocr_json = all_pic_ocr_pair[idx]
|
||
print(f"\n[DEBUG 帧{idx}] 提取: {repr(text)}")
|
||
for td in ocr_json['TextDetections']:
|
||
print(f" box: {repr(td['DetectedText'])} conf={td['Confidence']} watermark={td.get('is_watermark','?')}")
|
||
results.append((idx, text or ""))
|
||
|
||
non_empty = sum(1 for _, t in results if t)
|
||
print(f"\n共 {non_empty}/{len(results)} 帧有字幕")
|
||
|
||
# 生成 SRT
|
||
generate_srt_from_results(results, OUTPUT_SRT, interval_seconds=INTERVAL)
|
||
print(f"SRT 已输出: {OUTPUT_SRT}")
|
||
|
||
# 展示问题字幕附近的结果
|
||
with open(OUTPUT_SRT, 'r', encoding='utf-8') as f:
|
||
srt_content = f.readlines()
|
||
|
||
print("\n=== 修复验证（'曹小月从角落里拿出雨伞' 附近） ===")
|
||
found = False
|
||
for j, line in enumerate(srt_content):
|
||
if "曹小月从角落里拿出雨伞" in line:
|
||
start = max(0, j - 3)
|
||
end = min(len(srt_content), j + 10)
|
||
print("".join(srt_content[start:end]))
|
||
print("---")
|
||
found = True
|
||
break
|
||
|
||
if not found:
|
||
print("未在输出 SRT 中找到目标字幕，请检查提取结果。")
|