the model and the target language are parametrized

ros-ai · May 23, 2023 · b34c0ba · b34c0ba
1 parent fb8120a
commit b34c0ba
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 2 deletions.
diff --git a/launch/ros2_whisper.launch.py b/launch/ros2_whisper.launch.py
@@ -17,6 +17,7 @@ def generate_launch_description():
         package="ros2_whisper",
         executable="whisper_inference_node",
         output="screen",
+        parameters=[{"model": "base", "lanaguage": "english"}],
     )
 
     whisper_inference_node_event_handler = RegisterEventHandler(

diff --git a/ros2_whisper/whisper_inference_node.py b/ros2_whisper/whisper_inference_node.py
@@ -12,8 +12,13 @@
 class WhisperInferenceNode(Node):
     def __init__(self, node_name: str) -> None:
         super().__init__(node_name)
-        self.whisper_model_ = whisper.load_model("base")
-        self.whisper_options_ = whisper.DecodingOptions(language="english")
+        self.declare_parameter("model", "base")
+        self.declare_parameter("language", "english")
+        model = self.get_parameter("model").value
+        language = self.get_parameter("language").value
+
+        self.whisper_model_ = whisper.load_model(model)
+        self.whisper_options_ = whisper.DecodingOptions(language=language)
 
         self.declare_parameters(
             namespace="",