from osgeo import gdal
import os
os.environ['PROJ_LIB'] =r"C:/Soft/anaconda3/envs/yolo/Lib/site-packages/osgeo/data/proj"

def apply_reference_geotransform(reference_path, input_path, output_path):
    """
    将参考图像的地理变换应用到待转换图像上。

    :param reference_path: 参考图像路径
    :param input_path: 待转换图像路径
    :param output_path: 输出的已地理配准图像路径
    """
    # 打开源数据集
    ref_ds = gdal.Open(reference_path)
    if ref_ds is None:
        print(f"无法打开参考文件 {reference_path}")
        return

    # 获取参考图像的地理变换和投影信息
    geo_transform = ref_ds.GetGeoTransform()
    projection = ref_ds.GetProjection()

    # 关闭参考图像数据集
    ref_ds = None

    # 打开待转换图像数据集
    src_ds = gdal.Open(input_path)
    if src_ds is None:
        print(f"无法打开输入文件 {input_path}")
        return

    # 创建输出文件
    driver = gdal.GetDriverByName('GTiff')
    dst_ds = driver.CreateCopy(output_path, src_ds, strict=0)

    # 设置地理变换和投影信息
    dst_ds.SetGeoTransform(geo_transform)
    dst_ds.SetProjection(projection)

    # 关闭数据集
    dst_ds = None
    src_ds = None

    print(f"成功将 {input_path} 地理配准并保存至 {output_path}")

# 示例调用函数
input_tif = r'D:\zlc\drone\图像分割\栅格\Images\merged_image2.tif'
output_tif = r'D:\zlc\drone\图像分割\栅格\Images\merged_image3.tif'
reference_tif = r'D:\zlc\drone\图像分割\栅格\影像.tif'

try:
    apply_reference_geotransform(reference_tif, input_tif, output_tif)
except Exception as e:
    print(f"发生错误: {e}")

import cv2
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import ListedColormap, BoundaryNorm

def create_color_map(num_classes):
    """创建一个随机颜色映射表"""
    color_map = []
    for _ in range(num_classes):
        color = np.random.randint(0, 256, size=3).tolist()
        color_map.append(color)
    return color_map

def overlay_multi_class_mask_on_image(image_path, mask_path, alpha=0.5):
    image = cv2.imread(image_path)
    mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)

    if image is None or mask is None:
        print(f"无法读取图像或mask: {image_path} 或 {mask_path}")
        return

    # 确保mask和图像大小一致
    if image.shape[:2] != mask.shape:
        mask = cv2.resize(mask, (image.shape[1], image.shape[0]), interpolation=cv2.INTER_NEAREST)

    # 获取mask中存在的类别
    unique_classes = np.unique(mask)
    num_classes = len(unique_classes)

    # 创建颜色映射表
    color_map = create_color_map(num_classes)

    # 创建一个彩色mask图像
    color_mask = np.zeros_like(image)
    for i, class_id in enumerate(unique_classes):
        if class_id == 0:  # 跳过背景（假设0是背景）
            continue
        color_mask[mask == class_id] = color_map[i]

    # 将彩色mask与原始图像混合
    blended_image = cv2.addWeighted(image, 1 - alpha, color_mask, alpha, 0)

    return blended_image, color_map, unique_classes

def visualize_overlay_with_legend(image_path, mask_path, alpha=0.5):
    result_image, color_map, unique_classes = overlay_multi_class_mask_on_image(image_path, mask_path, alpha)

    if result_image is not None:
        plt.figure(figsize=(10, 5))

        # 显示带有标注的图像
        plt.imshow(cv2.cvtColor(result_image, cv2.COLOR_BGR2RGB))
        plt.title('Overlay Multi-Class Mask on Image')
        plt.axis('off')

        # 添加图例
        patches = [plt.Rectangle((0, 0), 1, 1, fc=np.array(color)/255) for color in color_map]
        classes = [f'Class {cls}' for cls in unique_classes if cls != 0]  # 排除背景类
        plt.legend(patches, classes, loc='upper right', bbox_to_anchor=(1.1, 1))
        plt.show()

# 测试函数
if __name__ == "__main__":
    image_path = r'D:\zlc\drone\drone-train\data\2021LoveDA\train\images\7.png'  # 替换为你的图像路径
    mask_path = r'D:\zlc\drone\drone-train\data\2021LoveDA\train\masks\7.png'    # 替换为你的mask路径
    visualize_overlay_with_legend(image_path, mask_path)