notebooks/dataset-visualize.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "# Dataset Visualization\n",
    "This notebooks allows to load and visualize a GEECO dataset using the same data loader function as used during model training."
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "'''imports'''\n",
    "import os\n",
    "from timeit import default_timer as timer\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "from mpl_toolkits.mplot3d import axes3d\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "\n",
    "from data.geeco_gym import pickplace_input_fn_v4\n",
    "from utils.plotting import create_image_grid"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "'''path setup'''\n",
    "root_path = os.environ['GEECO_ROOT']\n",
    "# dataset_dir = os.path.join(root_path, 'data', 'gym-push-pad2-cube2-v4')\n",
    "dataset_dir = os.path.join(root_path, 'data', 'gym-pick-pad2-cube2-v4')\n",
    "print(dataset_dir)"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "'''helper_fn'''\n",
    "def _frame_readout(d, step_idx=0, batch_idx=0):\n",
    "    step = d['step'][batch_idx][step_idx]\n",
    "    # arm: joint velocities\n",
    "    arm_jnt_qvel = [\n",
    "            'joint_qvel-robot0:shoulder_pan_joint',\n",
    "            'joint_qvel-robot0:shoulder_lift_joint',\n",
    "            'joint_qvel-robot0:upperarm_roll_joint',\n",
    "            'joint_qvel-robot0:elbow_flex_joint',\n",
    "            'joint_qvel-robot0:forearm_roll_joint',\n",
    "            'joint_qvel-robot0:wrist_flex_joint',\n",
    "            'joint_qvel-robot0:wrist_roll_joint',\n",
    "    ]\n",
    "    arm_qvel = np.zeros((len(arm_jnt_qvel), ))\n",
    "    for i, k in enumerate(arm_jnt_qvel):\n",
    "        arm_qvel[i] = d[k][batch_idx][step_idx]\n",
    "    # gripper: joint positions\n",
    "    gripper_jnt_qpos = [\n",
    "            'joint_qpos-robot0:r_gripper_finger_joint',\n",
    "            'joint_qpos-robot0:l_gripper_finger_joint',\n",
    "    ]\n",
    "    gripper_qpos = np.zeros((len(gripper_jnt_qpos), ))\n",
    "    for i, k in enumerate(gripper_jnt_qpos):\n",
    "        gripper_qpos[i] = d[k][batch_idx][step_idx]\n",
    "    # EE: qpos\n",
    "    ee_jnt_qpos = 'mocap_qpos-robot0:mocap'\n",
    "    ee_qpos = d[ee_jnt_qpos][batch_idx][step_idx][:3]  # only in xyz\n",
    "    # object: qpos\n",
    "    obj_jnt_qpos = 'object_qpos-object0:joint'\n",
    "    obj_qpos = d[obj_jnt_qpos][batch_idx][step_idx][:3]  # only in xyz\n",
    "    # data string\n",
    "#     data_str = \"#: %04d\\nvel: %s\\ngrp: %s\\nee: %s\\nobj: %s\" % (step, arm_qvel, gripper_qpos, ee_qpos, obj_qpos)\n",
    "    data_str = \"#: %04d\\ngrp: %s\\nee: %s\\nobj: %s\" % (step, gripper_qpos, ee_qpos, obj_qpos)\n",
    "    return data_str\n",
    "\n",
    "def _frame_readout_v2(d, step_idx=0, batch_idx=0):\n",
    "    step = d['step'][batch_idx][step_idx]\n",
    "    # arm: joint velocities\n",
    "    arm_qvel = d['vel_state'][batch_idx][step_idx]\n",
    "    # gripper: joint positions\n",
    "    gripper_qpos = d['grp_state'][batch_idx][step_idx]\n",
    "    # EE: qpos\n",
    "    ee_qpos = d['ee_state'][batch_idx][step_idx][:3]  # only in xyz\n",
    "    # object: qpos\n",
    "    obj_qpos = d['obj_state'][batch_idx][step_idx][:3]  # only in xyz\n",
    "    # targets / commands\n",
    "    grp_target = d['grp_target'][batch_idx][step_idx]\n",
    "    ee_target = d['ee_target'][batch_idx][step_idx][:3]  # only in xyz\n",
    "    # data string\n",
    "    data_str = \"#: %04d\\nx_grp: %s\\nx_ee: %s\\nx_obj: %s\\ny_grp: %s\\ny_ee: %s\" % (step, gripper_qpos, ee_qpos, obj_qpos, grp_target, ee_target)\n",
    "    return data_str\n",
    "\n",
    "def _feature_readout(feature, batch_idx=0, window_idx=0):\n",
    "    step = feature['step'][batch_idx][window_idx]\n",
    "    # state\n",
    "    vel_state = feature['vel_state'][batch_idx][window_idx]\n",
    "    grp_state = feature['grp_state'][batch_idx][window_idx]\n",
    "    ee_state = feature['ee_state'][batch_idx][window_idx][:3]  # only in xyz\n",
    "    obj_state = feature['obj_state'][batch_idx][window_idx][:3]  # only in xyz\n",
    "    # string readout\n",
    "    feature_str = \"#: %04d\\nx_grp: %s\\nx_ee: %s\\nx_obj: %s\" % (step, grp_state, ee_state, obj_state)\n",
    "    return feature_str\n",
    "\n",
    "def _feature_readout_v4(feature, batch_idx=0, window_idx=0):\n",
    "    step = feature['step'][batch_idx][window_idx]\n",
    "    # state\n",
    "    vel_state = feature['vel_state'][batch_idx][window_idx]\n",
    "    grp_state = feature['grp_state'][batch_idx][window_idx]\n",
    "    cmd = feature['cmd'][batch_idx][window_idx]\n",
    "    ee_state = feature['ee_state'][batch_idx][window_idx][:3]  # only in xyz\n",
    "    obj_state = feature['obj_state'][batch_idx][window_idx][:3]  # only in xyz\n",
    "    # string readout\n",
    "    feature_str = \"#: %04d\\ncmd: %s\\nx_ee: %s\\nx_obj: %s\" % (step, cmd, ee_state, obj_state)\n",
    "    return feature_str\n",
    "\n",
    "def _label_readout(label, batch_idx=0):\n",
    "    # targets / commands\n",
    "    grp_target = label['grp_target'][batch_idx]\n",
    "    ee_target = label['ee_target'][batch_idx][:3]  # only in xyz\n",
    "    # string readout\n",
    "    label_str = \"y_grp: %s\\ny_ee: %s\" % (grp_target, ee_target)\n",
    "    return label_str\n",
    "\n",
    "def _label_readout_v4(label, batch_idx=0):\n",
    "    # targets / commands\n",
    "    grp_target = label['grp_target'][batch_idx]\n",
    "    ee_target = label['ee_target'][batch_idx][:3]  # only in xyz\n",
    "    cmd = label['cmd'][batch_idx]\n",
    "    # string readout\n",
    "    label_str = \"cmd: %s\\ny_grp: %s\\ny_ee: %s\" % (cmd, grp_target, ee_target)\n",
    "    return label_str"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "'''input_fn'''\n",
    "split_name = 'debug'\n",
    "mode = 'eval'\n",
    "window_size = 4\n",
    "fetch_target = False\n",
    "shuffle_buffer = 32\n",
    "batch_size = 2\n",
    "dataset = pickplace_input_fn_v4(\n",
    "    dataset_dir=dataset_dir, split_name=split_name, mode=mode, window_size=window_size, fetch_target=fetch_target, shuffle_buffer=shuffle_buffer, batch_size=batch_size)\n",
    "iterator = dataset.make_one_shot_iterator()\n",
    "data = iterator.get_next()"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "'''session setup'''\n",
    "try:\n",
    "    sess.close()\n",
    "except:\n",
    "    pass\n",
    "sess = tf.InteractiveSession()"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "'''load a data batch'''\n",
    "np.set_printoptions(precision=2, suppress=True)\n",
    "t_start_fetch = timer()\n",
    "d = sess.run(data)\n",
    "t_end_fetch = timer()\n",
    "t_fetch = t_end_fetch - t_start_fetch\n",
    "print(\">>> Fetched one batch (size=%d, window=%d) in %.04f s!\" % (batch_size, window_size, t_fetch))\n",
    "feat, lbl = d\n",
    "print(\">>> Feature struct:\")\n",
    "for k, v in feat.items():\n",
    "    print(k, v.shape)\n",
    "print(\">>> Label struct:\")\n",
    "for k, v in lbl.items():\n",
    "    print(k, v.shape)"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "'''display a data batch'''\n",
    "if fetch_target:\n",
    "    grid_size = window_size + 1\n",
    "else:\n",
    "    grid_size = window_size\n",
    "\n",
    "for b_idx in range(batch_size):\n",
    "    \n",
    "    grid_rgb = create_image_grid(num_examples=grid_size, max_cols_per_row=10, tile_pad=1.2)  # window + target frame\n",
    "    for f_idx in range(window_size):  # display window\n",
    "        grid_rgb[f_idx].imshow(feat['rgb'][b_idx][f_idx])\n",
    "#         title = _feature_readout(feat, batch_idx=b_idx, window_idx=f_idx)\n",
    "        title = _feature_readout_v4(feat, batch_idx=b_idx, window_idx=f_idx)\n",
    "        grid_rgb[f_idx].set_title(title)\n",
    "    if fetch_target:\n",
    "        grid_rgb[window_size].imshow(feat['target_rgb'][b_idx])\n",
    "        grid_rgb[window_size].set_title('target frame\\n\\n\\n')\n",
    "    plt.show()\n",
    "    \n",
    "    grid_depth = create_image_grid(num_examples=grid_size, max_cols_per_row=10, tile_pad=1.2)  # window + target frame\n",
    "    for f_idx in range(window_size):\n",
    "        grid_depth[f_idx].imshow(np.squeeze(feat['depth'][b_idx][f_idx]))\n",
    "    if fetch_target:\n",
    "        grid_depth[window_size].imshow(np.squeeze(feat['target_depth'][b_idx]))\n",
    "    plt.show()\n",
    "    \n",
    "#     label_str = _label_readout(lbl, batch_idx=b_idx)\n",
    "    label_str = _label_readout_v4(lbl, batch_idx=b_idx)\n",
    "    print(label_str)"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [],
   "outputs": [],
   "metadata": {}
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}