import os
from osgeo import gdal
import numpy as np
from tqdm import tqdm
 
# 读取数据
def read_image(image_path, num_bands=None, selected_bands=None):
    dataset = gdal.Open(image_path)
    if dataset is None:
        print(f"Could not open image: {image_path}")
        return None
    else:
        if num_bands is None:
            num_bands = dataset.RasterCount
        image_data = []
        for i in range(1, num_bands + 1):
            if selected_bands is not None and i not in selected_bands:
                continue
            band = dataset.GetRasterBand(i)
            band_data = band.ReadAsArray()
            image_data.append(band_data)
        return np.array(image_data)
 
# 滑动裁剪
def sliding_crop(image, window_size=(512, 512), stride=256):
    height, width = image.shape[1], image.shape[2]
 
    # 计算需要填充的高度和宽度
    pad_height = 0
    pad_width = 0
    if height % window_size[0] != 0:
        pad_height = window_size[0] - (height % window_size[0])
    if width % window_size[1] != 0:
        pad_width = window_size[1] - (width % window_size[1])
 
    # 在图像右侧和下侧填充0值
    padded_image = np.pad(image, ((0, 0), (0, pad_height),
                          (0, pad_width)), mode='constant', constant_values=0)
 
    crops = []
    for y in range(0, height + pad_height - window_size[0] + 1, stride):
        for x in range(0, width + pad_width - window_size[1] + 1, stride):
            crop = padded_image[:, y:y+window_size[0], x:x+window_size[1]]
            crops.append(crop)
    return crops
 
# 保存数据
def save_crops(crops, output_dir):
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
 
    existing_files = os.listdir(output_dir)
    existing_indices = set()
    for filename in existing_files:
        if filename.startswith("crop_") and filename.endswith(".tif"):
            index_str = filename.split("_")[1].split(".")[0]
            existing_indices.add(int(index_str))
 
    start_index = max(existing_indices) + 1 if existing_indices else 0
 
    for i, crop in enumerate(crops):
        output_path = os.path.join(output_dir, f"crop_{start_index + i}.tif")
        save_image(crop, output_path)
 
def save_image(image_data, output_path):
    num_bands, height, width = image_data.shape
    driver = gdal.GetDriverByName("GTiff")
    dataset = driver.Create(output_path, width, height,
                            num_bands, gdal.GDT_Byte)
    for i in range(num_bands):
        dataset.GetRasterBand(i + 1).WriteArray(image_data[i])
    dataset.FlushCache()
 
 
if __name__ == "__main__":
    # 影像的路径
    image_path = r"xxxx.tif"
    # 标签的路径
    label_path = r"xxxx.tif"
    images = read_image(image_path)
    labels = read_image(label_path, 1)
    if images.shape[1:] == labels.shape[1:]:
        print("Images have same dimensions. Starting cropping...")
        # 设置滑动窗口的大小及步长，用于生成重叠的滑动裁剪块
        images_crops = sliding_crop(images, window_size=(512, 512), stride=256)
        labels_crops = sliding_crop(labels, window_size=(512, 512), stride=256)
        # 保存裁剪后的图像
        save_crops(images_crops,
                   r"H:\Images")# 影像块的文件夹
        save_crops(labels_crops,
                   r"H:\Labels")# 标签块的文件夹
        print("Cropping done and crops saved!")
    else:
        print("Images have different dimensions. Cannot proceed with cropping.")

import cv2
import os


path = r"D:\zlc\drone\图像分割\栅格\Labels"
files = os.listdir(path)
for file in files:
    name = file.split('.')[0]
    file_path = os.path.join(path,name+'.tif')
    img = cv2.imread(file_path)
    # img = cv2.imread(path)
    H,W=img.shape[0:2]
    print(H,W)

    #img1 = cv2.imread("F:/Deep_Learning/Model/YOLOv8_Seg/Dataset/images/20160222_080933_361_1.jpg")

    gray_img = cv2.cvtColor(img,cv2.COLOR_BGR2GRAY)
    ret,bin_img = cv2.threshold(gray_img,0,255,cv2.THRESH_BINARY+cv2.THRESH_OTSU)
    cnt,hit = cv2.findContours(bin_img,cv2.RETR_TREE,cv2.CHAIN_APPROX_TC89_KCOS)

    #cv2.drawContours(img1,cnt,-1,(0,255,0),5)

    cnt = list(cnt)
    f = open("D:\zlc\drone\图像分割\栅格\yolo-labels/{}.txt".format(file.split(".")[0]), "a+")
    for j in cnt:
        result = []
        pre = j[0]
        for i in j:
            if abs(i[0][0] - pre[0][0]) > 1 or abs(i[0][1] - pre[0][1]) > 1:# 在这里可以调整间隔点，我设置为1
                pre = i
                temp = list(i[0])
                temp[0] /= W
                temp[1] /= H
                result.append(temp)

                #cv2.circle(img1,i[0],1,(0,0,255),2)

        print(result)
        print(len(result))

        # if len(result) != 0:

        if len(result) != 0:
            f.write("0 ")
            for line in result:
                line = str(line)[1:-2].replace(",","")
                # print(line)
                f.write(line+" ")
            f.write("\n")
    f.close()

from ultralytics.utils import LOGGER
from pathlib import Path
import cv2
import numpy as np

