# coding=utf-8
import numpy as np
def xywh2xyxy(x):
y = np.copy(x)
y[..., 0] = x[..., 0] - x[..., 2] / 2 # top left x
y[..., 1] = x[..., 1] - x[..., 3] / 2 # top left y
y[..., 2] = x[..., 0] + x[..., 2] / 2 # bottom right x
y[..., 3] = x[..., 1] + x[..., 3] / 2 # bottom right y
return y
def cal_iou(det1, det2):
det1_x1, det1_y1 = det1[..., 0], det1[..., 1]
det1_x2, det1_y2 = det1[..., 2], det1[..., 3]
det2_x1, det2_y1 = det2[..., 0], det2[..., 1]
det2_x2, det2_y2 = det2[..., 2], det2[..., 3]
x1 = np.maximum(det1_x1, det2_x1)
y1 = np.maximum(det1_y1, det2_y1)
x2 = np.minimum(det1_x2, det2_x2)
y2 = np.minimum(det1_y2, det2_y2)
area_det1 = (det1_y2 - det1_y1 + 1) * (det1_x2 - det1_x1 + 1)
area_det2 = (det2_y2 - det2_y1 + 1) * (det2_x2 - det2_x1 + 1)
inter = np.maximum(0, (y2 - y1 + 1)) * np.maximum(0, (x2 - x1 + 1))
ious = inter / (area_det1 + area_det2 - inter)
return ious
def nms(detections, conf_thres=0.4, nms_thres=0.5):
outputs = []
detections = xywh2xyxy(detections) # 检测框格式转换 xywh -> xyxy
detections[..., 5:] *= detections[..., 4:5] # 类别的分数是用obj_conf * cls_conf
num_classes = detections.shape[2] - 5 # detections shape: (batch_size, num_bbox, 4 + 1 + num_classes)
candidates = detections[..., 4] > conf_thres # 用于后续的过滤筛选
for img_idx, dets in enumerate(detections):
output = []
dets = dets[candidates[img_idx]] # 根据conf_thres过滤
indexes = dets[..., 4].argsort()[::-1] # 根据obj_conf大小排序
dets = dets[indexes]
classes_dets = dets[..., 5:].argmax(axis=1) # 每个检测框的类别
for cls in range(num_classes): # 按类别循环遍历
dets_cls = dets[classes_dets == cls] # shape: (num_bbox, 4 + 1 + num_classes)
while len(dets_cls):
det_select = dets_cls[0]
# det_select shape: (x1, y1, x2, y2, score, label)
det_select = np.concatenate((det_select[..., :4], [det_select[..., 5:].max(), cls]))
output.append(det_select)
dets_cls = dets_cls[1:]
if len(dets_cls):
ious = cal_iou(det_select, dets_cls)
indexes = np.where(ious <= nms_thres)[0]
dets_cls = dets_cls[indexes]
outputs.append(output)
return np.array(outputs)
def scale_coords(output, ori_shape, ratio=1.0, top_pad=0, left_pad=0):
ori_h, ori_w = ori_shape
output[:, [0, 2]] -= left_pad
output[:, [1, 3]] -= top_pad
output[:, :4] /= ratio
output[:, [0, 2]] = output[:, [0, 2]].clip(0, ori_w) # x1, x2
output[:, [1, 3]] = output[:, [1, 3]].clip(0, ori_h) # y1, y2
完成。