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

# 文件夹路径
# folder_A = '/mnt/lustre/zhouyue1.vendor/geoground/data/exp_1103/mix_hbb/vis_seg_b8_val_box2mask'  # 替换为文件夹A的实际路径
folder_A = '/mnt/lustre/zhouyue1.vendor/geoground/data/exp_1104/mix_hbb/vis_seg_b8_test_box2mask'  # 替换为文件夹A的实际路径
folder_B = '/mnt/lustre/zhouyue1.vendor/geoground/data/masks/rrsisd'  # 替换为文件夹B的实际路径

# 定义计算IoU的函数
def calculate_iou(mask1, mask2):
    # mask1 和 mask2 应该是二值矩阵
    intersection = np.logical_and(mask1, mask2).sum()  # 交集
    union = np.logical_or(mask1, mask2).sum()  # 并集
    if union == 0:
        return 1.0  # 如果并集为0，说明两个图像都没有前景物体，IoU为1
    return intersection / union, intersection, union

# 获取文件夹A和文件夹B中的所有png文件
files_A = [f for f in os.listdir(folder_A) if f.endswith('.png')]
files_B = [f for f in os.listdir(folder_B) if f.endswith('.png')]

# 正则表达式提取下划线之前的数字部分
pattern = re.compile(r'^(\d+)_')

# 创建一个字典来存储B文件夹中的文件，键为下划线之前的数字，值为文件名列表
b_files_dict = {}

# 遍历文件夹B，提取文件名中的数字前缀并存入字典
for file_B in files_B:
    match_B = pattern.match(file_B)
    if match_B:
        prefix = match_B.group(1)
        if prefix in b_files_dict:
            b_files_dict[prefix].append(file_B)
        else:
            b_files_dict[prefix] = [file_B]

# 存储所有 IoU 的列表
ious = []
correct_count = 0  # 统计 IoU >= 0.5 的匹配对数量
total_count = 0    # 总的匹配对数量
cum_I, cum_U = 0, 0

# 匹配文件夹A中的文件，并计算对应文件的IoU
for file_A in files_A:
    match_A = pattern.match(file_A)
    if match_A:
        prefix = match_A.group(1)
        if prefix in b_files_dict:
            # 打开文件夹A和文件夹B中的匹配图片
            img_A_path = os.path.join(folder_A, file_A)
            img_B_path = os.path.join(folder_B, b_files_dict[prefix][0])  # 假设只取第一个匹配的文件

            # 打开图像并转换为二值化的numpy数组
            img_A = Image.open(img_A_path).convert('L')  # 转换为灰度图
            img_B = Image.open(img_B_path).convert('L')  # 转换为灰度图

            img_B_size = img_B.size  # (width, height)
            img_A = img_A.resize(img_B_size)

            # 将图像转换为二值数组，假设阈值为128
            mask_A = np.array(img_A) > 128
            mask_B = np.array(img_B) > 128

            # 计算 IoU
            iou, I, U = calculate_iou(mask_A, mask_B)
            ious.append(iou)

            cum_I += I
            cum_U += U

            # 统计总的匹配对数量
            total_count += 1

            # 统计 IoU >= 0.5 的匹配对
            if iou >= 0.5:
                correct_count += 1

            print(f"IoU for {file_A} and {b_files_dict[prefix][0]}: {iou}")
        else:
            print(f"No match found for {file_A} in folder B")

# 计算平均IoU
if ious:
    mean_iou = np.mean(ious)
    print(f"mIoU: {mean_iou}")
    print(f"oIoU: {cum_I / cum_U}")
else:
    mean_iou = 0
    print("No matches found to calculate IoU.")


# 计算 Acc@0.5
if total_count > 0:
    acc_0_5 = correct_count / total_count
    print(f"Acc@0.5: {acc_0_5}")
else:
    acc_0_5 = 0
    print("No matches found to calculate Acc@0.5.")