Go to the documentation of this file.00001 import numpy as np
00002
00003 from cpu_nms import cpu_nms
00004 from gpu_nms import gpu_nms
00005
00006 def py_nms_wrapper(thresh):
00007 def _nms(dets):
00008 return nms(dets, thresh)
00009 return _nms
00010
00011
00012 def cpu_nms_wrapper(thresh):
00013 def _nms(dets):
00014 return cpu_nms(dets, thresh)
00015 return _nms
00016
00017
00018 def gpu_nms_wrapper(thresh, device_id):
00019 def _nms(dets):
00020 return gpu_nms(dets, thresh, device_id)
00021 return _nms
00022
00023
00024 def nms(dets, thresh):
00025 """
00026 greedily select boxes with high confidence and overlap with current maximum <= thresh
00027 rule out overlap >= thresh
00028 :param dets: [[x1, y1, x2, y2 score]]
00029 :param thresh: retain overlap < thresh
00030 :return: indexes to keep
00031 """
00032 if dets.shape[0] == 0:
00033 return []
00034
00035 x1 = dets[:, 0]
00036 y1 = dets[:, 1]
00037 x2 = dets[:, 2]
00038 y2 = dets[:, 3]
00039 scores = dets[:, 4]
00040
00041 areas = (x2 - x1 + 1) * (y2 - y1 + 1)
00042 order = scores.argsort()[::-1]
00043
00044 keep = []
00045 while order.size > 0:
00046 i = order[0]
00047 keep.append(i)
00048 xx1 = np.maximum(x1[i], x1[order[1:]])
00049 yy1 = np.maximum(y1[i], y1[order[1:]])
00050 xx2 = np.minimum(x2[i], x2[order[1:]])
00051 yy2 = np.minimum(y2[i], y2[order[1:]])
00052
00053 w = np.maximum(0.0, xx2 - xx1 + 1)
00054 h = np.maximum(0.0, yy2 - yy1 + 1)
00055 inter = w * h
00056 ovr = inter / (areas[i] + areas[order[1:]] - inter)
00057
00058 inds = np.where(ovr <= thresh)[0]
00059 order = order[inds + 1]
00060
00061 return keep