1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114
| import torch from collections import Counter def mean_average_precision(pred_bboxes,true_boxes,iou_threshold,num_classes=20): average_precisions=[] epsilon=1e-6 for c in range(num_classes): detections=[] ground_truths=[] for detection in pred_bboxes: if detection[1]==c: detections.append(detection) for true_box in true_boxes: if true_box[1]==c: ground_truths.append(true_box) amount_bboxes=Counter(gt[0] for gt in ground_truths) for key,val in amount_bboxes.items(): amount_bboxes[key]=torch.zeros(val) detections.sort(key=lambda x:x[2],reverse=True) TP=torch.zeros(len(detections)) FP=torch.zeros(len(detections)) total_true_bboxes=len(ground_truths) if total_true_bboxes == 0: continue for detection_idx,detection in enumerate(detections): ground_truth_img=[bbox for bbox in ground_truths if bbox[0]==detections[0]] num_gts=len(ground_truth_img) best_iou=0 for idx,gt in emnumerate(ground_truth_img): iou=insert_over_union(torch.tensor(detection[3:]),torch.tensor(gt[3:])) if iou >best_iou: best_iou=iou best_gt_idx=idx if best_iou>iou_threshold: if amount_bboxes[detection[0]][best_gt_idx]==0: TP[detection_idx]=1 amount_bboxes[detection[0]][best_gt_idx]=1 else: FP[detection_idx]=1 else: FP[detection_idx]=1 TP_cumsum=torch.cumsum(TP,dim=0) FP_cumsum=torch.cumsum(FP,dim=0) recalls=TP_cumsum/(total_true_bboxes+epsilon) precisions=torch.divide(TP_cumsum,(TP_cumsu+FP_cumsum+epsilon)) precisions=torch.cat((torch.tensor([1]),precision)) recalls=torch.cat((torch.tensor([0]),recalls)) average_precisions.append(torch.trapz(precisions,recalls)) return sum(average_precisions)/len(average_precisions) def insert_over_union(boxes_preds,boxes_labels): box1_x1=boxes_preds[...,0:1] box1_y1=boxes_preds[...,1:2] box1_x2=boxes_preds[...,2:3] box1_y2=boxes_preds[...,3:4] box2_x1=boxes_labels[...,0:1] box2_y1=boxes_labels[...,1:2] box2_x2=boxes_labels[...,2:3] box2_y2=boxes_labels[...,3:4] x1=torch.max(box1_x1,box2_x1) y1=torch.max(box1_y1,box2_y1) x2=torch.min(box1_x2,box2_x2) y2=torch.min(box1_y2,box2_y2) intersection=(x2-x1).clamp(0)*(y2-y1).clamp(0) box1_area=abs((box1_x2-box1_x1)*(box1_y1-box1_y2)) box2_area=abs((box2_x2-box2_x1)*(box2_y1-box2_y2)) return intersection/(box1_area+box2_area-intersection+1e-6)
|