语义分割重要指标计算python程序网上有好些,各有特色,但最近在用paddle框架感觉还顺手,其计算分割的指标程序还可以,但原始的有些繁琐,而且没有对分割好的结果图和标记图做评估,在经过一段时间研读源码后,我在原始基础上做了一些修改和封装,最后形成如下代码,在此整理一下并分享出来。
在此先说明下语义分割指标是啥,感觉这几篇文章说的很好:
1 https://www.jeremyjordan.me/evaluating-image-segmentation-models/
2 https://www.cnblogs.com/sddai/p/14684335.html
代码简明扼要,不做过多解释。
#DQ 2020/12/31
import os
import sys
import cv2,time
import numpy as np
from scipy.sparse import csr_matrix
class ConfusionMatrix(object):
"""
Confusion Matrix for segmentation evaluation
"""
def __init__(self, num_classes=2, streaming=True):
self.confusion_matrix = np.zeros([num_classes, num_classes],
dtype='int64')
self.num_classes = num_classes
self.streaming = streaming
#pred:shape=(h,w);label:(h,w).val=clsnameid
def calculate(self, pred, label):
# If not in streaming mode, clear matrix everytime when call `calculate`
if not self.streaming:
self.zero_matrix()
pred=pred.flatten()
label=label.flatten()
one = np.ones_like(pred)
# Accumuate ([row=label, col=pred], 1) into sparse matrix
spm = csr_matrix((one, (label, pred)),
shape=(self.num_classes, self.num_classes))
spm = spm.todense()
self.confusion_matrix += spm
def zero_matrix(self):
""" Clear confusion matrix """
self.confusion_matrix = np.zeros([self.num_classes, self.num_classes],
dtype='int64')
def mean_iou(self):
iou_list = []
avg_iou = 0
# TODO: use numpy sum axis api to simpliy
vji = np.zeros(self.num_classes, dtype=int)
vij = np.zeros(self.num_classes, dtype=int)
for j in range(self.num_classes):
v_j = 0
for i in range(self.num_classes):
v_j += self.confusion_matrix[j][i]
vji[j] = v_j
for i in range(self.num_classes):
v_i = 0
for j in range(self.num_classes):
v_i += self.confusion_matrix[j][i]
vij[i] = v_i
for c in range(self.num_classes):
total = vji[c] + vij[c] - self.confusion_matrix[c][c]
if total == 0:
iou = 0
else:
iou = float(self.confusion_matrix[c][c]) / total
avg_iou += iou
iou_list.append(iou)
avg_iou = float(avg_iou) / float(self.num_classes)
return np.around(np.array(iou_list), decimals=3), np.around(avg_iou,decimals=3)
def accuracy(self):
total = self.confusion_matrix.sum()
total_right = 0
for c in range(self.num_classes):
total_right += self.confusion_matrix[c][c]
if total == 0:
avg_acc = 0
else:
avg_acc = float(total_right) / total
vij = np.zeros(self.num_classes, dtype=int)
for i in range(self.num_classes):
v_i = 0
for j in range(self.num_classes):
v_i += self.confusion_matrix[j][i]
vij[i] = v_i
acc_list = []
for c in range(self.num_classes):
if vij[c] == 0:
acc = 0
else:
acc = self.confusion_matrix[c][c] / float(vij[c])
acc_list.append(acc)
return np.around(np.array(acc_list),decimals=3), np.around(avg_acc,decimals=3)
def kappa(self):
vji = np.zeros(self.num_classes)
vij = np.zeros(self.num_classes)
for j in range(self.num_classes):
v_j = 0
for i in range(self.num_classes):
v_j += self.confusion_matrix[j][i]
vji[j] = v_j
for i in range(self.num_classes):
v_i = 0
for j in range(self.num_classes):
v_i += self.confusion_matrix[j][i]
vij[i] = v_i
total = self.confusion_matrix.sum()
# avoid spillovers
# TODO: is it reasonable to hard code 10000.0?
total = float(total) / 10000.0
vji = vji / 10000.0
vij = vij / 10000.0
tp = 0
tc = 0
for c in range(self.num_classes):
tp += vji[c] * vij[c]
tc += self.confusion_matrix[c][c]
tc = tc / 10000.0
pe = tp / (total * total)
po = tc / total
kappa = (po - pe) / (1 - pe)
return np.around(kappa,decimals=3)
class SegMetrics:
def __init__(self,predimdir,labelimdir,clsnum,pcolors=None):
self.predimdir=predimdir
self.labelimdir=labelimdir
self.clsnum=clsnum
self.pcolors=pcolors#pcolor2clsnameiddict=(np.array([0,0,0]),np.array([0,0,128]))
def get_confusion_matrix(self):
#彩色分割图转化为clsnameid,图,背景为0,第一个类灰度值为1,第二个类为2
def pcolorsegim2clsnameidim(im, pcolors): #pcolors=(np.array([b,g,r]),np.array([b,g,r]),..)
imhw = im.shape[:-1]
clsnameidim = np.zeros(imhw, np.int32) #must np.int32
for k,pcolor in enumerate(pcolors[1:]):
im1=im-pcolor
inds=np.argwhere(np.sum(im1,axis=2)==0)
clsnameidim[inds[:,0],inds[:,1]]=k+1
return clsnameidim
conf_mat = ConfusionMatrix(self.clsnum, streaming=True)
imnames=os.listdir(self.predimdir)
for imname in imnames:
print(imname)
predimpath=os.path.join(self.predimdir,imname)
predim=cv2.imread(predimpath)
labelimpath=os.path.join(self.labelimdir,imname)
labelim=cv2.imread(labelimpath)
if self.pcolors is not None:
predim = pcolorsegim2clsnameidim(predim, self.pcolors)
labelim = pcolorsegim2clsnameidim(labelim, self.pcolors)
conf_mat.calculate(predim, labelim)
_, iou = conf_mat.mean_iou()
_, acc = conf_mat.accuracy()
return conf_mat
def get_segmetrics(self):
conf_mat=self.get_confusion_matrix()
category_iou, avg_iou = conf_mat.mean_iou()
category_acc, avg_acc = conf_mat.accuracy()
kappa = conf_mat.kappa()
segmetrics={'avg_iou':avg_iou,'avg_acc':avg_acc,'category_iou':category_iou.tolist(),'category_acc':category_acc.tolist(),'kappa':kappa}
return segmetrics
if __name__ == '__main__':
predimdir='/data/predims'
labelimdir='/data/annotset'
clsnum=2
pcolors = (np.array([0, 0, 0]), np.array([0, 0, 128]))#伪彩色图表,第一个为背景,第二个第一个类,以此类推
segmet=SegMetrics(predimdir,labelimdir,clsnum,pcolors)
segmetrics=segmet.get_segmetrics()
print(segmetrics)