From 43838e3e9fc01af8853efce64b4807fc0e2b371a Mon Sep 17 00:00:00 2001
From: WenmuZhou <572459439@qq.com>
Date: Mon, 19 Sep 2022 11:39:22 +0000
Subject: [PATCH] add trt shape file to table and layout model

---
 deploy/cpp_infer/include/structure_layout.h | 3 +--
 deploy/cpp_infer/src/structure_layout.cpp   | 5 +++++
 deploy/cpp_infer/src/structure_table.cpp    | 5 +++++
 3 files changed, 11 insertions(+), 2 deletions(-)

diff --git a/deploy/cpp_infer/include/structure_layout.h b/deploy/cpp_infer/include/structure_layout.h
index 57f5f481..f3afb98f 100644
--- a/deploy/cpp_infer/include/structure_layout.h
+++ b/deploy/cpp_infer/include/structure_layout.h
@@ -88,7 +88,6 @@ private:
 
   // post-process
   PicodetPostProcessor post_processor_;
-
-}; // class StructureTableRecognizer
+};
 
 } // namespace PaddleOCR
\ No newline at end of file
diff --git a/deploy/cpp_infer/src/structure_layout.cpp b/deploy/cpp_infer/src/structure_layout.cpp
index 898ead7a..4f587d2b 100644
--- a/deploy/cpp_infer/src/structure_layout.cpp
+++ b/deploy/cpp_infer/src/structure_layout.cpp
@@ -120,6 +120,11 @@ void StructureLayoutRecognizer::LoadModel(const std::string &model_dir) {
         precision = paddle_infer::Config::Precision::kInt8;
       }
       config.EnableTensorRtEngine(1 << 20, 10, 3, precision, false, false);
+      if (!Utility::PathExists("./trt_layout_shape.txt")) {
+        config.CollectShapeRangeInfo("./trt_layout_shape.txt");
+      } else {
+        config.EnableTunedTensorRtDynamicShape("./trt_layout_shape.txt", true);
+      }
     }
   } else {
     config.DisableGpu();
diff --git a/deploy/cpp_infer/src/structure_table.cpp b/deploy/cpp_infer/src/structure_table.cpp
index 7df0ab94..1aa2e624 100644
--- a/deploy/cpp_infer/src/structure_table.cpp
+++ b/deploy/cpp_infer/src/structure_table.cpp
@@ -133,6 +133,11 @@ void StructureTableRecognizer::LoadModel(const std::string &model_dir) {
         precision = paddle_infer::Config::Precision::kInt8;
       }
       config.EnableTensorRtEngine(1 << 20, 10, 3, precision, false, false);
+      if (!Utility::PathExists("./trt_table_shape.txt")) {
+        config.CollectShapeRangeInfo("./trt_table_shape.txt");
+      } else {
+        config.EnableTunedTensorRtDynamicShape("./trt_table_shape.txt", true);
+      }
     }
   } else {
     config.DisableGpu();
-- 
GitLab