图像分割IOU计算代码

import os
os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE'
from PIL import Image
import numpy as np
from sklearn import metrics, neighbors


str1 = 'D:/program/tool/images/meaniou/driver_unet_01test_result.png'
im = Image.open(str1)
im = np.array(im)

str2 = 'D:/program/tool/images/meaniou/drive_unet_01_manual1.png'
im_gt = Image.open(str2)
im_gt = np.array(im_gt)

c = metrics.confusion_matrix(im.flatten(), im_gt.flatten())  # 混淆矩阵
TP = c[1][1]  # 预测为前景，GT为前景
TN = c[0][0]  # 预测为前景，GT为前景
FP = c[1][0]  # 预测为前景，GT为背景
FN = c[0][1]  # 预测为背景，GT为前景
print('IOU统计:\n'+str(c)+'\n\nFP为:'+str(FP)+'  FN为:'+str(FN)+'  TP为:'+str(TP)+'  TN:'+str(TN))

iou_p = TP/(TP+FN+FP)  # 前景的IOU
iou_n = TN/(TN+FN+FP)  # 背景的IOU
mean_iou = (iou_n+iou_p)/2

print('前景IOU:{:.3f}%\n背景IOU:{:.3f}%\nMeanIou:{:.3f}%'.format(iou_p*100, iou_n*100, mean_iou*100))