feat(traffic_light_classifier): add build only option for tensort eng…

…ine (#2985) Signed-off-by: yukke42 <yusuke.muramatsu@tier4.jp>
autowarefoundation · Mar 6, 2023 · 8c6797e · 8c6797e
1 parent 54d45e5
commit 8c6797e
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 8 deletions.
diff --git a/perception/traffic_light_classifier/README.md b/perception/traffic_light_classifier/README.md
@@ -54,14 +54,15 @@ These colors and shapes are assigned to the message as follows:
 
 #### cnn_classifier
 
-| Name              | Type | Description                          |
-| ----------------- | ---- | ------------------------------------ |
-| `model_file_path` | str  | path to the model file               |
-| `label_file_path` | str  | path to the label file               |
-| `precision`       | str  | TensorRT precision, `fp16` or `int8` |
-| `input_c`         | str  | the channel size of an input image   |
-| `input_h`         | str  | the height of an input image         |
-| `input_w`         | str  | the width of an input image          |
+| Name              | Type | Description                                      |
+| ----------------- | ---- | ------------------------------------------------ |
+| `model_file_path` | str  | path to the model file                           |
+| `label_file_path` | str  | path to the label file                           |
+| `precision`       | str  | TensorRT precision, `fp16` or `int8`             |
+| `input_c`         | str  | the channel size of an input image               |
+| `input_h`         | str  | the height of an input image                     |
+| `input_w`         | str  | the width of an input image                      |
+| `build_only`      | bool | shutdown node after TensorRT engin file is built |
 
 #### hsv_classifier
 

diff --git a/perception/traffic_light_classifier/src/cnn_classifier.cpp b/perception/traffic_light_classifier/src/cnn_classifier.cpp
@@ -47,6 +47,11 @@ CNNClassifier::CNNClassifier(rclcpp::Node * node_ptr) : node_ptr_(node_ptr)
 
   trt_ = std::make_shared<Tn::TrtCommon>(model_file_path, precision, input_name, output_name);
   trt_->setup();
+
+  if (node_ptr_->declare_parameter("build_only", false)) {
+    RCLCPP_INFO(node_ptr_->get_logger(), "TensorRT engine is built and shutdown node.");
+    rclcpp::shutdown();
+  }
 }
 
 bool CNNClassifier::getTrafficSignal(