AIStoryBoard/python/generate-anim/detect_comic_text_with_boxes.py

# -*- coding: utf-8 -*-
"""
使用 comic-text-detector-master 检测漫画页面的文字和位置
并绘制每个对话气泡的范围框
"""

import sys
import os
import json
from pathlib import Path
import cv2
import numpy as np
import warnings

# 抑制 pkg_resources 的弃用警告
warnings.filterwarnings('ignore', category=DeprecationWarning, message='.*pkg_resources.*')

# Windows编码修复
if sys.platform == 'win32':
    import io
    sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8', errors='replace')
    sys.stderr = io.TextIOWrapper(sys.stderr.buffer, encoding='utf-8', errors='replace')

# 添加comic-text-detector路径
project_root = Path(__file__).parent.parent
# 尝试两个可能的路径
comic_detector_path1 = project_root / 'comic-text-detector-master' / 'comic-text-detector-master'
comic_detector_path2 = project_root / 'comic-text-detector-master'
if comic_detector_path1.exists() and (comic_detector_path1 / 'inference.py').exists():
    comic_detector_path = comic_detector_path1
elif comic_detector_path2.exists() and (comic_detector_path2 / 'inference.py').exists():
    comic_detector_path = comic_detector_path2
else:
    comic_detector_path = comic_detector_path1  # 默认路径
sys.path.insert(0, str(comic_detector_path))

# 处理 wandb 可选依赖（comic-text-detector 需要但推理时不需要）
try:
    import wandb
except ImportError:
    # 创建一个假的 wandb 模块，避免导入错误
    class FakeWandb:
        @staticmethod
        def init(*args, **kwargs):
            pass
        @staticmethod
        def log(*args, **kwargs):
            pass
        @staticmethod
        def log_model(*args, **kwargs):
            pass
    sys.modules['wandb'] = FakeWandb()

try:
    from inference import TextDetector, REFINEMASK_INPAINT, REFINEMASK_ANNOTATION
    from utils.textblock import TextBlock
    from utils.io_utils import imread, imwrite
except ImportError as e:
    print(f"[ERROR] 无法导入comic-text-detector模块: {e}")
    print(f"[INFO] 请确保已安装依赖: pip install torch torchvision opencv-python numpy tqdm")
    sys.exit(1)


