leon 3 周之前
父節點
當前提交
bcd990403c
共有 2 個文件被更改,包括 3 次插入2 次删除
  1. 2 2
      src/infer/trt/classfier/classifier.cu
  2. 1 0
      src/infer/trt/classfier/classifier.hpp

+ 2 - 2
src/infer/trt/classfier/classifier.cu

@@ -170,6 +170,7 @@ virtual Result ClassifierModelImpl::forward(const tensor::Image &image, data::Bo
     
     uint8_t *image_device = preprocess_buffer_.gpu();
     uint8_t *image_host   = preprocess_buffer_.cpu();
+    size_t size_image     = image.width * image.height * 3;
     cudaStream_t stream_  = (cudaStream_t)stream;
     memcpy(image_host, image.bgrptr, size_image);
     checkRuntime(
@@ -218,10 +219,9 @@ virtual Result ClassifierModelImpl::forward(const tensor::Image &image, data::Bo
 
     for (int ib = 0; ib < num_image; ++ib) 
     {
-        float *output_array_cpu = output_array_.cpu() + ib * num_classes_;
         int *max_index = classes_indices_.cpu() + ib;
         int index = *max_index;
-        classfier_boxes_ptr[i]->label = class_names_[index];
+        classfier_boxes_ptr[ib]->label = class_names_[index];
     }
 }
 

+ 1 - 0
src/infer/trt/classfier/classifier.hpp

@@ -30,6 +30,7 @@ class ClassifierModelImpl : public Infer
 {
 
 public:
+    std::shared_ptr<TensorRT::Engine> trt_;
     std::vector<std::string> class_names_;
     std::string engine_file_;
     tensor::Memory<float> affine_matrix_;