提交 0611fe1a 编写于 作者: C Channingss

stylize cpp code

上级 a7474606
......@@ -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>();
......
......@@ -38,8 +38,9 @@ void Model::create_predictor(const std::string& model_dir,
config.SwitchSpecifyInputNames(true);
// 开启内存优化
config.EnableMemoryOptim();
if (use_trt){
config.EnableTensorRtEngine(1 << 20 /* workspace_size*/,
if (use_trt) {
config.EnableTensorRtEngine(
1 << 20 /* workspace_size*/,
32 /* max_batch_size*/,
20 /* min_subgraph_size*/,
paddle::AnalysisConfig::Precision::kFloat32 /* precision*/,
......@@ -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];
......
......@@ -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(
......
......@@ -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()
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册