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
|
""" @author: zj @file: iou-compute.py @time: 2020-01-19 """
import numpy as np import torch import cv2
def numpy_iou(target_boxes, pred_boxes): """ target_boxes和pred_boxes大小相同:[N, 4],其中N表示边框数目,4表示保存的是[xmin, ymin, xmax, ymax] """ xA = np.maximum(target_boxes[:, 0], pred_boxes[:, 0]) yA = np.maximum(target_boxes[:, 1], pred_boxes[:, 1]) xB = np.minimum(target_boxes[:, 2], pred_boxes[:, 2]) yB = np.minimum(target_boxes[:, 3], pred_boxes[:, 3]) intersection = np.maximum(0.0, xB - xA) * np.maximum(0.0, yB - yA) boxAArea = (target_boxes[:, 2] - target_boxes[:, 0]) * (target_boxes[:, 3] - target_boxes[:, 1]) boxBArea = (pred_boxes[:, 2] - pred_boxes[:, 0]) * (pred_boxes[:, 3] - pred_boxes[:, 1])
iou = intersection / (boxAArea + boxBArea - intersection) return iou
def pytorch_iou(target_boxes, pred_boxes): x_min = torch.max(target_boxes[:, 0], pred_boxes[:, 0]) y_min = torch.max(target_boxes[:, 1], pred_boxes[:, 1]) x_max = torch.min(target_boxes[:, 2], pred_boxes[:, 2]) y_max = torch.min(target_boxes[:, 3], pred_boxes[:, 3]) intersection = torch.max(torch.zeros(x_max.shape), x_max - x_min) \ * torch.max(torch.zeros(y_max.shape), y_max - y_min)
boxAArea = (target_boxes[:, 2] - target_boxes[:, 0]) * (target_boxes[:, 3] - target_boxes[:, 1]) boxBArea = (pred_boxes[:, 2] - pred_boxes[:, 0]) * (pred_boxes[:, 3] - pred_boxes[:, 1])
iou = intersection / (boxAArea + boxBArea - intersection) return iou
if __name__ == '__main__': target_boxes = np.array([[10, 10, 50, 50], [40, 270, 100, 380], [450, 300, 500, 500]]) pred_boxes = np.array([[20, 20, 40, 40], [30, 280, 200, 300], [400, 200, 450, 250]])
iou1 = numpy_iou(target_boxes, pred_boxes) iou2 = pytorch_iou(torch.Tensor(target_boxes), torch.Tensor(pred_boxes))
print('numpy:', iou1) print('pytorch:', iou2)
img = np.ones((650, 650, 3)) * 255 for item in target_boxes: xmin, ymin, xmax, ymax = item cv2.rectangle(img, (xmin, ymin), (xmax, ymax), (0, 0, 255), thickness=2) for item in pred_boxes: xmin, ymin, xmax, ymax = item cv2.rectangle(img, (xmin, ymin), (xmax, ymax), (0, 255, 0), thickness=2)
cv2.imshow('img', img) cv2.waitKey(0)
|