Skip to content

Commit

Permalink
add nms_cuda.c
Browse files Browse the repository at this point in the history
  • Loading branch information
Jianwei Yang committed Jan 20, 2018
1 parent 01e0f7c commit cfcc7fc
Showing 1 changed file with 19 additions and 0 deletions.
19 changes: 19 additions & 0 deletions lib/model/nms/src/nms_cuda.c
Original file line number Diff line number Diff line change
@@ -0,0 +1,19 @@
#include <THC/THC.h>
#include <stdio.h>
#include "nms_cuda_kernel.h"

// this symbol will be resolved automatically from PyTorch libs
extern THCState *state;

int nms_cuda(THCudaIntTensor *keep_out, THCudaTensor *boxes_host,
THCudaIntTensor *num_out, float nms_overlap_thresh) {

nms_cuda_compute(THCudaIntTensor_data(state, keep_out),
THCudaIntTensor_data(state, num_out),
THCudaTensor_data(state, boxes_host),
boxes_host->size[0],
boxes_host->size[1],
nms_overlap_thresh);

return 1;
}

0 comments on commit cfcc7fc

Please sign in to comment.