postprocess.h 2.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657
  1. #ifndef __CENTERPOINT_POSTPROCESS__
  2. #define __CENTERPOINT_POSTPROCESS__
  3. #include "buffers.h"
  4. #include "common.h"
  5. #include "config.h"
  6. #include <math.h>
  7. #include <stdint.h>
  8. #include <thrust/sort.h>
  9. #include <thrust/sequence.h>
  10. #include <thrust/execution_policy.h>
  11. #include <thrust/host_vector.h>
  12. struct Box{
  13. float x;
  14. float y;
  15. float z;
  16. float l;
  17. float h;
  18. float w;
  19. float velX;
  20. float velY;
  21. float theta;
  22. float score;
  23. int cls;
  24. bool isDrop; // for nms
  25. };
  26. int _raw_nms_gpu(const float* reg, const float* height, const float* dim , const float* rot,
  27. const int* indexs, long* dev_keep_data, unsigned long long* mask_cpu, unsigned long long* remv_gpu,
  28. int boxes_num, float nms_overlap_thresh);
  29. void _sort_by_key(float* keys, int* values,int size) ;
  30. void _gather_all(float* host_boxes, int* host_label,
  31. float* reg, float* height, float* dim, float* rot, float* sorted_score, int32_t* label,
  32. int* dev_indexs, long* host_keep_indexs, int boxSizeBef, int boxSizeAft) ;
  33. void _box_assign_launcher(float* reg, float* height , float* dim, float*rot, float* boxes, float*score, int* label, float* out_score, int*out_label,
  34. int* validIndexs ,int boxSize, int output_h, int output_w) ;
  35. void _index_assign_launcher(int* indexs, int output_h, int output_w) ;
  36. int _find_valid_score_num(float* score, float thre, int output_h, int output_w) ;
  37. // void _find_valid_score_num(float* score, float thre, int output_h, int output_w, int* box_size); //, thrust::host_vector<int> host_box_size);
  38. void postprocessGPU(samplesCommon::BufferManager * buffers,
  39. std::vector<Box>& predResult ,
  40. std::map<std::string, std::vector<std::string>>rpnOutputTensorNames,
  41. int* dev_score_indexs,
  42. unsigned long long* mask_cpu,
  43. unsigned long long* remv_cpu,
  44. int* host_score_indexs,
  45. long* host_keep_data,
  46. float* host_boxes,
  47. int* host_label);
  48. void postprocess(samplesCommon::BufferManager * buffers, std::vector<Box>& predResult);
  49. #endif