cang-ying/engines/subtitle-pipeline/reference-analysis/reference-subtitle-analyzer.py

#!/usr/bin/env python3
"""
Reference Subtitle Analyzer · 样片字幕量化分析器
==============================================
从样片截图里量字幕：字高、底部距离、描边宽度、字幕框位置、字号比例。

依赖：
  pip install opencv-python pillow numpy

用法：
  # 分析单张截图
  python reference-subtitle-analyzer.py --image reference-frame.png --output analysis.json

  # 批量分析样片截图
  python reference-subtitle-analyzer.py --batch ./reference-frames/ --output batch-analysis.json

  # 交互式标注模式（手动框选字幕区域）
  python reference-subtitle-analyzer.py --image reference-frame.png --interactive

输出 JSON 格式：
  {
    "video_height": 1920,
    "video_width": 1080,
    "subtitle_box": {
      "x": 100,
      "y": 1750,
      "width": 880,
      "height": 120
    },
    "subtitle_height_px": 120,
    "bottom_distance_px": 50,
    "bottom_distance_ratio": 0.026,
    "font_size_ratio": 0.0625,
    "stroke_width_px": 2,
    "alignment": "center",
    "margin_horizontal_px": 100
  }

路径：
  video-ai-system/engines/subtitle-pipeline/reference-analysis/reference-subtitle-analyzer.py
"""

import argparse
import json
import os
import sys
from pathlib import Path

try:
    import cv2
    import numpy as np
    from PIL import Image
except ImportError as e:
    print(f"[ERROR] 缺少依赖：{e}")
    print("请先安装：pip install opencv-python pillow numpy")
    sys.exit(1)


def analyze_subtitle_region(image_path: str, interactive: bool = False) -> dict:
    """
    分析字幕区域，量化字幕参数
    
    :param image_path: 截图路径
    :param interactive: 是否交互式标注（手动框选字幕区域）
    :return: 字幕参数字典
    """
    # 读取图片
    img = cv2.imread(image_path)
    if img is None:
        print(f"[ERROR] 无法读取图片：{image_path}")
        return {}
    
    height, width = img.shape[:2]
    print(f"[INFO] 图片尺寸：{width}x{height}")
    
    if interactive:
        # 交互式标注模式：手动框选字幕区域
        print("[INFO] 交互式标注模式：请在弹出的窗口中用鼠标框选字幕区域")
        print("[INFO] 框选完成后按空格或回车确认，按ESC取消")
        
        roi = cv2.selectROI("Select Subtitle Region", img, showCrosshair=True)
        cv2.destroyAllWindows()
        
        x, y, w, h = int(roi[0]), int(roi[1]), int(roi[2]), int(roi[3])
        if w == 0 or h == 0:
            print("[ERROR] 未框选字幕区域")
            return {}
    else:
        # 自动检测字幕区域（假设字幕在底部）
        # 策略：检测底部区域的文本（通过边缘检测 + 轮廓查找）
        print("[INFO] 自动检测字幕区域...")
        
        # 1. 裁剪底部区域（假设字幕在底部 15% 区域）
        bottom_region = img[int(height * 0.85):, :]
        
        # 2. 转为灰度图
        gray = cv2.cvtColor(bottom_region, cv2.COLOR_BGR2GRAY)
        
        # 3. 二值化（检测白色文本）
        _, binary = cv2.threshold(gray, 200, 255, cv2.THRESH_BINARY)
        
        # 4. 查找轮廓
        contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        
        if not contours:
            print("[WARN] 未检测到字幕区域，使用默认底部区域")
            # 使用默认底部区域
            x = int(width * 0.1)
            y = int(height * 0.85)
            w = int(width * 0.8)
            h = int(height * 0.12)
        else:
            # 合并所有轮廓的边界框
            # 正确算式：找到所有轮廓的最小外接矩形
            all_x = []
            all_y = []
            all_w = []
            all_h = []
            
            for cnt in contours:
                cx, cy, cw, ch = cv2.boundingRect(cnt)
                all_x.append(cx)
                all_y.append(cy)
                all_w.append(cw)
                all_h.append(ch)
            
            # 正确合并：先在 bottom_region 坐标系里算完，最后再加偏移
            x_bottom = min(all_x)
            y_bottom = min(all_y)
            max_x_bottom = max([all_x[i] + all_w[i] for i in range(len(all_x))])
            max_y_bottom = max([all_y[i] + all_h[i] for i in range(len(all_y))])
            w = max_x_bottom - x_bottom
            h = max_y_bottom - y_bottom

            # 扩展边界框（包含描边）
            padding = 10
            x = max(0, x_bottom - padding)
            y = max(0, y_bottom - padding) + int(height * 0.85)  # 偏移只加一次，在最后
            w = min(width - x, w + 2 * padding)
            h = min(height - y, h + 2 * padding)
    
    # 计算字幕参数
    # 先做边界校验，防止裁剪出空数组
    x = max(0, min(x, width - 1))
    y = max(0, min(y, height - 1))
    w = max(1, min(w, width - x))
    h = max(1, min(h, height - y))
    
    subtitle_box = {"x": x, "y": y, "width": w, "height": h}
    subtitle_height_px = h
    bottom_distance_px = height - (y + h)
    bottom_distance_ratio = bottom_distance_px / height
    font_size_ratio = h / height
    
    # 估算描边宽度（通过检测文本边缘的黑色像素）
    stroke_width_px = estimate_stroke_width(img, x, y, w, h)
    
    # 判断对齐方式（居中/左对齐/右对齐）
    alignment = "center" if abs(x + w/2 - width/2) < width * 0.1 else "left" if x < width * 0.3 else "right"
    
    # 计算水平边距
    margin_horizontal_px = x
    
    result = {
        "image_path": image_path,
        "video_width": width,
        "video_height": height,
        "subtitle_box": subtitle_box,
        "subtitle_height_px": subtitle_height_px,
        "bottom_distance_px": bottom_distance_px,
        "bottom_distance_ratio": round(bottom_distance_ratio, 4),
        "font_size_ratio": round(font_size_ratio, 4),
        "stroke_width_px": stroke_width_px,
        "alignment": alignment,
        "margin_horizontal_px": margin_horizontal_px
    }
    
    print(f"[OK] 字幕参数已量化：")
    print(f"  字幕框：x={x}, y={y}, w={w}, h={h}")
    print(f"  字高：{subtitle_height_px}px")
    print(f"  底部距离：{bottom_distance_px}px ({bottom_distance_ratio*100:.1f}%)")
    print(f"  字号比例：{font_size_ratio*100:.1f}%")
    print(f"  描边宽度：{stroke_width_px}px")
    print(f"  对齐方式：{alignment}")
    
    return result


