886 lines
40 KiB
Python
886 lines
40 KiB
Python
import os
|
||
import json
|
||
import cv2
|
||
import numpy as np
|
||
import random
|
||
import re
|
||
import string
|
||
import math
|
||
import tempfile
|
||
from collections import defaultdict
|
||
from concurrent.futures import ThreadPoolExecutor
|
||
from difflib import SequenceMatcher
|
||
from PIL import Image
|
||
from typing import Optional, List, Tuple, Dict, Any, Union
|
||
# 假设这些外部模块已正确导入
|
||
from tool.Ocr.tencentcloud_ocr import tencentcloud_ocr_api, img2base64
|
||
from tool.whisper_manager import whisper_manager
|
||
from tool.log.log_manager import logger
|
||
from config.config_utils import config
|
||
|
||
|
||
# ==================== 辅助函数 ====================
|
||
def natural_sort_key(filename: str) -> List[Union[int, str]]:
|
||
"""
|
||
自然排序键函数，将文件名中的数字部分转换为整数进行排序。
|
||
"""
|
||
return [int(part) if part.isdigit() else part.lower()
|
||
for part in re.split(r'(\d+)', filename)]
|
||
|
||
|
||
def has_chinese(text: str) -> bool:
|
||
"""文本是否含中文字符（CJK 统一表意文字）。"""
|
||
return any('一' <= ch <= '鿿' for ch in (text or ""))
|
||
|
||
|
||
def calculate_english_ratio(text: str) -> float:
|
||
"""计算文本中英文字符（字母、数字、常见标点）的占比。"""
|
||
if not text:
|
||
return 0.0
|
||
total_chars = len(text)
|
||
if total_chars == 0:
|
||
return 0.0
|
||
|
||
english_chars = 0
|
||
for ch in text:
|
||
if ch.isalpha() and ch.isascii():
|
||
english_chars += 1
|
||
elif ch.isdigit():
|
||
english_chars += 1
|
||
elif ch in ' .,!?;:\'"-()[]{}':
|
||
english_chars += 1
|
||
return english_chars / total_chars
|
||
|
||
|
||
def check_subtitle_quality_enhanced(best_subtitle: str, strict_mode: bool = True) -> bool:
|
||
"""
|
||
增强版字幕质量检测。
|
||
|
||
:param best_subtitle: 待检测的文本
|
||
:param strict_mode: 是否使用严格模式
|
||
:return: 若文本质量合格则返回 True
|
||
"""
|
||
if not best_subtitle or not best_subtitle.strip():
|
||
return True
|
||
|
||
subtitle = best_subtitle.strip()
|
||
total_chars = len(subtitle)
|
||
if total_chars == 0:
|
||
return True
|
||
|
||
digit_count = sum(1 for ch in subtitle if ch.isdigit())
|
||
symbol_count = sum(1 for ch in subtitle if ch in string.punctuation)
|
||
special_ratio = (digit_count + symbol_count) / total_chars
|
||
english_ratio = calculate_english_ratio(subtitle)
|
||
|
||
if strict_mode:
|
||
digit_ratio = digit_count / total_chars
|
||
if total_chars >= 10 and digit_ratio <= 0.9 and special_ratio <= 0.9:
|
||
return True
|
||
# 数字占比 ≤ 70% 且特殊字符占比 ≤ 80% 且英文占比 ≤ 90%
|
||
return digit_ratio <= 0.7 and special_ratio <= 0.8 and english_ratio <= 0.9
|
||
else:
|
||
# 普通模式仅检查特殊字符占比
|
||
return special_ratio <= 0.5
|
||
|
||
|
||
def detect_main_region(image_files: List[str],
|
||
input_dir: str,
|
||
sample_size: int = 10,
|
||
motion_threshold: int = 30) -> Tuple[int, int, int, int]:
|
||
"""
|
||
通过运动检测确定主画面区域。
|
||
|
||
:return: 边界框 (x, y, width, height)
|
||
"""
|
||
if not image_files:
|
||
raise ValueError("图片列表为空")
|
||
|
||
first_img_path = os.path.join(input_dir, image_files[0])
|
||
prev_frame = cv2.imread(first_img_path)
|
||
if prev_frame is None:
|
||
raise ValueError(f"无法读取第一帧: {first_img_path}")
|
||
prev_gray = cv2.cvtColor(prev_frame, cv2.COLOR_BGR2GRAY)
|
||
height, width = prev_gray.shape
|
||
|
||
# 随机采样中间帧（避开开头和结尾）
|
||
mid_start = len(image_files) // 4
|
||
mid_end = 3 * len(image_files) // 4
|
||
sample_indices = random.sample(range(mid_start, mid_end),
|
||
min(sample_size, mid_end - mid_start))
|
||
|
||
motion_accumulator = np.zeros((height, width), dtype=np.float32)
|
||
|
||
for idx in sample_indices:
|
||
img_path = os.path.join(input_dir, image_files[idx])
|
||
curr_frame = cv2.imread(img_path)
|
||
if curr_frame is None:
|
||
continue
|
||
curr_gray = cv2.cvtColor(curr_frame, cv2.COLOR_BGR2GRAY)
|
||
frame_diff = cv2.absdiff(curr_gray, prev_gray)
|
||
_, thresh = cv2.threshold(frame_diff, motion_threshold, 255, cv2.THRESH_BINARY)
|
||
motion_accumulator += thresh.astype(np.float32)
|
||
prev_gray = curr_gray
|
||
|
||
motion_accumulator = cv2.normalize(motion_accumulator, None, 0, 255, cv2.NORM_MINMAX)
|
||
_, motion_mask = cv2.threshold(motion_accumulator.astype(np.uint8), 100, 255, cv2.THRESH_BINARY)
|
||
|
||
contours, _ = cv2.findContours(motion_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
||
if not contours:
|
||
# 无运动区域时返回中心偏下区域
|
||
return (0, height // 4, width, height // 2)
|
||
|
||
main_contour = max(contours, key=cv2.contourArea)
|
||
x, y, w, h = cv2.boundingRect(main_contour)
|
||
|
||
# 合理性检查
|
||
if w < width * 0.1 or h < height * 0.1:
|
||
return (0, height // 4, width, height // 2)
|
||
|
||
# 考虑下方空间
|
||
bottom_space = height - (y + h)
|
||
if bottom_space > height * 0.1:
|
||
extend = int(0.8 * bottom_space)
|
||
h += extend
|
||
# 调整 y 使区域居中于原运动区域下方
|
||
y = y + h // 2
|
||
h //= 2
|
||
else:
|
||
h += bottom_space
|
||
|
||
return (x, y, w, h)
|
||
|
||
|
||
# ==================== 字幕区域检测 ====================
|
||
def _sample_indices_uniform(total: int, sample_size: int) -> List[int]:
|
||
"""均匀采样索引，确保帧间有间隔。"""
|
||
print(total, sample_size)
|
||
if sample_size <= 1:
|
||
return [total // 2]
|
||
step = total // (sample_size - 1)
|
||
indices = [i * step for i in range(sample_size)]
|
||
# 确保不超出范围
|
||
indices = [min(i, total - 1) for i in indices]
|
||
return indices
|
||
|
||
|
||
def _ocr_frames_concurrent(image_files: List[str], input_dir: str,
|
||
indices: List[int], workers: int) -> Dict[int, str]:
|
||
"""并发对指定帧索引做 OCR，返回 {idx: 原始响应字符串}。
|
||
在视频主线程内开子线程池，子线程只做网络请求、不写日志（满足日志线程隔离要求）。"""
|
||
results: Dict[int, str] = {}
|
||
if not indices:
|
||
return results
|
||
|
||
def _one(idx):
|
||
try:
|
||
img_path = os.path.join(input_dir, image_files[idx])
|
||
return idx, tencentcloud_ocr_api(img2base64(img_path=img_path))
|
||
except Exception:
|
||
return idx, None
|
||
|
||
with ThreadPoolExecutor(max_workers=max(1, min(workers, len(indices)))) as ex:
|
||
for idx, res in ex.map(_one, indices):
|
||
results[idx] = res
|
||
return results
|
||
|
||
|
||
def cv_locate_subtitle(image_files: List[str], input_dir: str,
|
||
height: int, width: int, sample: int = 80):
|
||
"""纯 CV(零 OCR)定位硬字幕带。读已抽好的 JPEG 帧(cv2.imread,与 AV1 解码无关)。
|
||
返回 (verdict, region, info):
|
||
'subtitle' (x,y,w,h) —— 找到下部动态字幕带
|
||
'none' None —— 整体静止+无动态文字带 = PPT/文档/截图,可直接跳过
|
||
'uncertain'None —— 拿不准,回退原 OCR 检测(不改变原行为)
|
||
安全准则: 只有【整体静止且无任何动态文字带】才判 none,绝不误杀真字幕视频。"""
|
||
n = len(image_files)
|
||
if n < 10:
|
||
return "uncertain", None, {"reason": "帧太少"}
|
||
idxs = np.linspace(n * 0.1, n * 0.9, min(sample, n)).astype(int)
|
||
grays = []
|
||
for i in idxs:
|
||
g = cv2.imread(os.path.join(input_dir, image_files[int(i)]), cv2.IMREAD_GRAYSCALE)
|
||
if g is not None:
|
||
grays.append(g)
|
||
if len(grays) < 8:
|
||
return "uncertain", None, {"reason": "读帧失败"}
|
||
H, W = grays[0].shape
|
||
x0, x1 = int(0.15 * W), int(0.85 * W) # 只看横向中间70%,避开角落logo
|
||
|
||
edge_acc = np.zeros(H); present = np.zeros(H)
|
||
for g in grays:
|
||
e = cv2.Canny(g, 80, 200)
|
||
row = e[:, x0:x1].sum(axis=1) / 255.0
|
||
edge_acc += row
|
||
present += (row > (x1 - x0) * 0.02)
|
||
edge_profile = edge_acc / len(grays)
|
||
presence = present / len(grays)
|
||
diff_acc = np.zeros(H)
|
||
for k in range(1, len(grays)):
|
||
diff_acc += cv2.absdiff(grays[k], grays[k - 1])[:, x0:x1].mean(axis=1)
|
||
temporal = diff_acc / max(1, len(grays) - 1)
|
||
|
||
TEMP_HI = 8.0 # 带变化≥此 = 字幕在翻动(动态)
|
||
TEMP_LO = 4.0 # 文字带变化<此 = 静止
|
||
STATIC_MED = 3.0 # 全画面 temporal 中位数<此 = 整体静止(PPT/文档)
|
||
PRES = 0.4
|
||
med = float(np.median(temporal))
|
||
|
||
cand = (presence > PRES) & (edge_profile > edge_profile.max() * 0.25)
|
||
rows = np.where(cand)[0]
|
||
if len(rows) == 0:
|
||
if med < STATIC_MED:
|
||
return "none", None, {"reason": "无文字带且画面静止", "med": round(med, 1)}
|
||
return "uncertain", None, {"reason": "无稳定文字带"}
|
||
|
||
bands = []
|
||
s = p = rows[0]
|
||
for r in rows[1:]:
|
||
if r - p > H * 0.03:
|
||
bands.append((s, p)); s = r
|
||
p = r
|
||
bands.append((s, p))
|
||
MIN_H = int(H * 0.012)
|
||
bands = [b for b in bands if (b[1] - b[0]) >= MIN_H]
|
||
|
||
def bchg(b):
|
||
return float(temporal[b[0]:b[1] + 1].mean())
|
||
|
||
dyn = [b for b in bands if bchg(b) >= TEMP_HI] if bands else []
|
||
if not dyn:
|
||
# 无动态字幕带: 整体静止 或 所有文字带都静止 → PPT跳过; 否则交OCR
|
||
if med < STATIC_MED:
|
||
return "none", None, {"reason": "无动态字幕带且画面静止(PPT)", "med": round(med, 1)}
|
||
if bands and max(bchg(b) for b in bands) < TEMP_LO:
|
||
return "none", None, {"reason": "所有文字带静止", "max_chg": round(max(bchg(b) for b in bands), 1)}
|
||
return "uncertain", None, {"reason": "拿不准(交OCR)", "med": round(med, 1)}
|
||
|
||
dyn.sort(key=lambda b: presence[b[0]:b[1] + 1].mean() * (0.4 + (b[0] + b[1]) / 2 / H), reverse=True)
|
||
ys, ye = dyn[0]
|
||
band_h = (ye - ys) / H; cy = (ys + ye) / 2 / H
|
||
# 若抽帧已裁到下半屏(crop_bottom_ratio),帧内"偏上"不再代表标题,放宽位置判定。
|
||
crop_r = config.get_splitframe_config().get('crop_bottom_ratio')
|
||
try:
|
||
crop_r = float(crop_r) if crop_r is not None else None
|
||
except (TypeError, ValueError):
|
||
crop_r = None
|
||
min_cy = 0.15 if (crop_r and 0.0 < crop_r < 1.0) else 0.45
|
||
# 有动态带=【不是PPT】。带太宽/偏上只代表不敢直接给区域 → 回退OCR,绝不判none。
|
||
if band_h > 0.35:
|
||
return "uncertain", None, {"reason": "动态带太宽,交OCR", "band_h": round(band_h, 2)}
|
||
if cy < min_cy:
|
||
return "uncertain", None, {"reason": "带偏上,交OCR", "cy": round(cy, 2)}
|
||
pad = int(H * 0.012)
|
||
y = max(0, ys - pad); h = min(H, ye + pad) - y
|
||
return "subtitle", (0, int(y), int(W), int(h)), {"cy": round(cy, 2), "chg": round(bchg((ys, ye)), 1)}
|
||
|
||
|
||
def detect_subtitle_region(image_files: List[str],
|
||
input_dir: str,
|
||
sample_size: int = 15,
|
||
keep_num: int = 9) -> Tuple[Optional[Tuple[int, int, int, int]],
|
||
List[int],
|
||
List[str],
|
||
List[Dict]]:
|
||
"""
|
||
使用 OCR 检测字幕区域。
|
||
|
||
:return: (subtitle_region, sample_indices, watermark_logs, candidate_regions)
|
||
subtitle_region 为 (x, y, width, height)；若无字幕则返回 None
|
||
"""
|
||
_sub_cfg = config.get_subtitle_config()
|
||
require_chinese = _sub_cfg.get('require_chinese', True) # 只要中文字幕
|
||
bands_thr = int(_sub_cfg.get('redundant_bands', 12)) # 冗余阈值(垂直位置簇数)，越大越宽松
|
||
|
||
# 读取第一张图片获取尺寸
|
||
first_img_path = os.path.join(input_dir, image_files[0])
|
||
img = cv2.imread(first_img_path)
|
||
if img is None:
|
||
raise ValueError(f"无法读取图片: {first_img_path}")
|
||
height, width = img.shape[:2]
|
||
|
||
# === CV 快速闸门(零 OCR,~1.4s)===
|
||
# 'none' → 静态/PPT,直接跳过(省掉对这类视频的几千次 OCR);
|
||
# 'subtitle' 且开启 cv_use_region → 直接用 CV 区域(跳过 OCR 检测);
|
||
# 否则 → 继续走下面的 OCR 检测(行为不变)。
|
||
if _sub_cfg.get('cv_detect', True):
|
||
try:
|
||
cv_verdict, cv_region, cv_info = cv_locate_subtitle(image_files, input_dir, height, width)
|
||
except Exception as _e:
|
||
cv_verdict, cv_region, cv_info = "uncertain", None, {"reason": f"CV异常:{_e}"}
|
||
if cv_verdict == "none":
|
||
logger.log(f"CV定位: 无动态字幕(静态/PPT),跳过 [{cv_info}]", level="warning", program="merge")
|
||
return None, [], [], []
|
||
if cv_verdict == "subtitle" and _sub_cfg.get('cv_use_region', False):
|
||
logger.log(f"CV定位: 字幕区域 {cv_region} (零OCR) [{cv_info}]", program="merge")
|
||
x, y, w, h = cv_region
|
||
cand = [{"x_min": x, "y_min": y, "x_max": x + w, "y_max": y + h, "width": w, "height": h}]
|
||
return cv_region, [], [], cand
|
||
logger.log(f"CV定位: {cv_verdict} [{cv_info}] → 走OCR检测", program="merge")
|
||
|
||
# 均匀采样帧（扩大到 10%-90% 区间，提高命中字幕帧的概率）
|
||
mid_start = len(image_files) // 10
|
||
mid_end = 9 * len(image_files) // 10
|
||
sample_range = mid_end - mid_start
|
||
sample_indices = [mid_start + idx for idx in _sample_indices_uniform(sample_range, sample_size)]
|
||
|
||
# 存储候选区域和水印信息
|
||
candidate_regions = [] # 每个元素为 dict: {x_min, y_min, x_max, y_max, width, height}
|
||
text_positions = defaultdict(list) # {text: [{'x':..., 'y':..., 'idx':...}]}
|
||
watermark_logs = []
|
||
|
||
# 并发预取采样帧的 OCR（原为逐帧串行，~15帧×2s 全在等网络）
|
||
ocr_workers = config.get_subtitle_config().get('ocr_workers', 10)
|
||
ocr_cache = _ocr_frames_concurrent(image_files, input_dir, sample_indices, ocr_workers)
|
||
|
||
for idx in sample_indices:
|
||
ocr_result = ocr_cache.get(idx)
|
||
if not ocr_result:
|
||
continue
|
||
ocr_data = json.loads(ocr_result)
|
||
if "TextDetections" not in ocr_data:
|
||
continue
|
||
|
||
for detection in ocr_data["TextDetections"]:
|
||
if "Polygon" not in detection:
|
||
continue
|
||
polygon = detection["Polygon"]
|
||
xs = [point["X"] for point in polygon]
|
||
ys = [point["Y"] for point in polygon]
|
||
x_min, x_max = min(xs), max(xs)
|
||
y_min, y_max = min(ys), max(ys)
|
||
w = x_max - x_min
|
||
h = y_max - y_min
|
||
center_x = (x_min + x_max) // 2
|
||
center_y = (y_min + y_max) // 2
|
||
|
||
# 过滤左右边缘区域的文字（可能是水印或logo）
|
||
if center_x > width * 0.85 or center_x < width * 0.15:
|
||
continue
|
||
if detection.get('Confidence', 0) <= 75:
|
||
continue
|
||
|
||
detected_text = detection.get("DetectedText", "")
|
||
if not check_subtitle_quality_enhanced(detected_text) and detection.get('Confidence', 0) <= 90:
|
||
continue
|
||
if len(detected_text) < 5:
|
||
continue
|
||
# 只要中文字幕：没有中文字符的(纯英文/数字)不当作候选字幕区域
|
||
if require_chinese and not has_chinese(detected_text):
|
||
continue
|
||
|
||
# 水印检测：文本相似度>0.7且位置相近
|
||
is_watermark = False
|
||
for existing_text, pos_list in text_positions.items():
|
||
similarity = SequenceMatcher(None, detected_text, existing_text).ratio()
|
||
if similarity > 0.7:
|
||
# 检查位置是否相近
|
||
for pos in pos_list:
|
||
if pos["idx"] != idx and abs(pos["y"] - center_y) < h:
|
||
is_watermark = True
|
||
watermark_logs.append(
|
||
f"第{idx}帧可能的水印: {detected_text} (位置: {center_x}, {center_y})")
|
||
logger.log(f"检测到可能的水印: {detected_text} (位置: {center_x}, {center_y})",
|
||
program="merge")
|
||
# 标记原文本为水印
|
||
pos["is_watermark"] = True
|
||
break
|
||
if is_watermark:
|
||
break
|
||
|
||
if is_watermark:
|
||
continue
|
||
|
||
# 记录文本位置
|
||
text_positions[detected_text].append({
|
||
"text": detected_text,
|
||
"x": center_x,
|
||
"y": center_y,
|
||
"is_watermark": False,
|
||
"idx": idx
|
||
})
|
||
|
||
# 存储候选区域
|
||
candidate_regions.append({
|
||
"x_min": x_min,
|
||
"y_min": y_min,
|
||
"x_max": x_max,
|
||
"y_max": y_max,
|
||
"width": w,
|
||
"height": h
|
||
})
|
||
|
||
# 如果没有候选区域，检查首尾帧是否有可能存在字幕
|
||
if not candidate_regions:
|
||
# 检查前 keep_num 张和后 keep_num 张是否有文字（并发）
|
||
keep_num = 9
|
||
n = len(image_files)
|
||
check_indices = sorted(set(
|
||
list(range(min(keep_num, n))) + list(range(max(0, n - keep_num), n))
|
||
))
|
||
fb_cache = _ocr_frames_concurrent(image_files, input_dir, check_indices, ocr_workers)
|
||
has_text = False
|
||
for idx in check_indices:
|
||
ocr_result = fb_cache.get(idx)
|
||
if ocr_result:
|
||
ocr_data = json.loads(ocr_result)
|
||
if "TextDetections" in ocr_data and ocr_data["TextDetections"]:
|
||
has_text = True
|
||
break
|
||
|
||
if not has_text:
|
||
logger.log("保留的前后9张图片均未检测到任何文字，认为该视频没有字幕", level="warning", program="merge")
|
||
return None, sample_indices, watermark_logs, candidate_regions
|
||
|
||
# 无候选区域但有文字：回退到主区域检测
|
||
main_region = detect_main_region(image_files, input_dir)
|
||
x, y, w, h = main_region
|
||
# 只取主区域下半部分
|
||
return (0, y + h // 2, width, h // 2), sample_indices, watermark_logs, candidate_regions
|
||
|
||
# 逆向清理：只清理画面上半部分（y < height*0.5）的水印候选区域
|
||
# 下半部分即便被误判为水印也保留，避免把字幕字幕一并删掉
|
||
confirmed_watermark_ys = {
|
||
pos["y"]
|
||
for pos_list in text_positions.values()
|
||
for pos in pos_list
|
||
if pos.get("is_watermark")
|
||
}
|
||
top_watermark_ys = {wy for wy in confirmed_watermark_ys if wy < height * 0.5}
|
||
if top_watermark_ys:
|
||
before = len(candidate_regions)
|
||
candidate_regions = [
|
||
r for r in candidate_regions
|
||
if not any(
|
||
abs((r["y_min"] + r["y_max"]) // 2 - wy) < 80
|
||
for wy in top_watermark_ys
|
||
)
|
||
]
|
||
removed = before - len(candidate_regions)
|
||
if removed:
|
||
logger.log(f"逆向移除 {removed} 个上半区水印候选区域", program="merge")
|
||
|
||
# 过滤过小的区域（高度<15）
|
||
candidate_regions = [r for r in candidate_regions if r["height"] >= 15]
|
||
|
||
# 水印移除/高度过滤后可能把候选全部清空 → 判无字幕(否则下面 min() 会因空序列崩成"拼接失败")
|
||
if not candidate_regions:
|
||
logger.log("候选字幕区域被过滤后为空(高度过小/全是水印)，判无字幕", level="warning", program="merge")
|
||
return None, sample_indices, watermark_logs, candidate_regions
|
||
|
||
# 冗余判定：用"不同垂直位置簇"的数量，而非候选框总数。
|
||
# 正常字幕会在多帧的同一位置重复出现，框数随采样帧数增长但集中在少数垂直位置；
|
||
# 纯文本/水印过多的视频才会把文字散布到很多不同的垂直位置。
|
||
region_cys = sorted((r["y_min"] + r["y_max"]) // 2 for r in candidate_regions)
|
||
distinct_bands = 1 if region_cys else 0
|
||
for _i in range(1, len(region_cys)):
|
||
if region_cys[_i] - region_cys[_i - 1] > 80:
|
||
distinct_bands += 1
|
||
if distinct_bands >= bands_thr:
|
||
logger.log(f"冗余文字过多（文字分散在{distinct_bands}个垂直位置，共{len(candidate_regions)}个候选区域），视频将被过滤",
|
||
level="warning", program="merge")
|
||
return None, sample_indices, watermark_logs, candidate_regions
|
||
|
||
# 合并所有候选区域
|
||
all_y_min = min(r["y_min"] for r in candidate_regions)
|
||
all_y_max = max(r["y_max"] for r in candidate_regions)
|
||
merged_h = all_y_max - all_y_min
|
||
# 扩展边界
|
||
expanded_y_min = max(0, int(all_y_min - 65))
|
||
expanded_y_max = min(height, int(all_y_max + 65))
|
||
final_h = expanded_y_max - expanded_y_min
|
||
|
||
# 如果合并区域过高（>60% 画面高度），用聚类找最大候选簇，避免退回全帧
|
||
if final_h >= height * 0.6 and len(candidate_regions) >= 4:
|
||
logger.log("合并区域过高，聚类寻找最大字幕区域簇", program="merge")
|
||
best_cluster = []
|
||
for region in candidate_regions:
|
||
cy = (region["y_min"] + region["y_max"]) // 2
|
||
cluster = [r for r in candidate_regions
|
||
if abs((r["y_min"] + r["y_max"]) // 2 - cy) < 200]
|
||
# 优先更大的簇；同等大小时选中心 y 更大的（靠近底部的字幕）
|
||
if len(cluster) > len(best_cluster) or (
|
||
len(cluster) == len(best_cluster) and best_cluster and
|
||
sum((r["y_min"] + r["y_max"]) // 2 for r in cluster) >
|
||
sum((r["y_min"] + r["y_max"]) // 2 for r in best_cluster)):
|
||
best_cluster = cluster
|
||
|
||
if best_cluster:
|
||
c_y_min = min(r["y_min"] for r in best_cluster)
|
||
c_y_max = max(r["y_max"] for r in best_cluster)
|
||
exp_y_min = max(0, c_y_min - 65)
|
||
exp_y_max = min(height, c_y_max + 65)
|
||
logger.log(f"聚类选定字幕区域: y={exp_y_min}-{exp_y_max}", program="merge")
|
||
return (0, exp_y_min, width, exp_y_max - exp_y_min), sample_indices, watermark_logs, candidate_regions
|
||
|
||
main_region = detect_main_region(image_files, input_dir)
|
||
x, y, w, h = main_region
|
||
return (0, y, width, h), sample_indices, watermark_logs, candidate_regions
|
||
|
||
return (0, expanded_y_min, width, final_h), sample_indices, watermark_logs, candidate_regions
|
||
|
||
|
||
# ==================== 拼图配置计算 ====================
|
||
def calculate_optimal_tile_configuration(subtitle_region: Tuple[int, int, int, int],
|
||
tile_width: int,
|
||
tile_height: int,
|
||
first_img_path: str,
|
||
logger,
|
||
max_base64_size: int = 10 * 1024 * 1024) -> Tuple[int, int]:
|
||
"""
|
||
计算最优的拼图配置（横向和纵向数量）。
|
||
|
||
:return: (tile_width, tile_height)
|
||
"""
|
||
x, y, w, h = subtitle_region
|
||
BORDER_SIZE = 40
|
||
tile_width_px = w + 2 * BORDER_SIZE
|
||
tile_height_px = h + 2 * BORDER_SIZE
|
||
|
||
# 估算单个图块的大小（基于第一张图片）
|
||
img = cv2.imread(first_img_path)
|
||
if img is None:
|
||
raise ValueError(f"无法读取图片: {first_img_path}")
|
||
img_h, img_w = img.shape[:2]
|
||
img_size = os.path.getsize(first_img_path)
|
||
single_tile_size = (tile_width_px * tile_height_px) / (img_w * img_h) * img_size * 3 # RGB估算
|
||
single_tile_base64 = single_tile_size * 4 // 3 # base64 增加 33%
|
||
|
||
# 每张合成图最多拼几块：硬上限(可配置) 与 base64 10MB 限制 取小。
|
||
# 加大上限 → OCR 请求数下降(每张装更多帧)，但单次 OCR 图更大、且文字条更密，
|
||
# 过大有 OCR 漏识别/超分辨率上限的风险，调大后务必抽查 SRT 质量。
|
||
hard_cap = config.get_subtitle_config().get('max_tiles_per_composite', 40)
|
||
try:
|
||
hard_cap = int(hard_cap)
|
||
except (TypeError, ValueError):
|
||
hard_cap = 40
|
||
max_tiles = min(hard_cap, int(max_base64_size // single_tile_base64))
|
||
max_tiles = max(1, max_tiles)
|
||
# 横向块数上限随总块数放宽(块多时仍想接近 0.5 的宽高比，避免拼成超细长条)
|
||
w_cap = max(7, int(round((max_tiles * 0.5) ** 0.5)) + 1)
|
||
logger.log(f"单个图块 base64 大小: {single_tile_base64 / 1024:.1f} KB", program="merge")
|
||
logger.log(f"最大允许图块数量: {max_tiles} (硬上限 {hard_cap})", program="merge")
|
||
|
||
# 寻找最佳宽高比（接近 0.5，即横向多于纵向）
|
||
best_config = (tile_width, tile_height) # 默认
|
||
best_ratio_diff = float('inf')
|
||
for w_tiles in range(1, min(max_tiles, w_cap) + 1):
|
||
h_tiles = int(max_tiles / w_tiles)
|
||
aspect = w_tiles / h_tiles
|
||
diff = abs(aspect - 0.5)
|
||
if diff < best_ratio_diff:
|
||
best_ratio_diff = diff
|
||
best_config = (w_tiles, h_tiles)
|
||
|
||
logger.log(f"推荐拼图配置: {best_config[0]} x {best_config[1]}", program="merge")
|
||
return best_config
|
||
|
||
|
||
# ==================== 音频提取（可复用） ====================
|
||
def extract_audio_text(video_path: str, logger) -> Dict[str, str]:
|
||
"""
|
||
从视频中提取开头和结尾的音频文字。
|
||
|
||
:return: {'first_3s': str, 'last_3s': str}
|
||
"""
|
||
audio_text = {}
|
||
if not video_path or not os.path.exists(video_path):
|
||
return audio_text
|
||
|
||
logger.log("开始提取视频中的音频文字...", program="merge")
|
||
# 从配置读取 whisper 模型大小（默认 small）
|
||
model_size = config.get_audio_config().get('model_size', 'small')
|
||
whisper_manager.load_model(model_size)
|
||
|
||
temp_audio_path = os.path.join(tempfile.gettempdir(),
|
||
f"temp_audio_{os.path.basename(video_path)}.wav")
|
||
try:
|
||
total_duration = whisper_manager.extract_audio_from_video(video_path, temp_audio_path, logger)
|
||
if total_duration is None:
|
||
logger.log("音频提取失败，跳过音频文字识别", program="merge")
|
||
return audio_text
|
||
|
||
if total_duration >= 5:
|
||
first_5s_text = whisper_manager.extract_audio_text_with_whisper(
|
||
temp_audio_path, 0, min(5, total_duration), logger
|
||
)
|
||
audio_text["first_3s"] = first_5s_text
|
||
logger.log(f"前五秒音频文字: {first_5s_text}", program="merge")
|
||
|
||
if total_duration >= 10:
|
||
last_5s_start = int(total_duration - 5)
|
||
last_5s_text = whisper_manager.extract_audio_text_with_whisper(
|
||
temp_audio_path, last_5s_start, total_duration, logger
|
||
)
|
||
audio_text["last_3s"] = last_5s_text
|
||
logger.log(f"后五秒音频文字: {last_5s_text}", program="merge")
|
||
finally:
|
||
try:
|
||
os.unlink(temp_audio_path)
|
||
except Exception:
|
||
pass
|
||
|
||
return audio_text
|
||
|
||
|
||
# ==================== 主处理函数 ====================
|
||
def process_images(input_dir, output_dir, logger, tile_width=3, tile_height=6, keep_num=9, interval_seconds=None,
|
||
video_path=None):
|
||
"""
|
||
处理图片拼接并生成元数据，自动计算最优拼图配置
|
||
|
||
参数:
|
||
input_dir: 输入图片目录
|
||
output_dir: 输出目录
|
||
keep_num: 保留的首尾图片数量
|
||
"""
|
||
# 创建输出目录
|
||
logger.start_program("merge")
|
||
os.makedirs(output_dir, exist_ok=True)
|
||
composite_dir = os.path.join(output_dir, "composite")
|
||
os.makedirs(composite_dir, exist_ok=True)
|
||
|
||
# 获取所有图片文件并排序
|
||
image_files = [f for f in os.listdir(input_dir) if f.lower().endswith(('.png', '.jpg', '.jpeg'))]
|
||
image_files.sort(key=natural_sort_key)
|
||
total_images = len(image_files)
|
||
duration = total_images / 3
|
||
# 确定字幕区域（通过OCR检测）
|
||
subtitle_region, sample_indices, watermark, candidate_regions = detect_subtitle_region(image_files, input_dir)
|
||
|
||
if subtitle_region is None:
|
||
logger.end_program("merge")
|
||
return duration, None, None, None, None, None, sample_indices, watermark, candidate_regions
|
||
logger.log(f"通过OCR检测确定字幕区域: {subtitle_region}", program="merge")
|
||
first_img_path = os.path.join(input_dir, image_files[0])
|
||
|
||
# 提取音频文字（如果提供了视频路径）
|
||
audio_text = {}
|
||
if video_path and os.path.exists(video_path):
|
||
logger.log("开始提取视频中的音频文字...", program="merge")
|
||
model_size = config.get_audio_config().get('model_size', 'small')
|
||
whisper_manager.load_model(model_size)
|
||
# 创建临时音频文件
|
||
temp_audio_path = os.path.join(tempfile.gettempdir(), f"temp_audio_{os.path.basename(video_path)}.wav")
|
||
total_duration = whisper_manager.extract_audio_from_video(video_path, temp_audio_path, logger)
|
||
# 从视频提取音频
|
||
if total_duration is not None:
|
||
|
||
# 提取前五秒文字
|
||
if total_duration >= 5:
|
||
first_5s_text = whisper_manager.extract_audio_text_with_whisper(temp_audio_path, 0,
|
||
min(5, total_duration), logger)
|
||
audio_text["first_3s"] = first_5s_text
|
||
logger.log(f"前五秒音频文字: {first_5s_text}", program="merge")
|
||
|
||
# 提取后五秒文字
|
||
if total_duration >= 10:
|
||
last_5s_start = int(total_duration - 5)
|
||
last_5s_text = whisper_manager.extract_audio_text_with_whisper(temp_audio_path, last_5s_start,
|
||
total_duration, logger)
|
||
audio_text["last_3s"] = last_5s_text
|
||
logger.log(f"后五秒音频文字: {last_5s_text}", program="merge")
|
||
|
||
# 清理临时音频文件
|
||
try:
|
||
os.unlink(temp_audio_path)
|
||
except OSError:
|
||
pass
|
||
else:
|
||
logger.log("音频提取失败，跳过音频文字识别", program="merge")
|
||
# 边框设置
|
||
BORDER_SIZE = 40
|
||
|
||
# 保留前keep_num张和最后keep_num张图片
|
||
kept_first = image_files[:keep_num]
|
||
kept_last = image_files[-keep_num:] if len(image_files) > 2 * keep_num else []
|
||
|
||
# 需要处理的图片（中间部分）
|
||
process_files = image_files[keep_num:-keep_num] if len(image_files) > 2 * keep_num else image_files[keep_num:]
|
||
|
||
# 归一化图块像素预算：高分辨率帧或过高的字幕区域会让单个图块的 base64 过大，
|
||
# 撑到 10MB 上限后每张合成图能塞的块数骤减（从~32 掉到 5~8），导致合成图（=OCR请求）
|
||
# 数量激增 3~6 倍。按基准像素量（约 1080p 单字幕条）等比缩小裁剪后的字幕条，使每张
|
||
# 合成图的块数稳定在~32，OCR 成本与源分辨率/区域高度无关；文字分辨率仍足够 OCR。
|
||
x0, y0, w0, h0 = subtitle_region
|
||
no_composite = config.get_subtitle_config().get('no_composite', False)
|
||
if no_composite:
|
||
# 不拼接：每个字幕条单独成图(tile 1×1)、不降采样(单条 base64 很小，全分辨率最准)。
|
||
# 下游 OCR 映射/网格搜索/去重完全复用，只是每张图=1 条、OCR 次数=帧数。
|
||
tile_scale = 1.0
|
||
w_scaled, h_scaled = w0, h0
|
||
tile_width, tile_height = 1, 1
|
||
logger.log("不拼接模式：每个字幕条单独 OCR(tile 1×1，全分辨率)", program="merge")
|
||
else:
|
||
# 归一化图块像素预算：高分辨率帧或过高的字幕区域会让单个图块的 base64 过大，
|
||
# 撑到 10MB 上限后每张合成图能塞的块数骤减（从~32 掉到 5~8），导致合成图（=OCR请求）
|
||
# 数量激增 3~6 倍。按基准像素量（约 1080p 单字幕条）等比缩小裁剪后的字幕条，使每张
|
||
# 合成图的块数稳定在~32，OCR 成本与源分辨率/区域高度无关；文字分辨率仍足够 OCR。
|
||
TARGET_TILE_AREA = 1920 * 240
|
||
tile_scale = min(1.0, math.sqrt(TARGET_TILE_AREA / float(max(1, w0 * h0))))
|
||
w_scaled = max(1, int(round(w0 * tile_scale)))
|
||
h_scaled = max(1, int(round(h0 * tile_scale)))
|
||
if tile_scale < 1.0:
|
||
logger.log(f"字幕条降采样: {w0}x{h0} → {w_scaled}x{h_scaled} (scale={tile_scale:.2f})", program="merge")
|
||
tile_width, tile_height = calculate_optimal_tile_configuration(
|
||
(x0, y0, w_scaled, h_scaled), tile_width, tile_height, first_img_path, logger)
|
||
|
||
# 初始化元数据
|
||
metadata = {
|
||
"tile_width": tile_width,
|
||
"tile_height": tile_height,
|
||
"subtitle_region": subtitle_region,
|
||
"border_size": BORDER_SIZE,
|
||
"interval_seconds": interval_seconds,
|
||
"audio_text": audio_text,
|
||
"kept_first": [],
|
||
"kept_last": [],
|
||
"composites": []
|
||
}
|
||
|
||
# 复制保留的图片到composite文件夹
|
||
kept_first_new_names = []
|
||
for i, img in enumerate(kept_first):
|
||
src = os.path.join(input_dir, img)
|
||
dst_filename = f"001_{i + 1:02d}_original_{img}"
|
||
dst = os.path.join(composite_dir, dst_filename)
|
||
|
||
if not os.path.exists(dst):
|
||
with Image.open(src) as im:
|
||
bordered_img = Image.new('RGB',
|
||
(im.width + 2 * BORDER_SIZE, im.height + 2 * BORDER_SIZE),
|
||
(255, 255, 255))
|
||
bordered_img.paste(im, (BORDER_SIZE, BORDER_SIZE))
|
||
|
||
if img.lower().endswith('.png'):
|
||
bordered_img.save(dst, format='PNG')
|
||
else:
|
||
bordered_img.save(dst, format='JPEG', quality=95)
|
||
|
||
kept_first_new_names.append(dst_filename)
|
||
|
||
metadata["kept_first"] = kept_first_new_names
|
||
|
||
# 处理中间图片进行拼接
|
||
if process_files:
|
||
x, y, w, h = subtitle_region # 原始帧上的裁剪坐标
|
||
ws, hs = w_scaled, h_scaled # 降采样后的图块尺寸（拼接/坐标用此）
|
||
|
||
# 计算图块和合成图尺寸（基于降采样后的尺寸）
|
||
tile_img_width = ws + 2 * BORDER_SIZE
|
||
tile_img_height = hs + 2 * BORDER_SIZE
|
||
composite_width = tile_width * tile_img_width
|
||
composite_height = tile_height * tile_img_height
|
||
|
||
# 分批：每批 = 一张合成图
|
||
tiles_per_composite = tile_width * tile_height
|
||
batches = []
|
||
for i in range(0, len(process_files), tiles_per_composite):
|
||
batches.append((len(batches) + 1, process_files[i:i + tiles_per_composite]))
|
||
|
||
def _build_composite(batch):
|
||
"""处理一张合成图：读源帧→裁剪→贴块→写盘。返回 (序号, composite_data)。
|
||
热点：逐帧读 1080p JPEG + 裁剪 + 写盘，原为单线程；这里并行化（PIL 编解码
|
||
在 C 层释放 GIL，线程池能真正并行，按核数提速）。"""
|
||
composite_num, batch_files = batch
|
||
composite_img = Image.new('RGB', (composite_width, composite_height), (255, 255, 255))
|
||
composite_filename = f"002_{composite_num:04d}_composite.jpg"
|
||
composite_data = {
|
||
"composite_id": f"composite_{composite_num:04d}",
|
||
"file_name": composite_filename,
|
||
"tiles": []
|
||
}
|
||
for pos in range(tiles_per_composite):
|
||
row = pos // tile_width
|
||
col = pos % tile_width
|
||
x_pos = col * tile_img_width
|
||
y_pos = row * tile_img_height
|
||
|
||
if pos < len(batch_files):
|
||
img_name = batch_files[pos]
|
||
img_path = os.path.join(input_dir, img_name)
|
||
with Image.open(img_path) as img:
|
||
subtitle_img = img.crop((x, y, x + w, y + h))
|
||
if tile_scale < 1.0:
|
||
subtitle_img = subtitle_img.resize((ws, hs))
|
||
# 合成图背景已是白色，直接把字幕条贴到带边框的偏移位置
|
||
composite_img.paste(subtitle_img, (x_pos + BORDER_SIZE, y_pos + BORDER_SIZE))
|
||
composite_data["tiles"].append({
|
||
"position": pos,
|
||
"x": x_pos + BORDER_SIZE,
|
||
"y": y_pos + BORDER_SIZE,
|
||
"width": ws,
|
||
"height": hs,
|
||
"source_image": img_name,
|
||
"is_blank": False
|
||
})
|
||
else:
|
||
# 空白块：合成图背景本就是白色，无需粘贴
|
||
composite_data["tiles"].append({
|
||
"position": pos,
|
||
"x": x_pos + BORDER_SIZE,
|
||
"y": y_pos + BORDER_SIZE,
|
||
"width": ws,
|
||
"height": hs,
|
||
"source_image": None,
|
||
"is_blank": True
|
||
})
|
||
|
||
output_path = os.path.join(composite_dir, composite_filename)
|
||
composite_img.save(output_path, format='JPEG', quality=95)
|
||
return composite_num, composite_data
|
||
|
||
# 并行处理所有合成图，结果按序号排序回填（顺序不能乱，元数据要对得上时间轴）
|
||
merge_workers = config.get_subtitle_config().get('merge_workers') or (os.cpu_count() or 8)
|
||
merge_workers = max(1, min(int(merge_workers), len(batches)))
|
||
results = []
|
||
with ThreadPoolExecutor(max_workers=merge_workers) as executor:
|
||
for r in executor.map(_build_composite, batches):
|
||
results.append(r)
|
||
results.sort(key=lambda t: t[0])
|
||
metadata["composites"] = [cd for _, cd in results]
|
||
|
||
# 处理最后保留的图片
|
||
kept_last_new_names = []
|
||
for i, img in enumerate(kept_last):
|
||
src = os.path.join(input_dir, img)
|
||
dst_filename = f"003_{i + 1:02d}_original_{img}"
|
||
dst = os.path.join(composite_dir, dst_filename)
|
||
|
||
if not os.path.exists(dst):
|
||
with Image.open(src) as im:
|
||
bordered_img = Image.new('RGB',
|
||
(im.width + 2 * BORDER_SIZE, im.height + 2 * BORDER_SIZE),
|
||
(255, 255, 255))
|
||
bordered_img.paste(im, (BORDER_SIZE, BORDER_SIZE))
|
||
|
||
if img.lower().endswith('.png'):
|
||
bordered_img.save(dst, format='PNG')
|
||
else:
|
||
bordered_img.save(dst, format='JPEG', quality=95)
|
||
|
||
kept_last_new_names.append(dst_filename)
|
||
|
||
metadata["kept_last"] = kept_last_new_names
|
||
|
||
# 保存元数据（保存到 composite_dir，避免被清理时删除）
|
||
metadata_path = os.path.join(composite_dir, "metadata.json")
|
||
with open(metadata_path, 'w') as f:
|
||
json.dump(metadata, f, indent=2)
|
||
|
||
logger.log(f"处理完成，拼图配置: {tile_width}x{tile_height}", program="merge")
|
||
logger.log(f"共生成 {len(metadata['composites'])} 张合成图", program="merge")
|
||
logger.log(f"元数据已保存到: {metadata_path}", program="merge")
|
||
logger.end_program("merge")
|
||
return duration, subtitle_region, tile_width, tile_height, len(
|
||
metadata['composites']), audio_text, sample_indices, watermark, candidate_regions
|
||
|
||
|
||
if __name__ == "__main__":
|
||
# 示例用法
|
||
input_directory = "input_images"
|
||
output_directory = "output"
|
||
from log.log_manager import logger # 假设已有全局 logger
|
||
|
||
result = process_images(input_directory, output_directory, logger)
|