| import numpy as np | |
| cimport numpy as np | |
| assert sizeof(int) == sizeof(np.int32_t) | |
| cdef extern from "gpu_nms.hpp": | |
| void _nms(np.int32_t*, int*, np.float32_t*, int, int, float, int) | |
| def gpu_nms(np.ndarray[np.float32_t, ndim=2] dets, np.float thresh, | |
| np.int32_t device_id=0): | |
| cdef int boxes_num = dets.shape[0] | |
| cdef int boxes_dim = dets.shape[1] | |
| cdef int num_out | |
| cdef np.ndarray[np.int32_t, ndim=1] \ | |
| keep = np.zeros(boxes_num, dtype=np.int32) | |
| cdef np.ndarray[np.float32_t, ndim=1] \ | |
| scores = dets[:, 4] | |
| cdef np.ndarray[np.int_t, ndim=1] \ | |
| order = scores.argsort()[::-1] | |
| cdef np.ndarray[np.float32_t, ndim=2] \ | |
| sorted_dets = dets[order, :] | |
| _nms(&keep[0], &num_out, &sorted_dets[0, 0], boxes_num, boxes_dim, thresh, device_id) | |
| keep = keep[:num_out] | |
| return list(order[keep]) | |