{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Rendering MultibodyPlant Tutorial\n",
    "For instructions on how to run these tutorial notebooks, please see the [index](./index.ipynb).\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 default 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 copy\n",
    "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.geometry import (\n",
    "    ClippingRange,\n",
    "    ColorRenderCamera,\n",
    "    DepthRange,\n",
    "    DepthRenderCamera,\n",
    "    MakeRenderEngineVtk,\n",
    "    RenderCameraCore,\n",
    "    RenderEngineVtkParams,\n",
    "    RenderLabel,\n",
    "    Role,\n",
    "    StartMeshcat,\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",
    ")\n",
    "from pydrake.visualization import (\n",
    "    AddDefaultVisualization,\n",
    "    ColorizeDepthImage,\n",
    "    ColorizeLabelImage,\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": "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": [
    "iiwa_url = (\n",
    "   \"package://drake_models/iiwa_description/sdf/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": [
    "(left_iiwa,) = Parser(plant, \"left\").AddModels(url=iiwa_url)\n",
    "plant.WeldFrames(\n",
    "    frame_on_parent_F=plant.world_frame(),\n",
    "    frame_on_child_M=plant.GetFrameByName(\"iiwa_link_0\", left_iiwa),\n",
    "    X_FM=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": [
    "(right_iiwa,) = Parser(plant, \"right\").AddModels(url=iiwa_url)\n",
    "plant.WeldFrames(\n",
    "    frame_on_parent_F=plant.world_frame(),\n",
    "    frame_on_child_M=plant.GetFrameByName(\"iiwa_link_0\", right_iiwa),\n",
    "    X_FM=xyz_rpy_deg([0, 0.5, 0], [0, 0, 0]),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Add a free body."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add a mesh from https://github.com/RobotLocomotion/models/ directly, without\n",
    "# using the SDFormat wrapper file.\n",
    "sugar_box_url = \"package://drake_models/ycb/meshes/004_sugar_box_textured.obj\"\n",
    "(sugar_box,) = Parser(plant).AddModels(url=sugar_box_url)\n",
    "sugar_box_body = plant.GetBodyByName(\"004_sugar_box_textured\", sugar_box)\n",
    "plant.SetDefaultFreeBodyPose(sugar_box_body, xyz_rpy_deg([0, 0, 0.5], [0, 0, 0]))"
   ]
  },
  {
   "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 depth and label colorizers.\n",
    "The systems convert a depth (or label) image into a color image, so we can easily preview it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "colorize_depth = builder.AddSystem(ColorizeDepthImage())\n",
    "colorize_label = builder.AddSystem(ColorizeLabelImage())\n",
    "colorize_label.background_color.set([0,0,0])\n",
    "builder.Connect(sensor.GetOutputPort(\"depth_image_32f\"),\n",
    "                colorize_depth.GetInputPort(\"depth_image_32f\"))\n",
    "builder.Connect(sensor.GetOutputPort(\"label_image\"),\n",
    "                colorize_label.GetInputPort(\"label_image\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Finalize the plant."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plant.Finalize()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Add visualization."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Adding visualization using default settings for both `Meshcat` and `DrakeVisualizer` can be done with a single call to `AddDefaultVisualization()`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "AddDefaultVisualization(builder=builder, meshcat=meshcat)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Build the diagram."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "diagram = builder.Build()\n",
    "diagram_context = diagram.CreateDefaultContext()"
   ]
  },
  {
   "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": [
    "## 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(\n",
    "    sensor.GetMyContextFromRoot(diagram_context)).data\n",
    "depth = colorize_depth.get_output_port().Eval(\n",
    "    colorize_depth.GetMyContextFromRoot(diagram_context)).data\n",
    "label = colorize_label.get_output_port().Eval(\n",
    "    colorize_label.GetMyContextFromRoot(diagram_context)).data\n",
    "\n",
    "fig, ax = plt.subplots(1, 3, figsize=(15, 10))\n",
    "ax[0].imshow(color)\n",
    "ax[1].imshow(depth)\n",
    "ax[2].imshow(label)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Change labels to model instance instead of body index."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll loop through each geometry item and change its label to reflect the model instance rather than body index."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "source_id = plant.get_source_id()\n",
    "scene_graph_context = scene_graph.GetMyMutableContextFromRoot(diagram_context)\n",
    "query_object = scene_graph.get_query_output_port().Eval(scene_graph_context)\n",
    "inspector = query_object.inspector()\n",
    "for geometry_id in inspector.GetAllGeometryIds():\n",
    "    properties = copy.deepcopy(inspector.GetPerceptionProperties(geometry_id))\n",
    "    if properties is None:\n",
    "        continue\n",
    "    frame_id = inspector.GetFrameId(geometry_id)\n",
    "    body = plant.GetBodyFromFrameId(frame_id)\n",
    "    new_label = int(body.model_instance())\n",
    "    properties.UpdateProperty(\"label\", \"id\", RenderLabel(new_label))\n",
    "    # TODO(#19123) Ideally we would use AssignRole(..., kReplace) here,\n",
    "    # but it is not yet supported for perception properties.\n",
    "    scene_graph.RemoveRole(scene_graph_context, source_id, geometry_id, Role.kPerception)\n",
    "    scene_graph.AssignRole(scene_graph_context, source_id, geometry_id, properties)\n",
    "\n",
    "label = colorize_label.get_output_port().Eval(\n",
    "    colorize_label.GetMyContextFromRoot(diagram_context)).data\n",
    "plt.figure(figsize=(5, 5))\n",
    "plt.imshow(label)"
   ]
  }
 ],
 "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
}