You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335
  1. # YOLOv5 🚀 by Ultralytics, GPL-3.0 license
  2. """
  3. Model validation metrics
  4. """
  5. import math
  6. import warnings
  7. from pathlib import Path
  8. import matplotlib.pyplot as plt
  9. import numpy as np
  10. import torch
  11. def fitness(x):
  12. # Model fitness as a weighted combination of metrics
  13. w = [0.0, 0.0, 0.1, 0.9] # weights for [P, R, mAP@0.5, mAP@0.5:0.95]
  14. return (x[:, :4] * w).sum(1)
  15. def ap_per_class(tp, conf, pred_cls, target_cls, plot=False, save_dir='.', names=()):
  16. """ Compute the average precision, given the recall and precision curves.
  17. Source: https://github.com/rafaelpadilla/Object-Detection-Metrics.
  18. # Arguments
  19. tp: True positives (nparray, nx1 or nx10).
  20. conf: Objectness value from 0-1 (nparray).
  21. pred_cls: Predicted object classes (nparray).
  22. target_cls: True object classes (nparray).
  23. plot: Plot precision-recall curve at mAP@0.5
  24. save_dir: Plot save directory
  25. # Returns
  26. The average precision as computed in py-faster-rcnn.
  27. """
  28. # Sort by objectness
  29. i = np.argsort(-conf)
  30. tp, conf, pred_cls = tp[i], conf[i], pred_cls[i]
  31. # Find unique classes
  32. unique_classes = np.unique(target_cls)
  33. nc = unique_classes.shape[0] # number of classes, number of detections
  34. # Create Precision-Recall curve and compute AP for each class
  35. px, py = np.linspace(0, 1, 1000), [] # for plotting
  36. ap, p, r = np.zeros((nc, tp.shape[1])), np.zeros((nc, 1000)), np.zeros((nc, 1000))
  37. for ci, c in enumerate(unique_classes):
  38. i = pred_cls == c
  39. n_l = (target_cls == c).sum() # number of labels
  40. n_p = i.sum() # number of predictions
  41. if n_p == 0 or n_l == 0:
  42. continue
  43. else:
  44. # Accumulate FPs and TPs
  45. fpc = (1 - tp[i]).cumsum(0)
  46. tpc = tp[i].cumsum(0)
  47. # Recall
  48. recall = tpc / (n_l + 1e-16) # recall curve
  49. r[ci] = np.interp(-px, -conf[i], recall[:, 0], left=0) # negative x, xp because xp decreases
  50. # Precision
  51. precision = tpc / (tpc + fpc) # precision curve
  52. p[ci] = np.interp(-px, -conf[i], precision[:, 0], left=1) # p at pr_score
  53. # AP from recall-precision curve
  54. for j in range(tp.shape[1]):
  55. ap[ci, j], mpre, mrec = compute_ap(recall[:, j], precision[:, j])
  56. if plot and j == 0:
  57. py.append(np.interp(px, mrec, mpre)) # precision at mAP@0.5
  58. # Compute F1 (harmonic mean of precision and recall)
  59. f1 = 2 * p * r / (p + r + 1e-16)
  60. names = [v for k, v in names.items() if k in unique_classes] # list: only classes that have data
  61. names = {i: v for i, v in enumerate(names)} # to dict
  62. if plot:
  63. plot_pr_curve(px, py, ap, Path(save_dir) / 'PR_curve.png', names)
  64. plot_mc_curve(px, f1, Path(save_dir) / 'F1_curve.png', names, ylabel='F1')
  65. plot_mc_curve(px, p, Path(save_dir) / 'P_curve.png', names, ylabel='Precision')
  66. plot_mc_curve(px, r, Path(save_dir) / 'R_curve.png', names, ylabel='Recall')
  67. i = f1.mean(0).argmax() # max F1 index
  68. return p[:, i], r[:, i], ap, f1[:, i], unique_classes.astype('int32')
  69. def compute_ap(recall, precision):
  70. """ Compute the average precision, given the recall and precision curves
  71. # Arguments
  72. recall: The recall curve (list)
  73. precision: The precision curve (list)
  74. # Returns
  75. Average precision, precision curve, recall curve
  76. """
  77. # Append sentinel values to beginning and end
  78. mrec = np.concatenate(([0.0], recall, [1.0]))
  79. mpre = np.concatenate(([1.0], precision, [0.0]))
  80. # Compute the precision envelope
  81. mpre = np.flip(np.maximum.accumulate(np.flip(mpre)))
  82. # Integrate area under curve
  83. method = 'interp' # methods: 'continuous', 'interp'
  84. if method == 'interp':
  85. x = np.linspace(0, 1, 101) # 101-point interp (COCO)
  86. ap = np.trapz(np.interp(x, mrec, mpre), x) # integrate
  87. else: # 'continuous'
  88. i = np.where(mrec[1:] != mrec[:-1])[0] # points where x axis (recall) changes
  89. ap = np.sum((mrec[i + 1] - mrec[i]) * mpre[i + 1]) # area under curve
  90. return ap, mpre, mrec
  91. class ConfusionMatrix:
  92. # Updated version of https://github.com/kaanakan/object_detection_confusion_matrix
  93. def __init__(self, nc, conf=0.25, iou_thres=0.45):
  94. self.matrix = np.zeros((nc + 1, nc + 1))
  95. self.nc = nc # number of classes
  96. self.conf = conf
  97. self.iou_thres = iou_thres
  98. def process_batch(self, detections, labels):
  99. """
  100. Return intersection-over-union (Jaccard index) of boxes.
  101. Both sets of boxes are expected to be in (x1, y1, x2, y2) format.
  102. Arguments:
  103. detections (Array[N, 6]), x1, y1, x2, y2, conf, class
  104. labels (Array[M, 5]), class, x1, y1, x2, y2
  105. Returns:
  106. None, updates confusion matrix accordingly
  107. """
  108. detections = detections[detections[:, 4] > self.conf]
  109. gt_classes = labels[:, 0].int()
  110. detection_classes = detections[:, 5].int()
  111. iou = box_iou(labels[:, 1:], detections[:, :4])
  112. x = torch.where(iou > self.iou_thres)
  113. if x[0].shape[0]:
  114. matches = torch.cat((torch.stack(x, 1), iou[x[0], x[1]][:, None]), 1).cpu().numpy()
  115. if x[0].shape[0] > 1:
  116. matches = matches[matches[:, 2].argsort()[::-1]]
  117. matches = matches[np.unique(matches[:, 1], return_index=True)[1]]
  118. matches = matches[matches[:, 2].argsort()[::-1]]
  119. matches = matches[np.unique(matches[:, 0], return_index=True)[1]]
  120. else:
  121. matches = np.zeros((0, 3))
  122. n = matches.shape[0] > 0
  123. m0, m1, _ = matches.transpose().astype(np.int16)
  124. for i, gc in enumerate(gt_classes):
  125. j = m0 == i
  126. if n and sum(j) == 1:
  127. self.matrix[detection_classes[m1[j]], gc] += 1 # correct
  128. else:
  129. self.matrix[self.nc, gc] += 1 # background FP
  130. if n:
  131. for i, dc in enumerate(detection_classes):
  132. if not any(m1 == i):
  133. self.matrix[dc, self.nc] += 1 # background FN
  134. def matrix(self):
  135. return self.matrix
  136. def plot(self, normalize=True, save_dir='', names=()):
  137. try:
  138. import seaborn as sn
  139. array = self.matrix / ((self.matrix.sum(0).reshape(1, -1) + 1E-6) if normalize else 1) # normalize columns
  140. array[array < 0.005] = np.nan # don't annotate (would appear as 0.00)
  141. fig = plt.figure(figsize=(12, 9), tight_layout=True)
  142. sn.set(font_scale=1.0 if self.nc < 50 else 0.8) # for label size
  143. labels = (0 < len(names) < 99) and len(names) == self.nc # apply names to ticklabels
  144. with warnings.catch_warnings():
  145. warnings.simplefilter('ignore') # suppress empty matrix RuntimeWarning: All-NaN slice encountered
  146. sn.heatmap(array, annot=self.nc < 30, annot_kws={"size": 8}, cmap='Blues', fmt='.2f', square=True,
  147. xticklabels=names + ['background FP'] if labels else "auto",
  148. yticklabels=names + ['background FN'] if labels else "auto").set_facecolor((1, 1, 1))
  149. fig.axes[0].set_xlabel('True')
  150. fig.axes[0].set_ylabel('Predicted')
  151. fig.savefig(Path(save_dir) / 'confusion_matrix.png', dpi=250)
  152. plt.close()
  153. except Exception as e:
  154. print(f'WARNING: ConfusionMatrix plot failure: {e}')
  155. def print(self):
  156. for i in range(self.nc + 1):
  157. print(' '.join(map(str, self.matrix[i])))
  158. def bbox_iou(box1, box2, x1y1x2y2=True, GIoU=False, DIoU=False, CIoU=False, eps=1e-7):
  159. # Returns the IoU of box1 to box2. box1 is 4, box2 is nx4
  160. box2 = box2.T
  161. # Get the coordinates of bounding boxes
  162. if x1y1x2y2: # x1, y1, x2, y2 = box1
  163. b1_x1, b1_y1, b1_x2, b1_y2 = box1[0], box1[1], box1[2], box1[3]
  164. b2_x1, b2_y1, b2_x2, b2_y2 = box2[0], box2[1], box2[2], box2[3]
  165. else: # transform from xywh to xyxy
  166. b1_x1, b1_x2 = box1[0] - box1[2] / 2, box1[0] + box1[2] / 2
  167. b1_y1, b1_y2 = box1[1] - box1[3] / 2, box1[1] + box1[3] / 2
  168. b2_x1, b2_x2 = box2[0] - box2[2] / 2, box2[0] + box2[2] / 2
  169. b2_y1, b2_y2 = box2[1] - box2[3] / 2, box2[1] + box2[3] / 2
  170. # Intersection area
  171. inter = (torch.min(b1_x2, b2_x2) - torch.max(b1_x1, b2_x1)).clamp(0) * \
  172. (torch.min(b1_y2, b2_y2) - torch.max(b1_y1, b2_y1)).clamp(0)
  173. # Union Area
  174. w1, h1 = b1_x2 - b1_x1, b1_y2 - b1_y1 + eps
  175. w2, h2 = b2_x2 - b2_x1, b2_y2 - b2_y1 + eps
  176. union = w1 * h1 + w2 * h2 - inter + eps
  177. iou = inter / union
  178. if GIoU or DIoU or CIoU:
  179. cw = torch.max(b1_x2, b2_x2) - torch.min(b1_x1, b2_x1) # convex (smallest enclosing box) width
  180. ch = torch.max(b1_y2, b2_y2) - torch.min(b1_y1, b2_y1) # convex height
  181. if CIoU or DIoU: # Distance or Complete IoU https://arxiv.org/abs/1911.08287v1
  182. c2 = cw ** 2 + ch ** 2 + eps # convex diagonal squared
  183. rho2 = ((b2_x1 + b2_x2 - b1_x1 - b1_x2) ** 2 +
  184. (b2_y1 + b2_y2 - b1_y1 - b1_y2) ** 2) / 4 # center distance squared
  185. if DIoU:
  186. return iou - rho2 / c2 # DIoU
  187. elif CIoU: # https://github.com/Zzh-tju/DIoU-SSD-pytorch/blob/master/utils/box/box_utils.py#L47
  188. v = (4 / math.pi ** 2) * torch.pow(torch.atan(w2 / h2) - torch.atan(w1 / h1), 2)
  189. with torch.no_grad():
  190. alpha = v / (v - iou + (1 + eps))
  191. return iou - (rho2 / c2 + v * alpha) # CIoU
  192. else: # GIoU https://arxiv.org/pdf/1902.09630.pdf
  193. c_area = cw * ch + eps # convex area
  194. return iou - (c_area - union) / c_area # GIoU
  195. else:
  196. return iou # IoU
  197. def box_iou(box1, box2):
  198. # https://github.com/pytorch/vision/blob/master/torchvision/ops/boxes.py
  199. """
  200. Return intersection-over-union (Jaccard index) of boxes.
  201. Both sets of boxes are expected to be in (x1, y1, x2, y2) format.
  202. Arguments:
  203. box1 (Tensor[N, 4])
  204. box2 (Tensor[M, 4])
  205. Returns:
  206. iou (Tensor[N, M]): the NxM matrix containing the pairwise
  207. IoU values for every element in boxes1 and boxes2
  208. """
  209. def box_area(box):
  210. # box = 4xn
  211. return (box[2] - box[0]) * (box[3] - box[1])
  212. area1 = box_area(box1.T)
  213. area2 = box_area(box2.T)
  214. # inter(N,M) = (rb(N,M,2) - lt(N,M,2)).clamp(0).prod(2)
  215. inter = (torch.min(box1[:, None, 2:], box2[:, 2:]) - torch.max(box1[:, None, :2], box2[:, :2])).clamp(0).prod(2)
  216. return inter / (area1[:, None] + area2 - inter) # iou = inter / (area1 + area2 - inter)
  217. def bbox_ioa(box1, box2, eps=1E-7):
  218. """ Returns the intersection over box2 area given box1, box2. Boxes are x1y1x2y2
  219. box1: np.array of shape(4)
  220. box2: np.array of shape(nx4)
  221. returns: np.array of shape(n)
  222. """
  223. box2 = box2.transpose()
  224. # Get the coordinates of bounding boxes
  225. b1_x1, b1_y1, b1_x2, b1_y2 = box1[0], box1[1], box1[2], box1[3]
  226. b2_x1, b2_y1, b2_x2, b2_y2 = box2[0], box2[1], box2[2], box2[3]
  227. # Intersection area
  228. inter_area = (np.minimum(b1_x2, b2_x2) - np.maximum(b1_x1, b2_x1)).clip(0) * \
  229. (np.minimum(b1_y2, b2_y2) - np.maximum(b1_y1, b2_y1)).clip(0)
  230. # box2 area
  231. box2_area = (b2_x2 - b2_x1) * (b2_y2 - b2_y1) + eps
  232. # Intersection over box2 area
  233. return inter_area / box2_area
  234. def wh_iou(wh1, wh2):
  235. # Returns the nxm IoU matrix. wh1 is nx2, wh2 is mx2
  236. wh1 = wh1[:, None] # [N,1,2]
  237. wh2 = wh2[None] # [1,M,2]
  238. inter = torch.min(wh1, wh2).prod(2) # [N,M]
  239. return inter / (wh1.prod(2) + wh2.prod(2) - inter) # iou = inter / (area1 + area2 - inter)
  240. # Plots ----------------------------------------------------------------------------------------------------------------
  241. def plot_pr_curve(px, py, ap, save_dir='pr_curve.png', names=()):
  242. # Precision-recall curve
  243. fig, ax = plt.subplots(1, 1, figsize=(9, 6), tight_layout=True)
  244. py = np.stack(py, axis=1)
  245. if 0 < len(names) < 21: # display per-class legend if < 21 classes
  246. for i, y in enumerate(py.T):
  247. ax.plot(px, y, linewidth=1, label=f'{names[i]} {ap[i, 0]:.3f}') # plot(recall, precision)
  248. else:
  249. ax.plot(px, py, linewidth=1, color='grey') # plot(recall, precision)
  250. ax.plot(px, py.mean(1), linewidth=3, color='blue', label='all classes %.3f mAP@0.5' % ap[:, 0].mean())
  251. ax.set_xlabel('Recall')
  252. ax.set_ylabel('Precision')
  253. ax.set_xlim(0, 1)
  254. ax.set_ylim(0, 1)
  255. plt.legend(bbox_to_anchor=(1.04, 1), loc="upper left")
  256. fig.savefig(Path(save_dir), dpi=250)
  257. plt.close()
  258. def plot_mc_curve(px, py, save_dir='mc_curve.png', names=(), xlabel='Confidence', ylabel='Metric'):
  259. # Metric-confidence curve
  260. fig, ax = plt.subplots(1, 1, figsize=(9, 6), tight_layout=True)
  261. if 0 < len(names) < 21: # display per-class legend if < 21 classes
  262. for i, y in enumerate(py):
  263. ax.plot(px, y, linewidth=1, label=f'{names[i]}') # plot(confidence, metric)
  264. else:
  265. ax.plot(px, py.T, linewidth=1, color='grey') # plot(confidence, metric)
  266. y = py.mean(0)
  267. ax.plot(px, y, linewidth=3, color='blue', label=f'all classes {y.max():.2f} at {px[y.argmax()]:.3f}')
  268. ax.set_xlabel(xlabel)
  269. ax.set_ylabel(ylabel)
  270. ax.set_xlim(0, 1)
  271. ax.set_ylim(0, 1)
  272. plt.legend(bbox_to_anchor=(1.04, 1), loc="upper left")
  273. fig.savefig(Path(save_dir), dpi=250)
  274. plt.close()