enable to change vocabulary runtime (#56)

CMakeLists.txt

@@ -15,6 +15,7 @@ add_message_files(
 add_service_files(
     FILES
     DeticSeg.srv
+    CustomVocabulary.srv
 )
 
 generate_messages(

node_script/node.py

@@ -11,6 +11,8 @@
 
 from detic_ros.msg import SegmentationInfo
 from detic_ros.srv import DeticSeg, DeticSegRequest, DeticSegResponse
+from detic_ros.srv import CustomVocabulary, CustomVocabularyRequest, CustomVocabularyResponse
+from std_srvs.srv import Empty, EmptyRequest, EmptyResponse
 
 
 class DeticRosNode:
@@ -36,6 +38,8 @@ def __init__(self, node_config: Optional[NodeConfig] = None):
 
         self.detic_wrapper = DeticWrapper(node_config)
         self.srv_handler = rospy.Service('~segment_image', DeticSeg, self.callback_srv)
+        self.vocab_srv_handler = rospy.Service('~custom_vocabulary', CustomVocabulary, self.custom_vocab_srv)
+        self.default_vocab_srv_handler = rospy.Service('~default_vocabulary', Empty, self.default_vocab_srv)
 
         if node_config.enable_pubsub:
             # As for large buff_size please see:
@@ -113,6 +117,17 @@ def callback_srv(self, req: DeticSegRequest) -> DeticSegResponse:
             resp.debug_image = debug_image
         return resp
 
+    def custom_vocab_srv(self, req: CustomVocabularyRequest) -> CustomVocabularyResponse:
+        rospy.loginfo("Change vocabulary to {}".format(req.vocabulary))
+        self.detic_wrapper.predictor.change_vocabulary(",".join(req.vocabulary))
+        res = CustomVocabularyResponse()
+        return res
+
+    def default_vocab_srv(self, req: EmptyRequest) -> EmptyResponse:
+        rospy.loginfo("Change to default vocabulary")
+        self.detic_wrapper.predictor.set_defalt_vocabulary()
+        res = EmptyResponse()
+        return res
 
 if __name__ == '__main__':
     rospy.init_node('detic_node', anonymous=True)

package.xml

@@ -11,6 +11,7 @@
   <build_depend>message_generation</build_depend>
   <exec_depend>message_runtime</exec_depend>
   <exec_depend>sensor_msgs</exec_depend>
+  <exec_depend>std_srvs</exec_depend>
   <exec_depend>jsk_recognition_msgs</exec_depend>
   <exec_depend>cv_bridge</exec_depend>
 

srv/CustomVocabulary.srv

@@ -0,0 +1,2 @@
+string[] vocabulary
+---