[audio_to_spectrogram] Enable publishing frequency vs amplitude plot

audio_to_spectrogram/README.md

@@ -84,3 +84,21 @@ roslaunch audio_to_spectrogram sample_audio_to_spectrogram.launch
     - `~publish_rate` (`Double`, default: `image_width / spectrogram_period`)
 
       Publish rate [Hz] of spectrogram topic.
+
+## spectrum_plot.py
+
+  A script to publish frequency vs amplitude plot image.
+
+![](docs/images/spectrum.jpg)
+
+  - ### Publishing topics
+
+    - `~output/viz` (`sensor_msgs/Image`)
+
+      Frequency vs amplitude plot image.
+
+  - ### Subscribing topics
+
+    - `~spectrum` (`jsk_recognition_msgs/Spectrum`)
+
+      Spectrum data calculated from audio by FFT.

audio_to_spectrogram/launch/audio_to_spectrogram.launch

@@ -51,10 +51,15 @@
     </rosparam>
   </node>
 
+  <!-- publish frequency vs amplitude image -->
+  <node name="spectrum_plot"
+        pkg="audio_to_spectrogram" type="spectrum_plot.py" >
+    <remap from="~spectrum" to="/audio_to_spectrum/spectrum_filtered" />
+  </node>
+
   <group if="$(arg gui)">
-    <!-- visualize spectrum -->
-    <node pkg="audio_to_spectrogram" type="spectrum_plot.py" name="spectrum_plot"  >
-      <remap from="~spectrum" to="/audio_to_spectrum/spectrum_filtered" />
+    <node pkg="image_view" type="image_view" name="spectrum_view"  >
+      <remap from="image" to="/spectrum_plot/output/viz" />
     </node>
 
     <!-- visualize spectrogram -->

audio_to_spectrogram/package.xml

@@ -14,6 +14,7 @@
   <exec_depend>cv_bridge</exec_depend>
   <exec_depend>image_view</exec_depend>
   <exec_depend>jsk_recognition_msgs</exec_depend>
+  <exec_depend>jsk_topic_tools</exec_depend>
   <exec_depend condition="$ROS_PYTHON_VERSION == 2">python-matplotlib</exec_depend>
   <exec_depend condition="$ROS_PYTHON_VERSION == 3">python3-matplotlib</exec_depend>
   <exec_depend>sensor_msgs</exec_depend>

audio_to_spectrogram/scripts/spectrum_plot.py

@@ -3,16 +3,22 @@
 
 from __future__ import division
 
+import cv_bridge
+from jsk_topic_tools import ConnectionBasedTransport
 import matplotlib.pyplot as plt
 import numpy as np
 import rospy
+import sensor_msgs.msg
 
 from jsk_recognition_msgs.msg import Spectrum
 
+from audio_to_spectrogram import convert_matplotlib_to_img
 
-class SpectrumPlot():
+
+class SpectrumPlot(ConnectionBasedTransport):
 
     def __init__(self):
+        super(SpectrumPlot, self).__init__()
         # Set matplotlib config
         self.fig = plt.figure(figsize=(8, 5))
         self.fig.suptitle('Spectrum plot', size=12)
@@ -23,10 +29,17 @@ def __init__(self):
         self.ax.set_xlabel('Frequency [Hz]', fontsize=12)
         self.ax.set_ylabel('Amplitude', fontsize=12)
         self.line, = self.ax.plot([0, 0], label='Amplitude of Spectrum')
-        # ROS subscriber
+        # ROS publisher and subscriber
+        self.pub_img = self.advertise(
+            '~output/viz', sensor_msgs.msg.Image, queue_size=1)
+
+    def subscribe(self):
         self.sub_spectrum = rospy.Subscriber(
             '~spectrum', Spectrum, self._cb, queue_size=1000)
 
+    def unsubscribe(self):
+        self.sub_spectrum.unregister()
+
     def _cb(self, msg):
         # Plot spectrum
         self.amp = np.array(msg.amplitude)
@@ -35,10 +48,15 @@ def _cb(self, msg):
         self.ax.set_xlim((self.freq.min(), self.freq.max()))
         self.ax.set_ylim((0.0, 20))
         self.ax.legend(loc='upper right')
+        if self.pub_img.get_num_connections() > 0:
+            bridge = cv_bridge.CvBridge()
+            img = convert_matplotlib_to_img(self.fig)
+            img_msg = bridge.cv2_to_imgmsg(img, encoding='rgb8')
+            img_msg.header = msg.header
+            self.pub_img.publish(img_msg)
 
 
 if __name__ == '__main__':
     rospy.init_node('spectrum_plot')
     SpectrumPlot()
-    while not rospy.is_shutdown():
-        plt.pause(.1)  # real-time plotting
+    rospy.spin()

audio_to_spectrogram/src/audio_to_spectrogram/__init__.py

@@ -0,0 +1,2 @@
+from .audio_buffer import AudioBuffer
+from .convert_matplotlib import convert_matplotlib_to_img

audio_to_spectrogram/src/audio_to_spectrogram/convert_matplotlib.py

@@ -0,0 +1,19 @@
+try:
+    # for python2.x
+    from StringIO import StringIO as BufIO
+except ModuleNotFoundError:
+    # for python3.x
+    from io import BytesIO as BufIO
+
+import numpy as np
+import PIL
+
+
+def convert_matplotlib_to_img(fig):
+    buf = BufIO()
+    fig.savefig(buf, format="png")
+    buf.seek(0)
+    img = np.array(
+        PIL.Image.open(buf), dtype=np.uint8)
+    img = img[..., :3]
+    return img