def detect_comic_text_with_boxes(image_path, output_dir, project_root=None, 
                                  input_size=1536, conf_thresh=0.4, nms_thresh=0.35, 
                                  mask_thresh=0.3, act='leaky', refine_mode=0, 
                                  keep_undetected_mask=False, erode_iterations=0, 
                                  invert_mask=False):
    """
    使用comic-text-detector检测漫画图片中的文字和位置
    只生成文字区域坐标JSON文件，不生成mask图片
    
    参数:
        image_path: 图片路径
        output_dir: 输出目录
        project_root: 项目根目录（可选）
        input_size: 输入尺寸（默认1536，越大精度越高但速度越慢，建议1024/1536/2048）
        conf_thresh: 置信度阈值（默认0.4，0-1，越高越严格，过滤低质量检测）
        nms_thresh: NMS阈值（默认0.35，0-1，越高保留越多重叠框）
        mask_thresh: Mask阈值（默认0.3，0-1，用于分割网络，影响mask精度）
        act: 激活函数（默认'leaky'，可选'leaky'或'relu'）
        refine_mode: 精炼模式（默认0=INPAINT填充模式，1=ANNOTATION标注模式）
        keep_undetected_mask: 是否保留未检测区域（默认False，True会保留更多区域但可能有噪声）
        erode_iterations: 腐蚀迭代次数（已废弃，不再使用）
        invert_mask: 是否反转mask（已废弃，不再使用）
    
    返回:
        检测结果字典，包含文字块的位置信息（不包含mask图片）
    """
    if project_root is None:
        project_root = Path(__file__).parent.parent
    else:
        project_root = Path(project_root)
    
    image_path = Path(image_path)
    output_dir = Path(output_dir)
    
    if not image_path.exists():
        raise FileNotFoundError(f"图片文件不存在: {image_path}")
    
    # 确保输出目录存在
    output_dir.mkdir(parents=True, exist_ok=True)
    
    print(f"📖 正在检测图片中的文字区域: {image_path.name}")
    
    # 设置模型路径
    possible_paths = [
        comic_detector_path / 'data' / 'comictextdetector.pt',
        comic_detector_path / 'data' / 'comictextdetector.pt.onnx',
        project_root / 'models' / 'comictextdetector.pt',
        project_root / 'models' / 'comictextdetector.pt.onnx',
    ]
    model_path = None
    for path in possible_paths:
        if path.exists():
            model_path = path
            break
    
    if model_path is None:
        raise FileNotFoundError(
            f"未找到模型文件。请下载模型并放到以下位置之一:\n" +
            "\n".join([f"  - {p}" for p in possible_paths])
        )
    
    # 初始化检测器
    device = 'cuda' if __import__('torch').cuda.is_available() else 'cpu'
    print(f"[INFO] 使用设备: {device}")
    
    try:
        detector = TextDetector(
            model_path=str(model_path),
            input_size=input_size,  # 可配置的输入尺寸
            device=device,
            conf_thresh=conf_thresh,  # 置信度阈值
            nms_thresh=nms_thresh,  # NMS阈值
            mask_thresh=mask_thresh,  # Mask阈值
            act=act  # 激活函数
        )
    except Exception as e:
        print(f"[ERROR] 初始化检测器失败: {e}")
        raise
    
    # 读取图片（处理中文路径）
    img_array = np.fromfile(str(image_path), dtype=np.uint8)
    img = cv2.imdecode(img_array, cv2.IMREAD_COLOR)
    
    if img is None:
        raise ValueError(f"无法读取图片文件: {image_path}")
    
    im_h, im_w = img.shape[:2]
    print(f"[INFO] 图片尺寸: {im_w}x{im_h}")
    
    # 执行检测
    print("[INFO] 正在检测文字区域...")
    print(f"[INFO] 检测参数: input_size={input_size}, conf_thresh={conf_thresh}, nms_thresh={nms_thresh}, refine_mode={refine_mode}")
    try:
        # 根据refine_mode选择模式
        refine_mode_enum = REFINEMASK_INPAINT if refine_mode == 0 else REFINEMASK_ANNOTATION
        mask, mask_refined, blk_list = detector(
            img,
            refine_mode=refine_mode_enum,  # 可配置的精炼模式
            keep_undetected_mask=keep_undetected_mask  # 可配置的未检测区域保留
        )
    except Exception as e:
        print(f"[ERROR] 检测失败: {e}")
        raise
    
    # 提取文字块信息（不绘制轮廓框，只提取数据）
    text_blocks = []
    
    for i, blk in enumerate(blk_list):
        # 获取边界框坐标
        x1, y1, x2, y2 = blk.xyxy
        x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
        
        # 获取文字行坐标（多边形轮廓）
        lines = blk.lines if hasattr(blk, 'lines') else []
        
        # 计算中心点
        center_x = (x1 + x2) / 2
        center_y = (y1 + y2) / 2
        
        # 计算宽度和高度
        width = x2 - x1
        height = y2 - y1
        
        # 处理lines，确保所有坐标都是可序列化的
        lines_data = []
        if lines:
            for line in lines:
                if isinstance(line, (list, tuple)):
                    line_points = []
                    for p in line:
                        if isinstance(p, (list, tuple)) and len(p) >= 2:
                            line_points.append([int(float(p[0])), int(float(p[1]))])
                    if line_points:
                        lines_data.append(line_points)
        
        text_block = {
            'index': i + 1,
            'bbox': {
                'x1': x1,
                'y1': y1,
                'x2': x2,
                'y2': y2,
                'width': int(width),
                'height': int(height),
                'center_x': int(center_x),
                'center_y': int(center_y)
            },
            'lines': lines_data,
            'language': str(getattr(blk, 'language', 'unknown')),
            'vertical': bool(getattr(blk, 'vertical', False))
        }
        text_blocks.append(text_block)
    
    print(f"[OK] 检测到 {len(text_blocks)} 个文字区域")
    
    # 保存结果
    image_name = image_path.stem
    result = {
        'image_file': image_path.name,
        'image_size': {
            'width': im_w,
            'height': im_h
        },
        'text_blocks': text_blocks,
        'total_count': len(text_blocks)
    }
    
    # 保存文字区域坐标JSON文件
    json_path = output_dir / f"{image_name}_text_regions.json"
    with open(json_path, 'w', encoding='utf-8') as f:
        json.dump(result, f, ensure_ascii=False, indent=2)
    print(f"[OK] 已保存文字区域坐标: {json_path}")
    
    # 同时保存简化版坐标文件（兼容OCR格式）
    ocr_compatible_result = {
        'dialogues': []
    }
    
    for block in text_blocks:
        # 转换为OCR兼容的格式
        bbox = block['bbox']
        # 构造四个角点坐标（左上、右上、右下、左下）
        bbox_points = [
            [bbox['x1'], bbox['y1']],  # 左上
            [bbox['x2'], bbox['y1']],  # 右上
            [bbox['x2'], bbox['y2']],  # 右下
            [bbox['x1'], bbox['y2']]   # 左下
        ]
        
        ocr_compatible_result['dialogues'].append({
            'bbox': bbox_points,
            'text': f'[文字区域{block["index"]}]',  # 占位符文字
            'confidence': 0.95,  # 高置信度，因为是专门的检测器
            'source': 'comic-text-detector',
            'region_info': {
                'width': bbox['width'],
                'height': bbox['height'],
                'center_x': bbox['center_x'],
                'center_y': bbox['center_y'],
                'vertical': block['vertical'],
                'language': block['language']
            }
        })
    
    # 不再生成mask图片，只生成JSON文件
    print(f"[INFO] OCR兼容数据将由Node.js处理，不生成中间文件")
    
    # 添加文件路径到返回结果（只包含JSON文件）
    result['output_files'] = {
        'text_regions_json': str(json_path)
    }
    
    # 添加OCR兼容数据供Node.js使用
    result['ocr_compatible_data'] = ocr_compatible_result
    
    return result


