pydrake drawing frames in director

director_client_drake_frames.h

/**
Example usage:

DrakeLcm lcm;
PublishFramesToLcm("DRAKE_DRAW_FRAMES", {
    {"X_WF", Isometry3d::Identity()},
    {"X_WG", Isometry3d::Identity()},
});

 */

// Needs other includes...
#include "drake/lcmt_viewer_draw.hpp"

void PublishFramesToLcm(
    const std::string& channel_name,
    const std::unordered_map<std::string, Eigen::Isometry3d>& name_to_frame_map,
    drake::lcm::DrakeLcmInterface* lcm) {
  std::vector<Eigen::Isometry3d> poses;
  std::vector<std::string> names;
  for (const auto& pair : name_to_frame_map) {
    poses.push_back(pair.second);
    names.push_back(pair.first);
  }
  PublishFramesToLcm(channel_name, poses, names, lcm);
}

void PublishFramesToLcm(const std::string& channel_name,
                        const std::vector<Eigen::Isometry3d>& poses,
                        const std::vector<std::string>& names,
                        drake::lcm::DrakeLcmInterface* dlcm) {
  DRAKE_DEMAND(poses.size() == names.size());
  drake::lcmt_viewer_draw frame_msg{};
  frame_msg.timestamp = 0;
  int32_t vsize = poses.size();
  frame_msg.num_links = vsize;
  frame_msg.link_name.resize(vsize);
  frame_msg.robot_num.resize(vsize, 0);

  for (size_t i = 0; i < poses.size(); i++) {
    Eigen::Isometry3f pose = poses[i].cast<float>();
    // Create a frame publisher
    Eigen::Vector3f goal_pos = pose.translation();
    Eigen::Quaternion<float> goal_quat =
        Eigen::Quaternion<float>(pose.linear());
    frame_msg.link_name[i] = names[i];
    frame_msg.position.push_back({goal_pos(0), goal_pos(1), goal_pos(2)});
    frame_msg.quaternion.push_back(
        {goal_quat.w(), goal_quat.x(), goal_quat.y(), goal_quat.z()});
  }

  const int num_bytes = frame_msg.getEncodedSize();
  const size_t size_bytes = static_cast<size_t>(num_bytes);
  std::vector<uint8_t> bytes(size_bytes);
  frame_msg.encode(bytes.data(), 0, num_bytes);
  dlcm->Publish(
      "DRAKE_DRAW_FRAMES_" + channel_name, bytes.data(), num_bytes, {});
}

director_client_drake_frames.py

"""Example usage:

X_WF = Isometry3.Identity()
X_WG = Isometry3.Identity()  # or something else
draw_frame_args(
  X_WF=X_WF,
  X_WG=W_WG,
)
"""

import os

import lcm
from lcm import LCM
import numpy as np

from robotlocomotion import viewer_draw_t

from drake import (
    lcmt_viewer_draw,
    lcmt_viewer_geometry_data,
    lcmt_viewer_link_data,
    lcmt_viewer_load_robot,
)
from pydrake.util.eigen_geometry import Isometry3

def draw_frames(names, X_WF_list, suffix="_PY", lcm=None):
    if lcm is None:
        lcm = LCM()
    msg = viewer_draw_t()
    msg.num_links = len(names)
    msg.link_name = names
    msg.robot_num = [0] * len(names)
    for X_WF in X_WF_list:
        X_WF = Isometry3(X_WF)
        msg.position.append(X_WF.translation())
        msg.quaternion.append(X_WF.quaternion().wxyz())
    lcm.publish("DRAKE_DRAW_FRAMES" + suffix, msg.encode())


def draw_frame_args(**kwargs):
    if "suffix" in kwargs:
        suffix = kwargs.pop("suffix")
    else:
        suffix = "_PY"
    names, X_WF_list = zip(*kwargs.items())
    draw_frames(names, X_WF_list, suffix)