Algorithm-Repo/SourceCode/MachineLearningUtils/TestTools/AIModelMetricsUtils/semantic_segmentation_metric.py

"""
语义分割：
[{'Label': 0, 'Confidence': 1.0, 'Image_size':[w, h], 'Contours': []}]
[{'Label': 1, 'Confidence': 1.0, 'Image_size':[w, h], 'Contours'：[ [ [1,2],[3,4] ] , [ [5,6],[7,8] ]  ]}]
[{'Label': 1, 'Confidence': 1.0, 'Image_size':[w, h], 'Contours'：[ [ [1,2],[3,4] ]  ]}]
[{'Label': 1, 'Confidence': 1.0, 'Image_size':[w, h], 'Contours'：[ [ [1,2],[3,4] ]  ]},{'Label': 2, 'Confidence': 1.0, 'Contours'：[ [ [1,2],[3,4] ]  ]}]
[]

"""
import json
import sys
import numpy as np
import cv2
from enum import Enum


class MethodAveragePrecision(Enum):
    """
    Class representing if the coordinates are relative to the
    image size or are absolute values.
    """
    EveryPointInterpolation = 1
    ElevenPointInterpolation = 2


class Evaluator_object_semamtic_segmentation(object):
    def __init__(self, iou_thres, method=MethodAveragePrecision.EveryPointInterpolation):
        """
        :param iou_thres: iou阈值
        :param method: 计算mAP的方法
        """
        self.method = method
        self.iou_thres = iou_thres
        self.wrong_file = {}
        self.background_images_all_indexs = []
        self.background_images_results = {}
        self.background_images_results_count = {}
        self.all_image_dict = {}
        self.all_no_background_images_pos_results = {}
        self.all_no_background_images_fp_results = {}
        self.all_no_background_images_tp_results = {}
        self.pred_results = []
        self.gt_results = []
        self.classes = []
        self.wrong_file['gt_wrong'] = []
        self.wrong_file['pred_wrong'] = []
        self.background_images_results_count['background_images_all_nums'] = 0
        self.all_image_dict['images_all_nums'] = 0

    def generate_metrics(self):
        classes = sorted(self.classes)
        # Precision x Recall is obtained individually by each class
        # Loop through by classes
        ret = []  # list containing metrics (precision, recall, average precision) of each class
        for c in classes:
            if c == 0:
                continue
            # Get only detection of class c
            dects = []
            [dects.append(d) for d in self.pred_results if d[1] == c]
            # Get only ground truths of class c, use filename as key
            gts = {}
            npos = 0
            for g in self.gt_results:
                if g[1] == c:
                    npos += 1
                    gts[g[0]] = gts.get(g[0], []) + [g]

                    if str(c) not in self.all_no_background_images_pos_results.keys():
                        self.all_no_background_images_pos_results[str(c)] = [g[0]]
                    else:
                        self.all_no_background_images_pos_results[str(c)].append(g[0])

            # sort detections by decreasing confidence
            dects = sorted(dects, key=lambda conf: conf[2], reverse=True)
            TP = np.zeros(len(dects))
            FP = np.zeros(len(dects))
            # create dictionary with amount of gts for each image
            det = {key: np.zeros(len(gts[key])) for key in gts}

            # print("Evaluating class: %s (%d detections)" % (str(c), len(dects)))
            # Loop through detections
            for d in range(len(dects)):
                # print('dect %s => %s' % (dects[d][0], dects[d][3],))
                # Find ground truth image
                gt = gts[dects[d][0]] if dects[d][0] in gts else []
                iouMax = sys.float_info.min
                for j in range(len(gt)):
                    # print('Ground truth gt => %s' % (gt[j][3],))
                    assert dects[d][4] == gt[j][4]
                    image_size = dects[d][4]
                    iou = Evaluator_object_semamtic_segmentation.iou(dects[d][3], gt[j][3], image_size)
                    if iou > iouMax:
                        iouMax = iou
                        jmax = j
                # Assign detection as true positive/don't care/false positive
                if iouMax >= self.iou_thres:
                    if det[dects[d][0]][jmax] == 0:
                        TP[d] = 1  # count as true positive
                        det[dects[d][0]][jmax] = 1  # flag as already 'seen'
                        # print("TP")

                        if str(c) not in self.all_no_background_images_tp_results.keys():
                            self.all_no_background_images_tp_results[str(c)] = [dects[d][0]]
                        else:
                            self.all_no_background_images_tp_results[str(c)].append(dects[d][0])

                    else:
                        FP[d] = 1  # count as false positive
                        if str(c) not in self.all_no_background_images_fp_results.keys():
                            self.all_no_background_images_fp_results[str(c)] = [dects[d][0]]
                        else:
                            self.all_no_background_images_fp_results[str(c)].append(dects[d][0])
                        # print("FP")
                # - A detected "cat" is overlaped with a GT "cat" with IOU >= IOUThreshold.
                else:
                    FP[d] = 1  # count as false positive
                    if str(c) not in self.all_no_background_images_fp_results.keys():
                        self.all_no_background_images_fp_results[str(c)] = [dects[d][0]]
                    else:
                        self.all_no_background_images_fp_results[str(c)].append(dects[d][0])
                    # print("FP")
            # compute precision, recall and average precision
            acc_FP = np.cumsum(FP)
            acc_TP = np.cumsum(TP)
            rec = acc_TP / npos
            prec = np.divide(acc_TP, (acc_FP + acc_TP))
            # Depending on the method, call the right implementation
            if self.method == MethodAveragePrecision.EveryPointInterpolation:
                [ap, mpre, mrec, ii] = Evaluator_object_semamtic_segmentation.CalculateAveragePrecision(rec, prec)
            else:
                [ap, mpre, mrec, _] = Evaluator_object_semamtic_segmentation.ElevenPointInterpolatedAP(rec, prec)
            # add class result in the dictionary to be returned
            r = {
                'class': c,
                'precision': prec,
                'recall': rec,
                'AP': ap,
                'interpolated precision': mpre,
                'interpolated recall': mrec,
                'total positives': npos,
                'total TP': np.sum(TP),
                'total FP': np.sum(FP)
            }
            ret.append(r)

        return ret

    def add_batch(self, gt_file, pred_file, image_index):
        if gt_file == []:
            self.wrong_file['gt_wrong'].append(image_index)
        elif pred_file == []:
            self.wrong_file['pred_wrong'].append(image_index)
        elif gt_file != [] and pred_file != []:
            # 判断gt为背景的图像：
            self.all_image_dict['images_all_nums'] += 1
            if len(gt_file) == 1 and gt_file[0]['Label'] == 0 and gt_file[0]['Contours'] == []:
                self.background_images_results_count['background_images_all_nums'] += 1
                if len(pred_file) == 1 and pred_file[0]['Label'] == 0 and gt_file[0]['Contours'] == []:
                    # if 'gtlabel_0_predlabel_{}'.format(0) not in self.background_images_results.keys():
                    #     self.background_images_results['gtlabel_0_predlabel_{}'.format(0)] = [image_index]
                    # else:
                    #     self.background_images_results['gtlabel_0_predlabel_{}'.format(0)].append(image_index)
                    if 'gtlabel_0_predlabel_{}'.format(0) not in self.background_images_results_count.keys():
                        self.background_images_results_count['gtlabel_0_predlabel_{}'.format(0)] = 1
                    else:
                        self.background_images_results_count['gtlabel_0_predlabel_{}'.format(0)] += 1
                else:
                    for i in range(len(pred_file)):
                        pred_label = pred_file[i]['Label']
                        if 'gtlabel_0_predlabel_{}'.format(pred_label) not in self.background_images_results.keys():
                            self.background_images_results['gtlabel_0_predlabel_{}'.format(pred_label)] = [image_index]
                        else:
                            self.background_images_results['gtlabel_0_predlabel_{}'.format(pred_label)].append(
                                image_index)
                        if 'gtlabel_0_predlabel_{}'.format(
                                pred_label) not in self.background_images_results_count.keys():
                            self.background_images_results_count['gtlabel_0_predlabel_{}'.format(pred_label)] = 1
                        else:
                            self.background_images_results_count['gtlabel_0_predlabel_{}'.format(pred_label)] += 1
            else:
                gt_len = len(gt_file)
                pred_len = len(pred_file)
                for i in range(gt_len):
                    each_gt_label = gt_file[i]['Label']
                    each_gt_counters = gt_file[i]['Contours']
                    each_image_size = gt_file[i]['Image_size']
                    self.gt_results.append([image_index, each_gt_label, 1.0, (each_gt_counters), each_image_size])
                    if each_gt_label not in self.classes:
                        self.classes.append(each_gt_label)
                for i in range(pred_len):
                    each_pred_label = pred_file[i]['Label']
                    each_pred_confidence = pred_file[i]['Confidence']
                    each_pred_counters = pred_file[i]['Contours']
                    each_pred_image_size = pred_file[i]['Image_size']
                    self.pred_results.append([image_index, each_pred_label, each_pred_confidence, (each_pred_counters),
                                              each_pred_image_size])
                    if each_pred_label not in self.classes:
                        self.classes.append(each_pred_label)

    @staticmethod
    def iou(contoursA, contoursB, image_size):

        pred_mask = Evaluator_object_semamtic_segmentation.contourtomask(contoursA, image_size)
        gt_mask = Evaluator_object_semamtic_segmentation.contourtomask(contoursB, image_size)
        iou = Evaluator_object_semamtic_segmentation.iou_score(pred_mask, gt_mask)
        assert iou >= 0
        return iou

    @staticmethod
    def iou_score(output, target):
        smooth = 1e-5
        intersection = (output & target).sum()
        union = (output | target).sum()
        return (intersection + smooth) / (union + smooth)

    @staticmethod
    def contourtomask(contours, image_size):
        mask = np.zeros((image_size[1], image_size[0]), dtype=np.int32)
        if len(contours) == 1:
            contours_cv = np.zeros((len(contours[0]), 1, 2), dtype=np.int32)
            for i in range(len(contours[0])):
                contours_cv[i] = [int(contours[0][i][0]), int(contours[0][i][1])]
            mask = cv2.drawContours(mask.copy(), [contours_cv], -1, 1, cv2.FILLED)

        # 暂时一个contours中，如果存放了两条轮廓，则必须是去差值
        # 后续有超过两条的情况，需要再添加相关代码
        if len(contours) == 2:
            contours_cv_1 = np.zeros((len(contours[0]), 1, 2), dtype=np.int32)
            for i in range(len(contours[0])):
                contours_cv_1[i] = [int(contours[0][i][0]), int(contours[0][i][1])]
            mask_1 = cv2.drawContours(mask.copy(), [contours_cv_1], -1, 1, cv2.FILLED)

            contours_cv_2 = np.zeros((len(contours[1]), 1, 2), dtype=np.int32)
            for i in range(len(contours[1])):
                contours_cv_2[i] = [int(contours[1][i][0]), int(contours[1][i][1])]
            mask_2 = cv2.drawContours(mask.copy(), [contours_cv_2], -1, 1, cv2.FILLED)

            mask = abs(mask_2 - mask_1)

        if len(contours) >= 3:
            raise IOError('Error: contours: %s could not be process.' % contours)

        return mask

    @staticmethod
    def CalculateAveragePrecision(rec, prec):
        mrec = []
        mrec.append(0)
        [mrec.append(e) for e in rec]
        mrec.append(1)
        mpre = []
        mpre.append(0)
        [mpre.append(e) for e in prec]
        mpre.append(0)
        for i in range(len(mpre) - 1, 0, -1):
            mpre[i - 1] = max(mpre[i - 1], mpre[i])
        ii = []
        for i in range(len(mrec) - 1):
            if mrec[1 + i] != mrec[i]:
                ii.append(i + 1)
        ap = 0
        for i in ii:
            ap = ap + np.sum((mrec[i] - mrec[i - 1]) * mpre[i])
        # return [ap, mpre[1:len(mpre)-1], mrec[1:len(mpre)-1], ii]
        return [ap, mpre[0:len(mpre) - 1], mrec[0:len(mpre) - 1], ii]

    @staticmethod
    # 11-point interpolated average precision
    def ElevenPointInterpolatedAP(rec, prec):
        # def CalculateAveragePrecision2(rec, prec):
        mrec = []
        # mrec.append(0)
        [mrec.append(e) for e in rec]
        # mrec.append(1)
        mpre = []
        # mpre.append(0)
        [mpre.append(e) for e in prec]
        # mpre.append(0)
        recallValues = np.linspace(0, 1, 11)
        recallValues = list(recallValues[::-1])
        rhoInterp = []
        recallValid = []
        # For each recallValues (0, 0.1, 0.2, ... , 1)
        for r in recallValues:
            # Obtain all recall values higher or equal than r
            argGreaterRecalls = np.argwhere(mrec[:] >= r)
            pmax = 0
            # If there are recalls above r
            if argGreaterRecalls.size != 0:
                pmax = max(mpre[argGreaterRecalls.min():])
            recallValid.append(r)
            rhoInterp.append(pmax)
        # By definition AP = sum(max(precision whose recall is above r))/11
        ap = sum(rhoInterp) / 11
        # Generating values for the plot
        rvals = []
        rvals.append(recallValid[0])
        [rvals.append(e) for e in recallValid]
        rvals.append(0)
        pvals = []
        pvals.append(0)
        [pvals.append(e) for e in rhoInterp]
        pvals.append(0)
        # rhoInterp = rhoInterp[::-1]
        cc = []
        for i in range(len(rvals)):
            p = (rvals[i], pvals[i - 1])
            if p not in cc:
                cc.append(p)
            p = (rvals[i], pvals[i])
            if p not in cc:
                cc.append(p)
        recallValues = [i[0] for i in cc]
        rhoInterp = [i[1] for i in cc]
        return [ap, rhoInterp, recallValues, None]


