leon 3 weeks ago
parent
commit
0e30d89c9c
2 changed files with 14 additions and 13 deletions
  1. 11 9
      src/infer/trt/classfier/classifier.cu
  2. 3 4
      src/infer/trt/classfier/classifier.hpp

+ 11 - 9
src/infer/trt/classfier/classifier.cu

@@ -91,11 +91,12 @@ bool ClassifierModelImpl::load(const std::string &engine_file, int gpu_id)
     return true;
 }
 
-void ClassifierModelImpl::preprocess(int ibatch, affine::CropResizeMatrix& matrix, int x, int y, int w, int h, void *stream)
+void ClassifierModelImpl::preprocess(int ibatch, const tensor::Image &image, affine::CropResizeMatrix& matrix, int x, int y, int w, int h, void *stream)
 {
     matrix.compute(
         std::make_tuple(w, h),
-        std::make_tuple(network_input_width_, network_input_height_));
+        std::make_tuple(network_input_width_, network_input_height_),
+        std::make_tuple(x, y));
     size_t input_numel = network_input_width_ * network_input_height_ * 3;
 
     float *input_device = input_buffer_.gpu() + ibatch * input_numel;
@@ -120,11 +121,11 @@ void ClassifierModelImpl::preprocess(int ibatch, affine::CropResizeMatrix& matri
 
 virtual Result ClassifierModelImpl::forward(const tensor::Image &image, void *stream)
 {
-    return;
+    return std::monostate{};
 }
 virtual Result ClassifierModelImpl::forward(const tensor::Image &image, int slice_width, int slice_height, float overlap_width_ratio, float overlap_height_ratio, void *stream)
 {
-    return;
+    return std::monostate{};
 }
 virtual Result ClassifierModelImpl::forward(const tensor::Image &image, data::BoxArray& boxes, void *stream)
 {
@@ -138,7 +139,7 @@ virtual Result ClassifierModelImpl::forward(const tensor::Image &image, data::Bo
         }
     }
     int num_image = classfier_boxes_ptr.size();
-    if (num_image == 0){ return; }
+    if (num_image == 0){ return std::monostate{}; }
 
     auto input_dims = trt_->static_dims(0);
     int infer_batch_size = input_dims[0];
@@ -151,7 +152,7 @@ virtual Result ClassifierModelImpl::forward(const tensor::Image &image, data::Bo
             if (!trt_->set_run_dims(0, input_dims)) 
             {
                 printf("Fail to set run dims\n");
-                return;
+                return std::monostate{};
             }
         } 
         else 
@@ -162,7 +163,7 @@ virtual Result ClassifierModelImpl::forward(const tensor::Image &image, data::Bo
                     "When using static shape model, number of images[%d] must be "
                     "less than or equal to the maximum batch[%d].",
                     num_image, infer_batch_size);
-                return;
+                return std::monostate{};
             }
         }
     }
@@ -183,7 +184,7 @@ virtual Result ClassifierModelImpl::forward(const tensor::Image &image, data::Bo
         int y = (int)box_ptr->top;
         int w = (int)box_ptr->right - x;
         int h = (int)box_ptr->bottom - y;
-        preprocess(i, crmatrix, x, y, w, h, stream);
+        preprocess(i, image, crmatrix, x, y, w, h, stream);
     }
 
     #ifdef TRT10
@@ -200,7 +201,7 @@ virtual Result ClassifierModelImpl::forward(const tensor::Image &image, data::Bo
     if (!trt_->forward(bindings, stream)) 
     {
         printf("Failed to tensorRT forward.");
-        return cv::Mat();
+        return std::monostate{};;
     }
     #endif
 
@@ -223,6 +224,7 @@ virtual Result ClassifierModelImpl::forward(const tensor::Image &image, data::Bo
         int index = *max_index;
         classfier_boxes_ptr[ib]->label = class_names_[index];
     }
+    return std::monostate{};
 }
 
 }

+ 3 - 4
src/infer/trt/classfier/classifier.hpp

@@ -51,8 +51,8 @@ public:
     affine::Norm normalize_;
 
 public:
-    YoloModelImpl() = default;
-    virtual ~YoloModelImpl() = default;
+    ClassifierModelImpl() = default;
+    virtual ~ClassifierModelImpl() = default;
 
     virtual int get_gpu_id() override
     {
@@ -65,7 +65,6 @@ public:
         size_t input_numel = network_input_width_ * network_input_height_ * 3;
         input_buffer_.gpu(batch_size * input_numel);
         output_buffer_.gpu(batch_size * num_classes_);
-        output_boxarray_.cpu(batch_size * num_classes_);
         classes_indices_.gpu(batch_size);
         classes_indices_.cpu(batch_size);
         preprocess_buffer_.gpu(image_width * image_height * 3);
@@ -73,7 +72,7 @@ public:
 
     bool load(const std::string &engine_file, int gpu_id);
 
-    void preprocess(int ibatch, affine::CropResizeMatrix& matrix, int x, int y, int w, int h, void *stream);
+    void preprocess(int ibatch, const tensor::Image &image, affine::CropResizeMatrix& matrix, int x, int y, int w, int h, void *stream);
     
     virtual Result forward(const tensor::Image &image, void *stream = nullptr);
     virtual Result forward(const tensor::Image &image, int slice_width, int slice_height, float overlap_width_ratio, float overlap_height_ratio, void *stream = nullptr);