Skip to content
Snippets Groups Projects
Unverified Commit d1d35451 authored by YouAnsheng's avatar YouAnsheng Committed by GitHub
Browse files

Update gpu_nms.pyx 

Fix bug for gpu_nms.
parent 64e310d5
No related branches found
No related tags found
No related merge requests found
Hide whitespace changes
Inline Side-by-side
mmdet/ops/nms/gpu_nms.pyx
+ 2
− 2
View file @ d1d35451
...@@ -19,7 +19,7 @@ memory_pool = {} ...@@ -19,7 +19,7 @@ memory_pool = {}
def gpu_nms(np.ndarray[np.float32_t, ndim=2] dets, np.float thresh, def gpu_nms(np.ndarray[np.float32_t, ndim=2] dets, np.float thresh,
np.int32_t device_id=0): np.int32_t device_id=0):
cdef int boxes_num = dets.shape[0] cdef int boxes_num = dets.shape[0]
cdef int boxes_dim = dets.shape[1] cdef int boxes_dim = 5
cdef int num_out cdef int num_out
cdef size_t base cdef size_t base
cdef np.ndarray[np.int32_t, ndim=1] \ cdef np.ndarray[np.int32_t, ndim=1] \
...@@ -29,7 +29,7 @@ def gpu_nms(np.ndarray[np.float32_t, ndim=2] dets, np.float thresh, ...@@ -29,7 +29,7 @@ def gpu_nms(np.ndarray[np.float32_t, ndim=2] dets, np.float thresh,
cdef np.ndarray[np.int_t, ndim=1] \ cdef np.ndarray[np.int_t, ndim=1] \
order = scores.argsort()[::-1] order = scores.argsort()[::-1]
cdef np.ndarray[np.float32_t, ndim=2] \ cdef np.ndarray[np.float32_t, ndim=2] \
sorted_dets = dets[order, :] sorted_dets = dets[order, :5]
cdef float cthresh = thresh cdef float cthresh = thresh
if device_id not in memory_pool: if device_id not in memory_pool:
with nogil: with nogil:
......
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment