generate_affordance.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy import ndimage\n",
    "import numpy as np\n",
    "import torch\n",
    "import os\n",
    "import sys\n",
    "if '/opt/ros/kinetic/lib/python2.7/dist-packages' in sys.path:\n",
    "    sys.path.remove('/opt/ros/kinetic/lib/python2.7/dist-packages')\n",
    "import cv2\n",
    "sys.path.append(os.path.join(os.getcwd(), 'src/grasp_suck/src'))\n",
    "from matplotlib import pyplot as plt\n",
    "from model_v2 import reinforcement_net"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Helper function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def preprocessing(color, depth):\n",
    "\t# Zoom 2 times\n",
    "\tcolor_img_2x = ndimage.zoom(color, zoom=[2, 2, 1], order=0)\n",
    "\tdepth_img_2x = ndimage.zoom(depth, zoom=[2, 2],    order=0)\n",
    "\t# Add extra padding to handle rotations inside network\n",
    "\tdiag_length = float(color_img_2x.shape[0])*np.sqrt(2)\n",
    "\tdiag_length = np.ceil(diag_length/32)*32 # Shrink 32 times in network\n",
    "\tpadding_width = int((diag_length - color_img_2x.shape[0])/2)\n",
    "\t# Convert BGR (cv) to RGB\n",
    "\tcolor_img_2x_b = np.pad(color_img_2x[:, :, 0], padding_width, 'constant', constant_values=0)\n",
    "\tcolor_img_2x_b.shape = (color_img_2x_b.shape[0], color_img_2x_b.shape[1], 1)\n",
    "\tcolor_img_2x_g = np.pad(color_img_2x[:, :, 1], padding_width, 'constant', constant_values=0)\n",
    "\tcolor_img_2x_g.shape = (color_img_2x_g.shape[0], color_img_2x_g.shape[1], 1)\n",
    "\tcolor_img_2x_r = np.pad(color_img_2x[:, :, 2], padding_width, 'constant', constant_values=0)\n",
    "\tcolor_img_2x_r.shape = (color_img_2x_r.shape[0], color_img_2x_r.shape[1], 1)\n",
    "\tcolor_img_2x = np.concatenate((color_img_2x_r, color_img_2x_g, color_img_2x_b), axis = 2)\n",
    "\tdepth_img_2x = np.pad(depth_img_2x, padding_width, 'constant', constant_values=0)\n",
    "\t# Normalize color image with ImageNet data\n",
    "\timage_mean = [0.485, 0.456, 0.406] # for sim: [0.20414721, 0.17816422, 0.15419899]\n",
    "\timage_std  = [0.229, 0.224, 0.225] # for sim: [0.1830081 , 0.16705943, 0.17520182]\n",
    "\tinput_color_img = color_img_2x.astype(float)/255 # np.uint8 to float\n",
    "\tfor c in range(3):\n",
    "\t\tinput_color_img[:, :, c] = (input_color_img[:, :, c] - image_mean[c]) / image_std[c]\n",
    "\t# Normalize depth image\n",
    "\tdepth_mean = 0.0909769548291 # for sim: 0.032723393\n",
    "\tdepth_std = 0.0397293901695 # for sim: 0.056900032\n",
    "\ttmp = depth_img_2x.astype(float)\n",
    "\ttmp = (tmp-depth_mean)/depth_std\n",
    "\t# Duplicate channel to DDD\n",
    "\ttmp.shape = (tmp.shape[0], tmp.shape[1], 1)\n",
    "\tinput_depth_img = np.concatenate((tmp, tmp, tmp), axis = 2)\n",
    "\t# Convert to tensor\n",
    "\t# H, W, C - > N, C, H, W\n",
    "\tinput_color_img.shape = (input_color_img.shape[0], input_color_img.shape[1], input_color_img.shape[2], 1)\n",
    "\tinput_depth_img.shape = (input_depth_img.shape[0], input_depth_img.shape[1], input_depth_img.shape[2], 1)\n",
    "\tinput_color_data = torch.from_numpy(input_color_img.astype(np.float32)).permute(3, 2, 0, 1)\n",
    "\tinput_depth_data = torch.from_numpy(input_depth_img.astype(np.float32)).permute(3, 2, 0, 1)\n",
    "\treturn input_color_data, input_depth_data, padding_width"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 95,
   "metadata": {},
   "outputs": [],
   "source": [
    "def vis_affordance(predictions):\n",
    "\ttmp = np.copy(predictions)\n",
    "\t# View the value as probability\n",
    "\ttmp[tmp<0] = 0\n",
    "\ttmp /= 5\n",
    "\ttmp[tmp>1] = 1\n",
    "\ttmp = (tmp*255).astype(np.uint8)\n",
    "\ttmp.shape = (tmp.shape[0], tmp.shape[1], 1)\n",
    "\theatmap = cv2.applyColorMap(tmp, cv2.COLORMAP_JET)\n",
    "\treturn heatmap"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Parameters Definition "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "########################### CHANGE HERE ###########################\n",
    "model_name = # TODO\n",
    "# input image should be the orthogonal projection from the point cloud along gravity direction\n",
    "# depth stands `height from bottom`\n",
    "color_name = # TODO\n",
    "depth_name = # TODO"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "net = reinforcement_net(use_cuda=True)\n",
    "net.load_state_dict(torch.load(model_name))\n",
    "net = net.cuda().eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 119,
   "metadata": {},
   "outputs": [],
   "source": [
    "color = cv2.imread(color_name)\n",
    "depth = np.load(depth_name)\n",
    "size = color.shape[0]\n",
    "# Preprocessing\n",
    "color_tensor, depth_tensor, pad = preprocessing(color, depth)\n",
    "color_tensor = color_tensor.cuda()\n",
    "depth_tensor = depth_tensor.cuda()\n",
    "prediction = net.forward(color_tensor, depth_tensor, is_volatile=True)\n",
    "# prediction: list with length 6\n",
    "# | index | tool |\n",
    "# | --- | --- |\n",
    "# | 0 | small suction cup |\n",
    "# | 1 | medium suction cup |\n",
    "# | 2 | gripper with -90 deg |\n",
    "# | 3 | gripper with -45 deg |\n",
    "# | 4 | gripper with 0 deg |\n",
    "# | 5 | gripper with 45 deg |\n",
    "# Only show small suction cup\n",
    "tool_1 = prediction[0][0, 0, pad//2:size+pad//2, pad//2:size+pad//2].detach().cpu().numpy() # small suction cup\n",
    "tool_1_cmap = vis_affordance(tool_1)\n",
    "combine = cv2.addWeighted(color, 1.0, tool_1_cmap, 0.8, 0.0)\n",
    "best = np.where(tool_1==np.max(tool_1))\n",
    "u, v = best[1][0], best[0][0]\n",
    "combine = cv2.circle(combine, (u, v), 3, (255, 255, 255), 2)\n",
    "plt.imshow(combine[:,:,::-1])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}