train_DSN.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train a Depth Seeding Network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys, os\n",
    "import json\n",
    "from time import time\n",
    "import glob\n",
    "\n",
    "import torch\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import scipy\n",
    "import scipy.io\n",
    "import cv2\n",
    "\n",
    "# My libraries\n",
    "import src.data_loader as data_loader\n",
    "import src.segmentation as segmentation\n",
    "import src.util.utilities as util_\n",
    "import src.util.flowlib as flowlib\n",
    "\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = \"0\" # TODO: Change this if you have more than 1 GPU"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def torch_to_numpy(torch_tensor, is_standardized_image = False):\n",
    "    \"\"\" Converts torch tensor (NCHW) to numpy tensor (NHWC) for plotting\n",
    "    \n",
    "        If it's an rgb image, it puts it back in [0,255] range (and undoes ImageNet standardization)\n",
    "    \"\"\"\n",
    "    np_tensor = torch_tensor.cpu().clone().detach().numpy()\n",
    "    if np_tensor.ndim == 4: # NCHW\n",
    "        np_tensor = np_tensor.transpose(0,2,3,1)\n",
    "    if is_standardized_image:\n",
    "        _mean=[0.485, 0.456, 0.406]; _std=[0.229, 0.224, 0.225]\n",
    "        for i in range(3):\n",
    "            np_tensor[...,i] *= _std[i]\n",
    "            np_tensor[...,i] += _mean[i]\n",
    "        np_tensor *= 255\n",
    "            \n",
    "    return np_tensor"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Example Dataset: TableTop Object Dataset (TOD)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "TOD_filepath = '...' # TODO: change this to the dataset you want to train on\n",
    "data_loading_params = {\n",
    "    \n",
    "    # Camera/Frustum parameters\n",
    "    'img_width' : 640, \n",
    "    'img_height' : 480,\n",
    "    'near' : 0.01,\n",
    "    'far' : 100,\n",
    "    'fov' : 45, # vertical field of view in degrees\n",
    "    \n",
    "    'use_data_augmentation' : True,\n",
    "\n",
    "    # Multiplicative noise\n",
    "    'gamma_shape' : 1000.,\n",
    "    'gamma_scale' : 0.001,\n",
    "    \n",
    "    # Additive noise\n",
    "    'gaussian_scale' : 0.005, # 5mm standard dev\n",
    "    'gp_rescale_factor' : 4,\n",
    "    \n",
    "    # Random ellipse dropout\n",
    "    'ellipse_dropout_mean' : 10, \n",
    "    'ellipse_gamma_shape' : 5.0, \n",
    "    'ellipse_gamma_scale' : 1.0,\n",
    "\n",
    "    # Random high gradient dropout\n",
    "    'gradient_dropout_left_mean' : 15, \n",
    "    'gradient_dropout_alpha' : 2., \n",
    "    'gradient_dropout_beta' : 5.,\n",
    "\n",
    "    # Random pixel dropout\n",
    "    'pixel_dropout_alpha' : 1., \n",
    "    'pixel_dropout_beta' : 10.,\n",
    "    \n",
    "}\n",
    "dl = data_loader.get_TOD_train_dataloader(TOD_filepath, data_loading_params, batch_size=4, num_workers=8, shuffle=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train Depth Seeding Network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dsn_params = {\n",
    "    \n",
    "    # Sizes\n",
    "    'feature_dim' : 64,\n",
    "    \n",
    "    # algorithm parameters\n",
    "    'lr' : 1e-2, # learning rate\n",
    "    'iter_collect' : 20, # Collect results every _ iterations\n",
    "    'max_iters' : 100000,\n",
    "    \n",
    "    # architecture parameters\n",
    "    'use_coordconv' : False,\n",
    "\n",
    "    # Loss function parameters\n",
    "    'lambda_fg' : 1,\n",
    "    'lambda_direction' : 1.,\n",
    "\n",
    "    # Hough Voting parameters\n",
    "    'skip_pixels' : 10, \n",
    "    'inlier_threshold' : 0.9, \n",
    "    'angle_discretization' : 100,\n",
    "    'inlier_distance' : 20,\n",
    "    'percentage_threshold' : 0.5, # this depends on skip_pixels, angle_discretization, inlier_distance. just gotta try it to see if it works\n",
    "    'object_center_kernel_radius' : 10,\n",
    "\n",
    "}\n",
    "depth_seeding_network = segmentation.DepthSeedingNetwork(dsn_params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Train the network for 1 epoch\n",
    "num_epochs = 1\n",
    "depth_seeding_network.train(num_epochs, dl)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plot some losses"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "fig = plt.figure(1, figsize=(15,3))\n",
    "total_subplots = 3\n",
    "starting_epoch = 0\n",
    "info_items = {k:v for (k,v) in depth_seeding_network.infos.items() if k > starting_epoch}\n",
    "\n",
    "plt.subplot(1,total_subplots,1)\n",
    "losses = [x['loss'] for (k,x) in info_items.items()]\n",
    "plt.plot(info_items.keys(), losses)\n",
    "plt.xlabel('Iteration')\n",
    "plt.title('Losses. {0}'.format(losses[-1]))\n",
    "\n",
    "plt.subplot(1,total_subplots,2)\n",
    "fg_losses = [x['FG loss'] for (k,x) in info_items.items()]\n",
    "plt.plot(info_items.keys(), fg_losses)\n",
    "plt.xlabel('Iteration')\n",
    "plt.title('Foreground Losses. {0}'.format(fg_losses[-1]))\n",
    "\n",
    "plt.subplot(1,total_subplots,3)\n",
    "direction_losses = [x['Direction loss'] for (k,x) in info_items.items()]\n",
    "plt.plot(info_items.keys(), direction_losses)\n",
    "plt.xlabel('Iteration')\n",
    "plt.title('Direction Losses. {0}'.format(direction_losses[-1]))\n",
    "\n",
    "print(\"Number of iterations: {0}\".format(depth_seeding_network.iter_num))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualize some stuff\n",
    "\n",
    "Run the network on a single batch, and plot the results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dl = data_loader.get_TOD_test_dataloader(TOD_filepath, data_loading_params, batch_size=8, num_workers=8, shuffle=True)\n",
    "dl_iter = dl.__iter__()\n",
    "\n",
    "batch = next(dl_iter)\n",
    "rgb_imgs = torch_to_numpy(batch['rgb'], is_standardized_image=True) # Shape: [N x H x W x 3]\n",
    "xyz_imgs = torch_to_numpy(batch['xyz']) # Shape: [N x H x W x 3]\n",
    "foreground_labels = torch_to_numpy(batch['foreground_labels']) # Shape: [N x H x W]\n",
    "direction_labels = torch_to_numpy(batch['direction_labels']) # Shape: [N x 2 x H x W]\n",
    "N, H, W = foreground_labels.shape[:3]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Number of images: {0}\".format(N))\n",
    "\n",
    "depth_seeding_network.eval_mode()\n",
    "\n",
    "### Compute segmentation masks ###\n",
    "st_time = time()\n",
    "seg_masks, direction_predictions, object_centers, initial_masks = depth_seeding_network.run_on_batch(batch)\n",
    "total_time = time() - st_time\n",
    "print('Total time taken for Segmentation: {0} seconds'.format(round(total_time, 3)))\n",
    "print('FPS: {0}'.format(round(N / total_time,3)))\n",
    "\n",
    "# Get results in numpy\n",
    "seg_masks = seg_masks.cpu().numpy()\n",
    "direction_predictions = direction_predictions.cpu().numpy().transpose(0,2,3,1)\n",
    "initial_masks = initial_masks.cpu().numpy()\n",
    "for i in range(N):\n",
    "    object_centers[i] = object_centers[i].cpu().numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig_index = 1\n",
    "for i in range(N):\n",
    "    \n",
    "    fig = plt.figure(fig_index); fig_index += 1\n",
    "    fig.set_size_inches(20,5)\n",
    "\n",
    "    # Plot image\n",
    "    plt.subplot(1,5,1)\n",
    "    plt.imshow(rgb_imgs[i,...].astype(np.uint8))\n",
    "    plt.title('Image {0}'.format(i+1))\n",
    "\n",
    "    # Plot Depth\n",
    "    plt.subplot(1,5,2)\n",
    "    plt.imshow(xyz_imgs[i,...,2])\n",
    "    plt.title('Depth')\n",
    "    \n",
    "    # Plot prediction\n",
    "    plt.subplot(1,5,3)\n",
    "    plt.imshow(util_.get_color_mask(seg_masks[i,...]))\n",
    "    plt.title(\"Predicted Masks\")\n",
    "    \n",
    "    # Plot Center Direction Predictions\n",
    "    plt.subplot(1,5,4)\n",
    "    fg_mask = np.expand_dims(seg_masks[i,...] == 2, axis=-1)\n",
    "    plt.imshow(flowlib.flow_to_image(direction_predictions[i,...] * fg_mask))\n",
    "    plt.title(\"Center Direction Predictions\")\n",
    "    \n",
    "    # Plot Initial Masks\n",
    "    plt.subplot(1,5,5)\n",
    "    plt.imshow(util_.get_color_mask(initial_masks[i,...]))\n",
    "    plt.title(f\"Initial Masks. #objects: {np.unique(initial_masks[i,...]).shape[0]-1}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:ssc]",
   "language": "python",
   "name": "conda-env-ssc-py"
  },
  "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
}