class BBox(object):
def __init__(self, img_name, bbox, conf):
self.name = img_name
self.bbox = bbox
self.conf = conf
@staticmethod
def IoU(bbox_a, bbox_b):
bbox_a = bbox_a.bbox
bbox_b = bbox_b.bbox
area1 = (bbox_a[0] - bbox_a[2]) * (bbox_a[1] - bbox_b[3])
area2 = (bbox_b[0] - bbox_b[2]) * (bbox_b[1] - bbox_b[3])
tx = min(bbox_a[0], bbox_b[0])
ty = min(bbox_a[1], bbox_b[1])
bx = max(bbox_a[2], bbox_b[2])
by = max(bbox_a[3], bbox_b[3])
inter_area = (bx - tx) * (by - ty)
return inter_area / (area1 + area2 - inter_area)
def compute_AP(prds, gts, IoU_threshold=0.5):
"""Given the prediction and ground truth to calculate the AP value
Args:
prds (List[]): Model prediction results
gts (_type_): True value
IoU_threshold (float, optional): _description_. Defaults to 0.5.
"""
# sort the prds using conf
prds = sorted(prds, key=lambda x: x.conf, reverse=True)
# collect the gts
name2gt_idx = dict()
for i, gt in enumerate(gts):
name2gt_idx[gt.name] = name2gt_idx.get(gt.name, []) + [i]
# compute the TP and FP
TP = np.zeros(len(prds))
FP = np.zeros(len(prds))
# initial used_gt
used_gt = np.zeros(len(gts))
for pred_idx, prd in enumerate(prds):
gt_idxs = name2gt_idx[prd.name]
# find the gt with max iou
max_iou_idx = -1
max_iou = 0
for gt_idx in gt_idxs:
gt = gts[gt_idx]
# If this GT bbox has been used, skip it
if used_gt[gt_idx]:
continue
iou = BBox.IoU(gt, prd)
if iou > IoU_threshold and iou > max_iou:
max_iou = iou
max_iou_idx = gt_idx
if max_iou_idx != -1:
TP[pred_idx] = 1
# Tag this GT bbox already in use, and it cannot match this GT bbox next time
used_gt[max_iou_idx] = 1
else:
FP[pred_idx] = 1
# compute the AP
acc_TP = np.cumsum(TP)
acc_FP = np.cumsum(FP)
n_total = len(gts)
Pr = acc_TP / (acc_TP + acc_FP)
Rc = acc_TP / n_total
return _compute_AP(Pr, Rc)
def _compute_AP(Pr, Rc):
Rc = np.concatenate(([0], Rc, [1]))
Pr = np.concatenate(([0], Pr, [0]))
# compute the real Pr
# Pr curve is not strictly decreasing, so the actual calculation only takes the maximum value on the right
for i in range(len(Pr) - 1, 0, -1):
Pr[i-1] = max(Pr[i-1], Pr[i])
# compute the AP
index = np.where(Rc[1:] != Rc[:-1])[0]
AP = np.sum(Pr[index + 1] * (Rc[index+1] - Rc[index]))
return AP
