File size: 3,513 Bytes
4963c36 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 | 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.") |