语义分割重要指标计算python程序

       语义分割重要指标计算python程序网上有好些，各有特色，但最近在用paddle框架感觉还顺手，其计算分割的指标程序还可以，但原始的有些繁琐，而且没有对分割好的结果图和标记图做评估，在经过一段时间研读源码后，我在原始基础上做了一些修改和封装，最后形成如下代码,在此整理一下并分享出来。

　　在此先说明下语义分割指标是啥，感觉这几篇文章说的很好：

１ 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)