Channingss 5 жил өмнө
parent
commit
0611fe1ab3

+ 1 - 1
deploy/cpp/include/paddlex/transforms.h

@@ -158,7 +158,7 @@ class Padding : public Transform {
         exit(-1);
       }
     }
-    if (item["target_size"].IsDefined()){
+    if (item["target_size"].IsDefined()) {
       if (item["target_size"].IsScalar()) {
         width_ = item["target_size"].as<int>();
         height_ = item["target_size"].as<int>();

+ 13 - 11
deploy/cpp/src/paddlex.cpp

@@ -38,13 +38,14 @@ void Model::create_predictor(const std::string& model_dir,
   config.SwitchSpecifyInputNames(true);
   // 开启内存优化
   config.EnableMemoryOptim();
-  if (use_trt){
-    config.EnableTensorRtEngine(1 << 20      /* workspace_size*/,
-                        32        /* max_batch_size*/,
-                        20                 /* min_subgraph_size*/,
-                        paddle::AnalysisConfig::Precision::kFloat32 /* precision*/,
-                        true             /* use_static*/,
-                        false             /* use_calib_mode*/);
+  if (use_trt) {
+    config.EnableTensorRtEngine(
+        1 << 20 /* workspace_size*/,
+        32 /* max_batch_size*/,
+        20 /* min_subgraph_size*/,
+        paddle::AnalysisConfig::Precision::kFloat32 /* precision*/,
+        true /* use_static*/,
+        false /* use_calib_mode*/);
   }
   predictor_ = std::move(CreatePaddlePredictor(config));
 }
@@ -294,20 +295,21 @@ bool Model::predict(const cv::Mat& im, SegResult* result) {
                      result->score_map.shape[3],
                      CV_32FC1,
                      result->score_map.data.data());
-  int idx=1;
+  int idx = 1;
   int len_postprocess = inputs_.im_size_before_resize_.size();
   for (std::vector<std::string>::reverse_iterator iter =
            inputs_.reshape_order_.rbegin();
-       iter != inputs_.reshape_order_.rend(); ++iter) {
+       iter != inputs_.reshape_order_.rend();
+       ++iter) {
     if (*iter == "padding") {
-      auto before_shape = inputs_.im_size_before_resize_[len_postprocess-idx];
+      auto before_shape = inputs_.im_size_before_resize_[len_postprocess - idx];
       inputs_.im_size_before_resize_.pop_back();
       auto padding_w = before_shape[0];
       auto padding_h = before_shape[1];
       mask_label = mask_label(cv::Rect(0, 0, padding_w, padding_h));
       mask_score = mask_score(cv::Rect(0, 0, padding_w, padding_h));
     } else if (*iter == "resize") {
-      auto before_shape = inputs_.im_size_before_resize_[len_postprocess-idx];
+      auto before_shape = inputs_.im_size_before_resize_[len_postprocess - idx];
       inputs_.im_size_before_resize_.pop_back();
       auto resize_w = before_shape[0];
       auto resize_h = before_shape[1];

+ 4 - 4
deploy/cpp/src/transforms.cpp

@@ -56,7 +56,7 @@ float ResizeByShort::GenerateScale(const cv::Mat& im) {
 }
 
 bool ResizeByShort::Run(cv::Mat* im, ImageBlob* data) {
-  data->im_size_before_resize_.push_back({im->rows,im->cols});
+  data->im_size_before_resize_.push_back({im->rows, im->cols});
   data->reshape_order_.push_back("resize");
 
   float scale = GenerateScale(*im);
@@ -87,7 +87,7 @@ bool CenterCrop::Run(cv::Mat* im, ImageBlob* data) {
 }
 
 bool Padding::Run(cv::Mat* im, ImageBlob* data) {
-  data->im_size_before_resize_.push_back({im->rows,im->cols});
+  data->im_size_before_resize_.push_back({im->rows, im->cols});
   data->reshape_order_.push_back("padding");
 
   int padding_w = 0;
@@ -121,7 +121,7 @@ bool ResizeByLong::Run(cv::Mat* im, ImageBlob* data) {
               << std::endl;
     return false;
   }
-  data->im_size_before_resize_.push_back({im->rows,im->cols});
+  data->im_size_before_resize_.push_back({im->rows, im->cols});
   data->reshape_order_.push_back("resize");
   int origin_w = im->cols;
   int origin_h = im->rows;
@@ -147,7 +147,7 @@ bool Resize::Run(cv::Mat* im, ImageBlob* data) {
               << std::endl;
     return false;
   }
-  data->im_size_before_resize_.push_back({im->rows,im->cols});
+  data->im_size_before_resize_.push_back({im->rows, im->cols});
   data->reshape_order_.push_back("resize");
 
   cv::resize(

+ 10 - 0
paddlex/command.py

@@ -64,6 +64,16 @@ def main():
         model = pdx.load_model(args.model_dir, fixed_input_shape)
         model.export_inference_model(args.save_dir)
 
+    if args.export_onnx:
+        assert args.model_dir is not None, "--model_dir should be defined while exporting onnx model"
+        assert args.save_dir is not None, "--save_dir should be defined to save onnx model"
+        fixed_input_shape = eval(args.fixed_input_shape)
+        assert len(
+            fixed_input_shape) == 2, "len of fixed input shape must == 2"
+
+        model = pdx.load_model(args.model_dir, fixed_input_shape)
+        model.export_onnx_model(args.save_dir)
+
 
 if __name__ == "__main__":
     main()