visualize_heatmap.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "id": "initial_id",
   "metadata": {},
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import torch\n",
    "\n",
    "import os\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = '0'\n",
    "\n",
    "from env.envs.hookScene import HookScene\n",
    "from env.envs.mugScene import MugScene\n",
    "from env.envs.spoonScene import SpoonScene\n",
    "from scripts.hooking import get_reference as get_hooking_reference\n",
    "from scripts.hanging import get_reference as get_hanging_reference\n",
    "from scripts.scooping import get_reference as get_scooping_reference\n",
    "from magic.match import match"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "327a7d33",
   "metadata": {},
   "source": [
    "task_name = 'hooking'\n",
    "tool_name = 'scissors'\n",
    "mug_id = tool_name\n",
    "spoon_id = tool_name\n",
    "sd_dino = False\n",
    "dift = False"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "ea7cd0dc2d4e1b5e",
   "metadata": {},
   "source": [
    "if task_name == 'hooking':\n",
    "    reference_img, reference_depth_img, reference_camera, reference_contact_center, reference_grasp_center, reference_scene, reference_init_pose, reference_pose = get_hooking_reference()\n",
    "elif task_name == 'hanging':\n",
    "    reference_img, reference_depth_img, reference_camera, reference_contact_center, reference_scene = get_hanging_reference()\n",
    "    reference_grasp_center = None\n",
    "elif task_name == 'scooping':\n",
    "    plane_origin = np.array([0, -0.35, 0])\n",
    "    plane_normal = np.array([0, -1, 0])\n",
    "    (reference_img, reference_contact_center, reference_collide_center, reference_grasp_center,\n",
    "     reference_pixel_to_3d_fn, reference_init_pose, reference_pcd_o3d) = get_scooping_reference(\n",
    "        manual_center=False, plane_origin=plane_origin, plane_normal=plane_normal\n",
    "    )\n",
    "else:\n",
    "    raise ValueError('Invalid task name')"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "58f921859dd3ed",
   "metadata": {},
   "source": [
    "if task_name == 'hooking':\n",
    "    target_scene = HookScene(\n",
    "        tool_name,\n",
    "        'box',\n",
    "    )\n",
    "\n",
    "    target_scene.hide_env_visual()\n",
    "    target_img, target_depth_img, target_camera = target_scene.get_picture(\n",
    "        direction='+z',\n",
    "        additional_translation=np.array([0.6, -0.6, 0.2-0.375*1.414]),\n",
    "        debug_viewer=False,\n",
    "        get_depth=True\n",
    "    )\n",
    "    target_scene.unhide_env_visual()\n",
    "elif task_name == 'hanging':\n",
    "    target_scene = MugScene(\n",
    "        mug_id,\n",
    "        add_robot=True,\n",
    "        fps=480\n",
    "    )\n",
    "\n",
    "    target_scene.hide_env_visual()\n",
    "    target_img, target_depth_img, target_camera = target_scene.get_picture(direction='+x', debug_viewer=False, get_depth=True)\n",
    "    target_scene.unhide_env_visual()\n",
    "else:\n",
    "    plane_origin = np.array([0, -0.35, 0])\n",
    "    plane_normal = np.array([0, -1, 0])\n",
    "    target_scene = SpoonScene(\n",
    "        spoon_id,\n",
    "        fps=480,\n",
    "        add_robot=True,\n",
    "        radius=0.035\n",
    "    )\n",
    "\n",
    "    target_img, target_pixel_to_3d_fn = target_scene.get_slice(plane_origin, plane_normal)\n",
    "\n",
    "results, resized_imgs, downsampled_images = match(reference_img, target_img, reference_contact_center, reference_grasp_center, only_compute_dino_feature=True, sd_dino=sd_dino, dift=dift)\n",
    "results = results.permute(0, 2, 3, 1)"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "8a60998b24cc5bcf",
   "metadata": {},
   "source": [
    "%matplotlib inline\n",
    "fig, axes = plt.subplots(3, 5, figsize=(2 * resized_imgs[0].size[0] / 100, resized_imgs[0].size[1] / 100), dpi=100)\n",
    "for i in range(3):\n",
    "    for j in range(5):\n",
    "        axes[i][j].axis('off')\n",
    "\n",
    "axes[0][0].imshow(resized_imgs[0])\n",
    "source_point = reference_contact_center\n",
    "resized_source_point = np.array(source_point) / reference_img.size[1] * resized_imgs[0].size[0]\n",
    "axes[0][0].scatter(resized_source_point[0], resized_source_point[1], c='r', s=50)\n",
    "\n",
    "\n",
    "import cv2\n",
    "from tqdm import tqdm \n",
    "\n",
    "for rotation_index in tqdm(range(0, 12)):\n",
    "    downsampled_source_point = (np.array(source_point) / reference_img.size[1] * downsampled_images[0].size[1]).astype(int)\n",
    "    source_feature = results[0][downsampled_source_point[1], downsampled_source_point[0]]\n",
    "\n",
    "    target_feature = results[rotation_index+1]\n",
    "    heatmap = torch.sum(target_feature * source_feature, dim=-1).cpu().numpy()\n",
    "\n",
    "    # resize heatmap to target image size\n",
    "    heatmap = cv2.resize(heatmap, (resized_imgs[0].size[0], resized_imgs[0].size[1]), interpolation=cv2.INTER_LINEAR)\n",
    "    # gaussian filter\n",
    "    heatmap = cv2.GaussianBlur(heatmap, (25, 25), 0)\n",
    "    heatmap = (heatmap - np.min(heatmap)) / (np.max(heatmap) - np.min(heatmap))\n",
    "\n",
    "    # boost the heatmap where the feature is high for visualization\n",
    "    if dift:\n",
    "        heatmap = np.power(heatmap, 3)\n",
    "    else:\n",
    "        heatmap = np.power(heatmap, 10)\n",
    "\n",
    "    # get the coordinates of the maximum value in the heatmap\n",
    "    max_index = np.unravel_index(heatmap.argmax(), heatmap.shape)\n",
    "    target_point = max_index\n",
    "\n",
    "    # set the size of the canvas to be the same as the img\n",
    "    i = rotation_index // 4\n",
    "    j = rotation_index % 4 + 1\n",
    "    axes[i][j].imshow(resized_imgs[rotation_index+1])\n",
    "    axes[i][j].imshow(heatmap, alpha=0.8)\n",
    "    axes[i][j].scatter(target_point[1], target_point[0], c='r', s=50, marker='*')\n",
    "\n",
    "plt.show()"
   ],
   "outputs": [],
   "execution_count": null
  }
 ],
 "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.10.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}