notebooks/joint_prediction_viewer.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Joint Prediction Viewer\n",
    "Visualize predictions for joint entities"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "import json\n",
    "from pathlib import Path\n",
    "import torch\n",
    "import numpy as np\n",
    "import meshplot as mp\n",
    "\n",
    "root_dir = Path().resolve().parent\n",
    "if str(root_dir) not in sys.path:\n",
    "    sys.path.append(str(root_dir))\n",
    "\n",
    "from joint.joint_prediction_set import JointPredictionSet\n",
    "from joint.joint_environment import JointEnvironment\n",
    "from datasets.joint_graph_dataset import JointGraphDataset\n",
    "from train import JointPrediction"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load the Network\n",
    "Load a pretrained checkpoint to use for inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_network(checkpoint_file):\n",
    "    \"\"\"Load the network\"\"\"\n",
    "    if not checkpoint_file.exists():\n",
    "        print(\"Checkpoint file does not exist\")\n",
    "        return None\n",
    "    model = JointPrediction.load_from_checkpoint(\n",
    "        checkpoint_file,\n",
    "        map_location=torch.device(\"cpu\")  # Just use the CPU\n",
    "    )\n",
    "    return model\n",
    "\n",
    "checkpoint_file = root_dir / \"pretrained/paper/last_run_0.ckpt\"\n",
    "model = load_network(checkpoint_file)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load the Data\n",
    "Load the dataset and create an instance of the JointPredictionSet class\n",
    "We assume that the joint json and mesh part files are in the same directory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Data cache loaded from: /Users/willisk/Autodesk/Code/Github/JoinABLe/data/tester/val.pickle\n"
     ]
    }
   ],
   "source": [
    "# Change to point to the Fusion 360 Gallery joint dataset\n",
    "# this directory should contain the joint json and obj part files\n",
    "data_dir = root_dir / \"data/tester\"\n",
    "\n",
    "dataset = JointGraphDataset(\n",
    "    root_dir=data_dir,\n",
    "    split=\"val\",\n",
    "    label_scheme=\"Joint,JointEquivalent\"\n",
    ")\n",
    "# Data sample in the dataset we want to visualize\n",
    "index = 1\n",
    "# Graphs for part one and two, and the densely connected joint graph\n",
    "g1, g2, joint_graph = dataset[index]\n",
    "# The joint file json\n",
    "joint_file = data_dir / dataset.files[index]\n",
    "\n",
    "# Load the prediction data\n",
    "jps = JointPredictionSet(\n",
    "    joint_file,\n",
    "    g1, g2, joint_graph,\n",
    "    model\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Assemble the Top-1 Prediction\n",
    "We use the JointEnvironment to calculate the transform that aligns the two parts together"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ed60c1e4bc6d4235bab2111fbc2c616d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Renderer(camera=PerspectiveCamera(children=(DirectionalLight(color='white', intensity=0.6, position=(2.1500000…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Get the transform to move body1, to align with the static body2\n",
    "# using default parameters, i.e. no offset, rotation, or flip\n",
    "transform = JointEnvironment.get_transform_from_parameters(\n",
    "    jps,\n",
    "    prediction_index=0,  # Top-1 prediction\n",
    "    offset=0,\n",
    "    rotation_in_degrees=0,\n",
    "    flip=False\n",
    ")\n",
    "\n",
    "# Render the meshes to visualize\n",
    "v, f, c, e, n, ni = jps.get_meshes(\n",
    "    apply_transform=True,\n",
    "    body_one_transform=transform,\n",
    "    show_joint_entity_colors=False,\n",
    "    show_joint_equivalent_colors=False,\n",
    "    return_vertex_normals=True\n",
    ")\n",
    "p = mp.plot(v, f, c=c);\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualize Entity Predictions\n",
    "Show the predictions as pink highlights on the joint bodies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "77aeb2010d6046dbafd69713372a55ca",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Renderer(camera=PerspectiveCamera(children=(DirectionalLight(color='white', intensity=0.6, position=(0.0, -0.1…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "body = 1\n",
    "v1, f1, _, _ = jps.get_mesh(\n",
    "    body=body,\n",
    "    apply_transform=False,\n",
    "    show_joint_entity_colors=False,\n",
    "    show_joint_equivalent_colors=False\n",
    ")\n",
    "c1, e1 = jps.get_joint_predictions(body=body, limit=1)\n",
    "p = mp.plot(v1, f1, c=c1, shading={\"colormap\": \"cool\", \"normalize\": [0, 1]})\n",
    "# TODO: Add support for edge colors\n",
    "if e1 is not None:\n",
    "    p.add_edges(v1, e1, shading={\"line_color\": \"red\"});"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e915e2dc9ee7455b998b641facd7f37c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Renderer(camera=PerspectiveCamera(children=(DirectionalLight(color='white', intensity=0.6, position=(2.1500000…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "body = 2\n",
    "v2, f2, _, _ = jps.get_mesh(\n",
    "    body=body,\n",
    "    apply_transform=False,\n",
    "    show_joint_entity_colors=False,\n",
    "    show_joint_equivalent_colors=False\n",
    ")\n",
    "c2, e2 = jps.get_joint_predictions(body=body, limit=1)\n",
    "p = mp.plot(v2, f2, c=c2, shading={\"colormap\": \"cool\", \"normalize\": [0, 1]})\n",
    "# TODO: Add support for edge colors\n",
    "if e2 is not None:\n",
    "    p.add_edges(v2, e2, shading={\"line_color\": \"red\"});"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualize Entity Axes\n",
    "Show the joint axes derived from the predicted B-Rep faces/edges"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "da6823e46dbe44e4adb1adbed27e1b3f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Renderer(camera=PerspectiveCamera(children=(DirectionalLight(color='white', intensity=0.6, position=(0.0, -0.1…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "body = 1\n",
    "p = mp.plot(v1, f1, c=c1, shading={\"colormap\": \"cool\", \"normalize\": [0, 1]})\n",
    "if e1 is not None:\n",
    "    p.add_edges(v1, e1, shading={\"line_color\": \"red\"});\n",
    "start_pts, end_pts = jps.get_joint_prediction_axis_lines(body=body, limit=1)\n",
    "p.add_lines(start_pts, end_pts, shading={\"line_color\": \"green\"});\n",
    "p.add_points(start_pts, shading={\"point_color\": \"green\", \"point_size\": 1});"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "00e6e864d27b413fae2a830426ab33bc",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Renderer(camera=PerspectiveCamera(children=(DirectionalLight(color='white', intensity=0.6, position=(2.1500000…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "body = 2\n",
    "p = mp.plot(v2, f2, c=c2, shading={\"colormap\": \"cool\", \"normalize\": [0, 1]})\n",
    "if e2 is not None:\n",
    "    p.add_edges(v2, e2, shading={\"line_color\": \"red\"});\n",
    "start_pts, end_pts = jps.get_joint_prediction_axis_lines(body=body, limit=1)\n",
    "p.add_lines(start_pts, end_pts, shading={\"line_color\": \"green\"});\n",
    "p.add_points(start_pts, shading={\"point_color\": \"green\", \"point_size\": 1});"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "4bf1c82db291c90cc4874b8785149ede8d32de9ceb77e89bfd1ded6bd42ea729"
  },
  "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.7.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}