if __name__ == "__main__":

    import os
    import glob

    currentPath = os.path.dirname(os.path.abspath(__file__))
    path_pred = os.path.join(currentPath, 'segment_test\\pred\\')
    path_gt = os.path.join(currentPath, 'segment_test\\gt\\')

    files = glob.glob(path_gt + "*.txt")
    files.sort()

    evaluator = Evaluator_object_semamtic_segmentation(iou_thres=0.5)

    for file in files:
        with open(file, "r", encoding='utf-8') as gtf:
            gt_anno = gtf.readlines()
            assert len(gt_anno) == 1, '每个txt对应一个json'
            gt_file = json.loads(gt_anno[0])

        pred_files_name = file.split('\\')[-1]
        pred_files = os.path.join(path_pred, pred_files_name)

        with open(pred_files, "r", encoding='utf-8') as predf:
            pred_anno = predf.readlines()
            assert len(pred_anno) == 1, '每个txt对应一个json'
            pred_file = json.loads(pred_anno[0])

        pred_len = len(pred_file)

        for i in range(pred_len):
            each_pred_label = pred_file[i]['Label']

        image_index = str(file.split('\\')[-1].split('.txt')[0])
        evaluator.add_batch(gt_file, pred_file, image_index)

    wrong_file = evaluator.wrong_file
    print(wrong_file)
    background_images_results = evaluator.background_images_results
    print(background_images_results)
    all_image_dict = evaluator.all_image_dict
    print(all_image_dict)

    metricsPerClass = evaluator.generate_metrics()
    for mc in metricsPerClass:
        c = mc['class']
        precision = mc['precision']
        recall = mc['recall']
        # Print AP per class

        # precision_all1 = precision[-1]
        # precision_all = mc['total TP'] / (mc['total TP'] + mc['total FP'])
        precision_all2 = 0 if (mc['total TP'] + mc['total FP']) == 0 else np.divide(mc['total TP'],
                                                                                    (mc['total TP'] + mc['total FP']))
        # recall_all1 = recall[-1]
        # recall_all = mc['total TP'] / mc['total positives']
        recall_all2 = 0 if mc['total positives'] == 0 else np.divide(mc['total TP'], mc['total positives'])
        print(c)
        print(precision_all2)
        print(recall_all2)
        print(mc['total TP'])
        print(mc['total FP'])
        print(mc['total positives'])
        print('------------------------------')