import cv2
import numpy as np
from pathlib import Path
from val_labels import load_yolo_polygon_labels,draw_mask_on_image

# 之前的load_yolo_polygon_labels和draw_mask_on_image函数保持不变

def visualize_labels_in_folder(src_img_dir, src_label_dir, dst_dir, alpha=0.5):
    """
    遍历源文件夹中的所有图片和对应标签，并在图片上绘制标签后保存至目标文件夹。
    :param src_img_dir: 源图像文件夹路径
    :param src_label_dir: 源标签文件夹路径
    :param dst_dir: 目标文件夹路径，用于保存可视化结果
    :param alpha: 叠加mask的透明度
    """
    # 确保目标文件夹存在
    Path(dst_dir).mkdir(parents=True, exist_ok=True)

    # 获取所有的图片文件
    img_paths = list(Path(src_img_dir).glob('*.jpg')) + list(Path(src_img_dir).glob('*.png'))

    for img_path in img_paths:
        txt_path = Path(src_label_dir) / f"{img_path.stem}.txt"
        if not txt_path.exists():
            print(f"⚠️ 跳过 {img_path.name}，未找到对应的标签文件")
            continue

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

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

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

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

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


if __name__ == "__main__":
    SRC_IMG_DIR = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/seg/resize1"  # 修改为你的源图像文件夹路径
    SRC_LABEL_DIR = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/seg/resize1"  # 修改为你的源标签文件夹路径
    DST_DIR = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/seg/resize_seg5"  # 修改为目标文件夹路径

    visualize_labels_in_folder(SRC_IMG_DIR, SRC_LABEL_DIR, DST_DIR)