00001
00002
00003 import argparse
00004 import os
00005 import sys
00006
00007 from cStringIO import StringIO
00008 import cv2
00009 from cv_bridge import CvBridge
00010 from distutils.version import LooseVersion
00011 from matplotlib.figure import Figure
00012 import numpy as np
00013 import python_qt_binding
00014 from python_qt_binding import loadUi
00015 from python_qt_binding.QtCore import Qt
00016 from python_qt_binding.QtCore import QTimer
00017 from python_qt_binding.QtCore import Slot
00018 from python_qt_binding.QtGui import QIcon
00019 import rospkg
00020 import rospy
00021 from rqt_gui_py.plugin import Plugin
00022 from rqt_plot.rosplot import ROSData as _ROSData
00023 from rqt_plot.rosplot import RosPlotException
00024 from rqt_py_common.topic_completer import TopicCompleter
00025 from sensor_msgs.msg import Image
00026 from sklearn import linear_model
00027
00028 from jsk_recognition_msgs.msg import PlotData
00029 from jsk_recognition_msgs.msg import PlotDataArray
00030
00031
00032 if LooseVersion(python_qt_binding.QT_BINDING_VERSION).version[0] >= 5:
00033 from python_qt_binding.QtWidgets import QSizePolicy
00034 from python_qt_binding.QtWidgets import QVBoxLayout
00035 from python_qt_binding.QtWidgets import QWidget
00036 try:
00037 from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg \
00038 as FigureCanvas
00039 except ImportError:
00040
00041 import thread
00042 sys.modules['_thread'] = thread
00043 from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg \
00044 as FigureCanvas
00045 try:
00046 from matplotlib.backends.backend_qt5agg import NavigationToolbar2QTAgg \
00047 as NavigationToolbar
00048 except ImportError:
00049 from matplotlib.backends.backend_qt5agg import NavigationToolbar2QT \
00050 as NavigationToolbar
00051 else:
00052 from python_qt_binding.QtGui import QSizePolicy
00053 from python_qt_binding.QtGui import QVBoxLayout
00054 from python_qt_binding.QtGui import QWidget
00055 try:
00056 from matplotlib.backends.backend_qt4agg import FigureCanvasQTAgg \
00057 as FigureCanvas
00058 except ImportError:
00059
00060 import thread
00061 sys.modules['_thread'] = thread
00062 from matplotlib.backends.backend_qt4agg import FigureCanvasQTAgg \
00063 as FigureCanvas
00064 try:
00065 from matplotlib.backends.backend_qt4agg import NavigationToolbar2QTAgg \
00066 as NavigationToolbar
00067 except ImportError:
00068 from matplotlib.backends.backend_qt4agg import NavigationToolbar2QT \
00069 as NavigationToolbar
00070
00071
00072
00073 class ROSData(_ROSData):
00074 def _get_data(self, msg):
00075 val = msg
00076 try:
00077 if not self.field_evals:
00078 return val
00079 for f in self.field_evals:
00080 val = f(val)
00081 return val
00082 except IndexError:
00083 self.error = RosPlotException(
00084 "{0} index error for: {1}".format(
00085 self.name, str(val).replace('\n', ', ')))
00086 except TypeError:
00087 self.error = RosPlotException(
00088 "{0} value was not numeric: {1}".format(
00089 self.name, val))
00090
00091
00092 def add_list(xss):
00093
00094 "xss = [[0, 1, 2, ...], [0, 1, 2, ...], ...]"
00095
00096 ret = []
00097 for xs in xss:
00098 ret.extend(xs)
00099 return ret
00100
00101
00102 class Plot2D(Plugin):
00103 def __init__(self, context):
00104 super(Plot2D, self).__init__(context)
00105 self.setObjectName('Plot2D')
00106 self._args = self._parse_args(context.argv())
00107 self._widget = Plot2DWidget(self._args.topics)
00108 self._widget.is_line = self._args.line
00109 self._widget.fit_line = self._args.fit_line
00110 self._widget.fit_line_ransac = self._args.fit_line_ransac
00111 outlier = self._args.fit_line_ransac_outlier
00112 self._widget.fit_line_ransac_outlier = outlier
00113 self._widget.xtitle = self._args.xtitle
00114 self._widget.ytitle = self._args.ytitle
00115 self._widget.no_legend = self._args.no_legend
00116 self._widget.sort_x = self._args.sort_x
00117 context.add_widget(self._widget)
00118
00119 def _parse_args(self, argv):
00120 parser = argparse.ArgumentParser(
00121 prog='rqt_histogram_plot', add_help=False)
00122 Plot2D.add_arguments(parser)
00123 args = parser.parse_args(argv)
00124 return args
00125
00126 @staticmethod
00127 def add_arguments(parser):
00128 group = parser.add_argument_group(
00129 'Options for rqt_histogram plugin')
00130 group.add_argument(
00131 'topics', nargs='?', default=[], help='Topics to plot')
00132 group.add_argument(
00133 '--line', action="store_true",
00134 help="Plot with lines instead of scatter")
00135 group.add_argument(
00136 '--fit-line', action="store_true",
00137 help="Plot line with least-square fitting")
00138 group.add_argument(
00139 '--fit-line-ransac', action="store_true",
00140 help="Plot line with RANSAC")
00141 group.add_argument(
00142 '--fit-line-ransac-outlier', type=float, default=0.1,
00143 help="Plot line with RANSAC")
00144 group.add_argument('--xtitle', help="Title in X axis")
00145 group.add_argument('--ytitle', help="Title in Y axis")
00146 group.add_argument('--no-legend', action="store_true")
00147 group.add_argument('--sort-x', action="store_true")
00148
00149
00150 class Plot2DWidget(QWidget):
00151 _redraw_interval = 10
00152
00153 def __init__(self, topics):
00154 super(Plot2DWidget, self).__init__()
00155 self.setObjectName('Plot2DWidget')
00156 rp = rospkg.RosPack()
00157 ui_file = os.path.join(
00158 rp.get_path('jsk_rqt_plugins'), 'resource', 'plot_histogram.ui')
00159 loadUi(ui_file, self)
00160 self.cv_bridge = CvBridge()
00161 self.subscribe_topic_button.setIcon(QIcon.fromTheme('add'))
00162 self.pause_button.setIcon(QIcon.fromTheme('media-playback-pause'))
00163 self.clear_button.setIcon(QIcon.fromTheme('edit-clear'))
00164 self.data_plot = MatPlot2D(self)
00165 self.data_plot_layout.addWidget(self.data_plot)
00166 self._topic_completer = TopicCompleter(self.topic_edit)
00167 self._topic_completer.update_topics()
00168 self.topic_edit.setCompleter(self._topic_completer)
00169 self.data_plot.dropEvent = self.dropEvent
00170 self.data_plot.dragEnterEvent = self.dragEnterEvent
00171 self._start_time = rospy.get_time()
00172 self._rosdata = None
00173 if len(topics) != 0:
00174 self.subscribe_topic(topics)
00175 self._update_plot_timer = QTimer(self)
00176 self._update_plot_timer.timeout.connect(self.update_plot)
00177 self._update_plot_timer.start(self._redraw_interval)
00178
00179 @Slot('QDropEvent*')
00180 def dropEvent(self, event):
00181 if event.mimeData().hasText():
00182 topic_name = str(event.mimeData().text())
00183 else:
00184 droped_item = event.source().selectedItems()[0]
00185 topic_name = str(droped_item.data(0, Qt.UserRole))
00186 self.subscribe_topic(topic_name)
00187
00188 @Slot()
00189 def on_topic_edit_returnPressed(self):
00190 "callback function when form is entered"
00191 if self.subscribe_topic_button.isEnabled():
00192 self.subscribe_topic(str(self.topic_edit.text()))
00193
00194 @Slot()
00195 def on_subscribe_topic_button_clicked(self):
00196 self.subscribe_topic(str(self.topic_edit.text()))
00197
00198 def subscribe_topic(self, topic_name):
00199 rospy.loginfo("subscribe topic")
00200 self.topic_with_field_name = topic_name
00201 self.pub_image = rospy.Publisher(
00202 topic_name + "/plot_image", Image, queue_size=1)
00203 if not self._rosdata:
00204 self._rosdata = ROSData(topic_name, self._start_time)
00205 else:
00206 if self._rosdata != topic_name:
00207 self._rosdata.close()
00208 self.data_plot.clear()
00209 self._rosdata = ROSData(topic_name, self._start_time)
00210 else:
00211 rospy.logwarn("%s is already subscribed", topic_name)
00212
00213 def enable_timer(self, enabled=True):
00214 if enabled:
00215 self._update_plot_timer.start(self._redraw_interval)
00216 else:
00217 self._update_plot_timer.stop()
00218
00219 @Slot()
00220 def on_clear_button_clicked(self):
00221 self.data_plot.clear()
00222
00223 @Slot(bool)
00224 def on_pause_button_clicked(self, checked):
00225 self.enable_timer(not checked)
00226
00227 def plot_one(self, msg, axes):
00228 concatenated_data = zip(msg.xs, msg.ys)
00229 if self.sort_x:
00230 concatenated_data.sort(key=lambda x: x[0])
00231 xs = [d[0] for d in concatenated_data]
00232 ys = [d[1] for d in concatenated_data]
00233 if self.is_line or msg.type is PlotData.LINE:
00234 axes.plot(xs, ys,
00235 label=msg.label or self.topic_with_field_name)
00236 else:
00237 axes.scatter(xs, ys,
00238 label=msg.label or self.topic_with_field_name)
00239 if msg.fit_line or self.fit_line:
00240 X = np.array(msg.xs)
00241 Y = np.array(msg.ys)
00242 A = np.array([X, np.ones(len(X))])
00243 A = A.T
00244 a, b = np.linalg.lstsq(A, Y, rcond=-1)[0]
00245 axes.plot(X, (a*X+b), "g--", label="{0} x + {1}".format(a, b))
00246 if msg.fit_line_ransac or self.fit_line_ransac:
00247 model_ransac = linear_model.RANSACRegressor(
00248 linear_model.LinearRegression(), min_samples=2,
00249 residual_threshold=self.fit_line_ransac_outlier)
00250 X = np.array(msg.xs).reshape((len(msg.xs), 1))
00251 Y = np.array(msg.ys)
00252 model_ransac.fit(X, Y)
00253 line_X = X
00254 line_y_ransac = model_ransac.predict(line_X)
00255 if len(model_ransac.estimator_.coef_) == 1:
00256 coef = model_ransac.estimator_.coef_[0]
00257 else:
00258 coef = model_ransac.estimator_.coef_[0][0]
00259 if not isinstance(model_ransac.estimator_.intercept_, list):
00260 intercept = model_ransac.estimator_.intercept_
00261 else:
00262 intercept = model_ransac.estimator_.intercept_[0]
00263 axes.plot(
00264 line_X, line_y_ransac, "r--",
00265 label="{0} x + {1}".format(coef, intercept))
00266
00267 def update_plot(self):
00268 if not self._rosdata:
00269 return
00270 try:
00271 data_x, data_y = self._rosdata.next()
00272 except RosPlotException, e:
00273 rospy.logerr("Exception in subscribing topic")
00274 rospy.logerr(e.message)
00275 return
00276 if len(data_y) == 0:
00277 return
00278 axes = self.data_plot._canvas.axes
00279 axes.cla()
00280
00281
00282 latest_msg = data_y[-1]
00283 min_x = None
00284 max_x = None
00285 min_y = None
00286 max_y = None
00287 if isinstance(latest_msg, PlotData):
00288 data = [latest_msg]
00289 legend_size = 8
00290 no_legend = False
00291 elif isinstance(latest_msg, PlotDataArray):
00292 data = latest_msg.data
00293 legend_size = latest_msg.legend_font_size or 8
00294 no_legend = latest_msg.no_legend
00295 if latest_msg.min_x != latest_msg.max_x:
00296 min_x = latest_msg.min_x
00297 max_x = latest_msg.max_x
00298 if latest_msg.min_y != latest_msg.max_y:
00299 min_y = latest_msg.min_y
00300 max_y = latest_msg.max_y
00301 else:
00302 rospy.logerr(
00303 "Topic should be jsk_recognition_msgs/PlotData",
00304 "or jsk_recognition_msgs/PlotDataArray")
00305 for d in data:
00306 self.plot_one(d, axes)
00307 xs = add_list([d.xs for d in data])
00308 ys = add_list([d.ys for d in data])
00309 if min_x is None:
00310 min_x = min(xs)
00311 if max_x is None:
00312 max_x = max(xs)
00313 if min_y is None:
00314 min_y = min(ys)
00315 if max_y is None:
00316 max_y = max(ys)
00317 axes.set_xlim(min_x, max_x)
00318 axes.set_ylim(min_y, max_y)
00319
00320 if not no_legend and not self.no_legend:
00321 axes.legend(prop={'size': legend_size})
00322 axes.grid()
00323 if self.xtitle:
00324 axes.set_xlabel(self.xtitle)
00325 if self.ytitle:
00326 axes.set_ylabel(self.ytitle)
00327 self.data_plot._canvas.draw()
00328 buffer = StringIO()
00329 self.data_plot._canvas.figure.savefig(buffer, format="png")
00330 buffer.seek(0)
00331 img_array = np.asarray(bytearray(buffer.read()), dtype=np.uint8)
00332 if LooseVersion(cv2.__version__).version[0] < 2:
00333 iscolor = cv2.CV_LOAD_IMAGE_COLOR
00334 else:
00335 iscolor = cv2.IMREAD_COLOR
00336 img = cv2.imdecode(img_array, iscolor)
00337 self.pub_image.publish(self.cv_bridge.cv2_to_imgmsg(img, "bgr8"))
00338
00339
00340 class MatPlot2D(QWidget):
00341
00342 class Canvas(FigureCanvas):
00343 def __init__(self, parent=None):
00344 super(MatPlot2D.Canvas, self).__init__(Figure())
00345 self.axes = self.figure.add_subplot(111)
00346 self.figure.tight_layout()
00347 self.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Expanding)
00348 self.updateGeometry()
00349
00350 def resizeEvent(self, event):
00351 super(MatPlot2D.Canvas, self).resizeEvent(event)
00352 self.figure.tight_layout()
00353
00354 def __init__(self, parent=None):
00355 super(MatPlot2D, self).__init__(parent)
00356 self._canvas = MatPlot2D.Canvas()
00357 self._toolbar = NavigationToolbar(self._canvas, self._canvas)
00358 vbox = QVBoxLayout()
00359 vbox.addWidget(self._toolbar)
00360 vbox.addWidget(self._canvas)
00361 self.setLayout(vbox)
00362
00363 def redraw(self):
00364 pass
00365
00366 def clear(self):
00367 self._canvas.axes.cla()
00368 self._canvas.draw()