graph_matching/GraphVisualizations.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "92d6a957",
   "metadata": {},
   "source": [
    "# Visualizing and Debugging Graphs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1a477f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from tbp.monty.frameworks.utils.logging_utils import load_stats\n",
    "from tbp.monty.frameworks.utils.plot_utils import plot_graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12f68b5e",
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib notebook"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "52299ed5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# exp data not really loaded, just for the load func to run\n",
    "# Needs to be an actual path if looking at intermediate models during training\n",
    "placeholder_exp_path = os.path.expanduser(\"~/tbp/tbp.monty/projects/monty_runs/logs/feature_pred_tests/\")\n",
    "\n",
    "sensor_pretrained_dict = os.path.expanduser(\"~/tbp/tbp.monty/projects/monty_runs/pretrained_models_ycb_high_res/supervised_pre_training/pretrained/\")\n",
    "_, _, _, lm_models_sensor = load_stats(placeholder_exp_path,\n",
    "                                load_train=False,\n",
    "                                load_eval=False,\n",
    "                                load_detailed=False,\n",
    "                                pretrained_dict=sensor_pretrained_dict,\n",
    "                               )\n",
    "# feature experiments on mesh models:\n",
    "mesh_pretrained_dict = os.path.expanduser(\"~/tbp/tbp.monty/projects/monty_runs/pretrained_models_from_mesh/ycb_mesh_models_low_res/\")\n",
    "mesh_pretrained_dict=os.path.expanduser(\"~/tbp/tbp.monty/projects/feature_eval_runs/pretrained_models_from_mesh/ycb4_3000p/\")\n",
    "_, _, _, lm_models_mesh = load_stats(placeholder_exp_path,\n",
    "                                load_train=False,\n",
    "                                load_eval=False,\n",
    "                                load_detailed=False,\n",
    "                                pretrained_dict=mesh_pretrained_dict,\n",
    "                               )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3cf684c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "pretrain_path = os.path.expanduser(\"~/tbp/results/monty/pretrained_models/\")\n",
    "pretrained_dict = pretrain_path + \"pretrained_ycb_v3/touch_1lm_10distinctobj/pretrained/\"\n",
    "\n",
    "sensor_pretrained_dict = os.path.expanduser(\"~/tbp/tbp.monty/projects/monty_runs/pretrained_models_ycb_high_res/supervised_pre_training/pretrained/\")\n",
    "_, _, _, lm_models_sensor = load_stats(placeholder_exp_path,\n",
    "                                load_train=False,\n",
    "                                load_eval=False,\n",
    "                                load_detailed=False,\n",
    "                                pretrained_dict=pretrained_dict,\n",
    "                               )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1e481752",
   "metadata": {},
   "source": [
    "### Scaling Curvature"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9e2e395a",
   "metadata": {},
   "outputs": [],
   "source": [
    "object_n = 'mug'\n",
    "lm_models = lm_models_sensor\n",
    "\n",
    "pc_ids_g = lm_models['pretrained'][0][object_n].feature_mapping['gaussian_curvature']\n",
    "gc = np.array(lm_models['pretrained'][0][object_n].x[:,pc_ids_g[0]:pc_ids_g[1]])\n",
    "pc_ids_m = lm_models['pretrained'][0][object_n].feature_mapping['mean_curvature']\n",
    "mc = np.array(lm_models['pretrained'][0][object_n].x[:,pc_ids_m[0]:pc_ids_m[1]])\n",
    "\n",
    "pc_ids_pc = lm_models['pretrained'][0][object_n].feature_mapping['principal_curvatures']\n",
    "pc = np.array(lm_models['pretrained'][0][object_n].x[:,pc_ids_pc[0]:pc_ids_pc[1]])\n",
    "\n",
    "def scale(to_scale):\n",
    "    negative = to_scale < 0\n",
    "    scaled = np.sqrt(np.abs(to_scale))\n",
    "    scaled[negative] = scaled[negative]*-1\n",
    "    return scaled\n",
    "\n",
    "def scale_clip(to_scale, clip):\n",
    "    to_scale = np.clip(to_scale, -clip, clip)\n",
    "    negative = to_scale < 0\n",
    "    scaled = np.sqrt(np.abs(to_scale))\n",
    "    scaled[negative] = scaled[negative]*-1\n",
    "    return scaled\n",
    "\n",
    "def log_sign(to_scale):\n",
    "    negative = to_scale < 0\n",
    "    scaled = np.log(np.abs(to_scale))\n",
    "    scaled[negative] = scaled[negative]*-1\n",
    "    return scaled\n",
    "\n",
    "gc = scale_clip(gc, 4096)\n",
    "mc = scale_clip(mc, 256)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5357de9e",
   "metadata": {},
   "outputs": [],
   "source": [
    "pcs = scale_clip(pc, 256)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f0decc7d",
   "metadata": {},
   "outputs": [],
   "source": [
    "pc_log = log_sign(pc)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6c6c2882",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "gc_range = (-64,64)\n",
    "mc_range = (-16,16)\n",
    "\n",
    "plt.figure(figsize=(10,7))\n",
    "plt.subplot(2,2,1)\n",
    "plt.hist(gc, bins=100, range=gc_range)\n",
    "plt.title('Gaussian Curvature')\n",
    "plt.subplot(2,2,2)\n",
    "plt.hist(mc, bins=100, range=mc_range)\n",
    "plt.title(\"Mean Curvature\")\n",
    "\n",
    "ax = plt.subplot(2,2,3,projection='3d')\n",
    "graph = lm_models['pretrained'][0][object_n]\n",
    "model_pos = graph.pos\n",
    "s = ax.scatter(\n",
    "    model_pos[:, 0],\n",
    "    model_pos[:, 1],\n",
    "    model_pos[:, 2],\n",
    "    s=5,\n",
    "    vmin=gc_range[0],\n",
    "    vmax=gc_range[1],\n",
    "    cmap='seismic',\n",
    "    c=gc\n",
    ")\n",
    "ax.set_box_aspect([1,1,1])\n",
    "ax.set_xticks([]), ax.set_yticks([]), ax.set_zticks([])\n",
    "\n",
    "ax2 = plt.subplot(2,2,4,projection='3d')\n",
    "graph = lm_models['pretrained'][0][object_n]\n",
    "model_pos = graph.pos\n",
    "s = ax2.scatter(\n",
    "    model_pos[:, 0],\n",
    "    model_pos[:, 1],\n",
    "    model_pos[:, 2],\n",
    "    s=5,\n",
    "    vmin=mc_range[0],\n",
    "    vmax=mc_range[1],\n",
    "    cmap='seismic',\n",
    "    c=mc\n",
    ")\n",
    "ax2.set_box_aspect([1,1,1])\n",
    "ax2.set_xticks([]), ax2.set_yticks([]), ax2.set_zticks([])\n",
    "plt.show()\n",
    "\n",
    "print(f\"gc in range ({max(gc)}, {min(gc)})\")\n",
    "print(f\"mc in range ({max(mc)}, {min(mc)})\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02671dae",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "plt.hist(pc)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c210a8d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "gc_range = (-256,256)\n",
    "mc_range = (-16,16)\n",
    "\n",
    "plt.figure(figsize=(10,7))\n",
    "plt.subplot(2,2,1)\n",
    "plt.hist(pc, bins=100, range=gc_range)\n",
    "plt.title('Principal Curvature')\n",
    "plt.subplot(2,2,2)\n",
    "plt.hist(pcs, bins=100, range=mc_range)\n",
    "plt.title(\"Principal Curvature Scaled\")\n",
    "\n",
    "ax = plt.subplot(2,2,3,projection='3d')\n",
    "graph = lm_models['pretrained'][0][object_n]\n",
    "model_pos = graph.pos\n",
    "s = ax.scatter(\n",
    "    model_pos[:, 0],\n",
    "    model_pos[:, 1],\n",
    "    model_pos[:, 2],\n",
    "    s=5,\n",
    "    vmin=gc_range[0],\n",
    "    vmax=gc_range[1],\n",
    "    cmap='seismic',\n",
    "    c=pc[:,0]\n",
    ")\n",
    "ax.set_box_aspect([1,1,1])\n",
    "ax.set_xticks([]), ax.set_yticks([]), ax.set_zticks([])\n",
    "\n",
    "ax2 = plt.subplot(2,2,4,projection='3d')\n",
    "graph = lm_models['pretrained'][0][object_n]\n",
    "model_pos = graph.pos\n",
    "s = ax2.scatter(\n",
    "    model_pos[:, 0],\n",
    "    model_pos[:, 1],\n",
    "    model_pos[:, 2],\n",
    "    s=5,\n",
    "    vmin=mc_range[0],\n",
    "    vmax=mc_range[1],\n",
    "    cmap='seismic',\n",
    "    c=pcs[:,0]\n",
    ")\n",
    "ax2.set_box_aspect([1,1,1])\n",
    "ax2.set_xticks([]), ax2.set_yticks([]), ax2.set_zticks([])\n",
    "plt.show()\n",
    "\n",
    "print(f\"gc in range ({max(gc)}, {min(gc)})\")\n",
    "print(f\"mc in range ({max(mc)}, {min(mc)})\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b8ee52f",
   "metadata": {},
   "outputs": [],
   "source": [
    "gc_range = (-256,256)\n",
    "mc_range = (-16,16)\n",
    "\n",
    "plt.figure(figsize=(10,7))\n",
    "plt.subplot(2,2,1)\n",
    "plt.hist(pc, bins=100, range=gc_range)\n",
    "plt.title('Principal Curvature')\n",
    "plt.subplot(2,2,2)\n",
    "plt.hist(pc_log, bins=100, range=mc_range)\n",
    "plt.title(\"Principal Curvature Logged\")\n",
    "\n",
    "ax = plt.subplot(2,2,3,projection='3d')\n",
    "graph = lm_models['pretrained'][0][object_n]\n",
    "model_pos = graph.pos\n",
    "s = ax.scatter(\n",
    "    model_pos[:, 0],\n",
    "    model_pos[:, 1],\n",
    "    model_pos[:, 2],\n",
    "    s=5,\n",
    "    vmin=mc_range[0],\n",
    "    vmax=mc_range[1],\n",
    "    cmap='seismic',\n",
    "    c=pc[:,0]\n",
    ")\n",
    "ax.set_box_aspect([1,1,1])\n",
    "ax.set_xticks([]), ax.set_yticks([]), ax.set_zticks([])\n",
    "\n",
    "ax2 = plt.subplot(2,2,4,projection='3d')\n",
    "graph = lm_models['pretrained'][0][object_n]\n",
    "model_pos = graph.pos\n",
    "s = ax2.scatter(\n",
    "    model_pos[:, 0],\n",
    "    model_pos[:, 1],\n",
    "    model_pos[:, 2],\n",
    "    s=5,\n",
    "    vmin=mc_range[0],\n",
    "    vmax=mc_range[1],\n",
    "    cmap='seismic',\n",
    "    c=pc_log[:,0]\n",
    ")\n",
    "ax2.set_box_aspect([1,1,1])\n",
    "ax2.set_xticks([]), ax2.set_yticks([]), ax2.set_zticks([])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c8e9337e",
   "metadata": {},
   "source": [
    "## Plot Features Stored in Models\n",
    "### Curvature"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d236d89",
   "metadata": {},
   "outputs": [],
   "source": [
    "object_name = 'mug'\n",
    "lm_models = lm_models_sensor\n",
    "\n",
    "fig = plt.figure()\n",
    "ax = plt.subplot(1,1,1,projection='3d')\n",
    "graph = lm_models['pretrained'][0][object_name]\n",
    "model_pos = graph.pos\n",
    "pc_ids = lm_models['pretrained'][0][object_name].feature_mapping['gaussian_curvature_sc']\n",
    "# pc_ids = lm_models['pretrained'][0][object_name].feature_mapping['principal_curvatures']\n",
    "curvs = graph.x[:,pc_ids[0]:pc_ids[1]]\n",
    "colors = curvs[:,0]\n",
    "s = ax.scatter(\n",
    "    model_pos[:, 0],\n",
    "    model_pos[:, 1],\n",
    "    model_pos[:, 2],\n",
    "    s=5,\n",
    "    vmin=-64,\n",
    "    vmax=64,\n",
    "    cmap='seismic',\n",
    "    c=colors\n",
    ")\n",
    "ax.set_xticks([]), ax.set_yticks([]), ax.set_zticks([])\n",
    "ax_range = 0.09\n",
    "ax.set_xlim(0-ax_range, 0+ax_range)\n",
    "# ax.set_ylim(1.5-ax_range,1.5+ax_range)\n",
    "ax.set_ylim(0-ax_range,0+ax_range)\n",
    "ax.set_zlim(0-ax_range, 0+ax_range)\n",
    "fig.colorbar(s)\n",
    "plt.title('Scaled Gaussian Curvature')\n",
    "# plt.title('Principal Curvature 1')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b912795f",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "object_name = 'mug'\n",
    "lm_models = lm_models_sensor\n",
    "\n",
    "graph = lm_models['pretrained'][0][object_name]\n",
    "model_pos = graph.pos\n",
    "fig = plt.figure(figsize=(10,7))\n",
    "\n",
    "ax = plt.subplot(2,2,1,projection='3d')\n",
    "feat_ids = lm_models['pretrained'][0][object_name].feature_mapping['gaussian_curvature']\n",
    "colors = graph.x[:,feat_ids[0]:feat_ids[1]]\n",
    "s = ax.scatter(\n",
    "    model_pos[:, 0],\n",
    "    model_pos[:, 1],\n",
    "    model_pos[:, 2],\n",
    "    s=5,\n",
    "    vmin=-1000,\n",
    "    vmax=1000,\n",
    "    cmap='seismic',\n",
    "    c=colors\n",
    ")\n",
    "ax.set_xticks([]), ax.set_yticks([]), ax.set_zticks([])\n",
    "ax.set_title('Gaussian Curvature')\n",
    "plt.colorbar(s)\n",
    "\n",
    "ax = plt.subplot(2,2,2,projection='3d')\n",
    "feat_ids = lm_models['pretrained'][0][object_name].feature_mapping['mean_curvature']\n",
    "colors = graph.x[:,feat_ids[0]:feat_ids[1]]\n",
    "s = ax.scatter(\n",
    "    model_pos[:, 0],\n",
    "    model_pos[:, 1],\n",
    "    model_pos[:, 2],\n",
    "    s=5,\n",
    "    vmin=-200,\n",
    "    vmax=200,\n",
    "    cmap='seismic',\n",
    "    c=colors\n",
    ")\n",
    "ax.set_xticks([]), ax.set_yticks([]), ax.set_zticks([])\n",
    "ax.set_title('Mean Curvature')\n",
    "plt.colorbar(s)\n",
    "\n",
    "ax = plt.subplot(2,2,3,projection='3d')\n",
    "feat_ids = lm_models['pretrained'][0][object_name].feature_mapping['gaussian_curvature_sc']\n",
    "colors = graph.x[:,feat_ids[0]:feat_ids[1]]\n",
    "s = ax.scatter(\n",
    "    model_pos[:, 0],\n",
    "    model_pos[:, 1],\n",
    "    model_pos[:, 2],\n",
    "    s=5,\n",
    "    vmin=-64,\n",
    "    vmax=64,\n",
    "    cmap='seismic',\n",
    "    c=colors\n",
    ")\n",
    "ax.set_xticks([]), ax.set_yticks([]), ax.set_zticks([])\n",
    "ax.set_title('Scaled Gaussian Curvature')\n",
    "plt.colorbar(s)\n",
    "\n",
    "ax = plt.subplot(2,2,4,projection='3d')\n",
    "feat_ids = lm_models['pretrained'][0][object_name].feature_mapping['mean_curvature_sc']\n",
    "colors = graph.x[:,feat_ids[0]:feat_ids[1]]\n",
    "s = ax.scatter(\n",
    "    model_pos[:, 0],\n",
    "    model_pos[:, 1],\n",
    "    model_pos[:, 2],\n",
    "    s=5,\n",
    "    vmin=-16,\n",
    "    vmax=16,\n",
    "    cmap='seismic',\n",
    "    c=colors\n",
    ")\n",
    "ax.set_xticks([]), ax.set_yticks([]), ax.set_zticks([])\n",
    "ax.set_title('Scaled Mean Curvature')\n",
    "plt.colorbar(s)\n",
    "\n",
    "plt.suptitle('From Sensor', fontsize=15)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ac81cd7",
   "metadata": {},
   "outputs": [],
   "source": [
    "object_name = 'potted_meat_can'\n",
    "\n",
    "graph_sensor = lm_models_sensor['pretrained'][object_name]\n",
    "model_pos_sensor = graph_sensor.pos\n",
    "graph_mesh = lm_models_mesh['pretrained'][object_name]\n",
    "model_pos_mesh = graph_mesh.pos\n",
    "fig = plt.figure(figsize=(10,7))\n",
    "\n",
    "ax = plt.subplot(2,2,1,projection='3d')\n",
    "feat_ids = lm_models_sensor['pretrained'][object_name].feature_mapping['gaussian_curvature_sc']\n",
    "colors = graph_sensor.x[:,feat_ids[0]:feat_ids[1]]\n",
    "s = ax.scatter(\n",
    "    model_pos_sensor[:, 0],\n",
    "    model_pos_sensor[:, 1],\n",
    "    model_pos_sensor[:, 2],\n",
    "    s=5,\n",
    "    vmin=-64,\n",
    "    vmax=64,\n",
    "    cmap='seismic',\n",
    "    c=colors\n",
    ")\n",
    "ax.set_xticks([]), ax.set_yticks([]), ax.set_zticks([])\n",
    "ax.set_title('Scaled Gaussian Curvature - Sensor')\n",
    "plt.colorbar(s)\n",
    "\n",
    "ax = plt.subplot(2,2,2,projection='3d')\n",
    "feat_ids = lm_models_sensor['pretrained'][object_name].feature_mapping['mean_curvature_sc']\n",
    "colors = graph_sensor.x[:,feat_ids[0]:feat_ids[1]]\n",
    "s = ax.scatter(\n",
    "    model_pos_sensor[:, 0],\n",
    "    model_pos_sensor[:, 1],\n",
    "    model_pos_sensor[:, 2],\n",
    "    s=5,\n",
    "    vmin=-16,\n",
    "    vmax=16,\n",
    "    cmap='seismic',\n",
    "    c=colors\n",
    ")\n",
    "ax.set_xticks([]), ax.set_yticks([]), ax.set_zticks([])\n",
    "ax.set_title('Scaled Mean Curvature - Sensor')\n",
    "plt.colorbar(s)\n",
    "\n",
    "ax = plt.subplot(2,2,3,projection='3d')\n",
    "feat_ids = lm_models_mesh['pretrained'][object_name].feature_mapping['gaussian_curvature_sc']\n",
    "colors = graph_mesh.x[:,feat_ids[0]:feat_ids[1]]\n",
    "s = ax.scatter(\n",
    "    model_pos_mesh[:, 0],\n",
    "    model_pos_mesh[:, 1],\n",
    "    model_pos_mesh[:, 2],\n",
    "    s=5,\n",
    "    vmin=-64,\n",
    "    vmax=64,\n",
    "    cmap='seismic',\n",
    "    c=colors\n",
    ")\n",
    "ax.set_xticks([]), ax.set_yticks([]), ax.set_zticks([])\n",
    "ax.set_title('Scaled Gaussian Curvature - Mesh')\n",
    "plt.colorbar(s)\n",
    "\n",
    "ax = plt.subplot(2,2,4,projection='3d')\n",
    "feat_ids = lm_models_mesh['pretrained'][object_name].feature_mapping['mean_curvature_sc']\n",
    "colors = graph_mesh.x[:,feat_ids[0]:feat_ids[1]]\n",
    "s = ax.scatter(\n",
    "    model_pos_mesh[:, 0],\n",
    "    model_pos_mesh[:, 1],\n",
    "    model_pos_mesh[:, 2],\n",
    "    s=5,\n",
    "    vmin=-16,\n",
    "    vmax=16,\n",
    "    cmap='seismic',\n",
    "    c=colors\n",
    ")\n",
    "ax.set_xticks([]), ax.set_yticks([]), ax.set_zticks([])\n",
    "ax.set_title('Scaled Mean Curvature - Mesh')\n",
    "plt.colorbar(s)\n",
    "\n",
    "plt.suptitle('Comparison Between Sensor and Mesh', fontsize=15)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d07bf269",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_ids_mesh = graph_mesh.feature_mapping['principal_curvatures']\n",
    "plot_ids_sensor = graph_sensor.feature_mapping['principal_curvatures']\n",
    "\n",
    "plt.figure(figsize=(10,4))\n",
    "plt.subplot(1,2,1)\n",
    "range = (-50,50)\n",
    "plt.hist(np.array(graph_sensor.x[:,plot_ids_sensor[0]]), alpha=0.6, bins=100, range=range, color='darkblue', density=True)\n",
    "plt.hist(np.array(graph_sensor.x[:,plot_ids_sensor[0]+1]), alpha=0.6, range=range, bins=100, color='lightblue', density=True)\n",
    "plt.hist(np.array(graph_mesh.x[:,plot_ids_mesh[0]]), alpha=0.6, bins=100, range=range, color=\"darkgreen\", density=True)\n",
    "plt.hist(np.array(graph_mesh.x[:,plot_ids_mesh[0] + 1]), alpha=0.6, bins=100, range=range, color=\"lightgreen\", density=True)\n",
    "plt.legend([\"sensor - k1\", \"sensor - k2\", \"mesh - k1\", \"mesh - k2\"])\n",
    "\n",
    "\n",
    "plt.subplot(1,2,2)\n",
    "plot_ids_mesh_g = graph_mesh.feature_mapping['gaussian_curvature_sc'][0]\n",
    "plot_ids_sensor_g = graph_sensor.feature_mapping['gaussian_curvature_sc'][0]\n",
    "range = (-64,64)\n",
    "plt.hist(np.array(graph_sensor.x[:,plot_ids_sensor_g]), alpha=0.6, bins=100, range=range, color='darkblue', density=True)\n",
    "plt.hist(np.array(graph_mesh.x[:,plot_ids_mesh_g]), alpha=0.6, bins=100, range=range, color=\"darkgreen\", density=True)\n",
    "plt.legend([\"sensor - gaussian c\", \"mesh - gaussian c\"])#, \"mesh - k1\", \"mesh - k2\"])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d4ac58b",
   "metadata": {},
   "outputs": [],
   "source": [
    "color_id_mesh = graph_mesh.feature_mapping['gaussian_curvature_sc'][0]\n",
    "color_id_sensor = graph_sensor.feature_mapping['gaussian_curvature_sc'][0]\n",
    "norm_len = 0.02\n",
    "\n",
    "fig = plt.figure()\n",
    "# vmin = -64\n",
    "# vmax = 64\n",
    "ax = plt.subplot(1,1,1,projection='3d')\n",
    "# s = ax.scatter(\n",
    "#     graph_mesh.pos[:, 0],\n",
    "#     graph_mesh.pos[:, 1],\n",
    "#     graph_mesh.pos[:, 2],\n",
    "#     s=5,\n",
    "#     alpha=0.5,\n",
    "#     c='green'#np.arctan(np.array(mesh_graph.x[:,color_id_mesh])/100)\n",
    "# )\n",
    "\n",
    "# ax.set_xticks([]), ax.set_yticks([]), ax.set_zticks([])\n",
    "# ax_range = 0.07\n",
    "# ax.set_xlim(0-ax_range, 0+ax_range)\n",
    "# ax.set_ylim(0.05-ax_range,0.05+ax_range)\n",
    "# ax.set_zlim(0-ax_range, 0+ax_range)\n",
    "# ax.set_title(\"Learned Locations and Point Normals\")\n",
    "\n",
    "align_offsets = [0.135, -1.475, 0.04]\n",
    "s2 = ax.scatter(\n",
    "    graph_sensor.pos[:, 0]+align_offsets[0],\n",
    "    graph_sensor.pos[:, 1]+align_offsets[1],\n",
    "    graph_sensor.pos[:, 2]+align_offsets[2],\n",
    "    s=5,\n",
    "    alpha=0.5,\n",
    "    c='blue'#np.arctan(np.array(sensor_graph.x[:,color_id_sensor])/100)\n",
    ")\n",
    "\n",
    "# s2 = ax.scatter(\n",
    "#     graph_mesh.pos[10:12, 0],\n",
    "#     graph_mesh.pos[10:12, 1],\n",
    "#     graph_mesh.pos[10:12, 2],\n",
    "#     s=10,\n",
    "#     alpha=1,\n",
    "#     c='red'#np.arctan(np.array(sensor_graph.x[:,color_id_sensor])/100)\n",
    "# )\n",
    "\n",
    "# for p_id, p in enumerate(np.array(graph_mesh.pos)):\n",
    "#     norm = graph_mesh.norm[p_id]\n",
    "#     ax.plot([p[0], p[0] + norm[0] * norm_len],\n",
    "#             [p[1], p[1] + norm[1] * norm_len],\n",
    "#             [p[2], p[2] + norm[2] * norm_len],\n",
    "#             c='lightgreen',\n",
    "#            )\n",
    "    \n",
    "for p_id, p in enumerate(np.array(graph_sensor.pos)):\n",
    "    p = p + align_offsets\n",
    "    norm = graph_sensor.norm[p_id]\n",
    "    ax.plot([p[0], p[0] + norm[0] * norm_len],\n",
    "            [p[1], p[1] + norm[1] * norm_len],\n",
    "            [p[2], p[2] + norm[2] * norm_len],\n",
    "            c='lightblue',\n",
    "           )\n",
    "    \n",
    "# plt.legend(['From Mesh','From Sensor'])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2ada01e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "node_distances = np.linalg.norm(\n",
    "                        graph_mesh.pos - graph_mesh.pos[1000],\n",
    "                        axis=1,\n",
    "                    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "361d4c04",
   "metadata": {},
   "outputs": [],
   "source": [
    "min(node_distances)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e2effe02",
   "metadata": {},
   "source": [
    "### HSV"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b448e69e",
   "metadata": {},
   "outputs": [],
   "source": [
    "object_name = 'mug'\n",
    "lm_models = lm_models_sensor\n",
    "\n",
    "fig = plt.figure()\n",
    "ax = plt.subplot(1,1,1,projection='3d')\n",
    "graph = lm_models['pretrained'][0][object_name]\n",
    "model_pos = graph.pos\n",
    "hsv_ids = lm_models['pretrained'][0][object_name].feature_mapping['hsv']\n",
    "hue = graph.x[:,hsv_ids[0]:hsv_ids[1]]\n",
    "colors = hue[:,0]\n",
    "s = ax.scatter(\n",
    "    model_pos[:, 0],\n",
    "    model_pos[:, 1],\n",
    "    model_pos[:, 2],\n",
    "    s=5,\n",
    "    vmin=0,\n",
    "    vmax=1,\n",
    "    cmap='hsv',\n",
    "    c=colors\n",
    ")\n",
    "norm_len = 0.01\n",
    "ax.set_xticks([]), ax.set_yticks([]), ax.set_zticks([])\n",
    "fig.colorbar(s)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c2cec48e",
   "metadata": {},
   "outputs": [],
   "source": [
    "object_n = \"mug\"\n",
    "norm_len=0.01\n",
    "plt.figure(figsize=(10,10))\n",
    "ax = plt.subplot(1,1,1,projection='3d')\n",
    "graph = lm_models_sensor['pretrained'][0][object_n]\n",
    "model_pos = graph.pos\n",
    "hsv_ids = lm_models['pretrained'][0][object_name].feature_mapping['hsv']\n",
    "hue = graph.x[:,hsv_ids[0]:hsv_ids[1]]\n",
    "colors = hue[:,0]\n",
    "curv_dir_ids = graph.feature_mapping['curvature_directions']\n",
    "\n",
    "s = ax.scatter(\n",
    "    model_pos[:, 0],\n",
    "    model_pos[:, 1],\n",
    "    model_pos[:, 2],\n",
    "    s=5,\n",
    "    cmap='hsv',\n",
    "    c=colors\n",
    ")\n",
    "# for p_id, p in enumerate(np.array(model_pos)):\n",
    "#     norm = graph.norm[p_id]\n",
    "#     ax.plot([p[0], p[0] + norm[0] * norm_len],\n",
    "#             [p[1], p[1] + norm[1] * norm_len],\n",
    "#             [p[2], p[2] + norm[2] * norm_len],\n",
    "#             c='royalblue',\n",
    "#            )\n",
    "#     cd1 = graph.x[p_id][curv_dir_ids[0]:curv_dir_ids[0]+3]\n",
    "#     cd2 = graph.x[p_id][curv_dir_ids[0]+3:curv_dir_ids[0]+6]\n",
    "#     ax.plot([p[0], p[0] + cd1[0] * norm_len],\n",
    "#             [p[1], p[1] + cd1[1] * norm_len],\n",
    "#             [p[2], p[2] + cd1[2] * norm_len],\n",
    "#             c='red',\n",
    "#            )\n",
    "#     ax.plot([p[0], p[0] + cd2[0] * norm_len * 0.5],\n",
    "#             [p[1], p[1] + cd2[1] * norm_len * 0.5],\n",
    "#             [p[2], p[2] + cd2[2] * norm_len * 0.5],\n",
    "#             c='orange',\n",
    "#            )\n",
    "ax.set_aspect('equal')\n",
    "ax.set_xticks([]), ax.set_yticks([]), ax.set_zticks([])\n",
    "ax.view_init(-45, 90)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b278d06a",
   "metadata": {},
   "source": [
    "### Locations (general plotting function)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "040c5028",
   "metadata": {},
   "outputs": [],
   "source": [
    "mug_graph = lm_models['pretrained'][0]['mug']\n",
    "bowl_graph = lm_models['pretrained'][0]['bowl']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99434f07",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_graph(mug_graph, show_edges=False, rotation=20)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9166ec8a",
   "metadata": {},
   "source": [
    "# Feature Change Models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2163ebc0",
   "metadata": {},
   "outputs": [],
   "source": [
    "log_path = os.path.expanduser(\"~/tbp/tbp.monty/projects/monty_runs/logs/\")\n",
    "placeholder_exp_path = log_path + \"feature_pred_tests/\"\n",
    "\n",
    "pretrain_path = os.path.expanduser(\"~/tbp/data/pretrained_models/\")\n",
    "# sensor_pretrained_dict_fc = pretrain_path + \"ycb_high_res/supervised_pre_training/pretrained/\"\n",
    "\n",
    "\n",
    "sensor_pretrained_dict_fc = pretrain_path + \"ycb_feature_change_S1_low_tol/supervised_pre_training_feature_change/pretrained/\"\n",
    "sensor_pretrained_dict_fc = pretrain_path + \"ycb_spiral_policy_s5/supervised_pre_training/pretrained\"\n",
    "_, _, _, lm_models_sensor_fc = load_stats(placeholder_exp_path,\n",
    "                                load_train=False,\n",
    "                                load_eval=False,\n",
    "                                load_detailed=False,\n",
    "                                pretrained_dict=sensor_pretrained_dict_fc,\n",
    "                               )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "db366e81",
   "metadata": {},
   "outputs": [],
   "source": [
    "object_name = 'master_chef_can'\n",
    "lm_models = lm_models_sensor_fc\n",
    "\n",
    "fig = plt.figure()\n",
    "ax = plt.subplot(1,1,1,projection='3d')\n",
    "graph = lm_models['pretrained'][0][object_name]\n",
    "model_pos = graph.pos\n",
    "hsv_ids = lm_models['pretrained'][0][object_name].feature_mapping['hsv']\n",
    "hue = graph.x[:,hsv_ids[0]:hsv_ids[0]+1]\n",
    "colors = hue[:,0]\n",
    "s = ax.scatter(\n",
    "    model_pos[:, 0],\n",
    "    model_pos[:, 1],\n",
    "    model_pos[:, 2],\n",
    "    s=5,\n",
    "    vmin=0,\n",
    "    vmax=1,\n",
    "    cmap='hsv',\n",
    "    c=colors\n",
    ")\n",
    "norm_len = 0.01\n",
    "ax.set_box_aspect([1,1,1])\n",
    "ax.set_xticks([]), ax.set_yticks([]), ax.set_zticks([])\n",
    "fig.colorbar(s)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a4a6c14",
   "metadata": {},
   "outputs": [],
   "source": [
    "object_name = 'mug'\n",
    "lm_models = lm_models_sensor_fc\n",
    "\n",
    "fig = plt.figure()\n",
    "ax = plt.subplot(1,1,1,projection='3d')\n",
    "graph = lm_models['pretrained'][0][object_name]\n",
    "model_pos = graph.pos\n",
    "# hsv_ids = lm_models['pretrained'][0][object_name].feature_mapping['hsv']\n",
    "curv_ids = lm_models['pretrained'][0][object_name].feature_mapping['principal_curvatures_log']\n",
    "curv = graph.x[:,curv_ids[0]:curv_ids[1]]\n",
    "colors = curv[:,0]\n",
    "s = ax.scatter(\n",
    "    model_pos[:, 0],\n",
    "    model_pos[:, 1],\n",
    "    model_pos[:, 2],\n",
    "    s=5,\n",
    "    vmin=-10,\n",
    "    vmax=10,\n",
    "    cmap='seismic',\n",
    "    c=colors\n",
    ")\n",
    "norm_len = 0.01\n",
    "ax.set_box_aspect([1,1,0.7])\n",
    "ax.set_xticks([]), ax.set_yticks([]), ax.set_zticks([])\n",
    "fig.colorbar(s)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9fd522fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_pos.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ac48c0a",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "plt.hist(np.array(curv))\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6b589276",
   "metadata": {},
   "source": [
    "## Augmented Sparse Models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8baa850b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# placeholder_exp_path = os.path.expanduser(\"~/tbp/tbp.monty/projects/monty_runs/evidence_tests/\")\n",
    "placeholder_exp_path = os.path.expanduser(\"~/tbp/results/monty/projects/monty_runs/evidence_tests/\")\n",
    "save_path = os.path.expanduser(\"~/tbp/results/monty/figures/models/\")\n",
    "\n",
    "pretrain_path = os.path.expanduser(\"~/tbp/results/monty/pretrained_models/\")\n",
    "\n",
    "pretrained_dict_0001 = pretrain_path + \"pretrained_ycb/supervised_pre_training_10obj_default_base/pretrained/\"\n",
    "_, _, _, lm_models_0001 = load_stats(placeholder_exp_path,\n",
    "                                load_train=False,\n",
    "                                load_eval=False,\n",
    "                                load_detailed=False,\n",
    "                                pretrained_dict=pretrained_dict_0001,\n",
    "                               )\n",
    "\n",
    "pretrained_dict_001 = pretrain_path + \"pretrained_ycb/supervised_pre_training_10obj_001dist_base/pretrained/\"\n",
    "_, _, _, lm_models_001 = load_stats(placeholder_exp_path,\n",
    "                                load_train=False,\n",
    "                                load_eval=False,\n",
    "                                load_detailed=False,\n",
    "                                pretrained_dict=pretrained_dict_001,\n",
    "                               )\n",
    "\n",
    "pretrained_dict_005 = pretrain_path + \"pretrained_ycb/supervised_pre_training_10obj_005dist_base/pretrained/\"\n",
    "_, _, _, lm_models_005 = load_stats(placeholder_exp_path,\n",
    "                                load_train=False,\n",
    "                                load_eval=False,\n",
    "                                load_detailed=False,\n",
    "                                pretrained_dict=pretrained_dict_005,\n",
    "                               )\n",
    "\n",
    "pretrained_dict_01 = pretrain_path + \"pretrained_ycb/supervised_pre_training_10obj_01dist_base/pretrained/\"\n",
    "_, _, _, lm_models_01 = load_stats(placeholder_exp_path,\n",
    "                                load_train=False,\n",
    "                                load_eval=False,\n",
    "                                load_detailed=False,\n",
    "                                pretrained_dict=pretrained_dict_01,\n",
    "                               )\n",
    "\n",
    "# ------ augmented\n",
    "pretrained_dict_aug_0001 = pretrain_path + \"pretrained_ycb/supervised_additional_training_touch_0001dist_10obj/pretrained/\"\n",
    "_, _, _, lm_models_aug_0001 = load_stats(placeholder_exp_path,\n",
    "                                load_train=False,\n",
    "                                load_eval=False,\n",
    "                                load_detailed=False,\n",
    "                                pretrained_dict=pretrained_dict_aug_0001,\n",
    "                               )\n",
    "\n",
    "pretrained_dict_aug_001 = pretrain_path + \"pretrained_ycb/supervised_additional_training_touch_001dist_10obj/pretrained/\"\n",
    "_, _, _, lm_models_aug_001 = load_stats(placeholder_exp_path,\n",
    "                                load_train=False,\n",
    "                                load_eval=False,\n",
    "                                load_detailed=False,\n",
    "                                pretrained_dict=pretrained_dict_aug_001,\n",
    "                               )\n",
    "\n",
    "pretrained_dict_aug_005 = pretrain_path + \"pretrained_ycb/supervised_additional_training_touch_005dist_10obj/pretrained/\"\n",
    "_, _, _, lm_models_aug_005 = load_stats(placeholder_exp_path,\n",
    "                                load_train=False,\n",
    "                                load_eval=False,\n",
    "                                load_detailed=False,\n",
    "                                pretrained_dict=pretrained_dict_aug_005,\n",
    "                               )\n",
    "\n",
    "pretrained_dict_aug_01 = pretrain_path + \"pretrained_ycb/supervised_additional_training_touch_10obj_01/pretrained/\"\n",
    "_, _, _, lm_models_aug_01 = load_stats(placeholder_exp_path,\n",
    "                                load_train=False,\n",
    "                                load_eval=False,\n",
    "                                load_detailed=False,\n",
    "                                pretrained_dict=pretrained_dict_aug_01,\n",
    "                               )\n",
    "\n",
    "pretrained_dict_aug_01on005 = pretrain_path + \"pretrained_ycb/supervised_additional_training_touch_01on005_10obj/pretrained/\"\n",
    "_, _, _, lm_models_aug_01on005 = load_stats(placeholder_exp_path,\n",
    "                                load_train=False,\n",
    "                                load_eval=False,\n",
    "                                load_detailed=False,\n",
    "                                pretrained_dict=pretrained_dict_aug_01on005,\n",
    "                               )\n",
    "\n",
    "pretrained_dict_touch_005 = pretrain_path + \"pretrained_ycb/only_touch_training_10obj_005dist/pretrained/\"\n",
    "_, _, _, lm_models_touch005 = load_stats(placeholder_exp_path,\n",
    "                                load_train=False,\n",
    "                                load_eval=False,\n",
    "                                load_detailed=False,\n",
    "                                pretrained_dict=pretrained_dict_touch_005,\n",
    "                               )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "be34c6d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "lm_models_aug_0001['0'].keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df39b5ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "lm_models_0001['pretrained'][0].keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a087034",
   "metadata": {},
   "outputs": [],
   "source": [
    "object_n = 'spoon'\n",
    "\n",
    "plt.figure(figsize=(15,12))\n",
    "ax = plt.subplot(3,4,1,projection='3d')\n",
    "graph = lm_models_0001['pretrained'][0][object_n]\n",
    "model_pos = graph.pos\n",
    "s = ax.scatter(\n",
    "    model_pos[:, 0],\n",
    "    model_pos[:, 1],\n",
    "    model_pos[:, 2],\n",
    "    s=1,\n",
    ")\n",
    "ax.set_aspect('equal')\n",
    "ax.set_xticks([]), ax.set_yticks([]), ax.set_zticks([])\n",
    "ax.view_init(-45, 90)\n",
    "plt.title(f\"mmd=0.0001\\nnum_points = {graph.x.shape[0]}\")\n",
    "\n",
    "ax = plt.subplot(3,4,2,projection='3d')\n",
    "graph = lm_models_001['pretrained'][0][object_n]\n",
    "model_pos = graph.pos\n",
    "s = ax.scatter(\n",
    "    model_pos[:, 0],\n",
    "    model_pos[:, 1],\n",
    "    model_pos[:, 2],\n",
    "    s=1,\n",
    ")\n",
    "ax.set_aspect('equal')\n",
    "ax.set_xticks([]), ax.set_yticks([]), ax.set_zticks([])\n",
    "ax.view_init(-45, 90)\n",
    "plt.title(f\"mmd=0.001\\nnum_points = {graph.x.shape[0]}\")\n",
    "\n",
    "ax = plt.subplot(3,4,3,projection='3d')\n",
    "graph = lm_models_005['pretrained'][0][object_n]\n",
    "model_pos = graph.pos\n",
    "s = ax.scatter(\n",
    "    model_pos[:, 0],\n",
    "    model_pos[:, 1],\n",
    "    model_pos[:, 2],\n",
    "    s=1,\n",
    ")\n",
    "ax.set_aspect('equal')\n",
    "ax.set_xticks([]), ax.set_yticks([]), ax.set_zticks([])\n",
    "ax.view_init(-45, 90)\n",
    "plt.title(f\"mmd=0.005\\nnum_points = {graph.x.shape[0]}\")\n",
    "\n",
    "ax = plt.subplot(3,4,4,projection='3d')\n",
    "graph = lm_models_01['pretrained'][0][object_n]\n",
    "model_pos = graph.pos\n",
    "s = ax.scatter(\n",
    "    model_pos[:, 0],\n",
    "    model_pos[:, 1],\n",
    "    model_pos[:, 2],\n",
    "    s=1,\n",
    ")\n",
    "ax.set_aspect('equal')\n",
    "ax.set_xticks([]), ax.set_yticks([]), ax.set_zticks([])\n",
    "ax.view_init(-45, 90)\n",
    "plt.title(f\"mmd=0.01\\nnum_points = {graph.x.shape[0]}\")\n",
    "\n",
    "# -------\n",
    "\n",
    "ax = plt.subplot(3,4,5,projection='3d')\n",
    "graph = lm_models_aug_0001['pretrained'][0][object_n]\n",
    "model_pos = graph.pos\n",
    "s = ax.scatter(\n",
    "    model_pos[:, 0],\n",
    "    model_pos[:, 1],\n",
    "    model_pos[:, 2],\n",
    "    s=1,\n",
    ")\n",
    "ax.set_aspect('equal')\n",
    "ax.set_xticks([]), ax.set_yticks([]), ax.set_zticks([])\n",
    "ax.view_init(-45, 90)\n",
    "plt.title(f\"mmd=0.0001\\nnum_points = {graph.x.shape[0]}\")\n",
    "\n",
    "ax = plt.subplot(3,4,6,projection='3d')\n",
    "graph = lm_models_aug_001['pretrained'][0][object_n]\n",
    "model_pos = graph.pos\n",
    "s = ax.scatter(\n",
    "    model_pos[:, 0],\n",
    "    model_pos[:, 1],\n",
    "    model_pos[:, 2],\n",
    "    s=1,\n",
    ")\n",
    "ax.set_aspect('equal')\n",
    "ax.set_xticks([]), ax.set_yticks([]), ax.set_zticks([])\n",
    "ax.view_init(-45, 90)\n",
    "plt.title(f\"models learned with vision + touch\\n\\n\"\n",
    "    f\"mmd=0.001\\nnum_points = {graph.x.shape[0]}\")\n",
    "\n",
    "ax = plt.subplot(3,4,7,projection='3d')\n",
    "graph = lm_models_aug_005['pretrained'][0][object_n]\n",
    "model_pos = graph.pos\n",
    "s = ax.scatter(\n",
    "    model_pos[:, 0],\n",
    "    model_pos[:, 1],\n",
    "    model_pos[:, 2],\n",
    "    s=1,\n",
    ")\n",
    "ax.set_aspect('equal')\n",
    "ax.set_xticks([]), ax.set_yticks([]), ax.set_zticks([])\n",
    "ax.view_init(-45, 90)\n",
    "plt.title(f\"mmd=0.005\\nnum_points = {graph.x.shape[0]}\")\n",
    "\n",
    "ax = plt.subplot(3,4,8,projection='3d')\n",
    "graph = lm_models_aug_01['pretrained'][0][object_n]\n",
    "model_pos = graph.pos\n",
    "s = ax.scatter(\n",
    "    model_pos[:, 0],\n",
    "    model_pos[:, 1],\n",
    "    model_pos[:, 2],\n",
    "    s=1,\n",
    ")\n",
    "ax.set_aspect('equal')\n",
    "ax.set_xticks([]), ax.set_yticks([]), ax.set_zticks([])\n",
    "ax.view_init(-45, 90)\n",
    "plt.title(f\"mmd=0.01\\nnum_points = {graph.x.shape[0]}\")\n",
    "\n",
    "#--------------------\n",
    "\n",
    "ax = plt.subplot(3,4,11,projection='3d')\n",
    "graph = lm_models_touch005['pretrained'][0][object_n]\n",
    "model_pos = graph.pos\n",
    "s = ax.scatter(\n",
    "    model_pos[:, 0],\n",
    "    model_pos[:, 1],\n",
    "    model_pos[:, 2],\n",
    "    s=1,\n",
    ")\n",
    "ax.set_aspect('equal')\n",
    "ax.set_xticks([]), ax.set_yticks([]), ax.set_zticks([])\n",
    "ax.view_init(-45, 90)\n",
    "plt.title(f\"mmd=0.005 just touch\\nnum_points = {graph.x.shape[0]}\")\n",
    "\n",
    "ax = plt.subplot(3,4,12,projection='3d')\n",
    "graph = lm_models_aug_01on005['pretrained'][0][object_n]\n",
    "model_pos = graph.pos\n",
    "s = ax.scatter(\n",
    "    model_pos[:, 0],\n",
    "    model_pos[:, 1],\n",
    "    model_pos[:, 2],\n",
    "    s=1,\n",
    ")\n",
    "ax.set_aspect('equal')\n",
    "ax.set_xticks([]), ax.set_yticks([]), ax.set_zticks([])\n",
    "ax.view_init(-45, 90)\n",
    "plt.title(f\"mmd=0.005 vision, 0.01 touch\\nnum_points = {graph.x.shape[0]}\")\n",
    "\n",
    "\n",
    "plt.suptitle(\"Models learned with vision policy\")\n",
    "# plt.show()\n",
    "plt.savefig(\n",
    "            save_path + f\"augmented_models_more2_{object_n}.png\",\n",
    "            bbox_inches=\"tight\",\n",
    "            dpi=300,\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eae854b5",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e023ef42",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_graph(lm_models_touch005['pretrained'][0]['spoon'])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5babc565",
   "metadata": {},
   "outputs": [],
   "source": [
    "object_n = \"bowl\"\n",
    "norm_len=0.005\n",
    "plt.figure()\n",
    "ax = plt.subplot(1,1,1,projection='3d')\n",
    "graph = lm_models_005['pretrained'][0][object_n]\n",
    "model_pos = graph.pos\n",
    "s = ax.scatter(\n",
    "    model_pos[:, 0],\n",
    "    model_pos[:, 1],\n",
    "    model_pos[:, 2],\n",
    "    s=1,\n",
    ")\n",
    "for p_id, p in enumerate(np.array(model_pos)):\n",
    "    norm = graph.norm[p_id]\n",
    "    ax.plot([p[0], p[0] + norm[0] * norm_len],\n",
    "            [p[1], p[1] + norm[1] * norm_len],\n",
    "            [p[2], p[2] + norm[2] * norm_len],\n",
    "            c='blue',\n",
    "           )\n",
    "    \n",
    "graph = lm_models_touch005['pretrained'][0][object_n]\n",
    "model_pos = graph.pos\n",
    "s = ax.scatter(\n",
    "    model_pos[:, 0],\n",
    "    model_pos[:, 1],\n",
    "    model_pos[:, 2],\n",
    "    s=1,\n",
    "    color='orange'\n",
    ")\n",
    "for p_id, p in enumerate(np.array(model_pos)):\n",
    "    norm = graph.norm[p_id]\n",
    "    ax.plot([p[0], p[0] + norm[0] * norm_len],\n",
    "            [p[1], p[1] + norm[1] * norm_len],\n",
    "            [p[2], p[2] + norm[2] * norm_len],\n",
    "            c='orange',\n",
    "           )\n",
    "ax.set_aspect('equal')\n",
    "ax.set_xticks([]), ax.set_yticks([]), ax.set_zticks([])\n",
    "ax.view_init(-45, 90)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a698bbd9",
   "metadata": {},
   "outputs": [],
   "source": [
    "pretrained_dict_touch_10u_005 = pretrain_path + \"pretrained_ycb/supervised_additional_training_10distinctobj_01on005/pretrained/\"\n",
    "_, _, _, lm_models_touch_10u_005 = load_stats(placeholder_exp_path,\n",
    "                                load_train=False,\n",
    "                                load_eval=False,\n",
    "                                load_detailed=False,\n",
    "                                pretrained_dict=pretrained_dict_touch_10u_005,\n",
    "                               )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25a4aa65",
   "metadata": {},
   "outputs": [],
   "source": [
    "lm_models_touch_10u_005['pretrained'][0].keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4bbf7fbb",
   "metadata": {},
   "outputs": [],
   "source": [
    "pretrained_dict_touch_10u_005 = pretrain_path + \"pretrained_ycb_v3/vision_1lm_10distinctobj/pretrained/\"\n",
    "_, _, _, lm_models_touch_10u_005 = load_stats(placeholder_exp_path,\n",
    "                                load_train=False,\n",
    "                                load_eval=False,\n",
    "                                load_detailed=False,\n",
    "                                pretrained_dict=pretrained_dict_touch_10u_005,\n",
    "                               )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4cd6b426",
   "metadata": {},
   "outputs": [],
   "source": [
    "lm_models_touch_10u_005['pretrained'][0].keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25be6d1f",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(15,7))\n",
    "for i, obj in enumerate(list(lm_models_touch_10u_005['pretrained'][0].keys())):\n",
    "    ax = plt.subplot(2,5,i+1,projection='3d')\n",
    "    graph = lm_models_touch_10u_005['pretrained'][0][obj]\n",
    "    model_pos = graph.pos\n",
    "    s = ax.scatter(\n",
    "        model_pos[:, 0],\n",
    "        model_pos[:, 1],\n",
    "        model_pos[:, 2],\n",
    "        s=1,\n",
    "    )\n",
    "    ax.set_aspect('equal')\n",
    "    ax.set_xticks([]), ax.set_yticks([]), ax.set_zticks([])\n",
    "    ax.view_init(-45, 90)\n",
    "    plt.title(f\"{obj}\\nnum_points = {graph.x.shape[0]}\")\n",
    "plt.suptitle(\"Models Learned with Vision (mmd=0.001, lower feat ths)\")\n",
    "plt.savefig(\n",
    "            save_path + f\"current_models_vision_distinct.png\",\n",
    "            bbox_inches=\"tight\",\n",
    "            dpi=300,\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d9c1a84b",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}