DataSet_Generation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Pick-And-Place Dataset Generation (Python2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## <div id=\"content\">Contents</div>\n",
    "1. <a href=\"#env\" style=\"color:green\">Setup Environment</a> \n",
    "2. <a href=\"#data\" style=\"color:green\">Dataset Generation</a>\n",
    "3. <a href=\"#process\" style=\"color:green\">Data Filtering and Processing</a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## <span id=\"env\">Setup Environment</span>\n",
    "<a href=\"#content\">Return to Contents</a>\n",
    "\n",
    "### <span id=\"setup-content\">Section Contents</span>\n",
    "\n",
    "1. <a href=\"#diagnostic\">Diagnostics</a>\n",
    "2. <a href=\"#packages\">Import Packages & Utilities</a>\n",
    "3. <a href=\"#launch\">Launch Simulation Environment</a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### <span id=\"diagnostic\"> Diagnostic Information</span>\n",
    "<span> <a href=\"#setup-content\">Return to Section Contents</a> </span>\n",
    "\n",
    "Notebook tested inside baxter.sh sudo environment. To access this use:\n",
    "```\n",
    "sudo ./baxter.sh sim\n",
    "cd ..\n",
    "jupyter notebook --allow-root\n",
    "```\n",
    "Otherwise Python path may not register properly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "%%bash\n",
    "if [ $(id -u) = 0 ];\n",
    "then\n",
    "    echo \"You are root\" # Should be root!\n",
    "else\n",
    "    echo \"You do not have the right priviliges, you must be root\"  \n",
    "fi\n",
    "python -V\n",
    "source /opt/ros/kinetic/setup.bash\n",
    "source ./workspace/devel/setup.bash"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Check that engines are running\n",
    "import ipyparallel as ipp # Run processes in parallel\n",
    "c = ipp.Client()\n",
    "print(\"Running Cluster with Engines: {}\".format(c.ids))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### <span id=\"packages\"> Python Packages </span>\n",
    "<span> <a href=\"#setup-content\">Return to Section Contents</a> </span>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import sys\n",
    "if sys.version_info[0] > 2:\n",
    "    raise Exception(\"Must be using Python 2\")\n",
    "\n",
    "from IPython.display import clear_output # Clear Output in cell programmatically\n",
    "from ipywidgets import IntProgress # Progress Bar\n",
    "import os, sys, shutil #Advanced File Manipulation\n",
    "from os import path\n",
    "import glob\n",
    "import matplotlib as mlt # Data Visualisation\n",
    "import matplotlib.pyplot as plt\n",
    "import cv2 # Interpret camera images\n",
    "import pandas, numpy # Data Manipulation\n",
    "from termcolor import colored # Colored Text\n",
    "\n",
    "import subprocess # Used to run scripts in the background\n",
    "import signal # Used to signal OS (e.g. to kill a process)\n",
    "import csv # easily read and write to CSV files (used for ground truths)\n",
    "import math # For trigonometry\n",
    "import json # Effeciently store data structures\n",
    "\n",
    "# Hyper-parameter Optimization Modules\n",
    "import types\n",
    "from types import * # Special Types (e.g. classes)\n",
    "import inspect # Inspect method/class signatures\n",
    "import pickle # Effeciently store classes to file\n",
    "\n",
    "\n",
    "import getpass # For sudo-based tasks (hides input)\n",
    "import re # Regular expressions\n",
    "import random # For variation\n",
    "import copy\n",
    "import time # Allow python to wait\n",
    "from enum import Enum\n",
    "\n",
    "# ROS Packages\n",
    "import rospy\n",
    "import rospkg\n",
    "import roslib\n",
    "import roslaunch\n",
    "import tf # Transformations\n",
    "from tf.transformations import *\n",
    "import std_msgs,gazebo_msgs,geometry_msgs,moveit_msgs\n",
    "from geometry_msgs.msg import Pose, Point, Quaternion\n",
    "from gazebo_msgs.msg import LinkStates\n",
    "from gazebo_msgs.srv import DeleteModel, SpawnModel, GetModelState, GetWorldProperties\n",
    "from std_msgs.msg import String\n",
    "from sensor_msgs.msg import Image # Listen to baxter Cameras\n",
    "import baxter_interface\n",
    "import moveit_commander\n",
    "import visual_interpreter\n",
    "import retina\n",
    "\n",
    "# >>> Utilities <<<\n",
    "def define_pose(x,y,z,**kwargs):\n",
    "    \"\"\"Returns a Pose based on x,y,z and optionally o_x,o_y,o_z and o_w\"\"\"\n",
    "    pose_target = geometry_msgs.msg.Pose()\n",
    "    pose_target.position.x = x\n",
    "    pose_target.position.y = y\n",
    "    pose_target.position.z = z\n",
    "\n",
    "    pose_target.orientation.x = kwargs.get(\"o_x\",0.0)\n",
    "    pose_target.orientation.y = kwargs.get(\"o_y\",1.0)\n",
    "    pose_target.orientation.z = kwargs.get(\"o_z\",0.0)\n",
    "    pose_target.orientation.w = kwargs.get(\"o_w\",0.0)\n",
    "    return pose_target\n",
    "\n",
    "def get_rpy_quaternion(roll,pitch,yaw):\n",
    "    \"\"\"Converts human readable roll (about x-axis), pitch (about y-axis)\n",
    "    and yaw (about z-axis) format to a 4D quaternion\"\"\"\n",
    "    origin, xaxis, yaxis, zaxis = (0, 0, 0), (1, 0, 0), (0, 1, 0), (0, 0, 1)\n",
    "\n",
    "    qz = quaternion_about_axis(yaw, zaxis)\n",
    "    qy = quaternion_about_axis(pitch, yaxis)\n",
    "    qx = quaternion_about_axis(roll, xaxis)\n",
    "\n",
    "    q = quaternion_multiply(qz, qy)\n",
    "    q = quaternion_multiply(q, qx)\n",
    "    \n",
    "    return q\n",
    "\n",
    "def isclose(a,b,rel_tol=1e-9,abs_tol=0.0):\n",
    "    \"\"\"Returns whether 2 floats are equal with a relative and absolute \n",
    "    tolerance to prevent problems from floating point precision\"\"\"\n",
    "    return abs(a-b) <= max( rel_tol * max(abs(a), abs(b)), abs_tol )\n",
    "\n",
    "class color:\n",
    "    \"\"\"Defines special characters for decorating print statements\"\"\"\n",
    "    PURPLE = '\\033[95m'\n",
    "    CYAN = '\\033[96m'\n",
    "    DARKCYAN = '\\033[36m'\n",
    "    BLUE = '\\033[94m'\n",
    "    GREEN = '\\033[92m'\n",
    "    YELLOW = '\\033[93m'\n",
    "    RED = '\\033[91m'\n",
    "    BOLD = '\\033[1m'\n",
    "    UNDERLINE = '\\033[4m'\n",
    "    END = '\\033[0m'\n",
    "    \n",
    "class label(Enum):\n",
    "    \"\"\"Defines labels for classifying pick-and-place images\"\"\"\n",
    "    GRASP = 0\n",
    "    BAD = 1\n",
    "    NOT_READY = 2\n",
    "    EMPTY = 3"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### <span id=\"launch\">Launch</span>\n",
    "<span> <a href=\"#setup-content\">Return to Section Contents</a> </span>\n",
    "\n",
    "Setup the appropriate tools for simulating Baxter"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Terminal 1** (Gazebo)\n",
    "> roslaunch baxter_gazebo baxter_world.launch\n",
    "\n",
    "Wait for message \"... Gravity Compensation was turned off\"\n",
    "\n",
    "**Terminal 2** (Joint Publishing)\n",
    "> rosrun baxter_tools enable_robot.py -e\n",
    "\n",
    "> rosrun baxter_interface joint_trajectory_action_server.py\n",
    "\n",
    "Note the joint trajectory action server may require sudo priviliges\n",
    "\n",
    "**Terminal 3** (Rviz)\n",
    "> roslaunch baxter_moveit_config demo_baxter.launch left_electric_gripper:=true right_electric_gripper:=true\n",
    "\n",
    "Motion planning will ALWAYS fail if electric grippers are not enabled\n",
    "\n",
    "**Terminal 4** (Gazebo2Rviz)\n",
    "> roslaunch gazebo2rviz gazebo2rviz.launch\n",
    "\n",
    "Continuosly transfers objects in Gazebo Simulation to Rviz"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "%%px --targets 0\n",
    "import rospkg, rospy, os, roslaunch \n",
    "\n",
    "gazebo_path = rospkg.RosPack().get_path('baxter_gazebo')\n",
    "launch_file_path = os.path.join(gazebo_path, \"launch\", \"baxter_world.launch\")\n",
    "\n",
    "uuid = roslaunch.rlutil.get_or_generate_uuid(None, False)\n",
    "roslaunch.configure_logging(uuid)\n",
    "launch = roslaunch.parent.ROSLaunchParent(uuid, [launch_file_path])\n",
    "launch.start()\n",
    "rospy.loginfo(\"Started Gazebo\")\n",
    "print(\"{}Gazebo{} now running on engine 0\".format(color.PURPLE,color.END))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "%%px --targets 1\n",
    "import rospkg, rospy, os, roslaunch\n",
    "\n",
    "package = 'baxter_tools'\n",
    "executable = 'enable_robot.py'\n",
    "node = roslaunch.core.Node(package, executable, args='-e')\n",
    "\n",
    "launch = roslaunch.scriptapi.ROSLaunch()\n",
    "launch.start()\n",
    "\n",
    "process = launch.launch(node)\n",
    "\n",
    "package = 'baxter_interface'\n",
    "executable = 'joint_trajectory_action_server.py'\n",
    "node = roslaunch.core.Node(package, executable)\n",
    "\n",
    "launch = roslaunch.scriptapi.ROSLaunch()\n",
    "launch.start()\n",
    "process = launch.launch(node)\n",
    "print(\"{}Joint Trajectory Action Server{} now running on engine 1\".format(color.PURPLE,color.END))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "%%px --targets 2\n",
    "import rospkg, rospy, os, roslaunch\n",
    "rviz_path = rospkg.RosPack().get_path('baxter_moveit_config')\n",
    "launch_file_path = os.path.join(rviz_path, \"launch\", \"demo_baxter.launch\")\n",
    "\n",
    "#rospy.init_node('rviz', anonymous=True)\n",
    "uuid = roslaunch.rlutil.get_or_generate_uuid(None, False)\n",
    "roslaunch.configure_logging(uuid)\n",
    "launch = roslaunch.parent.ROSLaunchParent(uuid, [launch_file_path])\n",
    "launch.start()\n",
    "rospy.loginfo(\"Started Rviz\")\n",
    "print(\"{}Rviz{} now running on engine 1\".format(color.PURPLE,color.END))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "%%px --targets 3\n",
    "# Note: Optional, transfers models from Gazebo to Rviz\n",
    "import rospkg, rospy, os, roslaunch\n",
    "\n",
    "gazebo2rviz_path = rospkg.RosPack().get_path('gazebo2rviz')\n",
    "launch_file_path = os.path.join(gazebo2rviz_path,\n",
    "                                \"launch\", \"gazebo2rviz.launch\")\n",
    "\n",
    "#rospy.init_node('gazebo2rviz', anonymous=True)\n",
    "uuid = roslaunch.rlutil.get_or_generate_uuid(None, False)\n",
    "roslaunch.configure_logging(uuid)\n",
    "launch = roslaunch.parent.ROSLaunchParent(uuid, [launch_file_path])\n",
    "\n",
    "launch.start()\n",
    "rospy.loginfo(\"Started Gazebo2Rviz\")\n",
    "print(\"{}Gazebo2Rviz{} now running on engine 1\".format(color.PURPLE,color.END))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "%%px --targets 0\n",
    "launch.shutdown()  # Shutdown Gazebo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "%%px --targets 1\n",
    "process.stop()  # Stop Joint Trajectory Action Server"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "%%px --targets 2\n",
    "launch.shutdown()  # Shutdown Rviz"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "%%px --targets 3 \n",
    "launch.shutdown()  # Shutdown Gazebo2Rviz"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## <span id=\"data\">Dataset Generation</span>\n",
    "<a href=\"#content\">Return to Contents</a>\n",
    "\n",
    "### <span id=\"dataset-content\">Section Contents</span>\n",
    "\n",
    "1. <a href=\"#spawn\">Object Spawning</a>\n",
    "    1. <a href=\"#spawn-table\">Spawn Table</a>\n",
    "    2. <a href=\"#spawn-block\">Spawn Block</a>\n",
    "    3. <a href=\"#spawn-custom\">Spawn Custom Object</a>\n",
    "2. <a href=\"#control\">Dataset Generation Master</a>\n",
    "    1. <a href=\"#control-manual\">Manual</a>\n",
    "    2. <a href=\"#control-automatic\">Automatic</a>\n",
    "3. <a href=\"#visual\">Visualiation Tools</a>\n",
    "    1. <a href=\"#visual-nodes\">Display ROS nodes</a>\n",
    "    2. <a href=\"#visual-camera\">Get live camera feed</a>\n",
    "\n",
    "In order to generate an appropriate dataset for Baxter to learn a generalized form of the pick-and-place task, a large set of examples is required. Since collecting data manually is resource-intensive, this has been automated as much as possible. Here hard-coded pick-and-place methods are used to move the robot to pick-up an object (of known location) in simulation. During this process, the video feed from Baxter's wrist ais recorded alongside useful metadata (pose, labels, etc.) under a clear naming scheme. \n",
    "\n",
    "The resulting dataset will contain sequences of images (ordered by their timestamp) where Baxter's wrist is seen to gradually approach the object until it is grasped. This data can be generated for a wide array of virtual objects to aid generalization, as the aim is to be able to pick-up unknown but gripper-compatible objects. \n",
    "\n",
    "The accompanying 'reward.csv' file contains all relevant metadata of frames from a specific iteration. Each iteration corresponds to 1 complete align-approach-grasp pipeline. The retract stage, which is used to determine whether the grasp was successful (i.e. does the object move along with the gripper), is not recorded as this information is redundant. The retino-cortical mapping will applied seperately in the next section to each of the filtered frames. Incorrect approaches, for instance where the motion planning fails, are also saved if these prove to make the final model more fault-tolerant.\n",
    "\n",
    "Note that in the context of this project, only rigid objects are considered. However, there is room to expand the dataset to include deformable and articulated objects."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### <span id=\"spawn\"> Object Spawning </span>\n",
    "In order to switch out different objects and avoid overlap, the ObjectSpawner class provides the methods needed to spawn any object from the ~/.gazebo/models directory or the baxter_sim_examples ROS package.\n",
    "\n",
    "It is possible to select any custom model via the private '_choose_model' interface, which is activated when no model_name is provided. This will recursively walk through the ~/.gazebo/models directory, allowing objects to be placed together in groups called 'categories'. Any selected models will be copied up to the root ~/.gazebo/models folder to allow Gazebo2Rviz to transfer the object from Gazebo to Rviz the object. An example model set can be found at [3DGems](http://data.nvision2.eecs.yorku.ca/3DGEMS/). \n",
    "\n",
    "BoundingBoxes are included to allow multiple objects to be spawned on the same table without overlap. However, since the dataset generator only spawns one at a time this can safely be ignored for most experiments.\n",
    "\n",
    "It is possible to use additional model directories by specifying their full path (using os.path). However, these must be included in the Gazebo_Model_Path environmental variable using .bashrc:\n",
    "\n",
    "```\n",
    "export GAZEBO_MODEL_PATH=[YOUR_NEW_PATH]/models:$GAZEBO_MODEL_PATH\n",
    "``` (See [Question](https://robotics.stackexchange.com/questions/1170/where-does-gazebo-set-the-gazebo-model-path-environment-variable))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "class BoundingBox(object):\n",
    "    \"\"\"Defines a cuboid volume to represent an object's bounding volume\n",
    "    If only x is supplied, constructs a x*x*x cube\n",
    "    If y is also supplied, constructs a x*y*x cuboid\n",
    "    If y and z are supplied, constructs a x*y*z cuboid\"\"\"\n",
    "    def __init__(self, x, y=None, z=None):  # x,y,z\n",
    "        self.x = x\n",
    "        self.y = x if y is None else y\n",
    "        self.z = x if z is None and y is None else (y if z is None else z)\n",
    "\n",
    "\n",
    "class ObjectSpawner(object):\n",
    "    \"\"\"Provides convenient API for spawning any objects in the Gazebo model directory\n",
    "    \n",
    "    Attributes:\n",
    "        -- Provided at initalization --\n",
    "        reference_frame: coordinates are used relative to ref frame [default: 'world']\n",
    "        bounding_box: defines a bounding volume for the object [default: 0.04*0.04*0.04]\n",
    "        model_name: if not provided, interactively choose model [default: None]\n",
    "        model_directory: if not provided, use default directory [default: None]\n",
    "        verbose: How much daignostic info to print [default: 1]\n",
    "        -- Generated --\n",
    "        nested_dir: Gazebo models can only be used if placed at root of Gazebo model directory\"\"\"\n",
    "    instances = {}\n",
    "\n",
    "    def __init__(self, reference_frame=\"world\", bounding_box=BoundingBox(0.04),\n",
    "                 model_name=None, model_directory=None,\n",
    "                 verbose=1):\n",
    "        self.reference_frame = reference_frame\n",
    "        self.bounding_box = bounding_box\n",
    "\n",
    "        self.default_dir = path.join(os.getcwd(), \"workspace\", \"src\",\n",
    "                                     \"baxter_simulator\",\n",
    "                                     \"baxter_sim_examples\", \"models\")\n",
    "        if model_directory is None:\n",
    "            self.model_dir = self.default_dir\n",
    "        else:\n",
    "            self.model_dir = model_directory\n",
    "\n",
    "        if model_name is None:\n",
    "            self.model_name = self._choose_model()\n",
    "        else:\n",
    "            self.model_name = model_name\n",
    "\n",
    "        # True if model is found in root model_path\n",
    "        if not hasattr(self, \"nested_dir\"):\n",
    "            self.nested_dir = self.model_dir\n",
    "\n",
    "        self.verbose = verbose\n",
    "\n",
    "    def spawn_model(self, pose):\n",
    "        \"\"\"Spawns a model at the target pose (position and orientation).\n",
    "        \n",
    "        If Spawning service failed: Returns None\n",
    "        Else: Returns name of generated object and its real pose\"\"\"\n",
    "        \n",
    "        model_path = path.join(self.nested_dir, self.model_name, \"model.sdf\")\n",
    "        with open(model_path, \"r\") as f:\n",
    "            xml = f.read()\n",
    "\n",
    "        if self.verbose >= 3:\n",
    "            print(\">> Loaded Model {} <<<\\n{}\".format(self.model_name, xml))\n",
    "\n",
    "        rospy.wait_for_service('/gazebo/spawn_sdf_model')\n",
    "        try:\n",
    "            spawn_sdf = rospy.ServiceProxy('/gazebo/spawn_sdf_model',\n",
    "                                           SpawnModel)\n",
    "\n",
    "            # Enforce unique names\n",
    "            if len(self.instances) == 0:\n",
    "                name = self.model_name\n",
    "            else:\n",
    "                \"{}_{}\".format(self.model_name, len(self.instances))\n",
    "\n",
    "            spawn_sdf(name, xml.replace(\"\\n\", \"\"),\n",
    "                      \"/\", pose, self.reference_frame)\n",
    "\n",
    "            self.instances[name] = pose\n",
    "            return name, pose\n",
    "        except rospy.ServiceException, e:\n",
    "            rospy.logerr(\"Spawn {} failed: {}\".format(self.model_name, e))\n",
    "            return None\n",
    "\n",
    "    def _choose_model(self):\n",
    "        \"\"\"Interactively select model by traversing Gazebo model subdirectories\n",
    "        Copies chosen model to top of Gazebo Model path so that it is accessible\n",
    "        Returns the chosen model name\"\"\"\n",
    "        # Only descend into N levels of file path tree\n",
    "        def walklevel(some_dir, level=1):\n",
    "            some_dir = some_dir.rstrip(os.path.sep)\n",
    "            assert os.path.isdir(some_dir)\n",
    "            num_sep = some_dir.count(os.path.sep)\n",
    "            for root, dirs, files in os.walk(some_dir):\n",
    "                yield root, dirs, files\n",
    "                num_sep_this = root.count(os.path.sep)\n",
    "                if num_sep + level <= num_sep_this:\n",
    "                    del dirs[:]\n",
    "\n",
    "        # Get all VALID folders with mesh options\n",
    "        def get_listdir(path, verbose=True):\n",
    "            valid_dirs = {}\n",
    "            count = 0\n",
    "            for dir_name, child_dirs, files in walklevel(path, level=1):\n",
    "                if child_dirs in [[], ['meshes'], ['materials', 'meshes']]:\n",
    "                    continue\n",
    "                else:\n",
    "                    if verbose: print(\"{}: {}\".format(count, dir_name))\n",
    "                    valid_dirs[count] = dir_name\n",
    "                    count += 1\n",
    "            return valid_dirs\n",
    "\n",
    "        # Ask user to pick directory by number\n",
    "        def select_by_number(options, option_desc):\n",
    "            while True:\n",
    "                resp = raw_input(\"Type no of desired {}\".format(option_desc))\n",
    "                if re.match(\"\\d+\", resp):\n",
    "                    if int(resp) in options.keys():\n",
    "                        return options[int(resp)]  # Selected option\n",
    "                print(\"Please pick desired {} by number\".format(option_desc))\n",
    "\n",
    "        if self.model_dir != self.default_dir:\n",
    "            valid_folders = get_listdir(self.model_dir)\n",
    "            category = select_by_number(valid_folders, \"Category\")\n",
    "            self.nested_dir = path.join(self.model_dir, category)\n",
    "            valid_models = os.listdir(self.nested_dir)\n",
    "        else:  # Default Directory does not have categories\n",
    "            valid_models = os.listdir(self.model_dir)\n",
    "\n",
    "        # Final Choice: One Specific Model\n",
    "        valid_models = dict(enumerate(valid_models, start=0))\n",
    "        for key, value in valid_models.items():\n",
    "            print(\"{}: {}\".format(key, value))\n",
    "        model_name = select_by_number(valid_models, \"Model\")\n",
    "\n",
    "        if self.verbose >= 1: print(\"You selected: {}\".format(model_name))\n",
    "\n",
    "        return model_name\n",
    "\n",
    "    def spawn_on_table(self, spawn=True, random_face=True,\n",
    "                       table_bbox = BoundingBox(0.32, 0.5, 0.7825),\n",
    "                       table_offset = Point(0.38, -0.1, 0.0)):\n",
    "        \"\"\"Spawns an object on a table's surface with random location and yaw. \n",
    "        \n",
    "        Attributes:\n",
    "            spawn: Whether to actually spawn the object [default:True]\n",
    "            random_face: Whether to vary which face is front-facing (24 possibilities) [default: True]\n",
    "            table_bbox: Defines table's bounding box, ensure this is reachable [default: 0.32*0.5*0.7825]\n",
    "            table_offset: Where table is placed relative to world origin (0,0,0)\n",
    "        Returns:\n",
    "            Unique name in simulation, overhead RPY and real pose\n",
    "        \"\"\"\n",
    "        x = random.uniform(table_offset.x, table_offset.x + table_bbox.x)  # Default: 0.38->0.7\n",
    "        y = random.uniform(table_offset.y, table_offset.y + table_bbox.y)  # Default: -0.1->0.4)\n",
    "        z = table_offset.z + table_bbox.z  # Default: 0.7825\n",
    "        fixed_rotations = numpy.arange(0, 2*numpy.pi, numpy.pi/2)\n",
    "        if random_face:\n",
    "            roll, pitch, yaw = random.choice(fixed_rotations), \\\n",
    "                               random.choice(fixed_rotations), \\\n",
    "                               random.uniform(0, 2*numpy.pi)\n",
    "        else:\n",
    "            roll, pitch, yaw = numpy.pi, 0, random.uniform(0, 2*numpy.pi)\n",
    "\n",
    "        q = get_rpy_quaternion(roll, pitch, yaw)\n",
    "\n",
    "        pose = define_pose(x=x, y=y, z=z, o_x=q[0],\n",
    "                           o_y=q[1], o_z=q[2], o_w=q[3])\n",
    "\n",
    "        # ToDo: Use bounding boxes to avoid overlap\n",
    "\n",
    "        name = None\n",
    "        if spawn:\n",
    "            name, _ = self.spawn_model(pose)\n",
    "\n",
    "            # Make model available at top level of Model directory\n",
    "            root_model_dir = path.join(self.model_dir, self.model_name)\n",
    "            if not path.exists(root_model_dir):\n",
    "                nested_model_dir = path.join(self.nested_dir, self.model_name)\n",
    "                shutil.copytree(nested_model_dir, root_model_dir)\n",
    "\n",
    "            if self.verbose >= 2:\n",
    "                print(\"Spawned {} at (x={},y={},z={})\".format(self.model_name,\n",
    "                                                              x, y, z))\n",
    "\n",
    "        return name if name is not None else \"EMPTY\", (numpy.pi, 0, yaw), pose\n",
    "\n",
    "    def delete_models(self, instance_name=None):\n",
    "        \"\"\"Use class-wide instances attribute to delete models. \n",
    "        instance_name: if specified, deletes specific model [default: None]\"\"\"\n",
    "        try:\n",
    "            delete_model = rospy.ServiceProxy(\"/gazebo/delete_model\",\n",
    "                                              DeleteModel)\n",
    "            if instance_name is None:\n",
    "                for instance_name in self.instances.keys():\n",
    "                    delete_model(instance_name)\n",
    "                    del self.instances[instance_name]\n",
    "            elif instance_name in self.instances:\n",
    "                delete_model(instance_name)\n",
    "                del self.instances[instance_name]\n",
    "        except rospy.ServiceException as e:\n",
    "            rospy.loginfo(\"Delete Model failed: {}\\n{}\".format(e, instances))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Object Spawner signature:\n",
    "* **reference_frame** (OPTIONAL): Reference Frame in which to spawn the object. Default: \"world\"\n",
    "* **bounding_box** (OPTIONAL): Defines overall size of the model using BoundingBox class. Default: BoundingBox(0.04)\n",
    "* **model_name** (OPTIONAL): General name of model, if not provided opens user selection wizard. default: none\n",
    "* **model_directory** (OPTIONAL): Root directory in ROS path where models can be found, uses WORKSPACE_PATH/ baxter_sim_examples/models if not provided. Default: None\n",
    "* **verbose** (OPTIONAL): Whether to print extra text. Default: False\n",
    "The ObjectSpawner keeps track of ALL spawned_objects across every ObjectSpawner instance generated. This means if you use delete_models() it will remove ALL objects in Gazebo that have been spawned using the Object Spawner.\n",
    "\n",
    "#### <span id=\"spawn-table\">Spawn Table</span>\n",
    "<span> <a href=\"#dataset-content\"> Return to Section Contents </a> </span>\n",
    "    \n",
    "If the table spawner fails, you may need to generate an sdf file first using [pysdf](http://wiki.ros.org/pysdf).\n",
    "\n",
    "The DataSet generator has a table_spawner built in, hence only use this cell for manual runs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "table_spawner = ObjectSpawner(reference_frame=\"world\",\n",
    "                              bounding_box=BoundingBox(0.913, 0.913, 0.7825),\n",
    "                              model_name=\"cafe_table\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "TableTop is 0.914m by 0.914m placed at x=0.8, y=0.0,z=0.7747. \n",
    "Hence valid area for spawning objects:\n",
    "* x = 0.343 m to x = 1.257 m\n",
    "* y = -0.257 m to x = 0.257 m\n",
    "* z 0.7825"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "table_spawner.spawn_model(Pose(position=Point(x=0.0, y=1.0, z=0.0)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "table_spawner.spawn_model(Pose(position=Point(x=0.8, y=0.0, z=0.0)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "table_spawner.delete_models()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### <span id=\"spawn-block\">Spawn Block</span>\n",
    "<span> <a href=\"#dataset-content\"> Return to Section Contents </a> </span>\n",
    "\n",
    "If the block spawner fails, you may need to generate an sdf file first using [pysdf](http://wiki.ros.org/pysdf).\n",
    "\n",
    "The DataSet generator has a spawner (default=block) built in, hence only use this cell for manual runs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "block_spawner = ObjectSpawner(reference_frame=\"world\",\n",
    "                              bounding_box=BoundingBox(0.04),\n",
    "                              model_name=\"block\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "block_spawner.spawn_model(Pose(position=Point(x=0.8, y=0.0, z=0.7825)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "block_spawner.delete_models()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### <span id=\"spawn-custom\"> Spawn Custom model Block </span>\n",
    "<span> <a href=\"#dataset-content\"> Return to Section Contents </a> </span>\n",
    "\n",
    "The DataSet generator has a custom spawner (default=block) built in, hence only use this cell for manual runs.\n",
    "\n",
    "Please follow instructions for [3DGEMS](http://data.nvision2.eecs.yorku.ca/3DGEMS/) under the Project_Setup manual to get the appropriate sdf file here. <br>\n",
    "Not all models will have their physics properties specified."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "custom_model_path = path.join(os.path.expanduser('~'), \".gazebo\", \"models\")\n",
    "custom_spawner = ObjectSpawner(reference_frame=\"world\",\n",
    "                               model_directory=custom_model_path)\n",
    "clear_output()\n",
    "print(\"Created {} spawner\".format(custom_spawner.model_name))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "custom_spawner.spawn_on_table()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "custom_spawner.delete_models()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### <span id=\"control\">Dataset Generation Master </span>\n",
    "<span> <a href=\"#dataset-content\">Return to Section Contents</a> </span>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "class MoveBaxter(object):\n",
    "    \"\"\"\n",
    "    MoveBaxter automates simple pick-and-place tasks using the Baxter Robot.\n",
    "\n",
    "    Attributes:\n",
    "        -- Provided at initialization --\n",
    "        id (string): Used for identifying the MoveBaxter instance when generating datasets\n",
    "        limb_name (string): Choice between 'left' and 'right' MoveGroup [default: 'left']\n",
    "        model_pattern (string): Name of spawnable model (must be in .gazebo/models or baxter_sim_examples/models) [default:'block']\n",
    "        max_attempts (int): Maximum number of attempts to try to find a motion plan [default: 100]\n",
    "        verbose (int): Print extra diagnostic information (0-None,1-Minimal,2-Moderate, 3-All)\n",
    "        planning_time (float): Time in seconds allowed per planning attempt [default: 1.0 seconds]\n",
    "        tolerance (float): Variation allowed in planning target joint, position and orientation of end effector\n",
    "\n",
    "        -- Generated --\n",
    "        spawner (ObjectSpawner): Class used to spawn Object of model_pattern on Table\n",
    "        gripper (baxter_interface): Used to open and close Baxter's gripper\n",
    "        display_trajectory_publisher (rospy Publisher): Publishes motion plans to Rviz for visualisation\n",
    "        get_model_state (rospy Service): Allows for pose retrieval of specific links\n",
    "        tf_listener (TransformListener): Provides reliable way to get pose of MoveGroup\n",
    "        robot (RobotComander): For diagnostics, allows for a complete recording of Robot Status\n",
    "        group (MoveGroupCommander): Target arm to be used for performing pick_and_place task\n",
    "        gt_path (string): Path to ground_truth csv file\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, id, limb_name=\"left\", model_pattern=\"block\", max_attempts=100,\n",
    "                 planning_time=1.0, tolerance=0.5, verbose=1):\n",
    "        self.id = id\n",
    "        self.limb_name = limb_name\n",
    "        self.verbose = verbose\n",
    "        self.max_attempts = max_attempts\n",
    "\n",
    "        # Which Model to Spawn on Table for picking\n",
    "        self.model_pattern = model_pattern\n",
    "        self.spawner = ObjectSpawner(reference_frame=\"world\",bounding_box=BoundingBox(0.04),\n",
    "                              model_name=model_pattern)\n",
    "        self.table_spawner = ObjectSpawner(reference_frame=\"world\",bounding_box=BoundingBox(0.913,0.913,0.7825),\n",
    "                              model_name=\"cafe_table\")\n",
    "\n",
    "        self.gripper = baxter_interface.Gripper(limb_name)\n",
    "\n",
    "        baxter_interface.RobotEnable(baxter_interface.CHECK_VERSION).state().enabled\n",
    "\n",
    "        if verbose >= 2: print(\"Enabled Robot\")\n",
    "\n",
    "        self.stage_pub = rospy.Publisher('pick_stage_publisher', String, queue_size=10)\n",
    "        self.display_trajectory_publisher = rospy.Publisher(\"/move_group/display_planned_path\",\n",
    "                                                            moveit_msgs.msg.DisplayTrajectory,\n",
    "                                                            queue_size=20)\n",
    "        self.get_model_state = rospy.ServiceProxy('/gazebo/get_model_state', GetModelState)\n",
    "        self.tf_listener = tf.TransformListener()\n",
    "\n",
    "        self.robot = moveit_commander.RobotCommander()\n",
    "\n",
    "        self.tolerance = tolerance\n",
    "        self.group = moveit_commander.MoveGroupCommander(\"left_arm\")\n",
    "        self.group.allow_replanning(True)\n",
    "        self.group.set_planning_time(planning_time) # Seconds\n",
    "        self.group.set_goal_tolerance(tolerance) # Joint, Position, Orientation \n",
    "        self.group.set_goal_position_tolerance(tolerance/10)\n",
    "\n",
    "        self.gt_path=os.path.join(\"workspace\",\"data\",\"{}_ground_truths.csv\".format(self.id))\n",
    "\n",
    "        if verbose >= 1: print(\"Successfully initialized MoveBaxter\")\n",
    "\n",
    "    def get_diagnostics(self):\n",
    "        \"\"\"\n",
    "        Provides a End Effector, Tolerance and (Optionally) Robot Status update\n",
    "        \n",
    "        Parameters:\n",
    "            verbose (bool): Override class verbosity to print complete Robot State [default: False]\n",
    "        \"\"\"\n",
    "        print(\"[MoveBaxter] Planning Frame: {}\".format(self.group.get_planning_frame()))\n",
    "        print(\"[MoveBaxter] End Effector: {}\\n{}\".format(self.group.get_end_effector_link(),\n",
    "                                                         self.get_current_pose()))\n",
    "        print(\"[MoveBaxter] Goal Tolerance:{}\".format(self.group.get_goal_tolerance()))\n",
    "        if self.verbose >= 3: print(\">>> Current Robot State <<<\\n{}\".format(self.robot.get_current_state()))\n",
    "\n",
    "    def get_object_state(self, model_name=None):\n",
    "        \"\"\"\n",
    "        Try to get status of a object in simulation.\n",
    "        \n",
    "        Parameters:\n",
    "            model_name (String): if None search all spawned instances [default:None]\n",
    "            \n",
    "        Returns:\n",
    "            GetModelState msg (if object state was found)\n",
    "            None (if failed to get object state)\"\"\"\n",
    "        rospy.wait_for_service(\"/gazebo/get_model_state\")\n",
    "        success = False\n",
    "        if model_name is None:\n",
    "            for i in range(len(self.spawner.instances)):\n",
    "                # If multiple instances of a model are spawned, Gazebo appends _NUM \n",
    "                model_name = self.model_pattern + (\"_{}\".format(i) if i > 0 else \"\")\n",
    "                result = self.get_model_state(model_name,\"\")\n",
    "                success = result.success\n",
    "                if success: break\n",
    "        else:\n",
    "            result = self.get_model_state(model_name, \"\")\n",
    "            if result is not None: success = result.success\n",
    "        if success is False:\n",
    "            return None\n",
    "        else:\n",
    "            return result # GetModelState Msg\n",
    "\n",
    "    def grasp_verification(self, original_pose, is_full):\n",
    "        \"\"\"\n",
    "        Checks whether target_object has changed its z-position as a result of a pick-attempt\n",
    "        \n",
    "        Parameters:\n",
    "            original_pose (Pose): Pose of object when it was originally spawned\n",
    "            \n",
    "        Returns:\n",
    "            None (if target link could not be found)\n",
    "            False (if object did not move vertically or fell)\n",
    "            True (if object did move vertically)\n",
    "            old_z, new_z (float, position along z-axis. Allows for cross-checking)\n",
    "        \"\"\"\n",
    "        new_z, old_z = None,original_pose.position.z\n",
    "        obj_z = self.spawner.bounding_box.z\n",
    "\n",
    "        # Empty state\n",
    "        if not is_full: return False, original_pose, original_pose\n",
    "\n",
    "        res = self.get_object_state()\n",
    "        if res is None: return None, None, None\n",
    "\n",
    "        # Compare original with new\n",
    "        new_pose = res.pose\n",
    "        new_z = new_pose.position.z\n",
    "        if isclose(old_z, new_z, rel_tol=1e-9, abs_tol=self.tolerance/10) or new_z < old_z + obj_z:\n",
    "            return False, original_pose, new_pose\n",
    "        else:\n",
    "            return True, original_pose, new_pose\n",
    "\n",
    "    def object_in_view(self, model_name = None, table_z = 0.7825):\n",
    "        \"\"\"\n",
    "        Checks whether a specific object is in the camera's viewing frustum.\n",
    "        Disclaimer: Not currently accurate for objects at distances <= 0.1\n",
    "        \n",
    "        Parameters:\n",
    "            model_name (String): If not provided, find any spawned instance [default: None]\n",
    "            table_z (float): Height of table [default: 0.7825]\n",
    "            \n",
    "        Returns:\n",
    "            True (if object is inside viewing frustum)\n",
    "            False (if object is outside of viewing frustum)\n",
    "            None (if failed to get object state)\n",
    "        \"\"\"\n",
    "\n",
    "        res = self.get_object_state(model_name=model_name)\n",
    "        if res is None: return None  # Object DNE\n",
    "\n",
    "        obj_pos = res.pose.position\n",
    "        bbox = move_baxter.spawner.bounding_box\n",
    "\n",
    "        # Get Camera Position\n",
    "        pos, _ = self.get_current_pose()\n",
    "        camera_offset = (0.03825, 0.012, 0.015355) # From urdf\n",
    "        # Adjust for Camera's offset from Gripper Position\n",
    "        pos[0] += camera_offset[0]; pos[1] += camera_offset[1]; pos[2] += camera_offset[2]\n",
    "\n",
    "        lim = pos[2] - table_z\n",
    "\n",
    "        if self.verbose >= 3:\n",
    "            print(\">>> Checking if {} Object in view <<<\".format(model_name if model_name is not None else \"\"))\n",
    "            print(\"Camera Position: {}\".format(pos))\n",
    "            print(\"Limit (+-): {}\".format(lim))\n",
    "            print(\"Object Position: {}+-{},{}+-{}\".format(obj_pos.x, bbox.x, obj_pos.y, bbox.y))\n",
    "            print(\"Table Height: {}\".format(table_z))\n",
    "\n",
    "        if (((pos[0] - lim < abs(obj_pos.x + bbox.x) < pos[0] + lim) or \n",
    "            (pos[0] - lim < abs(obj_pos.x - bbox.x) < pos[0] + lim)) and \n",
    "            ((pos[1] - lim < abs(obj_pos.y + bbox.y) < pos[1] + lim) or \n",
    "            (pos[1] - lim < abs(obj_pos.y - bbox.y) < pos[1] + lim))):\n",
    "            return True # In View\n",
    "        else:\n",
    "            return False # Not In View\n",
    "\n",
    "    def reset_models(self, spawn=True):\n",
    "        \"\"\"\n",
    "        Remove all object instances, but NOT the table\n",
    "        Then generate a new version of the object\n",
    "        \n",
    "        Parameters:\n",
    "            spawn (bool): Whether to spawn a new object afterwards\n",
    "        Returns:\n",
    "            Pose (to actual item if spawn=True, to empty spot if spawn=False)\n",
    "        \"\"\"\n",
    "        if self.verbose >= 1: print(\"Resetting all models with{} new spawn\".format(\"\" if spawn else \"out\"))\n",
    "        self.spawner.delete_models()\n",
    "        self.table_spawner.delete_models()\n",
    "        time.sleep(0.1)\n",
    "        self.table_spawner.spawn_model(Pose(position=Point(x=0.8, y=0.0, z=0.0)))\n",
    "        if spawn:\n",
    "            time.sleep(0.1)\n",
    "            return self.spawner.spawn_on_table() # Returns: name, RPY, Pose\n",
    "        else: \n",
    "            # Returns: 'EMPTY',RPY,Pose (no object at returned Pose!)\n",
    "            return self.spawner.spawn_on_table(spawn=False)\n",
    "\n",
    "    def loop_training(self,iterations=10, hover_distance=0.3, no_empty_state=False):\n",
    "        \"\"\"\n",
    "        Generates a image-sequence dataset with labels for a randomized pick task.\n",
    "        \n",
    "        Parameters:\n",
    "            iterations (int): How many times to pick up the object \n",
    "            hover_distance (float): What height to hover above the object (must clear object)\n",
    "        \"\"\"\n",
    "        self.return_to_home_pose(open_gripper=True)\n",
    "\n",
    "        for i in range(iterations):\n",
    "            session_id = \"{}_iteration_{}\".format(self.id, i)\n",
    "\n",
    "            # >>> Reset : New Iteration <<<\n",
    "            spawn_or_not = True if no_empty_state else False#bool(random.getrandbits(1))\n",
    "            object_name, RPY, spawned_pose = self.reset_models(spawn=spawn_or_not)\n",
    "            self.open_gripper()\n",
    "\n",
    "            hover_pose = define_pose(spawned_pose.position.x,\n",
    "                                    spawned_pose.position.y,\n",
    "                                    spawned_pose.position.z)\n",
    "            # Note: Spawned Pose quaternion will NOT match overhead gripper quaternion\n",
    "            q = get_rpy_quaternion(*RPY) # Hence we use the RPY of the upright-yaw component only\n",
    "            target_pose = define_pose(spawned_pose.position.x,\n",
    "                                      spawned_pose.position.y,\n",
    "                                      spawned_pose.position.z,\n",
    "                                      o_x=q[0], o_y=q[1], o_z=q[2], o_w=q[3])\n",
    "\n",
    "            # >>> Start Reward Publisher <<<\n",
    "            time_start = time.time()\n",
    "            reward_pub = subprocess.Popen(['bash', 'scripts/run_metadata_publisher.sh',\n",
    "                                           '-l', 'left',\n",
    "                                           '-o',object_name,\n",
    "                                           '-s', str(session_id),\n",
    "                                           '-e', 'false' if spawn_or_not else 'true'])\n",
    "            time.sleep(1) # Wait for Publisher to Initialize\n",
    "\n",
    "            # >>> Start Recording <<<\n",
    "            recorder = subprocess.Popen(['bash', 'scripts/run_recorder.sh',\n",
    "                                         'left', str(session_id)])\n",
    "            time.sleep(3) # Wait for Recorder to Initialize\n",
    "\n",
    "            # >>> Hover over Target <<<\n",
    "            self.stage_pub.publish(\"SEARCH\")\n",
    "            self.hover_over_target(hover_pose, hover_distance=hover_distance)\n",
    "            time.sleep(1)\n",
    "\n",
    "            # >>> Pick Target <<<\n",
    "            self.stage_pub.publish(\"APPROACH\")\n",
    "            self.move_to_target(target_pose, direct=True)\n",
    "            # Publish: In Graspable/Ungraspable Stage\n",
    "            self.stage_pub.publish(\"GRASP\")\n",
    "            time.sleep(1) # Ensure graspable position is captured\n",
    "            self.close_gripper()\n",
    "            time.sleep(1) # Ensure grasped object is captured\n",
    "\n",
    "            # >>> Stop Recording <<<\n",
    "            kill = subprocess.Popen(['bash', 'scripts/kill_recorder.sh', \n",
    "                                     session_id])\n",
    "            kill.wait(); recorder.wait() # Wait for node to shutdown cleanly\n",
    "            self.stage_pub.publish(\"IDLE\")\n",
    "\n",
    "            # >>> Stop Reward Publisher <<<\n",
    "            kill = subprocess.Popen(['bash', 'scripts/kill_metadata_publisher.sh', \n",
    "                                     session_id])\n",
    "            kill.wait(); reward_pub.wait()\n",
    "\n",
    "            duration = time.time() - time_start # in seconds\n",
    "\n",
    "            # >>> Feedback <<<\n",
    "            if self.verbose >= 3:\n",
    "                print(\"---- Iteration: {}, Session: {} ----\".format(self.id, i))\n",
    "                print(\">>> Spawned Pose <<<\\n{}\".format(spawned_pose))\n",
    "                print(\">>> Hover Pose <<<\\n{}\".format(hover_pose))\n",
    "\n",
    "            # >>> Check Grasp Success <<<\n",
    "            self.hover_over_target(hover_pose, \n",
    "                                   hover_distance=hover_distance)\n",
    "            rospy.sleep(0.5) # Ensure end position is captured\n",
    "\n",
    "            ground_truth_label, old_pose, new_pose = self.grasp_verification(spawned_pose, \n",
    "                                                                           spawn_or_not)\n",
    "            if self.verbose >= 1:\n",
    "                print(\"Grasp {}\".format(\"Successful\" if ground_truth_label else \"Unsuccessful\"))\n",
    "\n",
    "            # >>> Reset: Return Home <<< \n",
    "            self.return_to_home_pose(open_gripper=False)\n",
    "\n",
    "            # >>> Save Grasp Success and MetaData <<<\n",
    "            with open(self.gt_path, 'a') as f:\n",
    "                fieldnames = [\"data_id\", \"iteration\", \"grasp_success\", \"duration\",\n",
    "                              \"spawned_position\", \"end_position\",\n",
    "                              \"spawned_orientation\", \"end_orientation\"]\n",
    "                writer = csv.DictWriter(f, fieldnames=fieldnames)\n",
    "                if f.tell() == 0: # Header\n",
    "                    writer.writeheader()\n",
    "                writer.writerow({'data_id': self.id, 'iteration': i,\n",
    "                                 'grasp_success': ground_truth_label, 'duration': duration,\n",
    "                                 'spawned_position': old_pose.position,\n",
    "                                 'end_position': new_pose.position,\n",
    "                                 'spawned_orientation':old_pose.orientation,\n",
    "                                 'end_orientation':new_pose.orientation})\n",
    "\n",
    "        # Done\n",
    "        self.reset_models(spawn=False)\n",
    "        self.return_to_home_pose(open_gripper=True)\n",
    "\n",
    "    def display_trajectory(self, plan):\n",
    "        \"\"\"Publishes motion plan to Rviz for visualisation\"\"\"\n",
    "        display_trajectory = moveit_msgs.msg.DisplayTrajectory()\n",
    "        display_trajectory.trajectory_start = self.robot.get_current_state()\n",
    "        display_trajectory.trajectory.append(plan)\n",
    "        self.display_trajectory_publisher.publish(display_trajectory)\n",
    "\n",
    "    def get_current_pose(self,verbose=False):\n",
    "        \"\"\"Returns current pose (position/orientation) of gripper\"\"\"\n",
    "        current_pose = self.tf_listener.lookupTransform(\"base\",\n",
    "                                                        self.group.get_end_effector_link(),\n",
    "                                                        rospy.Time(0))\n",
    "        if verbose >= 2: print(\"Current End Effector Pose:\\n{}\".format(current_pose))\n",
    "        return current_pose\n",
    "\n",
    "    def move_to_target(self, target, direct=False):\n",
    "        \"\"\"\n",
    "        Attempt up to self.max_attempts to move the end effector to the target pose.\n",
    "        \n",
    "        Paramaters:\n",
    "            target (Pose): Desired position and orientation for end effector\n",
    "        \"\"\"\n",
    "        if not isinstance(target, geometry_msgs.msg.Pose):\n",
    "            raise TypeError(\"Target must be of type Pose\")\n",
    "        #self.group.clear_pose_targets()\n",
    "        self.group.set_pose_target(target)\n",
    "        result = False; fraction = 1.0\n",
    "        for attempt in range(0, self.max_attempts):\n",
    "            if direct:\n",
    "                (plan, fraction) = self.group.compute_cartesian_path(\n",
    "                                       [target],   # waypoints to follow\n",
    "                                       0.01,       # eef_step\n",
    "                                       0.0)        # jump_threshold\n",
    "            else:\n",
    "                plan = self.group.plan()\n",
    "\n",
    "            if plan is not None and len(plan.joint_trajectory.points)>0 and isclose(fraction, 1.0):\n",
    "                if self.verbose >= 2: print(\"Attempt {}, motion plan FOUND\".format(attempt))\n",
    "                self.display_trajectory(plan)\n",
    "                result = self.group.execute(plan,wait=True)\n",
    "                if result:\n",
    "                    if self.verbose >= 1: print(\"Motion Plan Execution Succeeded\")\n",
    "                    break\n",
    "            else:\n",
    "                if self.verbose >= 2: print(\"Attempt {}, motion plan NOT FOUND\".format(attempt))\n",
    "        if result and plan is not None:\n",
    "            rospy.sleep(1.0)\n",
    "            self.group.stop()\n",
    "            self.group.clear_pose_targets()\n",
    "        else:\n",
    "            if self.verbose >= 1: print(\"No Motion Plan Found in {} attempts\".format(self.max_attempts))\n",
    "                         \n",
    "        return result\n",
    "\n",
    "    def open_gripper(self):\n",
    "        \"\"\"Opens gripper and waits 1.0 seconds\"\"\"\n",
    "        if self.verbose >= 2: print(\"Opening {} gripper...\".format(self.limb_name))\n",
    "        self.gripper.open()\n",
    "        rospy.sleep(1.0)\n",
    "\n",
    "    def close_gripper(self):\n",
    "        \"\"\"Closes gripper and waits 1.0 seconds\"\"\"\n",
    "        if self.verbose >= 2: print(\"Closing {} gripper...\".format(self.limb_name))\n",
    "        self.gripper.close()\n",
    "        rospy.sleep(1.0)\n",
    "\n",
    "    def return_to_home_pose(self, open_gripper=True):\n",
    "        \"\"\"Returns arm to neutral home pose \n",
    "        \n",
    "        Parameters:\n",
    "            open_gripper (bool): Whether to open gripper\n",
    "        \"\"\"\n",
    "\n",
    "        if open_gripper: self.open_gripper()\n",
    "        pose_target = define_pose(x=0.46,\n",
    "                                  y=0.8 if self.limb_name == \"left\" else -0.8,\n",
    "                                  z=1.08)  # 0.69\n",
    "        success = self.move_to_target(pose_target)\n",
    "        if success and self.verbose >= 1:\n",
    "            print(\"Moved the {} arm to start pose...\".format(self.limb_name))\n",
    "        rospy.sleep(1.0)\n",
    "\n",
    "    def hover_over_target(self, supplied_pose, hover_distance=0.3):\n",
    "        \"\"\"Hovers gripper over object WITHOUT rotational alignment\"\"\"\n",
    "\n",
    "        if not isinstance(supplied_pose, geometry_msgs.msg.Pose):  # Pose\n",
    "            raise TypeError(\"Target was not of type Pose: {}\".format(type(target)))\n",
    "        target_pose = copy.deepcopy(supplied_pose)\n",
    "        target_pose.position.z += hover_distance\n",
    "        \n",
    "        success = self.move_to_target(target_pose)\n",
    "        if success and self.verbose >= 1: \n",
    "            print(\"Moving the {} arm over target\".format(self.limb_name))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Automatic Dataset Generation\n",
    "\n",
    "Automatically generate a custom dataset with 1000 iterations (approx 2GB). Will repeatedly spawn, align, approach, grasp and then retract. MoveBaxter will spawn appropriate publishers and subscribers to save information to the **PROJECT_DIR/workspace/data** directory.\n",
    "\n",
    "Parameters that may be customized:\n",
    "* **model_name**: Which object to spawn from ~/.gazebo/models or baxter_sim_examples directories\n",
    "* **tolerance**: How accurate motion planning should be (smaller is more accurate, but fails more often)\n",
    "* **limb_name**: Whether to use right or left planning group\n",
    "* **verbose**: How much diagnostic information to print during operation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(MoveBaxter.__doc__)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Note**: Gazebo, Rviz and JointTrajectoryActionServer must be running"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "run_name = raw_input(\"Enter a name for the current Run >>>\")\n",
    "\n",
    "moveit_commander.roscpp_initialize(sys.argv)\n",
    "rospy.init_node(\"MoveBaxter_Interface\", anonymous=True)\n",
    "move_baxter = MoveBaxter(run_name, limb_name=\"left\", tolerance=0.01, verbose=1)  \n",
    "move_baxter.open_gripper()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "move_baxter.loop_training(iterations=1000, no_empty_state=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Manual Data Generator\n",
    "MoveBaxter also allows manual control of the robot, this is primarily useful for testing new features or evaluating generated data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "moveit_commander.roscpp_initialize(sys.argv)\n",
    "rospy.init_node(\"MoveBaxter_Interface\", anonymous=True)\n",
    "move_baxter = MoveBaxter(\"Experimental\", limb_name=\"left\", tolerance=0.01, verbose=1)  \n",
    "move_baxter.open_gripper()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Reset**: Spawn 1 object (default=block) in a fresh environment. Then retrieve its pose for the alignment stage."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "name,RPY,spawned_pose = move_baxter.reset_models(spawn=True)\n",
    "hover_pose = define_pose(spawned_pose.position.x,\n",
    "                        spawned_pose.position.y,\n",
    "                        spawned_pose.position.z)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Target**: Set orientation (adjusted based on the yaw of the object) for the final pick action."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "q = get_rpy_quaternion(*RPY) # Hence we use the RPY of the upright-yaw component only\n",
    "target_pose = define_pose(spawned_pose.position.x,\n",
    "                          spawned_pose.position.y,\n",
    "                          spawned_pose.position.z,\n",
    "                          o_x=q[0],o_y=q[1],o_z=q[2],o_w=q[3])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Home**: Return to Home Pose, verify that object is not in view."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "move_baxter.return_to_home_pose()\n",
    "move_baxter.object_in_view()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Hover**: Move to align above the object (overhead orientation)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "move_baxter.hover_over_target(hover_pose)\n",
    "move_baxter.object_in_view()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Diagnostic**: Get gripper and object poses to deduce the distance between them (used to calculate reward)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "position,orientation=move_baxter.get_current_pose()\n",
    "print(type(Point(*position)),type(Quaternion(*orientation)))\n",
    "gripper_pose = Pose(Point(*position),Quaternion(*orientation))\n",
    "object_pose = move_baxter.get_model_state(\"block\",\"\").pose\n",
    "\n",
    "x2,y2,z2 = object_pose.position.x, object_pose.position.y, object_pose.position.z\n",
    "x1,y1,z1 = gripper_pose.position.x, gripper_pose.position.y, gripper_pose.position.z\n",
    "\n",
    "dx,dy,dz = x2-x1,y2-y1,z2-z1\n",
    "print(\"dx:{},dy:{},dz:{}\".format(dx,dy,dz))\n",
    "distance = math.sqrt(math.pow(dz,2)+math.pow(dy,2)+math.pow(dx,2))\n",
    "print(\"Distance: {}\".format(distance))\n",
    "print(type(object_pose),type(gripper_pose))\n",
    "#object_pose.inve"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Pick**: Move from hover position to move gripper around graspable points on spawned object. Use 'compute_cartesian_path' to plot a direct path."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "move_baxter.move_to_target(target_pose,direct=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Grasp**: Close grippers around object"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "move_baxter.close_gripper()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Retract & Verify**: Return to home pose (with object in grasp) and check whether object was grasped successfully"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "move_baxter.hover_over_target(hover_pose)\n",
    "grasp_successful = move_baxter.grasp_verification(original_pose=spawned_pose)\n",
    "print(\"Grasp {}\".format(\"successful\" if grasp_successful else \"unsuccessful\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### <span id=\"visual\">Visualisation Tools</span>\n",
    "\n",
    "#### <span id=\"visual-nodes\">Display ROS Nodes<span>\n",
    "<span> <a href=\"#dataset-content\"> Return to Section Contents </a> </span>\n",
    "     \n",
    "ROS employs a Publisher-Subscriber pattern for asynchronous communication between nodes. The \"rqt_graph\" node examines all published topics and their subscribers in order to illustrate the network of connections between different nodes. This can help diagnose what information a particular node can act on."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "rqt_graph = subprocess.Popen([\"rosrun\", \"rqt_graph\", \"rqt_graph\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "os.killpg(os.getpgid(rqt_graph.pid), signal.SIGINT)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### <span id=\"visual-camera\">Get Live Camera Feed</span>\n",
    "<span> <a href=\"#dataset-content\"> Return to Section Contents </a> </span>\n",
    "\n",
    "After startup in Gazebo, Baxter publishes a series of **/cameras/** topics for the 3 cameras on board Baxter:\n",
    "* 'head': head_camera\n",
    "* 'left': left_hand_camera\n",
    "* 'right': right_hand_camera\n",
    "* 'kinect': kinect raw bgr image \n",
    "\n",
    "These can be subscribed to by a different node to provide quasi-live feedback on Baxter's movements. In turn it is possible to superimpose visualisations, such as grasp points, to show how Baxter is operating. The feedback is not live because the asynchcronous communication between nodes does not gaurantee the receipt of data.\n",
    "\n",
    "For more information see [Baxter Wiki](https://github.com/RethinkRobotics/sdk-docs/wiki/Using-the-Cameras)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Open a Single Camera using Identifiers above:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "command = \"{}/scripts/open_single_camera.sh\".format(os.getcwd())\n",
    "visual_interpreter = subprocess.Popen([\"sudo\", \"bash\", command, \"left\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Open left, right and head camera simultaneously:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "command = \"{}/scripts/open_all_cameras.sh\".format(os.getcwd())\n",
    "baxter_view = subprocess.Popen([\"sudo\", \"bash\", command])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Force the camera listener to shutdown"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "%%bash\n",
    "rosnode kill /Baxter_Cam_Listener"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## <span id=\"process\">Data Processing</span>\n",
    "<a href=\"#content\">Return to Contents</a>\n",
    "\n",
    "\n",
    "### <span id=\"process-content\">Section Contents</span>\n",
    "\n",
    "1. <a href=\"#inspect\">Inspect Generated Data</a>\n",
    "1. <a href=\"#filter\">Filter Generated Data</a>\n",
    "    1. <a href=\"#filter-stage\">By Stage</a>\n",
    "    2. <a href=\"#filter-amount\">By Amount</a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### <span id=\"inspect\">Inspect Generated Data</span>\n",
    "<a href=\"#process-content\">Return to Section Contents</a>\n",
    "\n",
    "To help inform the filtering process, this section provides tools for inspecting the raw generated data. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_dir = path.join(os.getcwd(),\"workspace\",\"data\")\n",
    "result = glob.glob(path.join(data_dir,'*.{}'.format(\"csv\")))\n",
    "csvs = [r.split(data_dir+os.path.sep)[1] for r in result]\n",
    "run_ids = [csv.split(\"_\")[0] for csv in csvs]\n",
    "print(\"Choose a run ID: {}\".format(run_ids))\n",
    "run_id = input(\"\")\n",
    "while run_id not in run_ids:\n",
    "    run_id = input(\"\")\n",
    "print(\"Inpecting contents of Run: {}\".format(run_id))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Run Duration\n",
    "Inspect how much time it took to generate the data contained in a specific run. Note that this does not work if the data has been modified or copied elsewhere since creation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_run_duration(run_id,data_dir):\n",
    "    \"\"\"Looks up iteration creation dates for a particlar run\n",
    "    Returns difference between oldest and youngest iteration\"\"\"\n",
    "    run_iters = [path.join(data_dir,file) for file in os.listdir(data_dir) if file.startswith(\"session_\".format(run_id))]\n",
    "    start, *middle, end = sorted(run_iters, key=path.getctime)\n",
    "    time_dif = path.getctime(end)-path.getctime(start)\n",
    "    return datetime.timedelta(seconds=time_dif)\n",
    "\n",
    "print(\"Run Duration: {} (Hours:Minutes:Seconds)\".format(get_run_duration(run_id,data_dir)))\n",
    "total_time = datetime.timedelta(seconds=0)\n",
    "for run in run_ids:\n",
    "    total_time += get_run_duration(run,data_dir)\n",
    "print(\"Average Duration: {} (Hours:Minutes:Seconds)\".format(total_time/len(run_ids)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Run Size\n",
    "Inspect how much space the generated data occupies on disk. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def get_run_size(run_id,data_dir):\n",
    "    \"\"\"Walks iterations in a particular run and counts up their size\n",
    "    Returns total size of iterations in Bytes\"\"\"\n",
    "    byte_size = 0\n",
    "    for root, dirs, files in os.walk(data_dir,topdown=True):\n",
    "        dirs[:] = [d for d in dirs if d.startswith(\"session_{}\".format(run_id))]\n",
    "        for file in files:\n",
    "            byte_size += path.getsize(path.join(root,file))\n",
    "    return byte_size\n",
    "\n",
    "def convert_size(size_bytes):\n",
    "    \"\"\"Converts byte to human-readable form\n",
    "    Credit James Sapam \n",
    "    https://stackoverflow.com/questions/5194057/better-way-to-convert-file-sizes-in-python\"\"\"\n",
    "    if size_bytes == 0:\n",
    "        return \"0B\"\n",
    "    size_name = (\"B\", \"KB\", \"MB\", \"GB\", \"TB\", \"PB\", \"EB\", \"ZB\", \"YB\")\n",
    "    i = int(math.floor(math.log(size_bytes, 1024)))\n",
    "    p = math.pow(1024, i)\n",
    "    s = round(size_bytes / p, 2)\n",
    "    return \"%s %s\" % (s, size_name[i])\n",
    "\n",
    "print(\"Storage Space used: {}\".format(convert_size(get_run_size(run_id,data_dir))))\n",
    "\n",
    "total_storage = 0\n",
    "for run in run_ids:\n",
    "    total_storage += get_run_size(run,data_dir)\n",
    "print(\"Average Storage Space used: {}\".format(convert_size(total_storage/len(run_ids))))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Directory Structure\n",
    "Inspect the data directory structure. For effeciency only 1 file is listed per iteration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def list_files(startpath):\n",
    "    \"\"\"Walk through directory and print out its structure\n",
    "    Credit: Ellockie\n",
    "    https://stackoverflow.com/questions/9727673/list-directory-tree-structure-in-python/49620815#49620815\"\"\"\n",
    "    for root, dirs, files in os.walk(startpath,topdown=True):\n",
    "        #dirs[:] = [d for d in dirs if not d.endswith(\"session\")]\n",
    "        level = root.replace(startpath, '').count(os.sep)\n",
    "        indent = ' ' * 4 * (level)\n",
    "        print('{}{}/'.format(indent, os.path.basename(root)))\n",
    "        subindent = ' ' * 4 * (level + 1)\n",
    "        for f in files:\n",
    "            print('{}{}'.format(subindent, f))\n",
    "            break\n",
    "            \n",
    "list_files(path.join(os.getcwd(),\"workspace\",\"data\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### <span id=\"filter\">Filter Generated Data</span>\n",
    "<a href=\"#process-content\">Return to Section Contents</a>\n",
    "\n",
    "In order to train models on the data it is necessary to perform some pre-processing. To begin, we can filter the data so that we work with subsets only.\n",
    "\n",
    "#### <span id=\"filter-stage\">Filter By Stage</span>\n",
    "<a href=\"#process-content\">Return to Section Contents</a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Selecting Stages\n",
    "Interactively choose which stage(s) to filter by. <br>\n",
    "**Stages**: IDLE, SEARCH, APPROACH, GRASP, EMPTY"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# >>> Choose Which Stages to Include <<<\n",
    "stages = [\"IDLE\", \"SEARCH\", \"APPROACH\", \"GRASP\", \"EMPTY\"]\n",
    "print([\"{}:{}\".format(i,stage) for i,stage in enumerate(stages)])\n",
    "resp = raw_input(\"Pick stage to filter by:\")\n",
    "chosen_stages = []\n",
    "while resp:\n",
    "    try:\n",
    "        if 0 <= int(resp) < len(stages):\n",
    "            chosen_stages.append(stages[int(resp)])\n",
    "    except:\n",
    "        print(\"Enter a number between 0 and {}\")\n",
    "    resp = raw_input(\"Chosen: {}\\n ENTER to finish:\".format(chosen_stages,len(stages)))\n",
    "print(\"Final Selection: {}\".format(chosen_stages))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Filtered Data Dictionary\n",
    "Construct filtered data dictionary containing images and their MetaData"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# >>> Data Directory <<<\n",
    "data_dir = path.join(os.getcwd(),\"workspace\",\"data\")\n",
    "print(\"Data in directory: {}\".format(data_dir))\n",
    "\n",
    "# >>> Setup <<<\n",
    "data = {}\n",
    "counter = 0; img_count = 0; total_count = 0;\n",
    "\n",
    "for item in os.listdir(data_dir): \n",
    "    # >>> Images <<<\n",
    "    \n",
    "    if item.startswith(\"session\"): # Folder\n",
    "        # >>> Extract Identifiers <<<\n",
    "        match = re.match(\"^session_([^_]*)_iteration_(\\d+)\",item)\n",
    "        run,iteration = match.group(1),match.group(2)\n",
    "        \n",
    "        # >>> Match with MetaData <<<\n",
    "        csv_file = path.join(data_dir,item,\"reward.csv\")\n",
    "        with open(csv_file,'r') as f:\n",
    "            reader = csv.DictReader(f)\n",
    "            img_metadata = {}\n",
    "            for row in reader:\n",
    "                if row['stage'] in chosen_stages:\n",
    "                    img_metadata[row['img_name']] = row\n",
    "                    img_count += 1\n",
    "                total_count += 1\n",
    "            \n",
    "        # >>> Insert into Master Dictionary <<<\n",
    "        if run in data:\n",
    "            if iteration in data[run]:\n",
    "                data[run][iteration].update(img_metadata)\n",
    "            else:\n",
    "                data[run][iteration] = img_metadata          \n",
    "        else:\n",
    "            data[run] = {iteration:img_metadata}\n",
    "        #print(\"Run: {}, Session: {}, Iteration: {}\".format(run,item,iteration))\n",
    "        \n",
    "    # >>> Ground Truths <<<\n",
    "    elif item.endswith(\".csv\"): #File\n",
    "        gt_file = path.join(data_dir,item)\n",
    "        with open(gt_file,\"r\") as f:\n",
    "            reader = csv.DictReader(f)\n",
    "            for row in reader:\n",
    "                run,i = row[\"data_id\"],row[\"iteration\"]\n",
    "                if run in data:\n",
    "                    if i in data[run]:\n",
    "                        data[run][i][\"gt\"] = row\n",
    "                    else:\n",
    "                        data[run][i] = {\"gt\":row}\n",
    "                else: # EdgeCase\n",
    "                    data[run] = {i:{\"gt\":row}}\n",
    "print(\"Size of Filtered Dataset: {} out of {} total image ({:.2f}%)\".format(img_count,total_count,100.0*img_count/total_count))\n",
    "\n",
    "print(\">>>Data Dictionary Structure<<<\\nRunID\\n\\tIteration\\n\\t\\tImage\\n\\t\\t\\tImg_MetaData\\n\\t\\tGroundTruth\\n\\t\\t\\tGround_Truth_Data\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<b style=\"font-size:16pt\">Data Dictionary Structure</b> <br>\n",
    "Note: All keys and values are stored as Strings. \n",
    "\n",
    "* RunID\n",
    "    * Iteration\n",
    "        * Image\n",
    "            * Image MetaData\n",
    "        * Ground Truth\n",
    "            * GroundTruth MetaData\n",
    "            \n",
    "<b style=\"font-size:14pt\">Example</b>\n",
    "* 'Run0'\n",
    "    * '862'\n",
    "        * 'img_31_t_19227.125.jpg'\n",
    "            * **stage**: 'GRASP'\n",
    "            * **time**: '19227.125'\n",
    "            * **distance**: '0.0125309689858'\n",
    "            * **prev_distance**: '0.697887540807'\n",
    "            * **gripper_pos**: '[0.6335322702884931, 0.18987464581239905, 0.7978240862937698]'\n",
    "            * **object_pos**: '[0.5952855203224089, 0.17787252630049952, 0.7950000546788926]'\n",
    "            * **img_name**: 'img_31_t_19227.125.jpg'\n",
    "            * **reward**: '54.69302275010344'\n",
    "            * **visible**: 'False\n",
    "        * 'gt'\n",
    "            * **data_id**: 'Run0'\n",
    "            * **iteration**: '862'\n",
    "            * **duration**: '16.960762977600098'\n",
    "            * **grasp_success**: 'False'\n",
    "            * **spawned_position**: 'x: 0.572498505857\\ny: 0.186026069492\\nz: 0.7825'\n",
    "            * **end_position**: 'x: 0.572497609673\\ny: 0.186030828812\\nz: 0.795007721606'\n",
    "            * **spawned_orientation**: 'x: 0.0\\ny: 0.0\\nz: 0.051704243461168\\nw: -0.9986624410721113'\n",
    "            * **end_orientation**: 'x: -7.83541030138e-0\\ny: -0.000116963660073\\nz: -0.0517044027934\\nw: 0.9986624229'\n",
    "  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Filtered Dataset Name\n",
    "Pick a unique name for the Filtered Dataset, used for creating a new folder inside the **PROJECT_DIR/workspace/data** directory. Folder will contain retino-cortical images as well as copied MetaData."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_name = raw_input(\"Choose a name for the filtered dataset, this may NOT contain the word 'session'\")\n",
    "filtered_path = path.join(data_dir,dataset_name)\n",
    "if not path.exists(filtered_path):\n",
    "    os.makedirs(filtered_path)\n",
    "else:\n",
    "    raise NameError(\"{} already exists\".format(filtered_path))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Initialize Software Retina\n",
    "May take some time. Ensure machine has GPU and is running **CUDA 10.1** with Driver **418.87.01**. \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Add Cuda path to System if not there already:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output = subprocess.check_output(\"source scripts/switch_cuda.sh 10.1; env -0\", shell=True,\n",
    "                      executable=\"/bin/bash\")\n",
    "os.environ.update(line.partition('=')[::2] for line in output[46:].split('\\0'))\n",
    "print(output[:23])\n",
    "print(dict(os.environ)[\"CUDA_HOME\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Use 'nvcc --version' and 'nvidia-smi' in console to check status of GPU and CUDA."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "# Specified by CUDA installer\n",
    "nvcc --version\n",
    "# Specified by GPU\n",
    "nvidia-smi"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Start the retina**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Note: If Kernel crashes, logout of or reboot system\n",
    "print(\"Initiliazing Software Retina...\")\n",
    "retina = retina.Retina(800,800)\n",
    "print(\"Retina initialized\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Generate Retino-Cortical Images and MetaData\n",
    "Depending on GPU, will take a long time to complete. Creates compressed versions of the original raw image files and stores them in a seperate, filtered, dataset using the name provided above.\n",
    "\n",
    "Inside the MetaData for each image one can find the following information:\n",
    "\n",
    "    Visibility: True (visibile) OR False (not found) // Auto-generated label is often incorrect when we get too close to the object\n",
    "    Grasp: True (success) or False (failure) // Stored seperately in RUNID_ground_truths.csv\n",
    "    Stage: IDLE, SEARCH, APPROACH, GRASP\n",
    "    \n",
    "If the dataset is filtered to only include APPROACH and GRASP stages then the problem is constrained to the pick stage of a pick-and-place pipepline. IDLE and SEARCH arguably belong under an alignment framework and hence can be trained seperately.\n",
    "\n",
    "The \"GRASP\" label is published from the DataSet controller (moveBaxter) as soon as CloseGripper() is called. Hence it can accuractly correspond to the GRASPABLE state. \n",
    "\n",
    "However, whether the grasp succeeds or fails is determined by lifting the object back to the hover position. if the object moves with the gripper then the grasp is successful otherwise it fails. Consequently, since the GRASP state will only be active once the object is held firmly by the gripper it is difficult to assess which frames correspond to Grasp Success or Failure. One cannot easily derive this from the published GRASP stage because the amount of frames it takes to close the gripper may not be consistent. \n",
    "\n",
    "One rough solution to this is to conduct a manual analysis of a representative sample of iterations (see figure below). For each iteration one could record the amount of frames after the GRASP stage is published before the object is actually grasped. The maximum value of this analysis could then be used as a threshold to correctly label (most of the time) which images are in a GRASPED state. The images before this would then be in the GRASPABLE state.\n",
    "\n",
    "<div class=\"row\">\n",
    "    <div class=\"col1\" style=\"float:left\">\n",
    "        <img src=\"https://i.imgur.com/ASX0jF1.jpg\" height=200px width=200px/>\n",
    "        <b>GRASPABLE state</b>\n",
    "    </div>\n",
    "    <div class=\"col2\" style=\"float:left\">\n",
    "        <img src=\"https://i.imgur.com/nJGurpe.jpg\" height=200px width=200px/>\n",
    "        <b>GRASPED state</b>\n",
    "    </div>\n",
    "</div>\n",
    "\n",
    "Alternatively, it is possible to integrate haptic feedback into the dataset generation. That is, as soon as their is any resistance to motion in Baxter's gripper then the system is in a GRASPED state. However, this will likely lead to sparse data as it requires a 3rd subscriber. In order to match the frame, metadata and haptic feedback in an asynchronous environment it is necessary to match their timestamps. With 3 subscribers, it is likely that the timestamps will only rarely match and consequently the data may become too sparse. \n",
    "\n",
    "Hence, the optimal solution appears to be to merge the GRASPABLE state with the GRASPED state. This way the robot will only choose to close its grippers if the visual interpreter thinks it is likely to succeed. Otherwise the robot can either return to hover position and try again or reset the environment. \n",
    "\n",
    "To illustrate this problem, 4 different approaches to developing a Visual Interpreter solutuion are described below.\n",
    "\n",
    "<p style=\"font-size:16pt\"><b>Option A</b> (State Machine) <p>\n",
    "(image only has 1 label at a time) <br>\n",
    "    <b>Labels</b>: \"Ready\",\"NotReady\", \"Grasped\", \"Failed\", \"Empty\" <br>\n",
    "\n",
    "**Response**:\n",
    "* Ready: CloseGripper\n",
    "* NotReady: Approach (pick)\n",
    "* Grasped: Move onto Place Task\n",
    "* Empty: Search for Object\n",
    "* Failed: Reset\n",
    "\n",
    "<p style=\"font-size:16pt\"><b>Option B</b> (Merged Graspable/Grasped) <span style=\"color:red\">PREFERRED</span> <p>\n",
    "(image only has 1 label at a time) <br>\n",
    "    <b>Labels</b>: \"Grasp\",\"NotReady\", \"Bad\", \"Empty\" <br>\n",
    "\n",
    "**Response**:\n",
    "* Grasp: CloseGripper if not closed already and try to place\n",
    "* NotReady: Approach (pick)\n",
    "* Empty: Search for Object\n",
    "* Bad: Cannot grasp, reset\n",
    "\n",
    "<p style=\"font-size:16pt\"><b>Option C</b> (Paired Classes) <p>\n",
    "(image has 3 sets of labels) <br>\n",
    "    <b>Visibility</b>: Found versus Empty <br>\n",
    "    <b>Grasped</b>: Grasped or NotGrasped <br> \n",
    "    <b>Graspability</b>: Ready or NotReady <br>\n",
    "\n",
    "Main concern here is that the accuracy on Grasped versus NotGrasped will be poor because of the unreliable time taken to close the gripper.\n",
    "\n",
    "**Response** would be based on label combination (8 possible combinations):\n",
    "* Empty and NotReady and NotGrasped: Search\n",
    "* Empty and Ready and NotGrasped: Search (at least 1 incorrect label)\n",
    "* Empty and NotReady and Grasped: Reset (at least 1 incorrect label)\n",
    "* Empty and Ready and Grasped: Try to Place (at least 1 incorrect label)\n",
    "* Found and NotReady and Grasped: Try to Place (at least 1 incorrect label)\n",
    "* Found and Ready and NotGrasped: CloseGripper\n",
    "* Found and NotReady and NotGrasped: Approach (pick)\n",
    "* Found and Ready and Grasped: Try to Place\n",
    "\n",
    "<p style=\"font-size:16pt\"><b>Option D</b> (Merged Graspable/Grasped & Paired Classes) <p>\n",
    "(image has 2 sets of labels) <br>\n",
    "    <b>Visibility</b>: Found versus Empty <br>\n",
    "<b>Graspability</b>: Grasp versus NotReady versus Bad <br>\n",
    "\n",
    "**Response** would be based on label combination (6 possible combinations)\n",
    "* Found and Grasp: CloseGripper if not closed already and Try to Place\n",
    "* Found and NotReady: Approach (pick)\n",
    "* Found and Bad: Reset\n",
    "* Empty and Bad: Reset (at least 1 incorrect label)\n",
    "* Empty and NotReady: Search\n",
    "* Empty and Grasp: CloseGripper if not closed already and Try to Place (at least 1 incorrect label)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "from enum import Enum\n",
    "\n",
    "def get_iteration_path(run,iteration,data_dir=path.join(os.getcwd(),\"workspace\",\"data\")):\n",
    "    \"\"\"Returns absolute path to iteration directory using its identifying attributes\"\"\"\n",
    "    return path.join(data_dir,\"session_{}_iteration_{}\".format(run,iteration))\n",
    "    \n",
    "progress = IntProgress(min=0, max=img_count,description='Converting Dataset:',bar_style='success')\n",
    "display(progress)\n",
    "\n",
    "new_data = {}\n",
    "for session in data.keys():\n",
    "    session_data = data[session]\n",
    "    for iteration in session_data.keys():\n",
    "        iter_data = session_data[iteration]\n",
    "        iter_path = get_iteration_path(session,iteration)\n",
    "        \n",
    "        for img_name in iter_data.keys():\n",
    "            if img_name == 'gt': continue\n",
    "            img_data = iter_data[img_name]\n",
    "            if \"empty\" not in img_data: img_data[\"empty\"] = \"False\"\n",
    "            # Apply Retino-Cortical Transform\n",
    "            img_path = path.join(iter_path,img_name)\n",
    "            img = cv2.imread(img_path,cv2.IMREAD_COLOR)\n",
    "\n",
    "            if img is None: \n",
    "                print(\"ERROR at:\\n{},{},{}\".format(session,iteration,img_name)); continue\n",
    "                \n",
    "            retina_img = retina.sample(img)\n",
    "            new_img_name = \"{}_{}_{}\".format(session,iteration,img_name)\n",
    "            new_img_path = path.join(data_dir,dataset_name,new_img_name)\n",
    "            cv2.imwrite(new_img_path,retina_img)\n",
    "            \n",
    "            # >>> Get Label <<<\n",
    "            grasp_was_successful = True if iter_data['gt']['grasp_success'] == 'True' else False \n",
    "            if img_data[\"empty\"] in [\"True\",\"true\",True]: \n",
    "                grasp_label = label.EMPTY.value\n",
    "            elif img_data[\"stage\"] == \"GRASP\" and grasp_was_successful:\n",
    "                grasp_label = label.GRASP.value # Label: Grasp\n",
    "            elif img_data[\"stage\"] == \"GRASP\" and not grasp_was_successful:\n",
    "                grasp_label = label.BAD.value # Label: Bad\n",
    "            else:\n",
    "                grasp_label = label.NOT_READY.value # Label: NotReady\n",
    "            \n",
    "            with open(path.join(data_dir,dataset_name,\"metadata.csv\"),\"a\") as f:\n",
    "                field_names = img_data.keys() + [\"label\"]\n",
    "                writer = csv.DictWriter(f,fieldnames=field_names)\n",
    "                if f.tell() == 0:\n",
    "                    writer.writeheader()\n",
    "                written_data = copy.deepcopy(img_data)\n",
    "                written_data[\"img_name\"] = new_img_name\n",
    "                written_data.update({\"label\":grasp_label})\n",
    "                writer.writerow(written_data)\n",
    "            progress.value += 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### <span id=\"filter-amount\">Filter By Amount</span>\n",
    "<a href=\"#process-content\">Return to Section Contents</a>\n",
    "\n",
    "Filter to build a training/testing set of a fixed size."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_dir = os.path.join(os.getcwd(),\"workspace\",\"data\")\n",
    "iterations = [i for i in os.listdir(data_dir) if i.startswith(\"session\")]\n",
    "random.shuffle(iterations)\n",
    "resp = input(\"Select number of iterations out of {}\".format(len(iterations)))\n",
    "while True:\n",
    "    if resp.isdigit():\n",
    "        if 0 < int(resp) < len(iterations):\n",
    "            break\n",
    "    resp = input(\"Select number of iterations out of {}\".format(len(iterations)))\n",
    "    \n",
    "data = np.random.choice(iterations,size=int(resp),replace=False)\n",
    "print(data,len(data))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## <span id=\"play\">Experimental Playground</span>\n",
    "<a href=\"#content\">Return to Contents</a>\n",
    "\n",
    "This section is dedicated to experimental features. All code is tested here for robustness before it is used in the active environment. We retain examplars from successful tests here to aid understanding."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Examplars\n",
    "\n",
    "### Camera Recorder Factory\n",
    "Used to test the parallel launching and closing of a ROS node that records images from the 'left' hand camera approximately every 0.5 seconds. \n",
    "\n",
    "If **poll()** returns <b style=\"color:green\">None</b> then the node is still running. \n",
    "\n",
    "**wait()** ensures the next recorder is not initialized until the 1st has successfully terminated. \n",
    "\n",
    "**time.sleep()** simulates the execution of useful code, which may take some time to complete. In our case this would be the execution of Motion Plans for Baxter using MoveIt as we wish to record the Movement."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "class CameraRecorderFactory(object):\n",
    "    def __init__(self,limb,start,end):\n",
    "        self.limb = limb\n",
    "        self.start = start\n",
    "        self.end = end\n",
    "        self.cameras = []\n",
    "        \n",
    "        \n",
    "    def main_loop(self):\n",
    "        for session_id in range(self.start,self.end):\n",
    "            print(\"Starting {} recorder\".format(session_id))\n",
    "            recorder = subprocess.Popen(['bash', 'scripts/run_baxter_recorder.sh','left',str(session_id)])\n",
    "            print(\"Sleeping...\")\n",
    "            time.sleep(2) # Wait 2 seconds (approx 4 images)\n",
    "            print(\"Waking up!\")\n",
    "            print(recorder.poll())\n",
    "            print(\"Sleeping...\")\n",
    "            time.sleep(2)\n",
    "            print(\"Waking up!\")\n",
    "            kill = subprocess.Popen(['bash', 'scripts/kill_baxter_recorder.sh',str(session_id)])\n",
    "            kill.wait(); recorder.wait() # Wait for node to shutdown cleanly\n",
    "cam_factory = CameraRecorderFactory(\"left\",2,6):\n",
    "cam_factory.main_loop()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Gripper Offset\n",
    "When carrying out pick task a systematic offset in the gripper was noted. Here the appropriate corrections are found."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Wrist Target Pose**\n",
    "x=0.679, y=0.312, z=1.053, xo = -3.133, yo = 0.0 zo = -3.132\n",
    "\n",
    "**Block Target Pose**\n",
    "x=0.679, y=0.3100, z=0.820\n",
    "\n",
    "**Block Size**: \n",
    "    0.04 x 0.04 by 0.04\n",
    "    \n",
    "**Correcting Offset**\n",
    "x = x - 0.02\n",
    "y = y + 0.02\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "block_spawner.spawn_model(define_pose(x=0.7,y=0.4,z=0.7825))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "move_baxter.return_to_home_pose(open_gripper=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "move_baxter.get_diagnostics(verbose=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "move_baxter.hover_over_target(define_pose(x=0.7+offset,y=-0.1+offset,z=0.7825),pick=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "offset = 0.02\n",
    "move_baxter.hover_over_target(define_pose(x=0.6799-offset,y=0.310+offset,z=0.7825),pick=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "move_baxter.move_to_target(define_pose(x=0.6799-offset,y=0.310+offset,z=0.7825))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Frustum\n",
    "The frustum defines a 3D description for the Field of View of Camera. Only objects within the frustum are visible. \n",
    "![Frustum](https://upload.wikimedia.org/wikipedia/commons/thumb/8/8f/Square_frustum.png/200px-Square_frustum.png)\n",
    "The Frustum is a Trapezoid which we shall approximate as a Square Pyramid. This is a reasonable assumption since the gripper ensures objects will always be at a distance from the camera.\n",
    "\n",
    "A square pyramid can be described as follows: \n",
    "\n",
    "$|x| + |y| = k(h-z)$\n",
    "\n",
    "Since the pyramid is Square $|x| == |y|$. Therefore: \n",
    "\n",
    "$2|x| = k(h-z)$\n",
    "\n",
    "By placing Baxter's left arm at **x=0.0,y =1.0,z =0.64** and spawning blocks at the edges of the cameras FOV we find that: \n",
    "\n",
    "$|1.28|+|1.28| = k(0.64)$\n",
    "\n",
    "Therefore for our camera:\n",
    "$k = 4$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "pos,_ = move_baxter.get_current_pose()\n",
    "bbox = move_baxter.spawner.bounding_box\n",
    "resp_state = move_baxter.get_model_state(\"block\",\"\").pose\n",
    "camera_offset = (0.03825,0.012,0.015355) # From urdf\n",
    "pos[0] += camera_offset[0]\n",
    "pos[1] += camera_offset[1]\n",
    "pos[2] += camera_offset[2]\n",
    "cam_z = pos[2]; table_z = 0.7825\n",
    "x_lim = 2*(cam_z-table_z)\n",
    "y_lim = 2*(cam_z-table_z)\n",
    "print(\"Camera Pos:{}\".format(pos))\n",
    "print(\"x_lim: {}, y_lim: {}\".format(x_lim,y_lim))\n",
    "print(pos[0] - x_lim , abs(resp_state.position.x + bbox.x) , pos[0] + x_lim)\n",
    "print(pos[0] - x_lim < abs(resp_state.position.x + bbox.x) < pos[0] + x_lim)\n",
    "print(pos[0] - x_lim , abs(resp_state.position.x - bbox.x) , pos[0] + x_lim)\n",
    "print(pos[0] - x_lim < abs(resp_state.position.x - bbox.x) < pos[0] + x_lim)\n",
    "print(pos[1] - y_lim , abs(resp_state.position.y + bbox.y) , pos[1] + y_lim)\n",
    "print(pos[1] - y_lim < abs(resp_state.position.y + bbox.y) < pos[1] + y_lim)\n",
    "print(pos[1] - y_lim , abs(resp_state.position.y - bbox.y) , pos[1] + y_lim)\n",
    "print(pos[1] - y_lim < abs(resp_state.position.y - bbox.y) < pos[1] + y_lim)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Frustum Revisited\n",
    "After analysing the data, it appeared the viewing frustum was failing when objects came close to the gripper. As a result, an updated method is desired that also works well in the near-field.\n",
    "\n",
    "First it is necessary to test the object_in_view function directly, here we look at a particular case where the viewing frustum is producing incorrect results:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def object_in_view_test(grip_pos,obj_pos, table_z = 0.7825, bbox=BoundingBox(x=0.04)):\n",
    "\n",
    "    camera_offset = (0.03825,0.012,0.015355) # From urdf\n",
    "    grip_pos[0] += camera_offset[0]; grip_pos[1] += camera_offset[1]; grip_pos[2] += camera_offset[2]\n",
    "\n",
    "    lim = grip_pos[2] - table_z\n",
    "    print(\"Lower Bound: {:.5f}, Adjustment (+x): {:.5f}, Upper Bound: {:.5f}\".format(grip_pos[0],\n",
    "                      grip_pos[0]-lim,abs(obj_pos[0] + bbox.x),grip_pos[0]+lim))\n",
    "    print(\"Lower Bound: {:.5f}, Adjustment (+y): {:.5f}, Upper Bound: {:.5f}\\n\".format(grip_pos[1],\n",
    "                      grip_pos[1]-lim,abs(obj_pos[1] + bbox.y),grip_pos[1]+lim))\n",
    "    \n",
    "    print(\"Evaluate + Bounding Box Adjustment (x)\")\n",
    "    print((grip_pos[0] - lim < abs(obj_pos[0] + bbox.x) < grip_pos[0] + lim))\n",
    "    print(\"Evaluate - Bounding Box Adjustment (x)\")\n",
    "    print((grip_pos[0] - lim < abs(obj_pos[0] - bbox.x) < grip_pos[0] + lim))\n",
    "    print(\"Evaluate + Bounding Box Adjustment (y)\")\n",
    "    print((grip_pos[1] - lim < abs(obj_pos[1] + bbox.y) < grip_pos[1] + lim))\n",
    "    print(\"Evaluate - Bounding Box Adjustment (y)\")\n",
    "    print((grip_pos[1] - lim < abs(obj_pos[1] - bbox.y) < grip_pos[1] + lim))\n",
    "    \n",
    "    if (((grip_pos[0] - lim < abs(obj_pos[0] + bbox.x) < grip_pos[0] + lim) or \n",
    "        (grip_pos[0] - lim < abs(obj_pos[0] - bbox.x) < grip_pos[0] + lim)) and \n",
    "        ((grip_pos[1] - lim < abs(obj_pos[1] + bbox.y) < grip_pos[1] + lim) or \n",
    "        (grip_pos[1] - lim < abs(obj_pos[1] - bbox.y) < grip_pos[1] + lim))):\n",
    "        return True # In View\n",
    "    else:\n",
    "        return False # Not In View\n",
    "    \n",
    "test_data = data['Run1']['818']\n",
    "\n",
    "for img_name in sorted(test_data):\n",
    "    if img_name == \"gt\": continue\n",
    "    \n",
    "    img_data = test_data[img_name]\n",
    "    grip_pos = eval(img_data[\"gripper_pos\"])\n",
    "    obj_pos = eval(img_data[\"object_pos\"])\n",
    "\n",
    "    default_visiblity_label = True if img_data[\"visible\"] == 'True' else False\n",
    "    \n",
    "    print(img_name)\n",
    "    print(\"\\nEvaluated Positions:\\nGripper {}\\nObject {}\\n\".format(grip_pos,obj_pos))\n",
    "    print(\"Original Verdict: {}\\n\".format(img_data[\"visible\"]))\n",
    "    print(\"\\nNew Verdict: {}\".format(object_in_view_test(grip_pos,obj_pos)))\n",
    "    break"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Randomizing Cube Orientation\n",
    "![Cube Rotation](https://upload.wikimedia.org/wikipedia/en/7/74/Cube_rotation.gif)\n",
    "\n",
    "Our cube has 6 faces each with a different colour. If we stick to rotations of $\\pi/2$ then this means there are 24 unique ways of orienting the cube.\n",
    "\n",
    "To test this we spawn 14 cubes in a line along the surface of the table. Each one will have a random rotation about the:\n",
    "* X-axis (roll)\n",
    "* Y-axis (pitch)\n",
    "* Z-axis (yaw)\n",
    "\n",
    "Since a rotation of $2\\pi$ is the same as a rotation of $0$ radians, it follows that there are 4 choices for each axis: <br>\n",
    "\\[ $0,\\pi/2,\\pi,3\\pi/2$ \\]\n",
    "\n",
    "Spawning 14 of such random orientations at a time, we find most spawn correctly. However there are a few edge cases where the rotation is off-centre:\n",
    "* $3\\pi/2,\\pi/2,\\pi$\n",
    "* $\\pi,\\pi/2,3\\pi/2$\n",
    "* $\\pi/2,\\pi,3\\pi/2$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "#IEEE 754 are guaranteed to be accurate to within half a ULP\n",
    "for i in range(14):\n",
    "    fixed_rotations = numpy.arange(0,2*numpy.pi,numpy.pi/2)\n",
    "    roll,pitch,yaw=random.choice(fixed_rotations),random.choice(fixed_rotations),random.choice(fixed_rotations)#random.uniform(0,2*numpy.pi)\n",
    "    \n",
    "    q =tf.transformations.quaternion_from_euler(roll,pitch,yaw)\n",
    "\n",
    "    print(roll,pitch,yaw)\n",
    "    \n",
    "    bbox = move_baxter.spawner.bounding_box\n",
    "    move_baxter.spawner.spawn_model(define_pose(x=1.0,y=-0.4+1.4*i*bbox.y,z=0.7825,o_x=q[0],o_y=q[1],o_z=q[2],o_w=q[3]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "From examining the document for [tf.transformations](http://docs.ros.org/jade/api/tf/html/python/transformations.html) we find that this is due to **floating point precision** errors accumulating in the conversion from euler to quaternion. \n",
    "\n",
    "To circumvent this, we can use quaternions directly. However, 4-dimensional quaternions are not very intuitive for humans. Hence we use the property that quaternions are non-commutative, that is we need to multiply them in a specific order to get to the desired orientation. In this case we can use this by rotating about each axis  seperately and in the correct order (zyx).\n",
    "\n",
    "The result does not appear to generate any off-centre cubes.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Note: Floating Points (IEEE 754) are guaranteed to be accurate to within half a ULP\n",
    "for i in range(14):\n",
    "    fixed_rotations = numpy.arange(0,2*numpy.pi,numpy.pi/2)\n",
    "    roll,pitch,yaw=random.choice(fixed_rotations),random.choice(fixed_rotations),random.choice(fixed_rotations)#random.uniform(0,2*numpy.pi)\n",
    "    origin, xaxis, yaxis, zaxis = (0, 0, 0), (1, 0, 0), (0, 1, 0), (0, 0, 1)\n",
    "    \n",
    "    qz = quaternion_about_axis(yaw, zaxis)\n",
    "    qy = quaternion_about_axis(pitch, yaxis)\n",
    "    qx = quaternion_about_axis(roll, xaxis)\n",
    "    \n",
    "    \n",
    "    q = quaternion_multiply(qz, qy)\n",
    "    q = quaternion_multiply(q, qx)\n",
    "\n",
    "    print(roll,pitch,yaw)\n",
    "    \n",
    "    bbox = move_baxter.spawner.bounding_box\n",
    "    move_baxter.spawner.spawn_model(define_pose(x=1.0,y=-0.4+1.4*i*bbox.y,z=0.7825,o_x=q[0],o_y=q[1],o_z=q[2],o_w=q[3]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run Simulationusing Subprocess\n",
    "Unstable API for running Gazebo and Rviz as seperate processes (in the background). \n",
    "\n",
    "**Note**: All Rospy Log output is visible inside the terminal where Jupyter Notebook was launched."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "gazebo = Popen(['bash', 'scripts/run_gazebo.sh'])\n",
    "time.sleep(2)\n",
    "joint_process = Popen(['sudo','scripts/run_joint_trajectory_action_server.sh'], bufsize=1, universal_newlines=True)\n",
    "time.sleep(3)\n",
    "print(joint_process.poll())\n",
    "# stdout_value = joint_process.communicate()[0]\n",
    "# print 'Joint_Trajectory_Action_Server:\\n{}'.format(stdout_value)\n",
    "time.sleep(2)\n",
    "rviz = Popen(['bash', 'scripts/run_rviz.sh'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "gazebo2rviz = Popen(['sudo','scripts/run_model_transfer.sh'], bufsize=1, universal_newlines=True)\n",
    "time.sleep(1)\n",
    "print(gazebo2rviz.poll())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "os.killpg(os.getpgid(gazebo.pid), signal.SIGINT)\n",
    "time.sleep(1)\n",
    "os.killpg(os.getpgid(joint_process.pid), signal.SIGINT)\n",
    "time.sleep(1)\n",
    "os.killpg(os.getpgid(rviz.pid), signal.SIGINT)\n",
    "time.sleep(1)\n",
    "os.killpg(os.getpgid(gazebo2rviz.pid), signal.SIGINT)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Synchronising Image Feed with Meta Data\n",
    "In order to build a reliable dataset, it is necessary to associate a reward with each image from Baxter's camera feed. The challenge is that ROS's publisher-subscriber model is asynchronous, hence we cannot easily gaurantee that the image data and session metadata will be synchronised.\n",
    "\n",
    "To get around this we can use a TimeSychnroniser, from the [Message_Filters](http://wiki.ros.org/message_filters) package. This will filter out any incoming messages that do not have matching Timestamps in their header.\n",
    "\n",
    "However, the usual way of publishing MetaData - as a JSON parsable String - is not compatible with the message filters structure. This is because the message filter requires each of the data streams to have a header with a timestamp inside of it. To circumvent this we must define our own message type ('reward_message.msg'):\n",
    "\n",
    "```\n",
    "Header Header\n",
    "string data\n",
    "```\n",
    "For ROS to recognise this inside of our Publisher and Subscribers, we must also modify the build script (CMAKE) to use the message_generator and message_runtime when building our custom package. The details of this can be found [here](http://wiki.ros.org/msg).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from baxter_picker.msg import reward_message\n",
    "from std_msgs.msg import String, Header\n",
    "# JSON-parsable string\n",
    "reward_str ='{{\"stage\":\"{}\",\"reward\":{},\"visible\":{},\"pos\":{}}}'.format(\"IDLE\",10.2,\"true\" if True else \"false\",\n",
    "                                                              [3.46,2.97,6.52]) #,\"visible\":{},\"distance\":{},\"prev_distance\":{}}}'.format(\"IDLE\",10.2,True,\n",
    "                                                 #                           12.4,10.6)\n",
    "# Define message header: ID, TimeStamp, FrameID (not needed, hence using object_name)\n",
    "header = Header(24,rospy.Time.now(),\"block\") \n",
    "message = reward_message(header,reward_str)\n",
    "print(message)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import json\n",
    "print(message.data)\n",
    "json = json.loads(message.data)\n",
    "print(json)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Applying Retino-Cortical Transform\n",
    "The Software Retina provides a biologically-inspired mechanism for compressing images. Similar to the mammalian retina, the images output of a software retina is **foveated** - that is more detail is retained at the centre of the image than its peripheries. This is implemented using a set of overlapping 2D Gaussian curves and a log-polar mapping, whose net effect is to produce a distorted image:\n",
    "\n",
    "![Retina Image](https://i.imgur.com/Tub2cK3.jpg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "img_path = \"/home/xavierweiss/Automatic-Waste-Sorting/workspace/data/session_Run0_iteration_0/img_4_t_364.05.jpg\"\n",
    "print(\"Img Exists: {}\".format(os.path.isfile(img_path)))\n",
    "retina = retina.Retina(800,800)\n",
    "print(\"Retina initialized\")\n",
    "img = cv2.imread(img_path,cv2.IMREAD_COLOR)\n",
    "retina_img = retina.sample(img)\n",
    "cv2.imwrite(path.join(path.expanduser(\"~\"),\"Pictures\",\"retina_test.jpg\"), retina_img);\n",
    "cv2.imshow('image',retina_img)\n",
    "cv2.waitKey(0)\n",
    "cv2.destroyAllWindows()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Empty State\n",
    "In order to train the network to respond to objects being dynamically removed, it is necessary to train an EMPTY state. This visual state exists whenever an object is not Visible. The appropriate response is thus to search for the object until it is found or a user-specified timeout passses. \n",
    "\n",
    "As the default training set is quite robust in that objects are successfully grasped the majority of the time, it follows that few empty states will exist. Hence, to inject more empty states into into the pipeline we spawn 'imaginary' objects and follow the normal search-approach-grasp process. However, now the gripper will close on nothing (grasp will always 'fail') and the object will never be visible.\n",
    "\n",
    "The pick deep learning agent should not be trained with the empty state, as reward values will not be meaningful if there is no object to approach (empty state images use 0.0 as a placeholder reward).\n",
    "\n"
   ]
  }
 ],
 "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",
   "version": "3.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}