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(pred_masks, gt_masks):
    """Helper function calculate the IOU of masks

    Args:
        pred_masks (_type_): N x H x W, array of N masks 
        gt_masks (_type_): M x H x W, an array of M masks 

    Returns:
        iou: an array of NxM of IOU ([0,1]) 
            N rows - number of actual labels
            M columns - number of preds
    """
    if pred_masks.size == 0:
        return np.array([])

    # build function to take in two masks, compare them and see what their iou is.
    # similar to above but in mask.
    tp = np.sum(np.multiply(pred_masks[:, None], gt_masks), axis = (2,3))
    fp = np.sum(np.where(pred_masks[:, None] > gt_masks, 1, 0), axis = (2,3))
    fn = np.sum(np.where(pred_masks[:, None] < gt_masks, 1, 0), axis = (2,3))

    # print (f"tp: {tp}")
    # print (f"fp: {fp}")
    # print (f"fn: {fn}")
    iou = tp / (tp + fn + fp)

    return iou.T

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, task = "det"):
        '''
        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
        '''

        if task == 'det':
            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])
            want_idx = np.where(all_ious > self.IOU_THRESHOLD)
            
        elif task == 'seg':
            detections = [detection for detection in detections if detection[1] > self.CONF_THRESHOLD]
            gt_classes = np.array([label[0]for label in labels], dtype = np.int16) 
            detection_classes = np.array([detection[2] for detection in detections], dtype = np.int16)
            all_ious = mask_iou_calc(np.array([detection[0] for detection in detections]), np.array([label[1] for label in labels]))
            want_idx = np.where(all_ious > self.IOU_THRESHOLD)
        
        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
    
if __name__ == '__main__':
    # # test IOU for segmentation masks
    gtMasks = np.array([[[1, 1, 0], 
                        [0, 1, 0], 
                        [0, 0, 0]],
                        [[1, 1, 0], 
                        [0, 1, 1], 
                        [0, 0, 0]]])
    predMasks = np.array([[[1, 1, 0], 
                        [0, 1, 1], 
                        [0, 0, 0]],
                        [[1, 1, 0], 
                        [0, 1, 0], 
                        [0, 0, 0]]])
    
    # IOU is 0.75 
    IOU = mask_iou_calc(predMasks, gtMasks)
    print (IOU.shape)