AIStoryBoard/python/generate-anim/draw_red_panel_mask.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
创建透明PNG图片并绘制红线格子区域
"""
import cv2
import json
import sys
from pathlib import Path
import numpy as np

def draw_red_lines_on_transparent_image(origin_image_path, json_path, output_path):
    """
    根据原图片尺寸创建透明PNG，并根据JSON文件绘制红线格子区域
    """
    # 1. 从原图片获取真实尺寸
    try:
        # 使用cv2读取原图片获取尺寸（支持中文路径）
        img_data = np.fromfile(str(origin_image_path), dtype=np.uint8)
        origin_img = cv2.imdecode(img_data, cv2.IMREAD_COLOR)
        if origin_img is None:
            raise ValueError(f"无法读取原图片: {origin_image_path}")
        
        height, width = origin_img.shape[:2]
        print(f"[INFO] 从原图片获取尺寸: {width}x{height}")
        
    except Exception as e:
        print(f"[ERROR] 无法读取原图片: {e}")
        # 如果读取失败，使用默认尺寸
        width, height = 1334, 1940
        print(f"[INFO] 使用默认图片尺寸: {width}x{height}")
    
    # 2. 读取格子位置JSON
    with open(json_path, 'r', encoding='utf-8') as f:
        panel_data = json.load(f)
    
    panels = panel_data.get('panels', [])
    
    print(f"[INFO] 图片尺寸: {width}x{height}")
    print(f"[INFO] 需要绘制 {len(panels)} 个格子边框")
    
    # 创建透明背景图片 (RGBA格式)
    mask = np.zeros((height, width, 4), dtype=np.uint8)
    
    # 绘制每个格子的红色边框
    drawn_count = 0
    for i, panel in enumerate(panels):
        try:
            # 兼容两种格式：bbox数组 或 x,y,width,height字段
            if 'bbox' in panel:
                # 格式1: bbox数组 [x1, y1, x2, y2]
                bbox = panel['bbox']
                if isinstance(bbox, list) and len(bbox) >= 4:
                    x1, y1, x2, y2 = int(bbox[0]), int(bbox[1]), int(bbox[2]), int(bbox[3])
                else:
                    print(f"[WARN] 格子 {i+1} bbox格式无效: {bbox}")
                    continue
            elif 'x' in panel and 'y' in panel and 'width' in panel and 'height' in panel:
                # 格式2: x,y,width,height字段
                x = int(panel['x'])
                y = int(panel['y'])
                w = int(panel['width'])
                h = int(panel['height'])
                x1, y1, x2, y2 = x, y, x + w, y + h
            else:
                print(f"[WARN] 格子 {i+1} 缺少坐标信息")
                continue
            
            # 确保坐标在图片范围内
            x1 = max(0, min(x1, width-1))
            y1 = max(0, min(y1, height-1))
            x2 = max(0, min(x2, width-1))
            y2 = max(0, min(y2, height-1))
            
            if x2 > x1 and y2 > y1:
                # 绘制红色矩形框，线宽3，颜色为红色(0,0,255,255)
                cv2.rectangle(mask, (x1, y1), (x2, y2), (0, 0, 255, 255), 3)
                drawn_count += 1
                print(f"[INFO] 绘制格子 {drawn_count}: ({x1},{y1}) -> ({x2},{y2})")
            else:
                print(f"[WARN] 跳过无效格子 {i+1}: ({x1},{y1}) -> ({x2},{y2})")
        except Exception as e:
            print(f"[ERROR] 绘制格子 {i+1} 失败: {str(e)}")
    
    print(f"[INFO] 成功绘制 {drawn_count} 个格子边框")
    
    # 保存为PNG格式（支持透明度）
    try:
        # 先尝试简单保存
        output_path_str = str(output_path)
        print(f"[INFO] 尝试保存到: {output_path_str}")
        
        # 使用cv2.imencode处理中文路径
        success, encoded_img = cv2.imencode('.png', mask)
        if success:
            # 写入文件
            output_path.parent.mkdir(parents=True, exist_ok=True)  # 确保目录存在
            with open(output_path_str, 'wb') as f:
                f.write(encoded_img.tobytes())
            print(f"[SUCCESS] 已保存红线透明底遮罩图: {output_path}")
        else:
            raise RuntimeError(f"图片编码失败")
    except Exception as e:
        print(f"[ERROR] 详细错误信息: {str(e)}")
        import traceback
        traceback.print_exc()
        raise RuntimeError(f"保存图片失败: {output_path}, 错误: {str(e)}")

def main():
    if len(sys.argv) != 4:
        print("用法: python draw_red_transparent_image.py <原图片路径> <JSON文件路径> <输出图片路径>")
        sys.exit(1)
    
    origin_image_path = Path(sys.argv[1])
    json_path = Path(sys.argv[2])
    output_path = Path(sys.argv[3])
    
    try:
        draw_red_lines_on_transparent_image(origin_image_path, json_path, output_path)
    except Exception as e:
        print(f"[ERROR] 绘制失败: {e}")
        sys.exit(1)

if __name__ == "__main__":
    main()