src/confusion_matrix.py

import numpy as np
import pandas as pd


def box_iou_calc(boxes1, boxes2):
    # https://github.com/pytorch/vision/blob/master/torchvision/ops/boxes.py
    """
    Return intersection-over-union (Jaccard index) of boxes.
    Both sets of boxes are expected to be in (x1, y1, x2, y2) format.
    Arguments:
        boxes1 (Array[N, 4])
        boxes2 (Array[M, 4])
    Returns:
        iou (Array[N, M]): the NxM matrix containing the pairwise
            IoU values for every element in boxes1 and boxes2

    This implementation is taken from the above link and changed so that it only uses numpy.
    """

    def box_area(box):
        # box = 4xn
        return (box[2] - box[0]) * (box[3] - box[1])
    

    area1 = box_area(boxes1.T)
    area2 = box_area(boxes2.T)

    lt = np.maximum(boxes1[:, None, :2], boxes2[:, :2])  # [N,M,2]
    rb = np.minimum(boxes1[:, None, 2:], boxes2[:, 2:])  # [N,M,2]

    inter = np.prod(np.clip(rb - lt, a_min = 0, a_max = None), 2)
    return inter / (area1[:, None] + area2 - inter)  # iou = inter / (area1 + area2 - inter)

def mask_iou_calc(mask1, mask2):

    # build function to take in two masks, compare them and see what their iou is.
    # similar to above but in mask.

    return


class ConfusionMatrix:
    def __init__(self, num_classes, CONF_THRESHOLD = 0.2, IOU_THRESHOLD = 0.5):
        self.matrix = np.zeros((num_classes + 1, num_classes + 1))
        self.num_classes = num_classes
        self.CONF_THRESHOLD = CONF_THRESHOLD
        self.IOU_THRESHOLD = IOU_THRESHOLD
        self.got_tpfpfn = False
    
    def process_batch(self, detections, labels, return_matches=False):
        '''
        Return intersection-over-union (Jaccard index) of boxes.
        Both sets of boxes are expected to be in (x1, y1, x2, y2) format.
        Arguments:
            detections (Array[N, 6]), x1, y1, x2, y2, conf, class
            labels (Array[M, 5]), class, x1, y1, x2, y2
        Returns:
            None, updates confusion matrix accordingly
        '''
        detections = detections[detections[:, 4] > self.CONF_THRESHOLD]
        gt_classes = labels[:, 0].astype(np.int16)
        detection_classes = detections[:, 5].astype(np.int16)

        all_ious = box_iou_calc(labels[:, 1:], detections[:, :4])
        # print()
        # print('=== all_ious ===')
        # print(all_ious)
        want_idx = np.where(all_ious > self.IOU_THRESHOLD)
        # print('=== want_idx ===')
        # print(want_idx)
        # print()
        
        all_matches = []
        for i in range(want_idx[0].shape[0]):
            all_matches.append([want_idx[0][i], want_idx[1][i], all_ious[want_idx[0][i], want_idx[1][i]]])
        
        all_matches = np.array(all_matches)
        if all_matches.shape[0] > 0: # if there is match
            all_matches = all_matches[all_matches[:, 2].argsort()[::-1]]

            all_matches = all_matches[np.unique(all_matches[:, 1], return_index = True)[1]]

            all_matches = all_matches[all_matches[:, 2].argsort()[::-1]]

            all_matches = all_matches[np.unique(all_matches[:, 0], return_index = True)[1]]
        
        for i, label in enumerate(labels):
            if all_matches.shape[0] > 0 and all_matches[all_matches[:, 0] == i].shape[0] == 1:
                gt_class = gt_classes[i]
                detection_class = detection_classes[int(all_matches[all_matches[:, 0] == i, 1][0])]
                self.matrix[(gt_class), detection_class] += 1
            else:
                gt_class = gt_classes[i]
                self.matrix[(gt_class), self.num_classes] += 1
        
        for i, detection in enumerate(detections):
            if all_matches.shape[0] and all_matches[all_matches[:, 1] == i].shape[0] == 0:
                detection_class = detection_classes[i]
                self.matrix[self.num_classes ,detection_class] += 1
    
        if return_matches:
            return all_matches
        
    def get_tpfpfn(self):
        self.tp = np.diag(self.matrix).sum()
        fp = self.matrix.copy()
        np.fill_diagonal(fp, 0)
        self.fp = fp[:,:-1].sum()
        self.fn = self.matrix[:-1, -1].sum()
        self.got_tpfpfn = True

    def get_PR(self):
        if not self.got_tpfpfn:
            self.get_tpfpfn()
        # print (tp, fp, fn)
        self.precision = self.tp / (self.tp+self.fp)
        self.recall = self.tp/(self.tp+self.fn)

    def return_matrix(self):
        return self.matrix

    def process_full_matrix(self):
        """method to process matrix to something more readable
        """
        for idx, i in enumerate(self.matrix):
            i[0] = idx
        self.matrix = np.delete(self.matrix, 0, 0)
            
    def print_matrix_as_df(self):
        """method to print out processed matrix
        """
        df = pd.DataFrame(self.matrix)
        print (df.to_string(index=False))

    # def print_matrix(self):
    #     for i in range(self.num_classes + 1):
    #         print(' '.join(map(str, self.matrix[i])))

    def return_as_csv(self, csv_file_path):
        """method to print out processed matrix
        """
        df = pd.DataFrame(self.matrix)
        df.to_csv(csv_file_path, index = False)
        print (f"saved to: {csv_file_path}")

    def return_as_df(self):
        """method to print out processed matrix
        """
        df = pd.DataFrame(self.matrix)
        # df = df.set_index(0)
        # df.set_index(0)
        # print(df.columns)
        return df