AIStoryBoard/python/generate-anim/draw_panel_mask.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
根据格子检测结果绘制黑线白底遮罩图片
用于步骤5: 根据detectComicPanels()函数返回的结果，绘制出黑线白底遮罩图片
"""

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

def draw_panel_mask(json_path, output_dir):
    """
    根据格子检测JSON结果绘制遮罩图
    
    Args:
        json_path: 格子检测结果JSON文件路径
        output_dir: 输出目录
    """
    try:
        print(f"[INFO] 读取格子检测结果: {json_path}")
        
        # 读取格子检测结果
        if not os.path.exists(json_path):
            raise FileNotFoundError(f"JSON文件不存在: {json_path}")
        
        with open(json_path, 'r', encoding='utf-8') as f:
            panel_data = json.load(f)
        
        # 获取图片信息
        image_size = panel_data.get('image_size', {})
        width = image_size.get('width', 800)
        height = image_size.get('height', 600)
        panels = panel_data.get('panels', [])
        
        print(f"[INFO] 图片尺寸: {width}x{height}")
        print(f"[INFO] 格子数量: {len(panels)}")
        
        # 创建白色背景的遮罩图 (白底)
        mask = np.ones((height, width, 3), dtype=np.uint8) * 255
        
        # 绘制格子边框 (黑线)
        line_thickness = 2
        line_color = (0, 0, 0)  # 黑色
        
        for i, panel in enumerate(panels):
            bbox = panel.get('bbox', {})
            if not bbox:
                continue
                
            x1 = int(bbox.get('x1', 0))
            y1 = int(bbox.get('y1', 0))
            x2 = int(bbox.get('x2', width))
            y2 = int(bbox.get('y2', height))
            
            # 绘制矩形边框
            cv2.rectangle(mask, (x1, y1), (x2, y2), line_color, line_thickness)
            
            print(f"[DEBUG] 绘制格子 {i+1}: ({x1}, {y1}) -> ({x2}, {y2})")
        
        # 保存遮罩图
        json_filename = Path(json_path).stem
        # 移除 _panels 后缀，添加 _panel_mask 后缀
        if json_filename.endswith('_panels'):
            base_name = json_filename[:-8]  # 移除 '_panels'
        else:
            base_name = json_filename
            
        mask_filename = f"{base_name}_panel_mask.png"
        mask_path = os.path.join(output_dir, mask_filename)
        
        cv2.imwrite(mask_path, mask)
        
        print(f"[OK] 遮罩图已保存: {mask_path}")
        
        return {
            'mask_path': mask_path,
            'panels_count': len(panels),
            'image_size': {'width': width, 'height': height}
        }
        
    except Exception as e:
        print(f"[ERROR] 绘制遮罩图失败: {e}")
        import traceback
        traceback.print_exc()
        raise

def main():
    """主函数"""
    if len(sys.argv) < 3:
        print("用法: python draw_panel_mask.py <json_path> <output_dir>")
        print("示例: python draw_panel_mask.py panels.json ./output")
        sys.exit(1)
    
    json_path = sys.argv[1]
    output_dir = sys.argv[2]
    
    # 确保输出目录存在
    os.makedirs(output_dir, exist_ok=True)
    
    try:
        result = draw_panel_mask(json_path, output_dir)
        print(f"[SUCCESS] 成功绘制遮罩图，包含 {result['panels_count']} 个格子")
        sys.exit(0)
    except Exception as e:
        print(f"[ERROR] 程序执行失败: {e}")
        sys.exit(1)

if __name__ == "__main__":
    main()