AIStoryBoard/python/generate-anim/extract_panel_with_green_boxes.py

# -*- coding: utf-8 -*-
"""
从带绿色线框的完整图片中提取格子区域，并在格子图片上绘制对应的绿色线框
"""

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

# 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')


def extract_panel_with_green_boxes(full_image_path, panel_image_path, output_path, 
                                   panel_x, panel_y, panel_width, panel_height):
    """
    从完整图片中提取格子区域，并在格子图片上绘制对应的绿色线框
    
    参数:
        full_image_path: 带绿色线框的完整图片路径
        panel_image_path: 格子图片路径（原图）
        output_path: 输出图片路径（带绿色线框的格子图片）
        panel_x, panel_y, panel_width, panel_height: 格子在完整图片中的坐标
    """
    # 读取完整图片（带绿色线框）
    full_img_array = np.fromfile(str(full_image_path), dtype=np.uint8)
    full_img = cv2.imdecode(full_img_array, cv2.IMREAD_COLOR)
    
    if full_img is None:
        raise ValueError(f"无法读取完整图片: {full_image_path}")
    
    # 读取格子图片（原图）
    panel_img_array = np.fromfile(str(panel_image_path), dtype=np.uint8)
    panel_img = cv2.imdecode(panel_img_array, cv2.IMREAD_COLOR)
    
    if panel_img is None:
        raise ValueError(f"无法读取格子图片: {panel_image_path}")
    
    print(f"[INFO] 完整图片尺寸: {full_img.shape[1]}x{full_img.shape[0]}")
    print(f"[INFO] 格子图片尺寸: {panel_img.shape[1]}x{panel_img.shape[0]}")
    print(f"[INFO] 格子坐标: x={panel_x}, y={panel_y}, w={panel_width}, h={panel_height}")
    
    # 从完整图片中提取格子区域的绿色线框
    panel_region = full_img[panel_y:panel_y + panel_height, panel_x:panel_x + panel_width]
    
    # 检测绿色像素（BGR格式：绿色为 (0, 255, 0)）
    green_color_lower = np.array([0, 200, 0])
    green_color_upper = np.array([50, 255, 50])
    green_mask = cv2.inRange(panel_region, green_color_lower, green_color_upper)
    
    # 形态学操作，连接断开的绿色线条
    kernel = np.ones((3, 3), np.uint8)
    green_mask = cv2.morphologyEx(green_mask, cv2.MORPH_CLOSE, kernel, iterations=2)
    green_mask = cv2.dilate(green_mask, kernel, iterations=1)
    
    # 使用轮廓检测找到绿色框
    contours, _ = cv2.findContours(green_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    
    print(f"[INFO] 在格子区域检测到 {len(contours)} 个绿色轮廓")
    
    # 在格子图片上绘制绿色线框
    result_img = panel_img.copy()
    green_boxes = []
    
    for i, contour in enumerate(contours):
        # 计算边界框
        x, y, box_w, box_h = cv2.boundingRect(contour)
        
        # 过滤太小的区域
        min_area = panel_img.shape[0] * panel_img.shape[1] * 0.001
        if box_w * box_h < min_area:
            continue
        
        # 过滤太细的线
        if box_w < 5 or box_h < 5:
            continue
        
        # 绘制绿色矩形框
        cv2.rectangle(result_img, (x, y), (x + box_w, y + box_h), (0, 255, 0), 2)
        
        # 标注序号
        label = str(i + 1)
        font = cv2.FONT_HERSHEY_SIMPLEX
        font_scale = 0.6
        font_thickness = 1
        (text_width, text_height), baseline = cv2.getTextSize(label, font, font_scale, font_thickness)
        
        # 在左上角绘制文字背景（白色矩形）
        cv2.rectangle(result_img, (x, y - text_height - baseline - 2), 
                    (x + text_width, y), (255, 255, 255), -1)
        # 绘制文字（绿色）
        cv2.putText(result_img, label, (x, y - baseline - 2), 
                  font, font_scale, (0, 255, 0), font_thickness, cv2.LINE_AA)
        
        green_boxes.append({
            'x1': x,
            'y1': y,
            'x2': x + box_w,
            'y2': y + box_h
        })
    
    print(f"[INFO] 在格子图片上绘制了 {len(green_boxes)} 个绿色框")
    
    # 保存结果
    success, encoded_img = cv2.imencode('.png', result_img)
    if success:
        encoded_img.tofile(str(output_path))
        print(f"[OK] 已保存带绿色线框的格子图片: {output_path}")
        return True
    else:
        print(f"[ERROR] 保存图片失败: {output_path}")
        return False


if __name__ == '__main__':
    if len(sys.argv) < 8:
        print("用法: python extract_panel_with_green_boxes.py <full_image_path> <panel_image_path> <output_path> <panel_x> <panel_y> <panel_width> <panel_height>")
        print("示例: python extract_panel_with_green_boxes.py full.png panel.png output.png 0 0 750 746")
        sys.exit(1)
    
    full_image_path = sys.argv[1]
    panel_image_path = sys.argv[2]
    output_path = sys.argv[3]
    panel_x = int(sys.argv[4])
    panel_y = int(sys.argv[5])
    panel_width = int(sys.argv[6])
    panel_height = int(sys.argv[7])
    
    success = extract_panel_with_green_boxes(
        full_image_path, panel_image_path, output_path,
        panel_x, panel_y, panel_width, panel_height
    )
    
    sys.exit(0 if success else 1)