提交 24f8dfb7 编写于 作者: R root

Support fixed length TRT prediction

上级 1f99a635
...@@ -47,18 +47,20 @@ public: ...@@ -47,18 +47,20 @@ public:
class ResizeImgType0 { class ResizeImgType0 {
public: public:
virtual void Run(const cv::Mat &img, cv::Mat &resize_img, int max_size_len, 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 { class CrnnResizeImg {
public: public:
virtual void Run(const cv::Mat &img, cv::Mat &resize_img, float wh_ratio, 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}); const std::vector<int> &rec_image_shape = {3, 32, 320});
}; };
class ClsResizeImg { class ClsResizeImg {
public: public:
virtual void Run(const cv::Mat &img, cv::Mat &resize_img, 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}); const std::vector<int> &rec_image_shape = {3, 48, 192});
}; };
......
...@@ -77,7 +77,7 @@ int main(int argc, char **argv) { ...@@ -77,7 +77,7 @@ int main(int argc, char **argv) {
auto end = std::chrono::system_clock::now(); auto end = std::chrono::system_clock::now();
auto duration = auto duration =
std::chrono::duration_cast<std::chrono::microseconds>(end - start); std::chrono::duration_cast<std::chrono::microseconds>(end - start);
std::cout << "Cost" std::cout << "Cost "
<< double(duration.count()) * << double(duration.count()) *
std::chrono::microseconds::period::num / std::chrono::microseconds::period::num /
std::chrono::microseconds::period::den std::chrono::microseconds::period::den
......
...@@ -25,7 +25,7 @@ cv::Mat Classifier::Run(cv::Mat &img) { ...@@ -25,7 +25,7 @@ cv::Mat Classifier::Run(cv::Mat &img) {
int index = 0; int index = 0;
float wh_ratio = float(img.cols) / float(img.rows); 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->normalize_op_.Run(&resize_img, this->mean_, this->scale_,
this->is_scale_); this->is_scale_);
......
...@@ -61,7 +61,8 @@ void DBDetector::Run(cv::Mat &img, ...@@ -61,7 +61,8 @@ void DBDetector::Run(cv::Mat &img,
cv::Mat srcimg; cv::Mat srcimg;
cv::Mat resize_img; cv::Mat resize_img;
img.copyTo(srcimg); 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->normalize_op_.Run(&resize_img, this->mean_, this->scale_,
this->is_scale_); this->is_scale_);
......
...@@ -33,7 +33,7 @@ void CRNNRecognizer::Run(std::vector<std::vector<std::vector<int>>> boxes, ...@@ -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); 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->normalize_op_.Run(&resize_img, this->mean_, this->scale_,
this->is_scale_); this->is_scale_);
......
...@@ -60,7 +60,8 @@ void Normalize::Run(cv::Mat *im, const std::vector<float> &mean, ...@@ -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, 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 w = img.cols;
int h = img.rows; int h = img.rows;
...@@ -89,14 +90,19 @@ void ResizeImgType0::Run(const cv::Mat &img, cv::Mat &resize_img, ...@@ -89,14 +90,19 @@ void ResizeImgType0::Run(const cv::Mat &img, cv::Mat &resize_img,
resize_w = 32; resize_w = 32;
else else
resize_w = (resize_w / 32) * 32; resize_w = (resize_w / 32) * 32;
if (!use_tensorrt) {
cv::resize(img, resize_img, cv::Size(resize_w, resize_h)); cv::resize(img, resize_img, cv::Size(resize_w, resize_h));
ratio_h = float(resize_h) / float(h); ratio_h = float(resize_h) / float(h);
ratio_w = float(resize_w) / float(w); 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, void CrnnResizeImg::Run(const cv::Mat &img, cv::Mat &resize_img, float wh_ratio,
bool use_tensorrt,
const std::vector<int> &rec_image_shape) { const std::vector<int> &rec_image_shape) {
int imgC, imgH, imgW; int imgC, imgH, imgW;
imgC = rec_image_shape[0]; imgC = rec_image_shape[0];
...@@ -111,12 +117,27 @@ void CrnnResizeImg::Run(const cv::Mat &img, cv::Mat &resize_img, float wh_ratio, ...@@ -111,12 +117,27 @@ void CrnnResizeImg::Run(const cv::Mat &img, cv::Mat &resize_img, float wh_ratio,
resize_w = imgW; resize_w = imgW;
else else
resize_w = int(ceilf(imgH * ratio)); resize_w = int(ceilf(imgH * ratio));
if (!use_tensorrt) {
cv::resize(img, resize_img, cv::Size(resize_w, imgH), 0.f, 0.f, cv::resize(img, resize_img, cv::Size(resize_w, imgH), 0.f, 0.f,
cv::INTER_LINEAR); 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, void ClsResizeImg::Run(const cv::Mat &img, cv::Mat &resize_img,
bool use_tensorrt,
const std::vector<int> &rec_image_shape) { const std::vector<int> &rec_image_shape) {
int imgC, imgH, imgW; int imgC, imgH, imgW;
imgC = rec_image_shape[0]; imgC = rec_image_shape[0];
...@@ -130,12 +151,16 @@ void ClsResizeImg::Run(const cv::Mat &img, cv::Mat &resize_img, ...@@ -130,12 +151,16 @@ void ClsResizeImg::Run(const cv::Mat &img, cv::Mat &resize_img,
else else
resize_w = int(ceilf(imgH * ratio)); resize_w = int(ceilf(imgH * ratio));
if (!use_tensorrt) {
cv::resize(img, resize_img, cv::Size(resize_w, imgH), 0.f, 0.f, cv::resize(img, resize_img, cv::Size(resize_w, imgH), 0.f, 0.f,
cv::INTER_LINEAR); cv::INTER_LINEAR);
if (resize_w < imgW) { if (resize_w < imgW) {
cv::copyMakeBorder(resize_img, resize_img, 0, 0, 0, imgW - resize_w, cv::copyMakeBorder(resize_img, resize_img, 0, 0, 0, imgW - resize_w,
cv::BORDER_CONSTANT, cv::Scalar(0, 0, 0)); 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);
}
} }
} // namespace PaddleOCR } // namespace PaddleOCR
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册