mirror of
https://github.com/deepinsight/insightface.git
synced 2026-05-21 17:17:49 +00:00
refine repo structure
This commit is contained in:
67
detection/RetinaFace/rcnn/processing/nms.py
Normal file
67
detection/RetinaFace/rcnn/processing/nms.py
Normal file
@@ -0,0 +1,67 @@
|
||||
import numpy as np
|
||||
from ..cython.cpu_nms import cpu_nms
|
||||
try:
|
||||
from ..cython.gpu_nms import gpu_nms
|
||||
except ImportError:
|
||||
gpu_nms = None
|
||||
|
||||
|
||||
def py_nms_wrapper(thresh):
|
||||
def _nms(dets):
|
||||
return nms(dets, thresh)
|
||||
|
||||
return _nms
|
||||
|
||||
|
||||
def cpu_nms_wrapper(thresh):
|
||||
def _nms(dets):
|
||||
return cpu_nms(dets, thresh)
|
||||
|
||||
return _nms
|
||||
|
||||
|
||||
def gpu_nms_wrapper(thresh, device_id):
|
||||
def _nms(dets):
|
||||
return gpu_nms(dets, thresh, device_id)
|
||||
|
||||
if gpu_nms is not None:
|
||||
return _nms
|
||||
else:
|
||||
return cpu_nms_wrapper(thresh)
|
||||
|
||||
|
||||
def nms(dets, thresh):
|
||||
"""
|
||||
greedily select boxes with high confidence and overlap with current maximum <= thresh
|
||||
rule out overlap >= thresh
|
||||
:param dets: [[x1, y1, x2, y2 score]]
|
||||
:param thresh: retain overlap < thresh
|
||||
:return: indexes to keep
|
||||
"""
|
||||
x1 = dets[:, 0]
|
||||
y1 = dets[:, 1]
|
||||
x2 = dets[:, 2]
|
||||
y2 = dets[:, 3]
|
||||
scores = dets[:, 4]
|
||||
|
||||
areas = (x2 - x1 + 1) * (y2 - y1 + 1)
|
||||
order = scores.argsort()[::-1]
|
||||
|
||||
keep = []
|
||||
while order.size > 0:
|
||||
i = order[0]
|
||||
keep.append(i)
|
||||
xx1 = np.maximum(x1[i], x1[order[1:]])
|
||||
yy1 = np.maximum(y1[i], y1[order[1:]])
|
||||
xx2 = np.minimum(x2[i], x2[order[1:]])
|
||||
yy2 = np.minimum(y2[i], y2[order[1:]])
|
||||
|
||||
w = np.maximum(0.0, xx2 - xx1 + 1)
|
||||
h = np.maximum(0.0, yy2 - yy1 + 1)
|
||||
inter = w * h
|
||||
ovr = inter / (areas[i] + areas[order[1:]] - inter)
|
||||
|
||||
inds = np.where(ovr <= thresh)[0]
|
||||
order = order[inds + 1]
|
||||
|
||||
return keep
|
||||
Reference in New Issue
Block a user