color_histogram_visualizer.py
Go to the documentation of this file.
00001 #!/usr/bin/env python
00002 # -*- coding: utf-8 -*-
00003 # Copyright: Yuki Furuta <furushchev@jsk.imi.i.u-tokyo.ac.jp>
00004 
00005 import math
00006 import cv2
00007 from cv_bridge import CvBridge
00008 from dynamic_reconfigure.server import Server
00009 import numpy as np
00010 import matplotlib
00011 matplotlib.use('Agg')
00012 from matplotlib import pyplot as plt
00013 from matplotlib import gridspec
00014 import rospy
00015 from jsk_topic_tools import ConnectionBasedTransport
00016 from jsk_recognition_msgs.msg import ColorHistogram, ColorHistogramArray
00017 from sensor_msgs.msg import Image
00018 from jsk_pcl_ros.cfg import ColorHistogramVisualizerConfig as Config
00019 
00020 
00021 # matplotlib outputs image with size 800x600 by default
00022 IMG_WIDTH=800
00023 IMG_HEIGHT=600
00024 
00025 
00026 class ColorHistogramVisualizer(ConnectionBasedTransport):
00027     def __init__(self):
00028         super(ColorHistogramVisualizer, self).__init__()
00029 
00030         self.dyn_server = Server(Config, self.config_callback)
00031 
00032         self.cv_bridge = CvBridge()
00033         self.hsv_color_map = plt.cm.get_cmap('hsv')
00034         self.hsv_color_map_2d = self.get_hsv_color_map_2d()
00035 
00036         self.pub_image = self.advertise("~output", Image, queue_size=1)
00037 
00038     def config_callback(self, config, level):
00039         self.histogram_policy = config.histogram_policy
00040         self.histogram_index = config.histogram_index
00041         self.histogram_scale = config.histogram_scale
00042         return config
00043 
00044     def subscribe(self):
00045         self.sub_histogram = rospy.Subscriber("~input", ColorHistogram,
00046                                               self.callback, queue_size=1)
00047         self.sub_histogram_array = rospy.Subscriber("~input/array", ColorHistogramArray,
00048                                                     self.callback_array, queue_size=1)
00049 
00050     def unsubscribe(self):
00051         self.sub_histogram.unregister()
00052         self.sub_histogram_array.unregister()
00053 
00054     def get_hsv_color_map_2d(self):
00055         hsv_map = np.zeros((IMG_HEIGHT, IMG_WIDTH, 3), np.uint8)
00056         h, s = np.indices(hsv_map.shape[:2])
00057         hsv_map[:, :, 0] = h / float(IMG_HEIGHT) * 180.0
00058         hsv_map[:, :, 1] = 125.0
00059         hsv_map[:, :, 2] = s / float(IMG_WIDTH)  * 256.0
00060         hsv_map = cv2.cvtColor(hsv_map, cv2.COLOR_HLS2BGR)
00061         return hsv_map
00062 
00063     def callback(self, msg):
00064         if self.histogram_policy == Config.ColorHistogramVisualizer_HUE:
00065             img = self.plot_hist_hue(msg.histogram)
00066         elif self.histogram_policy == Config.ColorHistogramVisualizer_HUE_AND_SATURATION:
00067             img = self.plot_hist_hs(msg.histogram)
00068         else:
00069             rospy.logerr("Invalid histogram policy")
00070             return
00071         try:
00072             pub_msg = self.cv_bridge.cv2_to_imgmsg(img, "bgr8")
00073             self.pub_image.publish(pub_msg)
00074         except Exception as e:
00075             rospy.logerr("Failed to convert image: %s" % str(e))
00076 
00077     def callback_array(self, msg):
00078         if 0 <= self.histogram_index < len(msg.histograms):
00079             self.callback(msg.histograms[self.histogram_index])
00080         else:
00081             rospy.logerr("histogram index Out-of-index")
00082 
00083     def image_from_plot(self):
00084         fig = plt.gcf()
00085         fig.canvas.draw()
00086         w, h = fig.canvas.get_width_height()
00087         img = np.fromstring(fig.canvas.tostring_rgb(), dtype=np.uint8)
00088         fig.clf()
00089         img.shape = (h, w, 3)
00090         img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
00091         return img
00092 
00093     def plot_hist_hue(self, hist):
00094         gs = gridspec.GridSpec(1, 2, width_ratios=[4, 1])
00095         plt.subplot(gs[0], axisbg='silver')
00096         bin_size = len(hist) - 2
00097         bin_step = 360.0 / bin_size
00098         x = np.arange(360.0, step=bin_step)
00099         bars = plt.bar(x, hist[:-2], width=bin_step)
00100         cs = np.arange(0.0, 1.0, 1.0 / bin_size)
00101         for c, b in zip(cs, bars):
00102             b.set_facecolor(self.hsv_color_map(c))
00103         plt.xlim(0, 360.0)
00104         plt.ylim(ymin=0.0, ymax=1.0)
00105         ymin, ymax = plt.ylim()
00106         plt.subplot(gs[1], axisbg='silver')
00107         bars = plt.bar(range(2), hist[-2:], label=["white", "black"],
00108                        width=1.0, linewidth=2.0)
00109         bars[0].set_facecolor((1.0, 1.0, 1.0, 1.0))
00110         bars[1].set_facecolor((0.0, 0.0, 0.0, 1.0))
00111         plt.ylim(ymin, ymax)
00112         return self.image_from_plot()
00113 
00114     def plot_hist_saturation(self, hist):
00115         bin_size = len(hist)
00116         bin_step = 256.0 / bin_size
00117         x = np.arange(256.0, step=bin_step)
00118         plt.bar(x, hist, width=bin_step)
00119         plt.xlim(0, 256.0)
00120         return self.image_from_plot()
00121 
00122     def plot_hist_hs(self, hist):
00123         bin_size = int(math.sqrt(len(hist) - 2))
00124         white, black = hist[-2], hist[-1]
00125         hist = np.array(hist[:-2]).reshape(bin_size, bin_size).T
00126         rospy.loginfo("white: %f, black: %f" % (white, black))
00127         hist = np.clip(hist * 150 * self.histogram_scale, white, 1.0 - black)
00128         hist = hist[:, :, np.newaxis]
00129         hist = cv2.resize(hist, (IMG_WIDTH, IMG_HEIGHT))
00130         hist = hist[:, :, np.newaxis]
00131         img = self.hsv_color_map_2d * hist
00132         img = img.astype(np.uint8)
00133         return img
00134 
00135 if __name__ == '__main__':
00136     rospy.init_node("color_histogram_visualizer")
00137     v = ColorHistogramVisualizer()
00138     rospy.spin()


jsk_pcl_ros
Author(s): Yohei Kakiuchi
autogenerated on Sun Oct 8 2017 02:43:49