if __name__ == '__main__':
    if len(sys.argv) < 3:
        print("用法: python detect_comic_text_with_boxes.py <图片路径> <输出目录> [项目根目录] [input_size] [conf_thresh] [nms_thresh] [mask_thresh] [act] [refine_mode] [keep_undetected_mask] [erode_iterations]")
        print("可选参数:")
        print("  input_size: 输入尺寸（默认1536）")
        print("  conf_thresh: 置信度阈值（默认0.4）")
        print("  nms_thresh: NMS阈值（默认0.35）")
        print("  mask_thresh: Mask阈值（默认0.3）")
        print("  act: 激活函数（默认'leaky'）")
        print("  refine_mode: 精炼模式（默认0=INPAINT，1=ANNOTATION）")
        print("  keep_undetected_mask: 是否保留未检测区域（默认False，传入1为True）")
        print("  erode_iterations: 腐蚀迭代次数（默认0，值越大文字越细，建议0-3）")
        print("  invert_mask: 是否反转mask（默认False，传入1为True，输出黑字白底）")
        sys.exit(1)
    
    image_path = sys.argv[1]
    output_dir = sys.argv[2]
    project_root = sys.argv[3] if len(sys.argv) > 3 else None
    input_size = int(sys.argv[4]) if len(sys.argv) > 4 else 1536
    conf_thresh = float(sys.argv[5]) if len(sys.argv) > 5 else 0.4
    nms_thresh = float(sys.argv[6]) if len(sys.argv) > 6 else 0.35
    mask_thresh = float(sys.argv[7]) if len(sys.argv) > 7 else 0.3
    act = sys.argv[8] if len(sys.argv) > 8 else 'leaky'
    refine_mode = int(sys.argv[9]) if len(sys.argv) > 9 else 0
    keep_undetected_mask = bool(int(sys.argv[10])) if len(sys.argv) > 10 else False
    erode_iterations = int(sys.argv[11]) if len(sys.argv) > 11 else 0
    invert_mask = bool(int(sys.argv[12])) if len(sys.argv) > 12 else False
    
    try:
        detect_comic_text_with_boxes(
            image_path, output_dir, project_root,
            input_size, conf_thresh, nms_thresh, mask_thresh,
            act, refine_mode, keep_undetected_mask, erode_iterations, invert_mask
        )
    except Exception as e:
        print(f"[ERROR] 处理失败: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)