Notebook_tutorial.ipynb

{
 "nbformat": 4,
 "nbformat_minor": 2,
 "metadata": {
  "language_info": {
   "name": "python",
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "version": "3.6.10-final"
  },
  "orig_nbformat": 2,
  "file_extension": ".py",
  "mimetype": "text/x-python",
  "name": "python",
  "npconvert_exporter": "python",
  "pygments_lexer": "ipython3",
  "version": 3,
  "kernelspec": {
   "name": "python361064bitfraprojectenvcondaf4cac3dd88064b5783b04c582875a4bf",
   "display_name": "Python 3.6.10 64-bit ('fraproject_env': conda)"
  }
 },
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Introduction"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Localization is an essential task for robotics applications. To know the exact pose (position and orientation) of the agent it's essential for visualization, navigation, prediction, and planning."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### In the following I will show the whole pipeline to localize a drone which is based only on camera images and IMU data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import print_function, division\n",
    "\n",
    "import shutil\n",
    "\n",
    "import os\n",
    "import sys\n",
    "import numpy as np\n",
    "import glob\n",
    "import random\n",
    "from PIL import Image\n",
    "from tqdm import tqdm\n",
    "import argparse\n",
    "from bisect import bisect_right, bisect_left\n",
    "import pandas as pd\n",
    "from skimage import io, transform\n",
    "import time\n",
    "from datetime import datetime\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# Ignore warnings\n",
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "\n",
    "%matplotlib inline\n",
    "\n",
    "plt.ion()   # interactive mode\n",
    ""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from utils.training_fun import *\n",
    "from utils.utils import *\n",
    "from Dataset import mydataset\n",
    "from Model import  criterion, mymodel\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# The Dataset\n",
    "\n",
    "### Please, download the dataset from https://projects.asl.ethz.ch/datasets/doku.php?id=kmavvisualinertialdatasets\n",
    "### the data used in this work are Stereo Images (WVGA monochrome, 2×20 FPS) smd IMU (angular rate and acceleration, 200 Hz)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data Processing\n",
    "### First run dataprocessing.py to process the raw data from the EuRoC mav dataset and create the correct files"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Import the dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = NewDataset(csv_file='/home/francesca/euroc/cam0/file.csv', \n",
    "                         imu_data=new_imu_values,\n",
    "                           root_dir='/home/francesca/euroc/cam0/data/')\n",
    "\n",
    "dataloader = DataLoader(dataset, batch_size=32,\n",
    "                        shuffle=False, num_workers=4)\n",
    "\n",
    "print(dataloader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 32\n",
    "validation_split = .2\n",
    "shuffle_dataset = False\n",
    "random_seed= 42\n",
    "\n",
    "# Creating data indices for training and validation splits:\n",
    "dataset_size = len(dataset)\n",
    "indices = list(range(dataset_size))\n",
    "# split = int(np.floor(validation_split * dataset_size))\n",
    "split=320\n",
    "if shuffle_dataset :\n",
    "    np.random.seed(random_seed)\n",
    "    np.random.shuffle(indices)\n",
    "train_indices, val_indices = indices[split:], indices[:split]\n",
    "\n",
    "# Creating PT data samplers and loaders:\n",
    "train_sampler = SubsetRandomSampler(train_indices)\n",
    "valid_sampler = SubsetRandomSampler(val_indices)\n",
    "\n",
    "train_loader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, sampler=train_sampler, num_workers=4)\n",
    "validation_loader = torch.utils.data.DataLoader(dataset, batch_size=batch_size,sampler=valid_sampler, num_workers=4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Localization Task "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The localization task formulation consists into finding the current robot pose x_t given the current observation z_t and the previous state x_{t-1}.\n",
    "\n",
    "Our neural network regresses robot pose (translation and rotation in quaternions) from monocular image and Inertial measurements. Rotation is represented in quaternions because they do not suffer from a wrap around 2\\pi radians as Euler angles or axis-angle representations and more straightforward to deal than 3x3 rotation matrices."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### RNN-DNN-based regressor for camera pose on ResNet and modify it by adding a global average pooling layer after the last convolutional layer and introducing a fully-connected layer with 2048 neurons. Finally, it’s concluded with 6 DoF camera pose regressor for translation and rotation in quaternions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Device set up\n",
    "\n",
    "### determine the primary device "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = None\n",
    "if torch.cuda.is_available():\n",
    "    device = torch.device('cuda')\n",
    "else:\n",
    "    device = torch.device('cpu')\n",
    "print('device = {}'.format(device))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Train and Validate function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def train(train_loader, model, criterion, optimizer, epoch, max_epoch, log_freq=1, print_sum=True):\n",
    "    \n",
    "    position_loss=[]\n",
    "    orientation_loss = []\n",
    "    \n",
    "    model.train()\n",
    "    losses = AverageMeter()\n",
    "    epoch_time = time.time()\n",
    "    \n",
    "    gt_poses = np.empty((0, 7))\n",
    "    pred_poses = np.empty((0, 7))\n",
    "    \n",
    "    end = time.time()\n",
    "    \n",
    "    for idx, data in enumerate(train_loader):\n",
    "        if idx == 1920:\n",
    "            idx = 1919\n",
    "        batch_images= data[0].cuda()\n",
    "        imu_data= data[1].unsqueeze(0).cuda()\n",
    "    \n",
    "        batch_poses= data[2].cuda()\n",
    "        data_time = (time.time() - end)\n",
    "        batch_images = batch_images.to(device)\n",
    "        batch_poses = batch_poses.to(device)\n",
    "        batch_poses = torch.tensor(batch_poses, dtype=torch.float, device=device)\n",
    "#         imu_data = imu_data.to(device).double()\n",
    "        imu_data = torch.tensor(imu_data, dtype=torch.float, device=device)\n",
    "\n",
    "        out = model(data[0:2])\n",
    "        loss = criterion(out, batch_poses)\n",
    "        # Training step\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        \n",
    "        losses.update(loss.data[0], len(batch_images) * batch_images.size(0))\n",
    "        \n",
    "        # move data to cpu & numpy\n",
    "        bp = batch_poses.detach().cpu().numpy()\n",
    "        outp = out.detach().cpu().numpy()\n",
    "        gt_poses = np.vstack((gt_poses, bp))\n",
    "        pred_poses = np.vstack((pred_poses, outp))\n",
    "        \n",
    "        batch_time = (time.time() - end)\n",
    "        end = time.time()\n",
    "        \n",
    "        if log_freq != 0 and idx % log_freq == 0:\n",
    "            print('Epoch: [{}/{}]\\tBatch: [{}/{}]\\t'\n",
    "                  'Time: {batch_time:.3f}\\t'\n",
    "                  'Data Time: {data_time:.3f}\\t'\n",
    "                  'Loss: {losses.val:.3f}\\t'\n",
    "                  'Avg Loss: {losses.avg:.3f}\\t'.format(\n",
    "                   epoch, max_epoch - 1, idx, len(train_loader) - 1,\n",
    "                   batch_time=batch_time, data_time=data_time, losses=losses))\n",
    "            \n",
    "        t_loss = np.asarray([np.linalg.norm(p - t) for p, t in zip(pred_poses[:, :3], gt_poses[:, :3])])\n",
    "        q_loss = np.asarray([quaternion_angular_error(p, t) for p, t in zip(pred_poses[:, 3:], gt_poses[:, 3:])])\n",
    "    position_loss.append(t_loss)\n",
    "    orientation_loss.append(q_loss)\n",
    "\n",
    "    return losses\n",
    "    \n",
    "def validate(val_loader, model, criterion, epoch, log_freq=1, print_sum=True):\n",
    "    \n",
    "    losses = AverageMeter()\n",
    "    position_loss_testing=[]\n",
    "    orientation_loss_testing=[]\n",
    "    \n",
    "    # set model to evaluation\n",
    "    model.eval()\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        epoch_time = time.time()\n",
    "        end = time.time()\n",
    "        for idx, data in enumerate(val_loader):\n",
    "\n",
    "\n",
    "            batch_images= data[0].cuda()\n",
    "            imu_data= data[1].cuda()\n",
    "        \n",
    "            batch_poses= data[2].cuda()\n",
    "            data_time = time.time() - end\n",
    "            \n",
    "            batch_images = batch_images.to(device)\n",
    "            batch_poses = batch_poses.to(device)\n",
    "            batch_poses = torch.tensor(batch_poses, dtype=torch.float, device=device)\n",
    "            imu_data = imu_data.to(device)\n",
    "            imu_data = torch.tensor(imu_data, dtype=torch.float, device=device)\n",
    "            out = model(data[:2])\n",
    "            loss = criterion(out, batch_poses)           \n",
    "            losses.update(loss.data[0], len(data[0]) * data[0].size(0))\n",
    "\n",
    "\n",
    "\n",
    "            batch_time = time.time() - end\n",
    "            end = time.time()\n",
    "            if log_freq != 0 and idx % log_freq == 0:\n",
    "                print('Val Epoch: {}\\t'\n",
    "                      'Time: {batch_time:.3f}\\t'\n",
    "                      'Data Time: {data_time:.3f}\\t'\n",
    "                      'Loss: {losses.val:.3f}\\t'\n",
    "                      'Avg Loss: {losses.avg:.3f}'.format(\n",
    "                       epoch, batch_time=batch_time, data_time=data_time, losses=losses))\n",
    "\n",
    "    return losses\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create Model, Criterion and Optimizer\n",
    "\n",
    "### Objective: Regression for translation (3D pose) and rotation (quaternion)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# Create pretrained feature extractor\n",
    "feature_extractor = models.resnet18(pretrained=True)\n",
    "\n",
    "# Num features for the last layer before pose regressor\n",
    "num_features = 2048\n",
    "\n",
    "# Create model\n",
    "model = IMUNet(feature_extractor, num_features=num_features, pretrained=True)\n",
    "model = model.to(device)\n",
    "\n",
    "# Num features for tahe last layer before pose regressor\n",
    "num_features = 2048\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Net Criterion"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### As a loss function we use a learning approach to find an optimal weighting for translation and orientation with homoscedastic uncertainties that represent free scalar values that we learn through backpropagation with respect to the loss function. The effect is to decrease or increase the corresponding loss component automatically. To prevent the potential division by zero, we learn hte log which is more numerically stable. (check Alex Kendall paper for more detail)\n",
    "\n",
    "#### If learn_beta param is False it’s a simple weighted sum version of the loss, where beta balance the losses of two variables expressed in different units and of different scales."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lass NetCriterion(torch.nn.Module):\n",
    "    def __init__(self, beta = 512.0, learn_beta=True, sx=0.0, sq=3.0):\n",
    "        super(NetCriterion, self).__init__()\n",
    "        self.loss_fn1 = torch.nn.L1Loss()\n",
    "        self.loss_fn = torch.nn.MSELoss()\n",
    "        self.learn_beta = learn_beta\n",
    "        if not learn_beta:\n",
    "            self.beta = beta\n",
    "        else:\n",
    "            self.beta = 1.0\n",
    "        self.sx = torch.nn.Parameter(torch.Tensor([sx]), requires_grad=learn_beta)\n",
    "        self.sq = torch.nn.Parameter(torch.Tensor([sq]), requires_grad=learn_beta)\n",
    "\n",
    "    def forward(self, x, y):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            x: prediction \n",
    "            y: label\n",
    "        \"\"\"\n",
    "        \n",
    "        loss = 0\n",
    "        # Translation loss\n",
    "        loss = torch.exp(-self.sx) * (self.loss_fn(x[:, :3], y[:, :3])+0.01*self.loss_fn1(x[:, :3], y[:, :3])) \n",
    "        # Rotation loss\n",
    "        loss += torch.exp(-self.sq) * self.beta * self.loss_fn(x[:, 3:], y[:, 3:]+0.01*self.loss_fn1(x[:, 3:], y[:, 3:])) + self.sq\n",
    "\n",
    "        return loss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run Training"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### we use torch.optim.Adam optimizer with learning rate 1e-4, ResNet18 pretrained on ImageNet as a feature extractor and 2048 features on the last FC layer before pose regressors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "imu = next(iter(train_loader))[1]\n",
    "# Create model\n",
    "model = IMUNet(feature_extractor, num_features=num_features, pretrained=True)\n",
    "model = model.to(device)\n",
    "\n",
    "# Criterion\n",
    "criterion = NetCriterion(learn_beta=True)\n",
    "criterion = criterion.to(device)\n",
    "\n",
    "# Add all params for optimization\n",
    "param_list = [{'params': model.parameters()}]\n",
    "if criterion.learn_beta:\n",
    "    # Add sx and sq from loss function to optimizer params\n",
    "    param_list.append({'params': criterion.parameters()})\n",
    "\n",
    "# Create optimizer\n",
    "# optimizer = optim.Adam(params=param_list, lr=1e-5, weight_decay=0.0005)\n",
    "optimizer = optim.Adam(params=param_list, lr=0.0001, weight_decay=0.0001)\n",
    "\n",
    "\n",
    "# Epochs to train\n",
    "start_epoch = 0\n",
    "n_epochs = start_epoch + 620\n",
    "\n",
    "# Training\n",
    "\n",
    "train_loss=[]\n",
    "test_loss=[]\n",
    "print('Training ...')\n",
    "val_freq = 10\n",
    "for e in range(start_epoch, n_epochs):\n",
    "    losses = train(train_loader, model, criterion, optimizer, e, n_epochs, log_freq=100)\n",
    "    train_loss.append(losses)\n",
    "    if losses.val <= 0:\n",
    "        break\n",
    "    if e % val_freq == 0:\n",
    "#         end = time.time()\n",
    "        loss = validate(validation_loader, model, criterion, e, log_freq=100)\n",
    "        test_loss.append(loss)\n",
    "\n",
    "\n",
    "start_epoch = n_epochs\n",
    "\n",
    "print('n_epochs = {}'.format(n_epochs))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Save model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save(model, 'LSTM_pos_workin2_v2.torch')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Test Model Error on Validation and Train Datasets\n",
    "\n",
    "### Calculate translation and rotation error of the predicted poses on train and validation datasets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def model_results_pred_gt(model, dataloader):\n",
    "    \n",
    "    model.eval()\n",
    "\n",
    "    gt_poses = np.empty((0, 7))\n",
    "    pred_poses = np.empty((0, 7))\n",
    "\n",
    "    for idx, data in enumerate(dataloader):\n",
    "        batch_images = data[0].cuda()\n",
    "        batch_imu = data[1].cuda()\n",
    "        batch_poses  = data[2].cuda()\n",
    "        batch_images = batch_images.to(device)\n",
    "        batch_imu = batch_imu.to(device)\n",
    "        batch_imu = torch.tensor(batch_imu, dtype=torch.float, device=device)\n",
    "        batch_poses = batch_poses.to(device)\n",
    "        batch_poses = torch.tensor(batch_poses, dtype=torch.float, device=device)\n",
    "        out = model(data[0:2])        \n",
    "        loss = criterion(out, batch_poses)\n",
    "        bp = batch_poses.detach().cpu().numpy()\n",
    "        outp = out.detach().cpu().numpy()\n",
    "        gt_poses = np.vstack((gt_poses, bp))\n",
    "        pred_poses = np.vstack((pred_poses, outp))\n",
    "\n",
    "    gt_poses[:, :3] = gt_poses[:, :3] \n",
    "    pred_poses[:, :3] = pred_poses[:, :3] \n",
    "    \n",
    "    return pred_poses, gt_poses\n",
    "\n",
    "print('\\n=== Test Training Dataset ======')\n",
    "pred_poses, gt_poses = model_results_pred_gt(model, train_loader)\n",
    "\n",
    "print('gt_poses = {}'.format(gt_poses.shape))\n",
    "print('pred_poses = {}'.format(pred_poses.shape))\n",
    "t_loss = np.asarray([np.linalg.norm(p - t) for p, t in zip(pred_poses[:, :3], gt_poses[:, :3])])\n",
    "q_loss = np.asarray([quaternion_angular_error(p, t) for p, t in zip(pred_poses[:, 3:], gt_poses[:, 3:])])\n",
    "\n",
    "print('Translation(T) error in meters and Rotation(R) error in degrees:')\n",
    "print('T: median = {:.3f}, mean = {:.3f}'.format(np.median(t_loss), np.mean(t_loss)))\n",
    "print('R: median = {:.3f}, mean = {:.3f}'.format(np.median(q_loss), np.mean(q_loss)))\n",
    "\n",
    "# Save for later visualization\n",
    "pred_poses_train = pred_poses\n",
    "gt_poses_train = gt_poses\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('\\n=== Test Validation Dataset ======')\n",
    "pred_poses, gt_poses = model_results_pred_gt(model, validation_loader)\n",
    "\n",
    "print('gt_poses = {}'.format(gt_poses.shape))\n",
    "print('pred_poses = {}'.format(pred_poses.shape))\n",
    "t_loss = np.asarray([np.linalg.norm(p - t) for p, t in zip(pred_poses[:, :3], gt_poses[:, :3])])\n",
    "q_loss = np.asarray([quaternion_angular_error(p, t) for p, t in zip(pred_poses[:, 3:], gt_poses[:, 3:])])\n",
    "\n",
    "print('Translation(T) error in meters and Rotation(R) error in degrees:')\n",
    "print('T: median = {:.3f}, mean = {:.3f}'.format(np.median(t_loss), np.mean(t_loss)))\n",
    "print('R: median = {:.3f}, mean = {:.3f}'.format(np.median(q_loss), np.mean(q_loss)))\n",
    "\n",
    "# Save for later visualization\n",
    "pred_poses_val = pred_poses\n",
    "gt_poses_val = gt_poses"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualize Prediction and Ground Truth Poses\n",
    "#### (ground truth in `blue` and predictions in `red` colors)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from mpl_toolkits.mplot3d import Axes3D\n",
    "\n",
    "def calc_poses_params(poses, pose_format='full-mat'):\n",
    "    \"\"\"Calculates min, max, mean and std of translations of the poses\"\"\"\n",
    "\n",
    "    p = poses[0]\n",
    "    allp = extract_translation(p, pose_format)\n",
    "\n",
    "    for p in poses[1:]:\n",
    "        allp = np.vstack((allp, extract_translation(p, pose_format)))\n",
    "\n",
    "    p_min = np.min(allp, axis=0)\n",
    "    p_max = np.max(allp, axis=0)\n",
    "    p_mean = np.mean(allp, axis=0)\n",
    "    p_std = np.std(allp, axis=0)\n",
    "\n",
    "    return p_min, p_max, p_mean, p_std\n",
    "\n",
    "def set_3d_axes_limits(ax, poses, pose_format='quat'):\n",
    "    p_min, p_max, p_mean, p_std = calc_poses_params(poses, pose_format=pose_format)\n",
    "    ax.set_xlim(p_min[0], p_max[0])\n",
    "    ax.set_ylim(p_min[1], p_max[1])\n",
    "    ax.set_zlim(int(p_min[2] - 1), p_max[2])\n",
    "    return p_min, p_max, p_mean, p_std\n",
    "\n",
    "def extract_translation(p, pose_format='full-mat'):\n",
    "    if pose_format == 'full-mat':\n",
    "        return p[0:3, 3]\n",
    "    elif pose_format == 'quat':\n",
    "        return p[:3]\n",
    "    else:\n",
    "        warnings.warn(\"pose_format should be either 'full-mat' or 'quat'\")\n",
    "        return p\n",
    "\n",
    "\n",
    "def draw_pred_gt_poses(pred_poses, gt_poses):\n",
    "    fig = plt.figure(figsize=(8, 8))\n",
    "    ax = plt.axes(projection='3d')\n",
    "\n",
    "    ax.set_xlabel('$X$')\n",
    "    ax.set_ylabel('$Y$')\n",
    "    ax.set_zlabel('$Z$')\n",
    "    ax.view_init(50, 30)\n",
    "\n",
    "    all_poses = np.concatenate((pred_poses, gt_poses))\n",
    "    p_min, _, _, _ = set_3d_axes_limits(ax, all_poses, pose_format='quat')\n",
    "    \n",
    "    draw_poses(ax, pred_poses[:, :3], proj=False, proj_z=int(p_min[2] - 1), c='r', s=60)\n",
    "    draw_poses(ax, gt_poses[:, :3], proj=False, proj_z=int(p_min[2] - 1), c='b', s=60)\n",
    "    for i in range(pred_poses.shape[0]):\n",
    "        pp = pred_poses[i, :3]\n",
    "        gp = gt_poses[i, :3]\n",
    "        pps = np.vstack((pp, gp))\n",
    "        ax.plot(pps[:, 0], pps[:, 1], pps[:, 2], c=(0.7, 0.7, 0.7, 0.4))\n",
    "        \n",
    "    plt.draw()\n",
    "    \n",
    "def draw_poses(ax, poses, c='b', s=20, proj=False, proj_z=0, pose_format='quat'):\n",
    "    \"\"\"Draws the list of poses.\n",
    "    Args:\n",
    "        ax (Axes3D): 3D axes\n",
    "        poses (list): Poses list\n",
    "        c: matplotlib color\n",
    "        s: matplotlib size\n",
    "        proj (bool): True if draw projection of a path on z-axis\n",
    "        proj_z (float): Coord for z-projection\n",
    "    \"\"\"\n",
    "    coords = np.zeros((len(poses), 3))\n",
    "    for i, p in enumerate(poses):\n",
    "        # coords[i] = p[:3, 3]\n",
    "        # coords[i] = p\n",
    "        coords[i] = extract_translation(p, pose_format=pose_format)\n",
    "\n",
    "    # Draw projection\n",
    "    if proj:\n",
    "        if len(poses) > 1:\n",
    "            ax.plot(coords[:, 0], coords[:, 1], proj_z, c='g')\n",
    "        elif len(poses) == 1:\n",
    "            ax.scatter(coords[:, 0], coords[:, 1], proj_z, c=c)\n",
    "\n",
    "    # Draw path\n",
    "    ax.scatter(coords[:, 0], coords[:, 1], coords[:, 2], c=c, s=s)\n",
    "\n",
    "\n",
    "def draw_poses_list(ax, poses_list):\n",
    "    \"\"\"Draw list of lists of poses.\"\"\"\n",
    "    for poses in poses_list:\n",
    "        draw_poses(ax, poses)\n",
    "\n",
    "\n",
    "\n",
    "# Draw predicted vs ground truth poses\n",
    "draw_pred_gt_poses(pred_poses_train, gt_poses_train)\n",
    "plt.title('Train Dataset')\n",
    "plt.show()\n",
    "\n",
    "\n",
    "draw_pred_gt_poses(pred_poses_val, gt_poses_val)\n",
    "plt.title('Validation Dataset')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from mpl_toolkits.mplot3d import Axes3D\n",
    "\n",
    "def calc_poses_params(poses, pose_format='full-mat'):\n",
    "    \"\"\"Calculates min, max, mean and std of translations of the poses\"\"\"\n",
    "\n",
    "    p = poses[0]\n",
    "    allp = extract_translation(p, pose_format)\n",
    "\n",
    "    for p in poses[1:]:\n",
    "        allp = np.vstack((allp, extract_translation(p, pose_format)))\n",
    "\n",
    "    p_min = np.min(allp, axis=0)\n",
    "    p_max = np.max(allp, axis=0)\n",
    "    p_mean = np.mean(allp, axis=0)\n",
    "    p_std = np.std(allp, axis=0)\n",
    "\n",
    "    return p_min, p_max, p_mean, p_std\n",
    "\n",
    "def set_3d_axes_limits(ax, poses, pose_format='quat'):\n",
    "    p_min, p_max, p_mean, p_std = calc_poses_params(poses, pose_format=pose_format)\n",
    "    ax.set_xlim(p_min[0], p_max[0])\n",
    "    ax.set_ylim(p_min[1], p_max[1])\n",
    "    ax.set_zlim(int(p_min[2] - 1), p_max[2])\n",
    "    return p_min, p_max, p_mean, p_std\n",
    "\n",
    "def extract_translation(p, pose_format='full-mat'):\n",
    "    if pose_format == 'full-mat':\n",
    "        return p[0:3, 3]\n",
    "    elif pose_format == 'quat':\n",
    "        return p[:3]\n",
    "    else:\n",
    "        warnings.warn(\"pose_format should be either 'full-mat' or 'quat'\")\n",
    "        return p\n",
    "\n",
    "\n",
    "def draw_pred_gt_poses(pred_poses, gt_poses):\n",
    "    fig = plt.figure(figsize=(8, 8))\n",
    "    ax = plt.axes(projection='3d')\n",
    "\n",
    "    ax.set_xlabel('$X$')\n",
    "    ax.set_ylabel('$Y$')\n",
    "    ax.set_zlabel('$Z$')\n",
    "    ax.view_init(50, 30)\n",
    "\n",
    "    all_poses = np.concatenate((pred_poses, gt_poses))\n",
    "    p_min, _, _, _ = set_3d_axes_limits(ax, all_poses, pose_format='quat')\n",
    "    \n",
    "    draw_poses(ax, pred_poses[:, :3], proj=False, proj_z=int(p_min[2] - 1), c='r', s=60)\n",
    "    draw_poses(ax, gt_poses[:, :3], proj=False, proj_z=int(p_min[2] - 1), c='b', s=60)\n",
    "    for i in range(pred_poses.shape[0]):\n",
    "        pp = pred_poses[i, :3]\n",
    "        gp = gt_poses[i, :3]\n",
    "        pps = np.vstack((pp, gp))\n",
    "        ax.plot(pps[:, 0], pps[:, 1], pps[:, 2], c=(0.7, 0.7, 0.7, 0.4))\n",
    "        \n",
    "    plt.draw()\n",
    "    \n",
    "def draw_poses(ax, poses, c='b', s=20, proj=False, proj_z=0, pose_format='quat'):\n",
    "    \"\"\"Draws the list of poses.\n",
    "    Args:\n",
    "        ax (Axes3D): 3D axes\n",
    "        poses (list): Poses list\n",
    "        c: matplotlib color\n",
    "        s: matplotlib size\n",
    "        proj (bool): True if draw projection of a path on z-axis\n",
    "        proj_z (float): Coord for z-projection\n",
    "    \"\"\"\n",
    "    coords = np.zeros((len(poses), 3))\n",
    "    for i, p in enumerate(poses):\n",
    "        # coords[i] = p[:3, 3]\n",
    "        # coords[i] = p\n",
    "        coords[i] = extract_translation(p, pose_format=pose_format)\n",
    "\n",
    "    # Draw projection\n",
    "    if proj:\n",
    "        if len(poses) > 1:\n",
    "            ax.plot(coords[:, 0], coords[:, 1], proj_z, c='g')\n",
    "        elif len(poses) == 1:\n",
    "            ax.scatter(coords[:, 0], coords[:, 1], proj_z, c=c)\n",
    "\n",
    "    # Draw path\n",
    "    ax.scatter(coords[:, 0], coords[:, 1], coords[:, 2], c=c, s=s)\n",
    "\n",
    "\n",
    "def draw_poses_list(ax, poses_list):\n",
    "    \"\"\"Draw list of lists of poses.\"\"\"\n",
    "    for poses in poses_list:\n",
    "        draw_poses(ax, poses)\n",
    "\n",
    "\n",
    "# Draw predicted vs ground truth poses\n",
    "draw_pred_gt_poses(pred_poses_train, gt_poses_train)\n",
    "plt.title('Train Dataset')\n",
    "plt.show()\n",
    "\n",
    "\n",
    "draw_pred_gt_poses(pred_poses_val, gt_poses_val)\n",
    "plt.title('Testing Dataset')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Save model checkpoints"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "save_checkpoint(model, optimizer, criterion, 'nb_{}'.format(experiment_name), n_epochs)"
   ]
  }
 ]
}