Bounding box群のPrecisionとRecallを計算するための関数覚書

calc_precision_recall.py

def bb_intersection_over_union(boxA, boxB):
    """ calculate IOU score between boxA and boxB
    Args:
        boxA : [hori_lo,virt_lo,hori_hi,virt_hi]
        boxB : [hori_lo,virt_lo,hori_hi,virt_hi]
    Returns:
        iou_score
    """
    # determine the (x, y)-coordinates of the intersection rectangle
    xA = max(boxA[0], boxB[0])
    yA = max(boxA[1], boxB[1])
    xB = min(boxA[2], boxB[2])
    yB = min(boxA[3], boxB[3])
    # compute the area of intersection rectangle
    interArea = max(0, xB - xA + 1) * max(0, yB - yA + 1)
    # compute the area of both the prediction and ground-truth
    # rectangles
    boxAArea = (boxA[2] - boxA[0] + 1) * (boxA[3] - boxA[1] + 1)
    boxBArea = (boxB[2] - boxB[0] + 1) * (boxB[3] - boxB[1] + 1)
    # compute the intersection over union by taking the intersection
    # area and dividing it by the sum of prediction + ground-truth
    # areas - the interesection area
    iou = interArea / float(boxAArea + boxBArea - interArea)
    # return the intersection over union value
    return iou

import copy
def calc_precision_recall(pred_bbox_per_frame,gt_bbox_per_frame,iou_th=0.95):
    """ calculate precision and recall of one frame bbox set
    Args:
        pred_bbox_per_frame : [[hori_lo,virt_lo,hori_hi,virt_hi],...]
        gt_bbox_per_frame : [[hori_lo,virt_lo,hori_hi,virt_hi],...]
        iou_th : a threshould of correct or not correct
    Returns:
        sum_true_positive, sum_positive, sum_gt
        sum_true_positive : sum of true positive 
        sum_positive : length of pred_bbox_per_frame
        sum_gt : length of gt_bbox_per_frame
    """
    targets = copy.deepcopy(gt_bbox_per_frame)
    tp_count = 0
    for bb in pred_bbox_per_frame:
        for i, t_bb in enumerate(targets):
            iou_score = bb_intersection_over_union(bb, t_bb)
            if iou_score > iou_th:
                tp_count += 1
                del targets[i] #preserve the one-to-one correspondency
                break
                
    sum_pred = len(pred_bbox_per_frame)
    sum_gt = len(gt_bbox_per_frame)
    return tp_count, sum_pred, sum_gt

####### test code
import matplotlib.pyplot as plt
def draw_rect(bbox,color="red"):
    """ draw rectangle
    Args:
        box : [hori_lo,virt_lo,hori_hi,virt_hi]
    """
    [hori_lo,virt_lo,hori_hi,virt_hi] = bbox
    plt.plot(
        [hori_lo,hori_hi,hori_hi,hori_lo,hori_lo],
        [virt_lo,virt_lo,virt_hi,virt_hi,virt_lo],
        color = color,
    )
    
def test():
    gt_bboxes = [
    [0.09, -0.08, 0.99, 1.02],
    [1.19, 1.09, 2.19, 2.15],
    [1.45, 2.45, 2.95, 2.9]]
    pred_bboxes = [
        [0,0,1,1],
        [1.2,1.2,2.2,2.2],
        [1.5,2.5,3,3],
        [0.1,0.1,1.1,1.1]
    ]
    for bb in pred_bboxes:
        draw_rect(bb,color="blue")
    for bb in gt_bboxes:
        draw_rect(bb,color="red")
    plt.title("blue=pred,red=gt")
    
    tp,l_pred,l_gt = calc_precision_recall(
        pred_bboxes,
        gt_bboxes,
        iou_th = 0.9
    )
    print(f"TP={tp}, sum_pred={l_pred}, sum_gt={l_gt}")
    print(f"precisoin={tp/l_pred}, recall={tp/l_gt}")
test()

実行結果

TP=2, sum_pred=4, sum_gt=3
precisoin=0.5, recall=0.6666666666666666