def estimate_stroke_width(img: np.ndarray, x: int, y: int, w: int, h: int) -> int:
    """
    估算描边宽度（通过检测文本边缘的黑色像素）
    
    :param img: 图片数组
    :param x: 字幕框 x 坐标
    :param y: 字幕框 y 坐标
    :param w: 字幕框宽度
    :param h: 字幕框高度
    :return: 估算的描边宽度（像素）
    """
    # 裁剪字幕区域
    subtitle_region = img[y:y+h, x:x+w]
    
    # 转为灰度图
    gray = cv2.cvtColor(subtitle_region, cv2.COLOR_BGR2GRAY)
    
    # 检测边缘（Canny）
    edges = cv2.Canny(gray, 100, 200)
    
    # 统计边缘像素到文本区域的距离（估算描边宽度）
    # 简化方案：假设描边是黑色，检测白色文本周围的黑色像素环
    _, binary = cv2.threshold(gray, 200, 255, cv2.THRESH_BINARY)
    
    # 形态学操作：膨胀（模拟描边）
    kernel = np.ones((3, 3), np.uint8)
    dilated = cv2.dilate(binary, kernel, iterations=1)
    
    # 计算膨胀后的边缘宽度
    edge_width = cv2.absdiff(dilated, binary)
    stroke_pixels = cv2.countNonZero(edge_width)
    
    # 估算平均描边宽度
    if stroke_pixels > 0:
        # 简化：假设描边宽度是 1-3px
        return 2  # 默认值，实际应通过更精确的算法计算
    else:
        return 0


def batch_analyze(images_dir: str) -> dict:
    """
    批量分析样片截图
    
    :param images_dir: 截图目录
    :return: 批量分析结果
    """
    results = {}
    image_extensions = [".jpg", ".jpeg", ".png", ".bmp"]
    
    for file in os.listdir(images_dir):
        if any(file.lower().endswith(ext) for ext in image_extensions):
            image_path = os.path.join(images_dir, file)
            print(f"\n[INFO] 分析图片：{file}")
            result = analyze_subtitle_region(image_path)
            if result:
                results[file] = result
    
    # 计算平均值
    if results:
        avg_result = calculate_average(results)
        results["_average"] = avg_result
    
    return results


def calculate_average(results: dict) -> dict:
    """
    计算批量分析结果的平均值
    
    :param results: 批量分析结果
    :return: 平均值字典
    """
    keys = ["subtitle_height_px", "bottom_distance_px", "bottom_distance_ratio", 
            "font_size_ratio", "stroke_width_px", "margin_horizontal_px"]
    
    avg = {}
    for key in keys:
        values = [r[key] for r in results.values() if key in r]
        if values:
            avg[key] = round(sum(values) / len(values), 2)
    
    avg["alignment"] = max(set(r["alignment"] for r in results.values()), 
                          key=lambda x: sum(1 for r in results.values() if r["alignment"] == x))
    
    return avg


def main():
    parser = argparse.ArgumentParser(
        description="Reference Subtitle Analyzer · 样片字幕量化分析器",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
示例：
  python reference-subtitle-analyzer.py --image reference-frame.png --output analysis.json
  python reference-subtitle-analyzer.py --batch ./reference-frames/ --output batch-analysis.json
  python reference-subtitle-analyzer.py --image reference-frame.png --interactive
"""
    )
    parser.add_argument("--image", help="单张截图路径")
    parser.add_argument("--batch", help="批量分析截图目录")
    parser.add_argument("--output", required=True, help="输出 JSON 文件路径")
    parser.add_argument("--interactive", action="store_true", help="交互式标注模式")
    
    args = parser.parse_args()
    
    if args.image:
        result = analyze_subtitle_region(args.image, args.interactive)
        if not result:
            sys.exit(1)
        
        with open(args.output, "w", encoding="utf-8") as f:
            json.dump(result, f, ensure_ascii=False, indent=2)
        
        print(f"\n[OK] 分析结果已保存：{args.output}")
    
    elif args.batch:
        results = batch_analyze(args.batch)
        
        with open(args.output, "w", encoding="utf-8") as f:
            json.dump(results, f, ensure_ascii=False, indent=2)
        
        print(f"\n[OK] 批量分析结果已保存：{args.output}")
        if "_average" in results:
            print(f"[INFO] 平均字幕参数：{results['_average']}")
    
    else:
        print("[ERROR] 请指定 --image 或 --batch")
        sys.exit(1)


if __name__ == "__main__":
    main()