00001
00002
00003
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
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()