import numpy as np
import cv2
from matplotlib import pyplot as plt

def visualize_mask(points, size, save_path="mask.png"):
    """
    输入: 
        points    : 多边形坐标 [[x1,y1], [x2,y2], ...]
        size      : 图像尺寸 [height, width]
        save_path : 输出图片保存路径（默认当前目录mask.png）
    
    输出:
        展示并保存包含多边形顶点（红色）和掩码区域（半透明绿）的图片
    """
    # 数值有效性检查
    assert len(size) == 2, "size must be [height, width]"
    height, width = size
    assert height > 0 and width > 0, "invalid image size"
    
    # 1. 创建彩色画布（RGBA格式，支持半透明叠加）
    canvas = np.zeros((height, width, 4), dtype=np.uint8)
    
    # 2. 预处理坐标点
    polygon = []
    for x, y in points:
        # 将点限制在图像范围内（防止越界）
        x = int(np.clip(round(x), 0, width-1))
        y = int(np.clip(round(y), 0, height-1))
        polygon.append([x, y])
    polygon = np.array([polygon], dtype=np.int32)  # OpenCV要求形状为[N,1,2]
    
    # 3. 生成掩码区域（绿色半透明填充）
    mask = np.zeros((height, width), dtype=np.uint8)
    cv2.fillPoly(mask, polygon, color=1)
    canvas[..., 1] = 80 * mask  # 绿色通道
    canvas[..., 3] = 200 * mask  # 透明度通道（掩码区域半透明）

    # 4. 绘制多边形顶点（红色圆点标记）
    for idx, (x, y) in enumerate(polygon[0]):
        color = (0, 0, 255, 255)  # 纯红色（RGBA）
        cv2.circle(canvas, (x, y), radius=5, color=color, thickness=-1)
        cv2.putText(canvas, str(idx+1), (x+7, y+3),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255,255,255,255), 1)

    # 5. 绘制多边形边界线（蓝色）
    cv2.polylines(canvas, polygon, isClosed=True, 
                 color=(255, 0, 0, 255), thickness=2)  # 蓝色边界

    # 6. 保存和显示结果
    cv2.imwrite(save_path, canvas)
    
    # 使用matplotlib显示（确保透明天正确处理）
    plt.figure(figsize=(10, 6))
    plt.title(f"Mask Visualization (Saved to {save_path})")
    plt.imshow(cv2.cvtColor(canvas, cv2.COLOR_BGRA2RGBA))
    plt.axis('off')
    plt.show()

# ------------ 使用示例 ------------
if __name__ == "__main__":
    # 你的输入数据
    points = [
        [602.245, 290.396],  
        [585.264, 289.452],
        [568.283, 301.716],
        [557.905, 309.264],
        [546.584, 324.358],
        [543.754, 342.283],
        [544.698, 365.867],
        [552.245, 381.905],
        [571.113, 398.886],
        [583.377, 407.377],
        [597.528, 409.264],
        [607.905, 398.886],
        [626.773, 387.566],
        [636.208, 384.735],
        [646.584, 381.905],
        [660.735, 370.584],
        [661.679, 335.679],
        [654.132, 313.037],
        [639.981, 302.660],
        [618.283, 293.226]
    ]
    size = [960, 1280]  # 高度=488，宽度=640

    # 生成可视化图像
    visualize_mask(
        points=points,
        size=size,
        save_path="custom_mask_visualization.png"
    )