# debug_label_overlay.py
import cv2
import numpy as np
from pathlib import Path

def load_yolo_polygon_labels(txt_path, img_shape):
    """
    加载 YOLO 格式的多边形标签 (class x1 y1 x2 y2 ... xn yn)
    :param txt_path: .txt 标签路径
    :param img_shape: (H, W) 原图大小
    :return: 列表，元素为 [cls_id, polygon]，polygon 是 Nx2 的点阵
    """
    h, w = img_shape[:2]
    labels = []
    if not Path(txt_path).exists():
        print(f"⚠️ 标签文件不存在: {txt_path}")
        return labels

    with open(txt_path, 'r') as f:
        for line in f:
            line = line.strip()
            if not line or line.startswith('#'):
                continue
            parts = list(map(float, line.split()))
            cls_id = int(parts[0])
            coords = np.array(parts[1:]).reshape(-1, 2)  # (N, 2)
            # 反归一化到原图
            coords[:, 0] *= w
            coords[:, 1] *= h
            coords = coords.astype(np.int32)
            labels.append((cls_id, coords))
    return labels

def draw_mask_on_image(img, labels, alpha=0.5):
    """
    在图像上绘制半透明 mask 和边框
    """
    overlay = img.copy()
    colors = [
        (0, 255, 0),   # 绿
        (0, 0, 255),   # 蓝
        (255, 0, 0),   # 红
        (255, 255, 0),
        (255, 0, 255),
        (0, 255, 255),
        (128, 0, 128),
        (0, 128, 0),
        (0, 0, 128),
        (128, 128, 0)
    ]

    font = cv2.FONT_HERSHEY_SIMPLEX
    font_scale = 0.7
    thickness = 1

    for cls_id, poly in labels:
        color = colors[cls_id % len(colors)]

        # 填充半透明 mask
        cv2.fillPoly(overlay, [poly], color=color)

        # 绘制边缘轮廓
        cv2.polylines(overlay, [poly], isClosed=True, color=color, thickness=2)

        # 添加类别标签
        x, y = poly[0][0], poly[0][1]  # 使用第一个点作为标签位置
        label = f"Class {cls_id}"
        (text_w, text_h), _ = cv2.getTextSize(label, font, font_scale, thickness)
        cv2.rectangle(overlay, (x, y - text_h - 6), (x + text_w, y), color, -1)
        cv2.putText(overlay, label, (x, y - 4), font, font_scale,
                    (255, 255, 255), thickness, cv2.LINE_AA)

    # 混合叠加
    result = cv2.addWeighted(overlay, alpha, img, 1 - alpha, 0)
    return result

def main():
    # ====================== 配置区 ======================
    IMG_PATH = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/seg/yemian_seg_camera01/train/frame_20250805_121717_22320.jpg"       # 修改为你的某张测试图
    TXT_PATH = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/seg/yemian_seg_camera01/train/frame_20250805_121717_22320.txt"  # 对应的 .txt 文件
    OUTPUT_PATH = "/home/hx/yolo/output_masks2/debug_overlay.png"
    # ====================================================

    img_path = Path(IMG_PATH)
    txt_path = Path(TXT_PATH)
    output_path = Path(OUTPUT_PATH)
    output_path.parent.mkdir(parents=True, exist_ok=True)

    if not img_path.exists():
        print(f"❌ 图像文件不存在: {img_path}")
        return

    # 读取图像
    img = cv2.imread(str(img_path))
    if img is None:
        print(f"❌ 无法读取图像: {img_path}")
        return

    h, w = img.shape[:2]
    print(f"🖼️  图像大小: {w}x{h}")

    # 加载标签
    labels = load_yolo_polygon_labels(txt_path, img.shape)
    if len(labels) == 0:
        print(f"🟡 未找到有效标签: {txt_path}")
        return

    print(f"🎯 加载到 {len(labels)} 个对象")

    # 绘制叠加图
    result_img = draw_mask_on_image(img, labels, alpha=0.5)

    # 保存结果
    cv2.imwrite(str(output_path), result_img)
    print(f"✅ 已保存可视化结果: {output_path}")

    # 显示（可选）
    display_img = cv2.resize(result_img, (960, 540))  # 缩放便于显示
    cv2.imshow("Label Overlay Debug", display_img)
    print("📌 按任意键关闭窗口...")
    cv2.waitKey(0)
    cv2.destroyAllWindows()

if __name__ == "__main__":
    main()