AIStoryBoard/python/generate-anim/draw_green_boxes_on_mask.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
在遮罩图上绘制绿色文字区域边框
"""
import cv2
import json
import sys
from pathlib import Path
import numpy as np

def draw_green_boxes_on_mask(mask_path, regions_json_path, output_path):
    """
    在遮罩图上绘制绿色边框（支持中文路径）
    """
    # 读取遮罩图（支持中文路径）
    mask_data = np.fromfile(str(mask_path), dtype=np.uint8)
    mask = cv2.imdecode(mask_data, cv2.IMREAD_COLOR)
    if mask is None:
        raise ValueError(f"无法读取遮罩图: {mask_path}")
    
    print(f"[INFO] 遮罩图尺寸: {mask.shape[1]}x{mask.shape[0]}")
    
    # 读取文字区域JSON
    with open(regions_json_path, 'r', encoding='utf-8') as f:
        text_regions = json.load(f)
    
    print(f"[INFO] 需要绘制 {len(text_regions)} 个绿色边框")
    
    # 绘制每个文字区域的绿色边框
    for i, region in enumerate(text_regions):
        bbox = region['bbox']
        # bbox格式: [[x1,y1], [x2,y1], [x2,y2], [x1,y2]]
        x1, y1 = int(bbox[0][0]), int(bbox[0][1])
        x2, y2 = int(bbox[2][0]), int(bbox[2][1])
        
        # 绘制绿色矩形框
        cv2.rectangle(mask, (x1, y1), (x2, y2), (0, 255, 0), 2)  # 绿色，线宽2
        
        print(f"[INFO] 绘制区域 {i+1}: ({x1},{y1}) -> ({x2},{y2})")
    
    # 保存结果（支持中文路径）
    success, encoded_img = cv2.imencode('.png', mask)
    if success:
        encoded_img.tofile(str(output_path))
        print(f"[SUCCESS] 已保存带绿色边框的遮罩图: {output_path}")
    else:
        raise RuntimeError(f"保存图片失败: {output_path}")

def main():
    if len(sys.argv) != 4:
        print("用法: python draw_green_boxes_on_mask.py <遮罩图路径> <区域JSON路径> <输出图片路径>")
        sys.exit(1)
    
    mask_path = Path(sys.argv[1])
    regions_json_path = Path(sys.argv[2])
    output_path = Path(sys.argv[3])
    
    try:
        draw_green_boxes_on_mask(mask_path, regions_json_path, output_path)
    except Exception as e:
        print(f"[ERROR] 绘制失败: {e}")
        sys.exit(1)

if __name__ == "__main__":
    main()