director_script_util.py

"""Basic utilities for visualization scripts to be used via Director."""

from director.drakevisualizer import *


# TODO(eric.cousineau): Use RemoteTreeViewer or something like that.
class ModifiedDrakeVisualizer(DrakeVisualizer):
    """Extends DrakeVisualizer to subscribe to different channels."""

    # Ensure we are overriding the correct methods.
    assert hasattr(DrakeVisualizer, '_addSubscribers')
    assert hasattr(DrakeVisualizer, 'onViewerLoadRobot')

    def __init__(self, name, prefix, view):
        # See `director/.../mainwindowapp.py` for more information.
        # These must be set before base construction.
        self.name = name
        self._prefix = prefix
        DrakeVisualizer.__init__(self, view)

    def _addSubscribers(self):
        # Override drake visualizer's subscriptions.
        prefix = self._prefix
        self.subscribers.append(
            lcmUtils.addSubscriber(
                prefix + 'DRAKE_VIEWER_LOAD_ROBOT', lcmrl.viewer_load_robot_t,
                self.onViewerLoadRobot))
        self.subscribers.append(
            lcmUtils.addSubscriber(
                prefix + 'DRAKE_VIEWER_ADD_ROBOT', lcmrl.viewer_load_robot_t,
                self.onViewerAddRobot))
        self.subscribers.append(
            lcmUtils.addSubscriber(
                prefix + 'DRAKE_VIEWER_DRAW_ROBOT', lcmrl.viewer_draw_t,
                self.onViewerDraw))
        # N.B. Do not subscribe to point clouds or other things for now.

    def get_geom_pairs(self):
        """Returns link name and the corresponding geometry. For changing
        colors."""
        for robot in self.robots.itervalues():
            for link_name, link in robot.iteritems():
                if link_name != "world":
                    for geom in link.geometry:
                        yield (link_name, geom)


class RecoloredDrakeVisualizer(ModifiedDrakeVisualizer):
    """Visualizes traced detections, assigning a visually unique color to each
    individual object model."""
    def __init__(self, view):
        ModifiedDrakeVisualizer.__init__(
            self, "Tracked Models", prefix="TRACKED_MODELS_", view=view)

    def onViewerLoadRobot(self, msg):
        ModifiedDrakeVisualizer.onViewerLoadRobot(self, msg)
        # Override with dark green
        color = [0, 0, 0.8]
        for _, geom in self.get_geom_pairs():
            poly_data = geom.polyDataItem
            poly_data.setProperty('Color', color)