def convert_segment_masks_to_yolo_seg(masks_dir, output_dir, classes):
    """
    Converts a dataset of segmentation mask images to the YOLO segmentation format.

    This function takes the directory containing the binary format mask images and converts them into YOLO segmentation format.
    The converted masks are saved in the specified output directory.

    Args:
        masks_dir (str): The path to the directory where all mask images (png, jpg) are stored.
        output_dir (str): The path to the directory where the converted YOLO segmentation masks will be stored.
        classes (int): Total classes in the dataset i.e. for COCO classes=80

    Example:
        ```python
        from ultralytics.data.converter import convert_segment_masks_to_yolo_seg

        # The classes here is the total classes in the dataset, for COCO dataset we have 80 classes
        convert_segment_masks_to_yolo_seg("path/to/masks_directory", "path/to/output/directory", classes=80)

import os
import re
import numpy as np
import cv2

def merge_images_with_overlap(folder_path, row_overlap_width, col_overlap_height):
    # 正则表达式匹配 crop_{row}_{col}.tif 格式的文件名
    pattern = re.compile(r'crop_(\d+)_(\d+)\.tif$')

    images = {}
    missing_images = set()

    # 遍历文件夹中的所有文件并记录所有的 (row, col) 键
    for filename in os.listdir(folder_path):
        match = pattern.match(filename)
        if match:
            try:
                row, col = map(int, match.groups())
                images[(row, col)] = os.path.join(folder_path, filename)
            except ValueError as e:
                print(f"转换行或列编号时出错: {e}，文件名: {filename}")
                continue

    if not images:
        print("没有找到匹配的图像文件")
        return None

    # 获取实际存在的行和列集合
    rows = {key[0] for key in images.keys()}
    cols = {key[1] for key in images.keys()}

    # 计算 max_row 和 max_col
    max_row = max(rows) if rows else -1
    max_col = max(cols) if cols else -1

    # 检查是否有缺失的图像文件
    for row in range(max_row + 1):
        for col in range(max_col + 1):
            if (row, col) not in images:
                missing_images.add((row, col))

    if missing_images:
        print("警告：以下图像文件缺失:")
        for missing in missing_images:
            print(f"行 {missing[0]} 列 {missing[1]}")

    # 假设所有的图片尺寸是一致的
    sample_image_path = list(images.values())[0]
    try:
        sample_image = cv2.imread(sample_image_path, cv2.IMREAD_COLOR)
        image_height, image_width, _ = sample_image.shape
    except Exception as e:
        print(f"无法打开样本图像: {sample_image_path}, 错误: {e}")
        return None

    result_width = (max_col + 1) * image_width - max_col * col_overlap_height
    result_height = (max_row + 1) * image_height - max_row * row_overlap_width

    # 创建一个新的空白图像来存放最终结果
    final_result = np.zeros((result_height, result_width, 3), dtype=np.uint8)

    # 处理每一行每一列
    for row in range(max_row + 1):
        for col in range(max_col + 1):
            path = images.get((row, col))
            if not path:
                print(f"跳过缺失的行 {row} 列 {col} 的图像")
                continue

            try:
                img = cv2.imread(path, cv2.IMREAD_COLOR)

                left = col * (image_width - col_overlap_height)
                upper = row * (image_height - row_overlap_width)
                right = left + image_width
                lower = upper + image_height

                # 如果不是第一列，则需要处理左边的重叠区域
                if col > 0 and (row, col-1) in images:
                    prev_img_path = images.get((row, col-1))
                    if prev_img_path:
                        prev_img = cv2.imread(prev_img_path, cv2.IMREAD_COLOR)
                        blended_part = cv2.addWeighted(
                            prev_img[:, -(col_overlap_height):, :], 0.5,
                            img[:, :col_overlap_height, :], 0.5, 0
                        )
                        final_result[upper:lower, left:left+col_overlap_height, :] = blended_part

                # 如果不是第一行，则需要处理上面的重叠区域
                if row > 0 and (row-1, col) in images:
                    prev_img_path = images.get((row-1, col))
                    if prev_img_path:
                        prev_img = cv2.imread(prev_img_path, cv2.IMREAD_COLOR)
                        blended_part = cv2.addWeighted(
                            prev_img[-row_overlap_width:, :, :], 0.5,
                            img[:row_overlap_width, :, :], 0.5, 0
                        )
                        final_result[upper:upper+row_overlap_width, left:right, :] = blended_part

                # 粘贴非重叠部分
                non_overlap_box = (
                    left + (col_overlap_height if col > 0 else 0),
                    upper + (row_overlap_width if row > 0 else 0),
                    right,
                    lower
                )
                final_result[non_overlap_box[1]:non_overlap_box[3], non_overlap_box[0]:non_overlap_box[2], :] = \
                    img[non_overlap_box[1]-upper:non_overlap_box[3]-upper, non_overlap_box[0]-left:non_overlap_box[2]-left, :]
            except Exception as e:
                print(f"处理图像时发生错误: {path}, 错误: {e}")

    # 保存最终结果图像
    output_path = os.path.join(folder_path, 'merged_image.tif')
    try:
        cv2.imwrite(output_path, final_result)
        print(f"合并后的图像已保存至: {output_path}")
    except Exception as e:
        print(f"保存结果图像时出错: {e}")

# 使用函数，指定文件夹路径、行重叠宽度和列重叠高度
folder_path = r'D:\zlc\drone\图像分割\栅格\Images'
row_overlap_width = 256  # 行之间重叠的宽度（像素）
col_overlap_height = 256  # 列之间重叠的高度（像素）

merge_images_with_overlap(folder_path, row_overlap_width, col_overlap_height)