visualize_policy_paper.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "code_folding": [],
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Visualizing policy using method from https://arxiv.org/abs/1711.00138\n",
    "# GitHub repo: https://github.com/greydanus/visualize_atari\n",
    "# Imports\n",
    "import os\n",
    "import sys\n",
    "import pdb\n",
    "import json\n",
    "import torch\n",
    "import random\n",
    "import argparse\n",
    "import torchvision\n",
    "import tensorboardX\n",
    "import torch.optim as optim\n",
    "import torchvision.utils as vutils\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "\"\"\"\n",
    "TODO: Set the absolute repository path.\n",
    "\"\"\"\n",
    "repo_path = \"/scratch/cluster/srama/Research/LearningToLookAround/SidekickPolicyLearning\"\n",
    "sys.path.append(os.path.join(repo_path, 'misc/'))\n",
    "\n",
    "from utils import *\n",
    "from models import *\n",
    "from torch.optim import lr_scheduler\n",
    "from tensorboardX import SummaryWriter\n",
    "from scipy.ndimage.filters import gaussian_filter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "code_folding": [
     9
    ],
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Refer https://github.com/greydanus/visualize_atari/blob/master/saliency.py\n",
    "\n",
    "def occlude(I, mask):\n",
    "    num_channels = I.shape[2]\n",
    "    I_out = np.zeros(I.shape)\n",
    "    for chn in range(num_channels):\n",
    "        I_out[:, :, chn] = I[:, :, chn]*(1-mask) + gaussian_filter(I[:, :, chn], sigma=5)*mask \n",
    "    return I_out\n",
    "    \n",
    "def get_mask(center, size, r):\n",
    "    y,x = np.ogrid[-center[0]:size[0]-center[0], -center[1]:size[1]-center[1]]\n",
    "    keep = x*x + y*y <= 1\n",
    "    mask = np.zeros(size) ; mask[keep] = 1 # select a circle of pixels\n",
    "    mask = gaussian_filter(mask, sigma=r) # blur the circle of pixels. this is a 2D Gaussian for r=r^2=1\n",
    "    return mask/mask.max()\n",
    "\n",
    "def saliency_on_image(saliency, image, fudge_factor, channels=[2], sigma=0):\n",
    "    # Refer https://github.com/greydanus/visualize_atari/blob/master/saliency.py\n",
    "    pmax = saliency.max()\n",
    "    S = saliency\n",
    "    S = S if sigma == 0 else gaussian_filter(S, sigma=sigma)\n",
    "    S -= S.min()\n",
    "    S = S.clip(0, 1)\n",
    "    S = 255.0*fudge_factor*pmax*S/(S.max() + 1e-7)\n",
    "    I = np.copy(image.astype('float32'))\n",
    "    if image.shape[2] == 1:\n",
    "        I[:, :, 0] += S.astype('float32')\n",
    "    else:\n",
    "        for channel in channels:\n",
    "            I[:, :, channel] += S.astype('float32')\n",
    "    I = I.clip(0, 255.0).astype('uint8')\n",
    "    return I\n",
    "\n",
    "def get_belief_saliency(belief, others, agent, iscuda, lr=1e-4, weight_decay=1e-1, iters=200):\n",
    "    \"\"\"\n",
    "    belief: current hidden state of the aggregator\n",
    "    others: dictionary of elevation('elev'), azimuth('azim'), time('time'), proprioception change('pro')\n",
    "    NOTE: Assuming batch size = 1\n",
    "          Ensure that the agent's parameters have requires_grad=False\n",
    "    \"\"\"\n",
    "    M = agent.M\n",
    "    N = agent.N\n",
    "    C = agent.C\n",
    "    d_belief_orig = Variable(torch.randn(1, 256)*1e-3, requires_grad=True)\n",
    "    optimizer = optim.SGD([d_belief_orig], lr=lr, weight_decay=weight_decay, momentum=0.9)\n",
    "    \n",
    "    act_input_1 = torch.cat([belief.view(1, -1), others['pro'][:, :2]], dim=1) \n",
    "    \n",
    "    if 'elev' in others:\n",
    "        xe = others['elev']\n",
    "        act_input_1 = torch.cat([act_input_1, xe], dim=1)\n",
    "    if 'azim' in others:\n",
    "        xa = others['azim']\n",
    "        act_input_1 = torch.cat([act_input_1, xa], dim=1)\n",
    "    if 'time' in others:\n",
    "        xt = others['time']\n",
    "        act_input_1 = torch.cat([act_input_1, xt], dim=1)\n",
    "    \n",
    "    activations_1 = agent.policy.act(act_input_1)\n",
    "    # Solve for d_belief that maximises change in policy activations\n",
    "    for i in range(iters):\n",
    "        if iscuda:\n",
    "            d_belief = d_belief_orig.cuda()\n",
    "        else:\n",
    "            d_belief = d_belief_orig\n",
    "            \n",
    "        act_input_2 = torch.cat([d_belief + belief.view(1, -1), others['pro'][:, :2]], dim=1)\n",
    "        \n",
    "        if 'elev' in others:\n",
    "            xe = others['elev']\n",
    "            act_input_2 = torch.cat([act_input_2, xe], dim=1)\n",
    "        if 'azim' in others:\n",
    "            xa = others['azim']\n",
    "            act_input_2 = torch.cat([act_input_2, xa], dim=1)\n",
    "        if 'time' in others:\n",
    "            xt = others['time']\n",
    "            act_input_2 = torch.cat([act_input_2, xt], dim=1)\n",
    "\n",
    "        activations_2 = agent.policy.act(act_input_2)\n",
    "        loss = -((activations_1 - activations_2)**2).sum() + weight_decay*(d_belief**2).sum()\n",
    "        loss.backward()\n",
    "        #pdb.set_trace()\n",
    "        optimizer.step()\n",
    "        if torch.norm(d_belief_orig.data) >= 0.75 * torch.norm(belief.data):\n",
    "            break\n",
    "    diff = -((activations_1 - activations_2)**2).sum()\n",
    "    diff_probs = -((F.softmax(activations_1, dim=1) - F.softmax(activations_2, dim=1))**2).sum()\n",
    "    \n",
    "    # Get the saliency\n",
    "    decoded_2 = agent.policy.decode(F.normalize(d_belief + belief.view(1, -1), p=1, dim=1))\n",
    "    decoded_2 = decoded_2.view(1, N, M, C, 32, 32)\n",
    "    decoded_1 = agent.policy.decode(F.normalize(belief.view(1, -1), p=1, dim=1))\n",
    "    decoded_1 = decoded_1.view(1, N, M, C, 32, 32)\n",
    "    saliency = np.sum(0.5*((decoded_1 - decoded_2).data.cpu().numpy())**2, axis=3)\n",
    "    \n",
    "    return saliency, diff.data[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "code_folding": []
   },
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "TODO: Specify the path to model file to be visualized\n",
    "\"\"\"\n",
    "model_path = os.path.join(repo_path, 'models/sun360/ltla.net')\n",
    "\n",
    "loaded_state = torch.load(model_path)\n",
    "opts = loaded_state['opts']\n",
    "\n",
    "# Unset all the unnecessary variables\n",
    "opts.critic_full_obs = False\n",
    "opts.act_full_obs = False\n",
    "opts.batch_size = 50\n",
    "opts.trajectories_type = 'utility_maps'\n",
    "opts.expert_trajectories = False\n",
    "opts.expert_rewards = False\n",
    "\n",
    "if opts.dataset == 0:\n",
    "    opts.h5_path = os.path.join(repo_path, 'data/sun360/sun360_processed.h5')\n",
    "    opts.utility_h5_path = os.path.join(repo_path, 'scores/sun360/ours-demo-scores.h5')\n",
    "else:\n",
    "    opts.h5_path = os.path.join(repo_path, 'data/modelnet_hard/modelnet30_processed.h5')    \n",
    "    opts.utility_h5_path = os.path.join(repo_path, 'scores/modelnet_hard/ours-demo-scores.h5')\n",
    "    \n",
    "if not hasattr(opts, 'h5_path_unseen'):\n",
    "    if opts.dataset == 0:\n",
    "        opts.h5_path_unseen = ''\n",
    "    else:\n",
    "        opts.h5_path_unseen = os.path.join(repo_path, 'data/modelnet_hard/modelnet10_processed.h5')\n",
    "        \n",
    "if opts.expert_trajectories:\n",
    "    opts.T_sup = 3\n",
    "\n",
    "if opts.expert_trajectories:\n",
    "    agent = AgentSupervised(opts)\n",
    "else:\n",
    "    agent = Agent(opts)\n",
    "    \n",
    "agent.policy.load_state_dict(loaded_state['state_dict'])\n",
    "agent.policy.eval()\n",
    "\n",
    "if opts.expert_rewards:\n",
    "    from DataLoader import DataLoaderExpert as DataLoader\n",
    "elif opts.expert_trajectories:\n",
    "    from DataLoader import DataLoaderExpertPolicy as DataLoader\n",
    "else:\n",
    "    from DataLoader import DataLoaderSimple as DataLoader\n",
    "\n",
    "loader = DataLoader(opts)\n",
    "\n",
    "set_random_seeds(opts.seed)\n",
    "\n",
    "if opts.dataset == 0:\n",
    "    opts.num_channels = 3\n",
    "    if opts.mean_subtract:\n",
    "        # R, G, B means and stds\n",
    "        opts.mean = [119.16, 107.68, 95.12]\n",
    "        opts.std = [61.88, 61.72, 67.24]\n",
    "    else:\n",
    "        opts.mean = [0, 0, 0]\n",
    "        opts.std = [1, 1, 1]\n",
    "elif opts.dataset == 1:\n",
    "    opts.num_channels = 1\n",
    "    if opts.mean_subtract:\n",
    "        # R, G, B means and stds\n",
    "        opts.mean = [193.0162338615919]\n",
    "        opts.std = [37.716024486312811]\n",
    "    else:\n",
    "        opts.mean = [0]\n",
    "        opts.std = [1]\n",
    "else:\n",
    "    raise ValueError('Dataset %d does not exist!'%(opts.dataset))\n",
    "    \n",
    "# Avoid computing gradients for the agent in backward\n",
    "for parameter in agent.policy.parameters():\n",
    "    parameter.requires_grad = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "code_folding": [],
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "saliency_belief_all = []\n",
    "images_all = []\n",
    "pano_all = []\n",
    "decoded_all = []\n",
    "decoded_saliency_all = []\n",
    "decoded_all_raw = []\n",
    "\n",
    "num_iters = 6\n",
    "for iters in range(num_iters):\n",
    "    if opts.expert_rewards:\n",
    "        pano, _, depleted = loader.next_batch_test()\n",
    "        pano_maps = None\n",
    "        pano_rewards = None\n",
    "    elif opts.expert_trajectories:\n",
    "        pano, _, _, depleted = loader.next_batch_test()\n",
    "        pano_rewards = None \n",
    "        pano_maps = None\n",
    "    else:\n",
    "        pano, _, depleted = loader.next_batch_test()\n",
    "        pano_rewards = None\n",
    "        pano_maps = None\n",
    "    \n",
    "    batch_size = opts.batch_size\n",
    "    start_idx = get_starts(opts.N, opts.M, batch_size, opts.start_view)\n",
    "    state_object = State(pano, pano_rewards, start_idx, opts)\n",
    "    hidden = None\n",
    "    images_all.append([])\n",
    "    decoded_all.append([])\n",
    "    saliency_belief_all.append([])\n",
    "    decoded_all_raw.append([])\n",
    "    visited_idxes = []\n",
    "    decoded_saliency_all.append([])    \n",
    "    \n",
    "    pano_all.append([])\n",
    "    \n",
    "    for t in range(opts.T):\n",
    "        # Get the original image's action distribution\n",
    "        im_np, pro = state_object.get_view() # im_np - BxCx32x32\n",
    "        im, pro = preprocess(im_np, pro)\n",
    "        \n",
    "        # Store the visited locations\n",
    "        visited_idxes.append(state_object.idx)\n",
    "        \n",
    "        C = im.shape[1]\n",
    "        W = im.shape[2]\n",
    "        H = im.shape[3]\n",
    "        \n",
    "        # Get the panoramas with all views filled in till the current time step\n",
    "        pano_full_view = np.copy(state_object.views)\n",
    "        if opts.num_channels == 1:\n",
    "            pano_full_view_3chn = np.zeros((batch_size, opts.N, opts.M, 3, 32, 32))\n",
    "            for c in range(3):\n",
    "                pano_full_view_3chn[:, :, :, c, :, :] = pano_full_view[:, :, :, 0, :, :]\n",
    "            pano_full_view = pano_full_view_3chn\n",
    "\n",
    "        for i in range(batch_size):\n",
    "            for t_prime in range(t+1):\n",
    "                \n",
    "                pano_full_view[i, visited_idxes[t_prime][i][0], visited_idxes[t_prime][i][1], 0, :3, :] = 255\n",
    "                pano_full_view[i, visited_idxes[t_prime][i][0], visited_idxes[t_prime][i][1], 0, :, :3] = 255\n",
    "                pano_full_view[i, visited_idxes[t_prime][i][0], visited_idxes[t_prime][i][1], 0, -3:, :] = 255\n",
    "                pano_full_view[i, visited_idxes[t_prime][i][0], visited_idxes[t_prime][i][1], 0, :, -3:] = 255\n",
    "\n",
    "                pano_full_view[i, visited_idxes[t_prime][i][0], visited_idxes[t_prime][i][1], 1:, :3, :] = 0\n",
    "                pano_full_view[i, visited_idxes[t_prime][i][0], visited_idxes[t_prime][i][1], 1:, :, :3] = 0\n",
    "                pano_full_view[i, visited_idxes[t_prime][i][0], visited_idxes[t_prime][i][1], 1:, -3:, :] = 0\n",
    "                pano_full_view[i, visited_idxes[t_prime][i][0], visited_idxes[t_prime][i][1], 1:, :, -3:] = 0\n",
    "\n",
    "        pano_full_view = pano_full_view.transpose((0, 3, 1, 4, 2, 5)).reshape(\\\n",
    "                                    batch_size, 1, 3, opts.N*32, opts.M*32)\n",
    "        pano_all[iters].append(pano_full_view)\n",
    "        \n",
    "        policy_input = {'im': im, 'pro': pro}\n",
    "        # Assuming only actOnElev\n",
    "        policy_input['elev'] = torch.Tensor([[state_object.idx[i][0]] for i in range(batch_size)])\n",
    "        policy_input['time'] = torch.Tensor([[t] for i in range(batch_size)])\n",
    "        if opts.iscuda:\n",
    "            for var in policy_input:\n",
    "                policy_input[var] = policy_input[var].cuda()\n",
    "        \n",
    "        for var in policy_input:\n",
    "            policy_input[var] = Variable(policy_input[var])\n",
    "        \n",
    "        probs, decoded, hidden_new, value = agent.policy.forward(policy_input, hidden)\n",
    "\n",
    "        # Create the decoded image with saliency modified view filled in\n",
    "        decoded_images = decoded.data.cpu().numpy()*255.0\n",
    "        # Add the means\n",
    "        for chn in range(len(opts.mean)):\n",
    "            decoded_images[:, :, :, chn] += opts.mean[chn] \n",
    "\n",
    "        decoded_images_shifted = np.zeros(decoded_images.shape)\n",
    "        decoded_images_shifted_wo_sal = np.zeros(decoded_images.shape)\n",
    "\n",
    "        for i in range(batch_size):\n",
    "            decoded_images_shifted[i] = np.copy(np.roll(decoded_images[i], state_object.start_idx[i][1], axis=1))\n",
    "            decoded_images_shifted_wo_sal[i] = np.copy(np.roll(decoded_images[i], state_object.start_idx[i][1], axis=1))\n",
    "        if opts.num_channels == 1:\n",
    "            decoded_images_shifted_3chn = np.zeros((batch_size, opts.N, opts.M, 3, 32, 32))\n",
    "            decoded_images_shifted_wo_sal_3chn = np.zeros((batch_size, opts.N, opts.M, 3, 32, 32))\n",
    "            for c in range(3):\n",
    "                decoded_images_shifted_3chn[:, :, :, c, :, :] = np.copy(decoded_images_shifted[:, :, :, 0, :, :])\n",
    "                decoded_images_shifted_wo_sal_3chn[:, :, :, c, :, :] = np.copy(decoded_images_shifted_wo_sal[:, :, :, 0, :, :])\n",
    "            decoded_images_shifted = decoded_images_shifted_3chn\n",
    "            decoded_images_shifted_wo_sal = decoded_images_shifted_wo_sal_3chn\n",
    "            \n",
    "        for t_view, visited_view in enumerate(visited_idxes):\n",
    "            for i in range(batch_size):                \n",
    "                curr_view = np.copy(state_object.views[i, visited_view[i][0], visited_view[i][1]])\n",
    "                curr_view = np.transpose(curr_view, (1, 2, 0))\n",
    "                if opts.num_channels == 1:\n",
    "                    curr_view_3chn = np.zeros((32, 32, 3))\n",
    "                    for c in range(3):\n",
    "                        curr_view_3chn[:, :, c] = np.copy(curr_view[: ,:, 0])\n",
    "                    curr_view = curr_view_3chn\n",
    "                \n",
    "                # Modify the current view to have some borders (for easy visualization)\n",
    "                if t_view == t:\n",
    "                    curr_view[:3, :, 0] = 255\n",
    "                    curr_view[-3:, :, 0] = 255\n",
    "                    curr_view[:, :3, 0] = 255\n",
    "                    curr_view[:, -3:, 0] = 255\n",
    "                    curr_view[:3, :, 1:] = 0 \n",
    "                    curr_view[-3:, :, 1:] = 0\n",
    "                    curr_view[:, :3, 1:] = 0\n",
    "                    curr_view[:, -3:, 1:] = 0\n",
    "\n",
    "                decoded_images_shifted[i, visited_view[i][0], visited_view[i][1]] = \\\n",
    "                    np.transpose(curr_view, (2, 0, 1))\n",
    "                    \n",
    "                decoded_images_shifted_wo_sal[i, visited_view[i][0], visited_view[i][1]] = \\\n",
    "                    np.transpose(curr_view, (2, 0, 1))\n",
    "        \n",
    "        # Convert BxNxMxCx32x32 -> BxCx(32*N)x(32*M)\n",
    "        decoded_image_full_view = decoded_images_shifted.transpose((0, 3, 1, 4, 2, 5)).reshape(\\\n",
    "                                            batch_size, 1, 3, opts.N*32, opts.M*32)\n",
    "        decoded_image_full_view_wo_sal = decoded_images_shifted_wo_sal.transpose((0, 3, 1, 4, 2, 5)).reshape(\\\n",
    "                                            batch_size, 1, 3, opts.N*32, opts.M*32)\n",
    "        \n",
    "        decoded_all[iters].append(decoded_image_full_view)\n",
    "        decoded_all_raw[iters].append(decoded_image_full_view_wo_sal)\n",
    "        \n",
    "        # ============ Get the belief state saliency ============ \n",
    "        saliency_belief = []\n",
    "        for i in range(batch_size):\n",
    "            policy_input_curr = {}\n",
    "            # Get only current batch element\n",
    "            for k, v in policy_input.iteritems():\n",
    "                policy_input_curr[k] = v[i:i+1]\n",
    "            saliency_belief_curr, _ = get_belief_saliency(hidden_new[0][:, i:i+1], policy_input_curr, agent, opts.iscuda)\n",
    "            # roll saliency_belief_curr to reflect the unknown azimuth\n",
    "            saliency_belief_curr = np.roll(saliency_belief_curr, state_object.start_idx[i][1], axis=2)\n",
    "            saliency_belief_curr = saliency_belief_curr.transpose((0, 1, 3, 2, 4)).reshape(1, opts.N*32, opts.M*32)\n",
    "            saliency_belief.append(saliency_belief_curr)\n",
    "        \n",
    "        saliency_belief = np.concatenate(saliency_belief, axis=0)\n",
    "        saliency_belief_all[iters].append(saliency_belief)\n",
    "                        \n",
    "        # Act greedily\n",
    "        _, act = probs.max(dim=1)\n",
    "        act = act.data.view(-1, 1)\n",
    "        # Rotate the view\n",
    "        _ = state_object.rotate(act[:, 0])\n",
    "        \n",
    "        # Set hidden\n",
    "        hidden = hidden_new\n",
    "    \n",
    "    # Normalize the saliency values\n",
    "    for i in range(batch_size):        \n",
    "        saliency_max = 0\n",
    "        for t in range(opts.T):\n",
    "            saliency_max = saliency_belief_all[iters][t][i].max()\n",
    "            saliency_belief_all[iters][t][i] /= (saliency_max + 1e-9)\n",
    "    \n",
    "    # Impose saliency on decoded views\n",
    "    for t in range(opts.T):\n",
    "\n",
    "        decoded_saliency_full_view = np.zeros((batch_size, 1, 3, opts.N*32, opts.M*32))\n",
    "        for i in range(batch_size):\n",
    "            curr_saliency = saliency_belief_all[iters][t][i]\n",
    "            curr_view = decoded_all_raw[iters][t][i, 0].transpose(1, 2, 0)\n",
    "            decoded_saliency_full_view[i, 0] = saliency_on_image(curr_saliency, curr_view, 1, channels=[0, 2], sigma=2).transpose(2, 0, 1)\n",
    "        \n",
    "        decoded_saliency_all[iters].append(decoded_saliency_full_view)    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "code_folding": [],
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Write outputs to tensorboard\n",
    "save_path = os.path.join(repo_path, 'visualizations/sun360/')\n",
    "writer = SummaryWriter(log_dir=save_path)\n",
    "images_count = 0\n",
    "for iters in range(num_iters):\n",
    "    batch_size = decoded_all[iters][0].shape[0]\n",
    "    for i in range(batch_size):\n",
    "        images_count += 1\n",
    "        outputs_all = []\n",
    "        #display_image = np.transpose(pano_all[iters][i], (0, 2, 3, 1)).astype('float32')\n",
    "        #display_image = (display_image - display_image.min())/(display_image.max()-display_image.min())\n",
    "        #outputs_all.append(display_image)\n",
    "        for t in range(opts.T):\n",
    "            display_image = np.transpose(decoded_all[iters][t][i], (0, 2, 3, 1))\n",
    "            display_image = (display_image - display_image.min())/(display_image.max()-display_image.min())\n",
    "            outputs_all.append(display_image)\n",
    "            \n",
    "        outputs_all = np.concatenate(outputs_all, axis=0).transpose((0, 3, 1, 2))\n",
    "        x = vutils.make_grid(torch.Tensor(outputs_all), padding=5, normalize=True, scale_each=True, nrow=opts.T+1, pad_value=1.0)\n",
    "        writer.add_image('#%d Decoded Viewgrid'%(images_count), x, 0)\n",
    "        \n",
    "        outputs_all = []\n",
    "        for t in range(opts.T):\n",
    "            display_image = np.transpose(pano_all[iters][t][i], (0, 2, 3, 1)).astype('float32')\n",
    "            display_image = (display_image - display_image.min())/(display_image.max()-display_image.min())\n",
    "            outputs_all.append(display_image)\n",
    "        \n",
    "        outputs_all = np.concatenate(outputs_all, axis=0).transpose((0, 3, 1, 2))\n",
    "        x = vutils.make_grid(torch.Tensor(outputs_all), padding=5, normalize=False, scale_each=False, nrow=opts.T+1, pad_value=1.0)\n",
    "        writer.add_image('#%d GT Viewgrid'%(images_count), x, 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "code_folding": [],
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Visualize belief saliency in numpy\n",
    "for iters in range(1):\n",
    "    batch_size = decoded_saliency_all[iters][0].shape[0]\n",
    "    for i in range(batch_size):\n",
    "        for t in range(opts.T):\n",
    "            display_image = np.transpose(decoded_saliency_all[iters][t][i][0], (1, 2, 0)).astype('float32')\n",
    "            display_image = (display_image - display_image.min())/(display_image.max()-display_image.min())\n",
    "            fig = plt.figure(figsize=(12, 16))\n",
    "            if display_image.shape[2] == 1:\n",
    "                plt.imshow(display_image[:, :, 0], cmap='gray')\n",
    "            else:\n",
    "                plt.imshow(display_image)\n",
    "            plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "code_folding": [],
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Write outputs to tensorboard\n",
    "images_count = 0\n",
    "for iters in range(num_iters):\n",
    "    batch_size = decoded_saliency_all[iters][0].shape[0]\n",
    "    for i in range(batch_size):\n",
    "        images_count += 1\n",
    "        outputs_all = []\n",
    "        #display_image = pano_all[iters][i].astype('float32')\n",
    "        #display_image = (display_image - display_image.min())/(display_image.max()-display_image.min())\n",
    "        #outputs_all.append(display_image)\n",
    "        for t in range(opts.T):\n",
    "            display_image = decoded_saliency_all[iters][t][i].astype('float32')\n",
    "            display_image = (display_image - display_image.min())/(display_image.max()-display_image.min())\n",
    "            outputs_all.append(display_image)\n",
    "        \n",
    "        outputs_all = np.concatenate(outputs_all, axis=0)\n",
    "        x = vutils.make_grid(torch.Tensor(outputs_all), padding=5, normalize=True, scale_each=True, nrow=opts.T+1, pad_value=1.0)\n",
    "        writer.add_image('#%d Belief Saliency'%(images_count), x, 0)\n",
    "\n",
    "writer.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}