import os
import numpy as np
from PIL import Image
import torch

def resize_to_image_8():
    # 原始图片文件夹路径
    input_folder = "D:\XMU\mac\hujie\\3D\DUSt3R\dust3r\data\llff(sanerf-hq)\colinepiano\images"
    # 新的保存缩放图片的文件夹路径
    output_folder = "D:\XMU\mac\hujie\\3D\DUSt3R\dust3r\data\llff(sanerf-hq)\colinepiano\images_8"

    # 创建输出文件夹（如果不存在）
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)

    # 遍历输入文件夹中的所有文件
    for filename in os.listdir(input_folder):
        # 获取文件路径
        file_path = os.path.join(input_folder, filename)

        # 确保是文件且是图片文件（这里假设图片格式为jpg或png）
        if os.path.isfile(file_path) and filename.lower().endswith(('.jpg', '.jpeg', '.png')):
            # 打开图片
            original_image = Image.open(file_path)

            # 获取原始尺寸
            original_width, original_height = original_image.size

            # 计算缩放后的尺寸（1/8）
            new_width = original_width // 8
            new_height = original_height // 8

            # 进行缩放
            resized_image = original_image.resize((new_width, new_height), Image.LANCZOS)

            # 生成新的文件名并保存到输出文件夹
            new_file_path = os.path.join(output_folder, filename)
            resized_image.save(new_file_path)
            print(f"Image {filename} resized and saved to {new_file_path}")

    print("All images have been resized and saved.")


def RGB_to_mask():

    # 输入文件夹路径
    input_folder = "D:\XMU\mac\hujie\\3D\DUSt3R\dust3r\data\llff(sanerf-hq)\\piano\gt"
    # 输出文件夹路径
    output_folder = "D:\XMU\mac\hujie\\3D\DUSt3R\dust3r\data\llff(sanerf-hq)\\piano\gt_masks"

    # 创建输出文件夹（如果不存在）
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)

    # 遍历输入文件夹中的所有文件
    for filename in os.listdir(input_folder):
        # 构建完整的文件路径
        file_path = os.path.join(input_folder, filename)

        # 打开图片
        image = Image.open(file_path).convert('RGB')

        # 转换为灰度图像
        gray_image = image.convert('L')

        # 转换为NumPy数组
        gray_array = np.array(gray_image)

        # 二值化（假设阈值为128，可根据需要调整）
        threshold = 10
        binary_array = (gray_array > threshold).astype(np.uint8) * 255

        # 转换为Image对象
        binary_image = Image.fromarray(binary_array, 'L')  # 'L'代表灰度模式

        base_name, extension = os.path.splitext(filename)

        # 生成新的文件名
        new_filename = base_name + "_mask" + extension

        new_file_path = os.path.join(output_folder, new_filename)

        # 保存新的图片
        binary_image.save(new_file_path)
        print(f"Converted {filename} to {new_filename} and saved.")

    print("All images have been processed.")

# 读取mask文件夹下的所有ground truth masks，不再需要经过SAM生成mask
def get_gt_masks(folder_path):

    from dust3r.utils.image import load_images, rgb
    imgs_mask = load_images(folder_path, 512)
    # 定义保存布尔mask的列表
    bool_masks = []

    for mask in imgs_mask:
        image_array = mask['img'].squeeze(0).numpy()
        # 将RGB图像转换为灰度图像
        # 使用简单的加权方法转换为灰度: Y = 0.299*R + 0.587*G + 0.114*B
        gray_image = 0.299 * image_array[0] + 0.587 * image_array[1] + 0.114 * image_array[2]

        # 将灰度图像转换为布尔数组（前景为True，背景为False）
        bool_array = gray_image > 0

        # 将布尔数组添加到列表中
        bool_masks.append(bool_array)

    # 输出布尔mask的数量
    print(f"Total number of mask images processed: {len(bool_masks)}")
    return bool_masks

if __name__ == "__main__":
    # 输入文件夹路径
    folder_path = "D:\XMU\mac\hujie\\3D\DUSt3R\dust3r\data\llff(sanerf-hq)\cecread\gt_masks"

    RGB_to_mask()