AIStoryBoard/python/generate-anim/cut_dialog_blocks.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
简单的绿色线框检测和文字块切割脚本
"""
import cv2
import numpy as np
import sys
from pathlib import Path

def detect_green_boxes_and_cut(image_path, output_dir):
    """检测绿色线框并切割文字块"""
    # 读取图片
    img_data = np.fromfile(str(image_path), dtype=np.uint8)
    img = cv2.imdecode(img_data, cv2.IMREAD_COLOR)
    if img is None:
        raise ValueError(f"无法读取图片: {image_path}")
    
    print(f"[INFO] 图片尺寸: {img.shape[1]}x{img.shape[0]}")
    
    # 检测绿色线框
    # 绿色范围 (BGR格式)
    lower_green = np.array([0, 100, 0])
    upper_green = np.array([100, 255, 100])
    
    # 创建绿色掩码
    mask = cv2.inRange(img, lower_green, upper_green)
    
    # 寻找轮廓
    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    
    if len(contours) == 0:
        print("[INFO] 未检测到绿色线框")
        return
    
    # 获取边界矩形
    boxes = []
    for contour in contours:
        x, y, w, h = cv2.boundingRect(contour)
        if w > 10 and h > 10:  # 过滤小框
            boxes.append({'x': x, 'y': y, 'w': w, 'h': h})
    
    if len(boxes) == 0:
        print("[INFO] 未找到有效的文字框")
        return
    
    print(f"[INFO] 检测到 {len(boxes)} 个绿色文字框")
    
    # 按从右到左、从上到下排序
    boxes.sort(key=lambda b: (b['y'], -b['x']))
    
    # 确保输出目录存在
    Path(output_dir).mkdir(parents=True, exist_ok=True)
    
    # 切割并保存每个文字块
    for i, box in enumerate(boxes, 1):
        x, y, w, h = box['x'], box['y'], box['w'], box['h']
        
        # 添加少量边距
        padding = 5
        x = max(0, x - padding)
        y = max(0, y - padding)
        w = min(w + 2 * padding, img.shape[1] - x)
        h = min(h + 2 * padding, img.shape[0] - y)
        
        # 切割区域
        dialog_img = img[y:y+h, x:x+w]
        
        # 保存文件
        output_file = Path(output_dir) / f"dialog_{i}.png"
        success, encoded = cv2.imencode('.png', dialog_img)
        if success:
            encoded.tofile(str(output_file))
            print(f"[{i}/{len(boxes)}] 保存: {output_file.name} ({w}x{h})")
    
    print(f"✅ 文字块切割完成: {len(boxes)} 个")

if __name__ == '__main__':
    if len(sys.argv) != 3:
        print("用法: python cut_dialog_blocks.py <图片路径> <输出目录>")
        sys.exit(1)
    
    try:
        detect_green_boxes_and_cut(sys.argv[1], sys.argv[2])
    except Exception as e:
        print(f"[ERROR] {e}")
        sys.exit(1)