{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Rendering MultibodyPlant Tutorial\n",
    "For instructions on how to run these tutorial notebooks, please see the [README](https://github.com/RobotLocomotion/drake/blob/master/tutorials/README.md).\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This shows examples of:\n",
    "* Adding a `MultibodyPlant` and `SceneGraph` to a diagram\n",
    "* Adding two separate IIWAs to the `MultibodyPlant`\n",
    "* Adding `Meshcat` visualization\n",
    "* Adding a camera with a VTK renderer\n",
    "* Rendering color and label images (at zero configuration)\n",
    "* Using `SceneGraphInspector` to query `SceneGraph` geometries\n",
    "* Associating `SceneGraph` geometries with `MultibodyPlant` bodies\n",
    "* Extracting `RenderLabel`s from given geometries\n",
    "* Remapping labels to only distinguish by `ModelInstanceIndex`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Necessary Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib as mpl\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pydrake.common import FindResourceOrThrow\n",
    "from pydrake.geometry import MeshcatVisualizerCpp, StartMeshcat\n",
    "from pydrake.geometry.render import (\n",
    "    ClippingRange,\n",
    "    ColorRenderCamera,\n",
    "    DepthRange,\n",
    "    DepthRenderCamera,\n",
    "    RenderCameraCore,\n",
    "    RenderLabel,\n",
    "    MakeRenderEngineVtk,\n",
    "    RenderEngineVtkParams,\n",
    ")\n",
    "from pydrake.math import RigidTransform, RollPitchYaw\n",
    "from pydrake.multibody.parsing import Parser\n",
    "from pydrake.multibody.plant import AddMultibodyPlantSceneGraph\n",
    "from pydrake.multibody.tree import BodyIndex\n",
    "from pydrake.systems.analysis import Simulator\n",
    "from pydrake.systems.framework import DiagramBuilder\n",
    "from pydrake.systems.sensors import (\n",
    "    CameraInfo,\n",
    "    RgbdSensor,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll also start a `Meshcat` instance here.  It is typical to do this once at the beginning of the notebook, and use the same instance throughout.  You should open a second browser window to the URL that is displayed in the cell output."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "meshcat = StartMeshcat()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define helper methods"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def xyz_rpy_deg(xyz, rpy_deg):\n",
    "    \"\"\"Shorthand for defining a pose.\"\"\"\n",
    "    rpy_deg = np.asarray(rpy_deg)\n",
    "    return RigidTransform(RollPitchYaw(rpy_deg * np.pi / 180), xyz)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "reserved_labels = [\n",
    "    RenderLabel.kDoNotRender,\n",
    "    RenderLabel.kDontCare,\n",
    "    RenderLabel.kEmpty,\n",
    "    RenderLabel.kUnspecified,\n",
    "]\n",
    "\n",
    "def colorize_labels(image):\n",
    "    \"\"\"Colorizes labels.\"\"\"\n",
    "    # TODO(eric.cousineau): Revive and use Kuni's palette.\n",
    "    cc = mpl.colors.ColorConverter()\n",
    "    color_cycle = plt.rcParams[\"axes.prop_cycle\"]\n",
    "    colors = np.array([cc.to_rgb(c[\"color\"]) for c in color_cycle])\n",
    "    bg_color = [0, 0, 0]\n",
    "    image = np.squeeze(image)\n",
    "    background = np.zeros(image.shape[:2], dtype=bool)\n",
    "    for label in reserved_labels:\n",
    "        background |= image == int(label)\n",
    "    color_image = colors[image % len(colors)]\n",
    "    color_image[background] = bg_color\n",
    "    return color_image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def colorize_depth(depth, invalid_color=(100, 0, 0)):\n",
    "    \"\"\"\n",
    "    Colorizes depth (in meters) as an RGB image.\n",
    "\n",
    "    In general, it'd be better to use a matplotlib colormap, but\n",
    "    this simplifies handling for invalid depth values.\n",
    "    \"\"\"\n",
    "    assert depth.dtype == np.float32\n",
    "    invalid = (depth <= 0) | ~np.isfinite(depth)\n",
    "    scale_min = np.min(depth[~invalid])\n",
    "    scale_max = np.max(depth[~invalid])\n",
    "    # Normalize.\n",
    "    depth_scaled = (depth - scale_min) / (scale_max - scale_min)\n",
    "    h, w = depth_scaled.shape\n",
    "    # Convert to rgb.\n",
    "    depth_rgb = np.tile(depth_scaled.reshape((h, w, 1)), (1, 1, 3))\n",
    "    # Invert coloring (white is up front).\n",
    "    depth_rgb = 1 - depth_rgb\n",
    "    # Recolor.\n",
    "    depth_rgb = (depth_rgb * 255).astype(np.uint8)\n",
    "    depth_rgb[invalid] = np.array(invalid_color)\n",
    "    return depth_rgb"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create diagram builder with plant and scene graph."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "builder = DiagramBuilder()\n",
    "plant, scene_graph = AddMultibodyPlantSceneGraph(builder, 0.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "parser = Parser(plant)\n",
    "iiwa_file = FindResourceOrThrow(\n",
    "   \"drake/manipulation/models/iiwa_description/sdf/\"\n",
    "   \"iiwa14_no_collision.sdf\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Add first IIWA on left side."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "iiwa_1 = parser.AddModelFromFile(iiwa_file, model_name=\"iiwa_1\")\n",
    "plant.WeldFrames(\n",
    "    frame_on_parent_P=plant.world_frame(),\n",
    "    frame_on_child_C=plant.GetFrameByName(\"iiwa_link_0\", iiwa_1),\n",
    "    X_PC=xyz_rpy_deg([0, -0.5, 0], [0, 0, 0]),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Add second IIWA on right side."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "iiwa_2 = parser.AddModelFromFile(iiwa_file, model_name=\"iiwa_2\")\n",
    "plant.WeldFrames(\n",
    "    frame_on_parent_P=plant.world_frame(),\n",
    "    frame_on_child_C=plant.GetFrameByName(\"iiwa_link_0\", iiwa_2),\n",
    "    X_PC=xyz_rpy_deg([0, 0.5, 0], [0, 0, 0]),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Add renderer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "renderer_name = \"renderer\"\n",
    "scene_graph.AddRenderer(\n",
    "    renderer_name, MakeRenderEngineVtk(RenderEngineVtkParams()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Add camera with same color and depth properties."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# N.B. These properties are chosen arbitrarily.\n",
    "intrinsics = CameraInfo(\n",
    "    width=640,\n",
    "    height=480,\n",
    "    fov_y=np.pi/4,\n",
    ")\n",
    "core = RenderCameraCore(\n",
    "    renderer_name,\n",
    "    intrinsics,\n",
    "    ClippingRange(0.01, 10.0),\n",
    "    RigidTransform(),\n",
    ")\n",
    "color_camera = ColorRenderCamera(core, show_window=False)\n",
    "depth_camera = DepthRenderCamera(core, DepthRange(0.01, 10.0))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Because we add this to `plant.world_body()`, this is a *scene-fixed* camera.\n",
    "\n",
    "To make it a moving camera, simply add it to a body that is not anchored\n",
    "against the world."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "world_id = plant.GetBodyFrameIdOrThrow(plant.world_body().index())\n",
    "X_WB = xyz_rpy_deg([2, 0, 0.75], [-90, 0, 90])\n",
    "sensor = RgbdSensor(\n",
    "    world_id,\n",
    "    X_PB=X_WB,\n",
    "    color_camera=color_camera,\n",
    "    depth_camera=depth_camera,\n",
    ")\n",
    "\n",
    "builder.AddSystem(sensor)\n",
    "builder.Connect(\n",
    "    scene_graph.get_query_output_port(),\n",
    "    sensor.query_object_input_port(),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Add `Meshcat` Visualizer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "meshcat_vis = MeshcatVisualizerCpp.AddToBuilder(builder, scene_graph, meshcat)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Finalize plant and build diagram."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plant.Finalize()\n",
    "diagram = builder.Build()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Publish initial visualization message with default context."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO(eric.cousineau): Replace this with `diagram.Publish(diagram_context)`\n",
    "# once all visualizers no longer use initialization events.\n",
    "Simulator(diagram).Initialize()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create context and get subsystem contexts."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "diagram_context = diagram.CreateDefaultContext()\n",
    "sensor_context = sensor.GetMyMutableContextFromRoot(diagram_context)\n",
    "sg_context = scene_graph.GetMyMutableContextFromRoot(diagram_context)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Render color and label images using matplotlib\n",
    "Note that this uses the default labeling scheme, using `body.index()`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "color = sensor.color_image_output_port().Eval(sensor_context).data\n",
    "depth = sensor.depth_image_32F_output_port().Eval(sensor_context).data.squeeze(2)\n",
    "label = sensor.label_image_output_port().Eval(sensor_context).data\n",
    "\n",
    "fig, ax = plt.subplots(1, 3, figsize=(15, 10))\n",
    "ax[0].imshow(color)\n",
    "ax[1].imshow(colorize_depth(depth))\n",
    "ax[2].imshow(colorize_labels(label))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Change labels to model instance by direct remapping."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will loop through each geometry item, get it's corresponding body, and then\n",
    "remap the original label to a label that is only distinct for model instances.\n",
    "\n",
    "**NOTE**: If the labels in the given plant merge the two model instances together,\n",
    "then this will not fix that. The correct behavior would be to update the RenderLabel\n",
    "values in the geometry instances themselves.\n",
    "\n",
    "*TODO(eric.cousineau)*: Add an example of updating only the RenderLabel."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "query_object = scene_graph.get_query_output_port().Eval(sg_context)\n",
    "inspector = query_object.inspector()\n",
    "\n",
    "label_by_model = label.copy()\n",
    "for geometry_id in inspector.GetAllGeometryIds():\n",
    "    body = plant.GetBodyFromFrameId(inspector.GetFrameId(geometry_id))\n",
    "    geometry_label = inspector.GetPerceptionProperties(\n",
    "        geometry_id).GetProperty(\"label\", \"id\")\n",
    "    # WARNING: If you do not cast the `geometry_label` to `int`, this\n",
    "    # comparison will take a long time since NumPy will do\n",
    "    # element-by-element comparison using `RenderLabel.__eq__`.\n",
    "    mask = (label == int(geometry_label))\n",
    "    label_by_model[mask] = int(body.model_instance())\n",
    "\n",
    "plt.figure(figsize=(5, 5))\n",
    "plt.imshow(colorize_labels(label_by_model))"
   ]
  }
 ],
 "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}