import cv2
import numpy as np
import os

# 输入：已保存的黑白掩码图像路径
mask_path = 'composite_mask.png'

# 输出：可视化结果保存路径
output_dir = 'output_masks'
os.makedirs(output_dir, exist_ok=True)
result_image_path = os.path.join(output_dir, 'mask_with_rotated_rects.jpg')

# 读取掩码图像（灰度图）
mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
if mask is None:
    raise FileNotFoundError(f"❌ 无法读取掩码图像: {mask_path}")

# 二值化（确保是 0 和 255）
_, binary = cv2.threshold(mask, 127, 255, cv2.THRESH_BINARY)

# 查找轮廓
contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

# 创建三通道图像用于可视化（白色背景，黑色轮廓）
vis_image = np.ones((*mask.shape, 3), dtype=np.uint8) * 255  # 白底

# 存储所有旋转矩形结果
rotated_rects = []

for i, cnt in enumerate(contours):
    if len(cnt) < 2:
        continue  # 至少需要两个点才能拟合

    # 计算最小外接旋
    rect = cv2.minAreaRect(cnt)  # (center, (width, height), angle)
    rotated_rects.append(rect)

    # 获取四个角
    box = cv2.boxPoints(rect)
    box = np.float16(box)

    # 在可视化图像上绘制旋转矩形（绿色）
    cv2.drawContours(vis_image, [box], 0, (0, 180, 0), 2)

    # 可选：标注序号
    cx, cy = int(rect[0][0]), int(rect[0][1])
    cv2.putText(vis_image, str(i), (cx, cy), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 0), 1)

    # 打印信息
    center = rect[0]
    size = rect[1]
    angle = rect[2]
    print(f"物体 {i}:")
    print(f"  中心点: ({center[0]:.1f}, {center[1]:.1f})")
    print(f"  尺寸: 宽={size[0]:.1f}, 高={size[1]:.1f}")
    print(f"  旋转角度: {angle:.2f}°")
    print(f"  四个角点: {box.tolist()}")
    print("-" * 40)

# 保存可视化结果
cv2.imwrite(result_image_path, vis_image)
print(f"✅ 旋转矩形可视化已保存至: {result_image_path}")