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