Det为预测框 Gt为真实框
import numpy as np
from functools import reduce
def iou_cal(boxa, boxb):
x1, y1, x2, y2 = boxa
x3, y3, x4, y4 = boxb
left_max = max(x1, x3)
right_min = min(x2, x4)
top_max = max(y1, y3)
bottom_min = min(y2, y4)
inter_area = max(0, right_min-left_max) * max(0, bottom_min-top_max)
a_area = (x2 - x1) * (y2 - y1)
b_area = (x4 - x3) * (y4 - y3)
iou = inter_area / (a_area + b_area - inter_area)
return iou
cls_num = 3
pre_boxes = []
gt_boxes = []
cls_pre_boxes = {i: [] for i in range(cls_num)}
cls_gt_boxes = {j: [] for j in range(cls_num)}
for i in range(50):
x1 = np.random.randint(10, 50)
x2 = x1 + np.random.randint(20, 50)
y1 = np.random.randint(20, 60)
y2 = y1 + np.random.randint(30, 80)
cls = np.random.randint(0, 3)
conf = np.random.rand()
pre_boxes.append([x1, y1, x2, y2, cls, conf])
gt_boxes.append([x1, y1, x2, y2, cls, 0])
for pre_box in pre_boxes:
cls_pre_boxes[pre_box[4]].append(pre_box)
for gt_box in gt_boxes:
cls_gt_boxes[gt_box[4]].append(gt_box)
res = []
for pre_cls, pre_boxes in cls_pre_boxes.items():
gt_cls = pre_cls
gt_boxes = cls_gt_boxes[gt_cls]
pre_boxes = sorted(pre_boxes, key=lambda x: x[5], reverse=True)
TP = np.zeros(len(pre_boxes))
FP = np.zeros(len(pre_boxes))
for ind, pre_box in enumerate(pre_boxes):
iou_max = 0
for indx, gt_box in enumerate(gt_boxes):
iou = IOU_calculate.iou_cal(pre_box[:4], gt_box[:4])
if iou > iou_max:
iou_max = iou
jmax = indx
if iou_max > 0.5 and gt_boxes[jmax][5] == 0:
TP[ind] = 1
gt_boxes[jmax][5] = 1
else:
FP[ind] = 1
acc_TP = np.cumsum(TP)
acc_FP = np.cumsum(FP)
rec = acc_TP / len(gt_boxes)
acc = np.divide(acc_TP, acc_TP + acc_FP)
ap = np.trapz(acc, rec)
res.append((pre_cls, ap))
print(res)
mAP = reduce(lambda x, y: x + y, [re[1] for re in res]) /cls_num
print("mAP = ", mAp)
|