Merge pull request #1696 from LDOUBLEV/trt

Support fixed length TRT prediction

deploy/cpp_infer/include/preprocess_op.h

@@ -47,18 +47,20 @@ public:
 class ResizeImgType0 {
 public:
   virtual void Run(const cv::Mat &img, cv::Mat &resize_img, int max_size_len,
-                   float &ratio_h, float &ratio_w);
+                   float &ratio_h, float &ratio_w, bool use_tensorrt);
 };
 
 class CrnnResizeImg {
 public:
   virtual void Run(const cv::Mat &img, cv::Mat &resize_img, float wh_ratio,
+                   bool use_tensorrt = false,
                    const std::vector<int> &rec_image_shape = {3, 32, 320});
 };
 
 class ClsResizeImg {
 public:
   virtual void Run(const cv::Mat &img, cv::Mat &resize_img,
+                   bool use_tensorrt = false,
                    const std::vector<int> &rec_image_shape = {3, 48, 192});
 };
 

deploy/cpp_infer/src/main.cpp

@@ -77,7 +77,7 @@ int main(int argc, char **argv) {
   auto end = std::chrono::system_clock::now();
   auto duration =
       std::chrono::duration_cast<std::chrono::microseconds>(end - start);
-  std::cout << "Cost"
+  std::cout << "Cost  "
             << double(duration.count()) *
                    std::chrono::microseconds::period::num /
                    std::chrono::microseconds::period::den

deploy/cpp_infer/src/ocr_cls.cpp

@@ -25,7 +25,7 @@ cv::Mat Classifier::Run(cv::Mat &img) {
   int index = 0;
   float wh_ratio = float(img.cols) / float(img.rows);
 
-  this->resize_op_.Run(img, resize_img, cls_image_shape);
+  this->resize_op_.Run(img, resize_img, this->use_tensorrt_, cls_image_shape);
 
   this->normalize_op_.Run(&resize_img, this->mean_, this->scale_,
                           this->is_scale_);

deploy/cpp_infer/src/ocr_det.cpp

@@ -61,7 +61,8 @@ void DBDetector::Run(cv::Mat &img,
   cv::Mat srcimg;
   cv::Mat resize_img;
   img.copyTo(srcimg);
-  this->resize_op_.Run(img, resize_img, this->max_side_len_, ratio_h, ratio_w);
+  this->resize_op_.Run(img, resize_img, this->max_side_len_, ratio_h, ratio_w,
+                       this->use_tensorrt_);
 
   this->normalize_op_.Run(&resize_img, this->mean_, this->scale_,
                           this->is_scale_);

deploy/cpp_infer/src/ocr_rec.cpp

@@ -33,7 +33,7 @@ void CRNNRecognizer::Run(std::vector<std::vector<std::vector<int>>> boxes,
 
     float wh_ratio = float(crop_img.cols) / float(crop_img.rows);
 
-    this->resize_op_.Run(crop_img, resize_img, wh_ratio);
+    this->resize_op_.Run(crop_img, resize_img, wh_ratio, this->use_tensorrt_);
 
     this->normalize_op_.Run(&resize_img, this->mean_, this->scale_,
                             this->is_scale_);

deploy/cpp_infer/src/preprocess_op.cpp

@@ -60,7 +60,8 @@ void Normalize::Run(cv::Mat *im, const std::vector<float> &mean,
 }
 
 void ResizeImgType0::Run(const cv::Mat &img, cv::Mat &resize_img,
-                         int max_size_len, float &ratio_h, float &ratio_w) {
+                         int max_size_len, float &ratio_h, float &ratio_w,
+                         bool use_tensorrt) {
   int w = img.cols;
   int h = img.rows;
 
@@ -89,14 +90,19 @@ void ResizeImgType0::Run(const cv::Mat &img, cv::Mat &resize_img,
     resize_w = 32;
   else
     resize_w = (resize_w / 32) * 32;
-
-  cv::resize(img, resize_img, cv::Size(resize_w, resize_h));
-
-  ratio_h = float(resize_h) / float(h);
-  ratio_w = float(resize_w) / float(w);
+  if (!use_tensorrt) {
+    cv::resize(img, resize_img, cv::Size(resize_w, resize_h));
+    ratio_h = float(resize_h) / float(h);
+    ratio_w = float(resize_w) / float(w);
+  } else {
+    cv::resize(img, resize_img, cv::Size(640, 640));
+    ratio_h = float(640) / float(h);
+    ratio_w = float(640) / float(w);
+  }
 }
 
 void CrnnResizeImg::Run(const cv::Mat &img, cv::Mat &resize_img, float wh_ratio,
+                        bool use_tensorrt,
                         const std::vector<int> &rec_image_shape) {
   int imgC, imgH, imgW;
   imgC = rec_image_shape[0];
@@ -111,12 +117,27 @@ void CrnnResizeImg::Run(const cv::Mat &img, cv::Mat &resize_img, float wh_ratio,
     resize_w = imgW;
   else
     resize_w = int(ceilf(imgH * ratio));
-
-  cv::resize(img, resize_img, cv::Size(resize_w, imgH), 0.f, 0.f,
-             cv::INTER_LINEAR);
+  if (!use_tensorrt) {
+    cv::resize(img, resize_img, cv::Size(resize_w, imgH), 0.f, 0.f,
+               cv::INTER_LINEAR);
+    cv::copyMakeBorder(resize_img, resize_img, 0, 0, 0,
+                       int(imgW - resize_img.cols), cv::BORDER_CONSTANT,
+                       {127, 127, 127});
+  } else {
+    int k = int(img.cols * 32 / img.rows);
+    if (k >= 100) {
+      cv::resize(img, resize_img, cv::Size(100, 32), 0.f, 0.f,
+                 cv::INTER_LINEAR);
+    } else {
+      cv::resize(img, resize_img, cv::Size(k, 32), 0.f, 0.f, cv::INTER_LINEAR);
+      cv::copyMakeBorder(resize_img, resize_img, 0, 0, 0, int(100 - k),
+                         cv::BORDER_CONSTANT, {127, 127, 127});
+    }
+  }
 }
 
 void ClsResizeImg::Run(const cv::Mat &img, cv::Mat &resize_img,
+                       bool use_tensorrt,
                        const std::vector<int> &rec_image_shape) {
   int imgC, imgH, imgW;
   imgC = rec_image_shape[0];
@@ -130,11 +151,15 @@ void ClsResizeImg::Run(const cv::Mat &img, cv::Mat &resize_img,
   else
     resize_w = int(ceilf(imgH * ratio));
 
-  cv::resize(img, resize_img, cv::Size(resize_w, imgH), 0.f, 0.f,
-             cv::INTER_LINEAR);
-  if (resize_w < imgW) {
-    cv::copyMakeBorder(resize_img, resize_img, 0, 0, 0, imgW - resize_w,
-                       cv::BORDER_CONSTANT, cv::Scalar(0, 0, 0));
+  if (!use_tensorrt) {
+    cv::resize(img, resize_img, cv::Size(resize_w, imgH), 0.f, 0.f,
+               cv::INTER_LINEAR);
+    if (resize_w < imgW) {
+      cv::copyMakeBorder(resize_img, resize_img, 0, 0, 0, imgW - resize_w,
+                         cv::BORDER_CONSTANT, cv::Scalar(0, 0, 0));
+    }
+  } else {
+    cv::resize(img, resize_img, cv::Size(100, 32), 0.f, 0.f, cv::INTER_LINEAR);
   }
 }