From 24f8dfb7e17a7b4d4354432a835659fc5dfb8840 Mon Sep 17 00:00:00 2001 From: root Date: Fri, 8 Jan 2021 09:34:29 +0000 Subject: [PATCH] Support fixed length TRT prediction --- deploy/cpp_infer/include/preprocess_op.h | 4 +- deploy/cpp_infer/src/main.cpp | 2 +- deploy/cpp_infer/src/ocr_cls.cpp | 2 +- deploy/cpp_infer/src/ocr_det.cpp | 3 +- deploy/cpp_infer/src/ocr_rec.cpp | 2 +- deploy/cpp_infer/src/preprocess_op.cpp | 53 +++++++++++++++++------- 6 files changed, 47 insertions(+), 19 deletions(-) diff --git a/deploy/cpp_infer/include/preprocess_op.h b/deploy/cpp_infer/include/preprocess_op.h index bda7d123..ab4c1400 100644 --- a/deploy/cpp_infer/include/preprocess_op.h +++ b/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 &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 &rec_image_shape = {3, 48, 192}); }; diff --git a/deploy/cpp_infer/src/main.cpp b/deploy/cpp_infer/src/main.cpp index f40e5edf..49c68f55 100644 --- a/deploy/cpp_infer/src/main.cpp +++ b/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(end - start); - std::cout << "Cost" + std::cout << "Cost " << double(duration.count()) * std::chrono::microseconds::period::num / std::chrono::microseconds::period::den diff --git a/deploy/cpp_infer/src/ocr_cls.cpp b/deploy/cpp_infer/src/ocr_cls.cpp index 3aeda2ed..9199e082 100644 --- a/deploy/cpp_infer/src/ocr_cls.cpp +++ b/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_); diff --git a/deploy/cpp_infer/src/ocr_det.cpp b/deploy/cpp_infer/src/ocr_det.cpp index 3678f37d..489940f0 100644 --- a/deploy/cpp_infer/src/ocr_det.cpp +++ b/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_); diff --git a/deploy/cpp_infer/src/ocr_rec.cpp b/deploy/cpp_infer/src/ocr_rec.cpp index 27cfe4c9..fa1e44ff 100644 --- a/deploy/cpp_infer/src/ocr_rec.cpp +++ b/deploy/cpp_infer/src/ocr_rec.cpp @@ -33,7 +33,7 @@ void CRNNRecognizer::Run(std::vector>> 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_); diff --git a/deploy/cpp_infer/src/preprocess_op.cpp b/deploy/cpp_infer/src/preprocess_op.cpp index cdb20c31..87d8dbbd 100644 --- a/deploy/cpp_infer/src/preprocess_op.cpp +++ b/deploy/cpp_infer/src/preprocess_op.cpp @@ -60,7 +60,8 @@ void Normalize::Run(cv::Mat *im, const std::vector &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 &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 &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); } } -- GitLab