pusht/qualitative-visualize.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Qualitative Visualization  \n",
    "\n",
    "This notebook demonstrates how to create the qualitative visualization of Figure 9 and Figure 10 in paper. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step-1: Generate Data\n",
    "\n",
    "You can use `pusht_human.py` to find a preferable layout and save both the images and pointer locations by pressing `s` key. \n",
    "\n",
    "Then, you can draw the trajectory with `draw_human_trajectory.ipynb`, just use your image as the background instead of the default \"000.jpg\".\n",
    "\n",
    "You can download our data from the `datasets/visualization` folder (in box). We assume that you already download the `datasets/visualization` folder (in box) and put it under `outputs` (local, in `pusht` folder)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "vscode": {
     "languageId": "plaintext"
    }
   },
   "source": [
    "## Step-2: Estimate Likelihood and conditioned Trajectories"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, import libraries and define some functions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from omegaconf import OmegaConf\n",
    "import os.path as osp\n",
    "import torch\n",
    "torch.set_grad_enabled(False)\n",
    "import os\n",
    "import hydra\n",
    "from copy import deepcopy\n",
    "from torchvision.transforms.functional import to_pil_image, to_tensor\n",
    "from diffusion_policy.common.pytorch_util import dict_apply\n",
    "import torch.nn.functional as F\n",
    "import pathlib\n",
    "from scipy import interpolate\n",
    "from PIL import Image, ImageDraw\n",
    "import numpy as np\n",
    "import json\n",
    "from tqdm.auto import tqdm\n",
    "\n",
    "def load_json(file_path):\n",
    "    with open(file_path, 'r') as f:\n",
    "        data = json.load(f)\n",
    "    return data\n",
    "\n",
    "def smooth_curve(points, window_size=3):\n",
    "    points = np.array(points)\n",
    "    smoothed_points = []\n",
    "    for i in range(len(points)):\n",
    "        start = max(0, i - window_size // 2)\n",
    "        end = min(len(points), i + window_size // 2 + 1)\n",
    "        avg_point = np.mean(points[start:end], axis=0)\n",
    "        smoothed_points.append(avg_point)\n",
    "    return smoothed_points\n",
    "\n",
    "def add_circle(draw, end_point, radius=5, fill=(255, 0, 0)):\n",
    "    x, y = end_point\n",
    "    bounding_box = [x - radius, y - radius, x + radius, y + radius]\n",
    "    draw.ellipse(bounding_box, fill=fill)\n",
    "\n",
    "def draw_smooth_line_with_dot(img, points, line_color=(255, 0, 0), line_width=3, dot_radius=5, dot_color=(255, 0, 0)):\n",
    "    img = img.convert('RGBA')\n",
    "    draw = ImageDraw.Draw(img)\n",
    "    smooth_points = smooth_curve(points)\n",
    "    smooth_points_tuple = [tuple(map(int, point)) for point in smooth_points]\n",
    "    draw.line(smooth_points_tuple, fill=line_color, width=line_width)\n",
    "    if len(smooth_points_tuple) > 1:\n",
    "        add_circle(draw, smooth_points_tuple[-1], radius=dot_radius, fill=dot_color)\n",
    "    return img"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = 0 # selecting the device"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we load two ARP models. They are simpler than the one in the main result: they use a flat action sequence. One of them treats every action as a continuous value, and the another one discretize the action space."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "============= Initialized Observation Utils with Obs Spec =============\n",
      "\n",
      "using obs modality: rgb with keys: ['image']\n",
      "using obs modality: depth with keys: []\n",
      "using obs modality: scan with keys: []\n",
      "using obs modality: low_dim with keys: []\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/common/users/xz653/anaconda3/envs/rvt/lib/python3.8/site-packages/torchvision/models/_utils.py:208: UserWarning: The parameter 'pretrained' is deprecated since 0.13 and may be removed in the future, please use 'weights' instead.\n",
      "  warnings.warn(\n",
      "/common/users/xz653/anaconda3/envs/rvt/lib/python3.8/site-packages/torchvision/models/_utils.py:223: UserWarning: Arguments other than a weight enum or `None` for 'weights' are deprecated since 0.13 and may be removed in the future. The current behavior is equivalent to passing `weights=None`.\n",
      "  warnings.warn(msg)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "============= Initialized Observation Utils with Obs Spec =============\n",
      "\n",
      "using obs modality: rgb with keys: ['image']\n",
      "using obs modality: depth with keys: []\n",
      "using obs modality: scan with keys: []\n",
      "using obs modality: low_dim with keys: []\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cfg1 = OmegaConf.load(f'outputs/visualization/flat.yaml')\n",
    "\n",
    "P1 = hydra.utils.instantiate(cfg1.policy)\n",
    "P1.load_state_dict(torch.load('outputs/visualization/flat_epoch=0500-test_mean_score=0.822.ckpt')['state_dicts']['model'])\n",
    "\n",
    "\n",
    "cfg2 = OmegaConf.load(f'outputs/visualization/flat_dis.yaml')\n",
    "P2 = hydra.utils.instantiate(cfg2.policy)\n",
    "P2.load_state_dict(torch.load('outputs/visualization/flat_dis_epoch=1500-test_mean_score=0.815.ckpt')['state_dicts']['model'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Likelihood Inference\n",
    "\n",
    "First, load all human data into `all_batches`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Data Preparation\n",
    "all_batches = []\n",
    "\n",
    "for i in range(4):\n",
    "    fig = Image.open(f'./outputs/visualization/human-data/{i:03d}.jpg')\n",
    "    pic = to_tensor(fig)\n",
    "    state = load_json(f'./outputs/visualization/human-data/{i:03d}.json')\n",
    "    action = load_json(f'./outputs/visualization/human-data/{i:03d}.action.json')\n",
    "\n",
    "    agent_pos = torch.as_tensor(state['pos_agent']).reshape(1, 2)\n",
    "\n",
    "    action_tensor = []\n",
    "    for l in action:\n",
    "        l = F.interpolate(torch.as_tensor(l).float().reshape(1, -1, 2).permute(0, 2, 1), size=(16,), mode='linear').permute(0, 2, 1)\n",
    "        l = l / 256 * 512\n",
    "        l = l.clamp(0, 511)\n",
    "        action_tensor.append(l)\n",
    "    action_tensor = torch.cat(action_tensor, 0)\n",
    "\n",
    "    batch = {\n",
    "        'image': F.interpolate(pic[None, ...], size=(96, 96), mode='bilinear', align_corners=False)[None, ...],\n",
    "        'origin_image': fig,\n",
    "        'agent_pos': agent_pos,\n",
    "        'action': action_tensor\n",
    "    }\n",
    "    all_batches.append(batch)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following code will estimate the likelihood of human-drawn trajectories, and save them to `./outputs/visualization/regenerated_likelihood.{images, states}`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "73d2f2bb543d459dbf4f01d374046ad2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/44 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "total = 0\n",
    "for bi in range(4):\n",
    "     for self_name, self in [('P1', P1), ('P2', P2)]:\n",
    "          batch = deepcopy(all_batches[bi])\n",
    "          for action_i, action in enumerate(batch['action']):\n",
    "               total += 1\n",
    "                    \n",
    "with tqdm(total=total) as pbar:\n",
    "     for bi in range(4):\n",
    "          for self_name, self in [('P1', P1), ('P2', P2)]:\n",
    "               batch = deepcopy(all_batches[bi])\n",
    "               for action_i, action in enumerate(batch['action']):\n",
    "                    action = batch['action'][[action_i]]\n",
    "\n",
    "                    self.eval()\n",
    "                    nobs = {\n",
    "                         'image': batch['image'],\n",
    "                         'agent_pos': batch['agent_pos']\n",
    "                    }\n",
    "\n",
    "                    nobs['image'] -= 0.5\n",
    "                    nobs['image'] /= 0.5\n",
    "\n",
    "                    nobs['agent_pos'] -= 256\n",
    "                    nobs['agent_pos'] /= 256.\n",
    "\n",
    "                    action -= 256\n",
    "                    action /= 256\n",
    "\n",
    "                    batch_size = 1\n",
    "                         \n",
    "                    future_tk_types = ['x', 'y'] * self.horizon\n",
    "                    future_chk_ids = list(range(self.horizon * 2))\n",
    "\n",
    "\n",
    "                    this_nobs = dict_apply(nobs, \n",
    "                         lambda x: x[:,:self.n_obs_steps,...].reshape(-1,*x.shape[2:]))\n",
    "                    nobs_features = self.obs_encoder(this_nobs)\n",
    "                    nobs_features = self.obs_feat_linear(nobs_features)\n",
    "                    nobs_features = nobs_features.reshape(batch_size, self.n_obs_steps, self.policy.cfg.n_embd)\n",
    "\n",
    "                    tk_vals = action.flatten(1).unsqueeze(-1)\n",
    "                    tk_names = future_tk_types\n",
    "                    tk_types = torch.as_tensor([self.policy.token_name_2_ids[tname] for tname in tk_names]).reshape(1, -1, 1).repeat(batch_size, 1, 1)\n",
    "                    seq = torch.cat([tk_vals, tk_types], dim=-1)\n",
    "                    loss_dict, log_prob = self.policy.compute_loss(seq, contexts={'visual-token': nobs_features}, log_prob=True) \n",
    "                    log_prob = log_prob.sum().item()\n",
    "\n",
    "                    pts = (action * 256 + 256) / 512 * 256\n",
    "                    fig = batch['origin_image'].copy()\n",
    "                    fig = draw_smooth_line_with_dot(fig, pts[0], line_color=(255, 0, 0, 208), line_width=4, dot_radius=4, dot_color=(255, 0, 0, 208))\n",
    "                    \n",
    "                    output_folder_images = f'./outputs/visualization/regenerated_likelihood.images/{self_name}/{bi}/'\n",
    "                    output_folder_states = f'./outputs/visualization/regenerated_likelihood.states/{self_name}/{bi}/'\n",
    "                    os.makedirs(output_folder_images, exist_ok=True)\n",
    "                    os.makedirs(output_folder_states, exist_ok=True)\n",
    "                    fig.save(osp.join(output_folder_images, f'{action_i}_{log_prob:.02f}.png'))\n",
    "                    torch.save({'pts': pts, 'log_prob': log_prob}, osp.join(output_folder_states, f'{action_i}.pt'))\n",
    "                    pbar.update()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Predict with Human Guidance"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following code will generate trajectories conditioned on first half of the human trajectory,  and save them to `./outputs/visualization/regenerated_guide.{images, states}`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "359dc0d98e1f4b90b714c6f3920846b4",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/440 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "error at 3 6, possibly get an NaN on out-of-distribution data\n",
      "error at 3 6, possibly get an NaN on out-of-distribution data\n",
      "error at 3 6, possibly get an NaN on out-of-distribution data\n",
      "error at 3 6, possibly get an NaN on out-of-distribution data\n",
      "error at 3 6, possibly get an NaN on out-of-distribution data\n",
      "error at 3 6, possibly get an NaN on out-of-distribution data\n",
      "error at 3 6, possibly get an NaN on out-of-distribution data\n",
      "error at 3 6, possibly get an NaN on out-of-distribution data\n",
      "error at 3 6, possibly get an NaN on out-of-distribution data\n"
     ]
    }
   ],
   "source": [
    "total = 0\n",
    "for bi in range(4):\n",
    "     for self_name, self in [('P1', P1), ('P2', P2)]:\n",
    "          batch = deepcopy(all_batches[bi])\n",
    "          for action_i, action in enumerate(batch['action']):\n",
    "               for sample_i in range(10):\n",
    "                    total += 1\n",
    "               \n",
    "               \n",
    "with tqdm(total=total) as pbar:       \n",
    "     for bi in range(4):\n",
    "          for self_name, self in [('P1', P1), ('P2', P2)]:\n",
    "               batch = deepcopy(all_batches[bi])\n",
    "               for action_i, action in enumerate(batch['action']):\n",
    "                    for sample_i in range(10):\n",
    "                         action = batch['action'][[action_i]]\n",
    "\n",
    "                         self.eval()\n",
    "                         nobs = {\n",
    "                              'image': batch['image'],\n",
    "                              'agent_pos': batch['agent_pos']\n",
    "                         }\n",
    "\n",
    "                         nobs['image'] -= 0.5\n",
    "                         nobs['image'] /= 0.5\n",
    "\n",
    "                         nobs['agent_pos'] -= 256\n",
    "                         nobs['agent_pos'] /= 256.\n",
    "\n",
    "                         action -= 256\n",
    "                         action /= 256\n",
    "\n",
    "                         batch_size = 1\n",
    "                              \n",
    "                         future_tk_types = ['x', 'y'] * (self.horizon // 2)   \n",
    "                         future_chk_ids = list(range(self.horizon, self.horizon * 2))\n",
    "                         future_tk_chk_ids = [{'tk_id': self.policy.token_name_2_ids[tk_type], 'chk_id': chk_id} for chk_id, tk_type in zip(future_chk_ids, future_tk_types)]       \n",
    "\n",
    "\n",
    "                         this_nobs = dict_apply(nobs, \n",
    "                              lambda x: x[:,:self.n_obs_steps,...].reshape(-1,*x.shape[2:]))\n",
    "                         nobs_features = self.obs_encoder(this_nobs)\n",
    "                         nobs_features = self.obs_feat_linear(nobs_features)\n",
    "                         nobs_features = nobs_features.reshape(batch_size, self.n_obs_steps, self.policy.cfg.n_embd)\n",
    "\n",
    "\n",
    "                         seq = torch.zeros(batch_size, self.horizon, 2)\n",
    "\n",
    "                         seq[:, :, 1] = torch.as_tensor([0, 1] * (self.horizon // 2))\n",
    "                         seq[:, :, 0] = action[0, :self.horizon//2, :].flatten()\n",
    "\n",
    "                         try:\n",
    "                              action_pred = self.policy.generate(seq, future_tk_chk_ids, contexts={'visual-token': nobs_features}, sample=True)\n",
    "                         except:\n",
    "                              tqdm.write(f'error at {bi} {action_i}, possibly get an NaN on out-of-distribution data')\n",
    "                              continue\n",
    "                         action_pred = action_pred[..., 0].reshape(-1, self.horizon, 2)\n",
    "\n",
    "                         action_pred *= 256.\n",
    "                         action_pred += 256.\n",
    "                         action_pred.clamp_(0, 511)\n",
    "\n",
    "\n",
    "                         pts = action_pred[0] / 512 * 256\n",
    "                         fig =  batch['origin_image'].copy()\n",
    "                         # draw_keypoints(fig, pts, radius=1)\n",
    "                         fig = draw_smooth_line_with_dot(fig, pts[:self.horizon//2+1], line_color=(255, 0, 0, 245), line_width=4, dot_radius=2, dot_color=(255, 0, 0, 208))\n",
    "                         fig = draw_smooth_line_with_dot(fig, pts[self.horizon//2:], line_color=(0, 0, 255, 208), line_width=4, dot_radius=4, dot_color=(0, 0, 255, 208))\n",
    "\n",
    "                         output_folder_images = f'./outputs/visualization/regenerated_guide.images/{self_name}/{bi}/'\n",
    "                         output_folder_states = f'./outputs/visualization/regenerated_guide.states/{self_name}/{bi}/'\n",
    "                         os.makedirs(output_folder_images, exist_ok=True)\n",
    "                         os.makedirs(output_folder_states, exist_ok=True)\n",
    "\n",
    "                         fig.save(osp.join(output_folder_images, f'{action_i}-{sample_i}.png'))\n",
    "                         torch.save(pts, osp.join(output_folder_states, f'{action_i}-{sample_i}.pt'))\n",
    "                         pbar.update()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step-3: Visualization"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Unzip the `outputs/visualization/visualized_data.zip` in the `outputs/visualization` folder. It contains our likelihood and prediction with guidance results (we use the discrete model at a different epoch).  Here we recreate the visualization in paper.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Creating Colormap\n",
    "\n",
    "We use the `fast` colormap from https://www.kennethmoreland.com/color-advice/. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "from colormath.color_objects import *\n",
    "from colormath.color_conversions import convert_color\n",
    "\n",
    "file_descriptor = open('outputs/visualization/fast.colormap.json', 'r')\n",
    "raw_color_data = json.load(file_descriptor)[0]\n",
    "import pandas\n",
    "\n",
    "scalar = []\n",
    "rgb_values = []\n",
    "for i in range(0, len(raw_color_data['RGBPoints']), 4):\n",
    "    scalar.append(raw_color_data['RGBPoints'][i+0])\n",
    "    rgb_values.append(sRGBColor(\n",
    "        raw_color_data['RGBPoints'][i+1],\n",
    "        raw_color_data['RGBPoints'][i+2],\n",
    "        raw_color_data['RGBPoints'][i+3]\n",
    "    ))\n",
    "\n",
    "data = pandas.DataFrame({'scalar': scalar, 'rgb_values': rgb_values})\n",
    "data['lab_values'] = data['rgb_values'].apply(lambda rgb: convert_color(rgb, LabColor))\n",
    "\n",
    "\n",
    "def color_lookup_sRGBColor(x):\n",
    "    if x < 0:\n",
    "        return sRGBColor(0, 0, 0)\n",
    "    for index in range(0, data.index.size-1):\n",
    "        low_scalar = data['scalar'][index]\n",
    "        high_scalar = data['scalar'][index+1]\n",
    "        if (x > high_scalar):\n",
    "            continue\n",
    "        low_lab = data['lab_values'][index]\n",
    "        high_lab = data['lab_values'][index+1]\n",
    "        interp = (x-low_scalar)/(high_scalar-low_scalar)\n",
    "        mid_lab = LabColor(interp*(high_lab.lab_l-low_lab.lab_l) + low_lab.lab_l,\n",
    "                           interp*(high_lab.lab_a-low_lab.lab_a) + low_lab.lab_a,\n",
    "                           interp*(high_lab.lab_b-low_lab.lab_b) + low_lab.lab_b,\n",
    "                           observer=low_lab.observer,\n",
    "                           illuminant=low_lab.illuminant)\n",
    "        return convert_color(mid_lab, sRGBColor)\n",
    "    return sRGBColor(1, 1, 1)\n",
    "\n",
    "def color_lookup(x):\n",
    "    return color_lookup_sRGBColor(x).get_value_tuple()\n",
    "\n",
    "def color_lookup_upscaled(x):\n",
    "    return color_lookup_sRGBColor(x).get_upscaled_value_tuple()\n",
    "\n",
    "def to_color(v):\n",
    "    r = color_lookup(v)[:3]\n",
    "    return tuple([int(i * 255) for i in r])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Predict with Human Guidance\n",
    "\n",
    "The following dicts are used to select which image / trajectory to visualization. You always want to choose the good looking ones. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "assignments ={\n",
    "0: {\n",
    "    'bad': [0,1,2],\n",
    "    'good': [3,4]\n",
    "},\n",
    "1: {\n",
    "    'bad': [2,3,4,5],\n",
    "    'good': [0,1]\n",
    "},\n",
    "2: {\n",
    "    'good': [0,1],\n",
    "    'bad': [2,3]\n",
    "},\n",
    "3: {\n",
    "    'good': [0, 1, 2, 3, 6],\n",
    "    'bad': [4, ]\n",
    "}\n",
    "}\n",
    "\n",
    "guide_assignments ={\n",
    "0: {\n",
    "    'bad': [(0, 0), (2, 0)],\n",
    "    'good': [(3, 4), (4, 0)]\n",
    "},\n",
    "1: {\n",
    "    'bad': [(2, 8), (3, 5) , (4, 9), (5, 1)],\n",
    "    'good': [(0, 0), (1, 7)]\n",
    "},\n",
    "2: {\n",
    "    'good': [(0, 0), (1, 9), (2, 3)],\n",
    "    'bad': [(3, 9)]\n",
    "},\n",
    "3: {\n",
    "    'good': [(0, 9), (1, 9), (2, 5), (5, 9)],\n",
    "    'bad': [(3, 3), (4, 2), (6, 5)]\n",
    "}\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following code generate the prediction with guidance visualization images, and save them to `./outputs/visualization/results/guide`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "dict_keys([0, 2])\n",
      "dict_keys([2, 3, 4, 5])\n",
      "dict_keys([3])\n",
      "dict_keys([3, 4, 6])\n"
     ]
    }
   ],
   "source": [
    "model_v = 'P1'\n",
    "horizon = 16\n",
    "\n",
    "os.makedirs('outputs/visualization/results/guide/unroll', exist_ok=True)\n",
    "\n",
    "for img_id in range(4):\n",
    "    origin_fig = Image.open(f'./outputs/visualization/human-data/{str(img_id).zfill(3)}.jpg')\n",
    "    guide_pts = {}\n",
    "\n",
    "    for k in guide_assignments[img_id]:\n",
    "        guide_pts[k] = {}\n",
    "        for aid, sid in guide_assignments[img_id][k]:\n",
    "            state = torch.load(f'./outputs/visualization/guide.states/{model_v}/{img_id}/{aid}-{sid}.pt')\n",
    "            guide_pts[k][aid] = state\n",
    "            \n",
    "    fig = origin_fig.copy()\n",
    "    for cid, aid in enumerate(sorted(guide_pts['good'])):\n",
    "        c1, c2 = (255, 0, 0, 200), (0, 0, 255, 200)\n",
    "        fig = draw_smooth_line_with_dot(fig, guide_pts['good'][aid][:horizon//2+1], line_color=c1, line_width=4, dot_radius=2, dot_color=c1)\n",
    "        fig = draw_smooth_line_with_dot(fig, guide_pts['good'][aid][horizon//2:], line_color=c2, line_width=4, dot_radius=4, dot_color=c2)\n",
    "\n",
    "        _fig = origin_fig.copy()\n",
    "        _fig = draw_smooth_line_with_dot(_fig, guide_pts['good'][aid][:horizon//2+1], line_color=c1, line_width=4, dot_radius=2, dot_color=c1)\n",
    "        _fig = draw_smooth_line_with_dot(_fig, guide_pts['good'][aid][horizon//2:], line_color=c2, line_width=4, dot_radius=4, dot_color=c2)\n",
    "        _fig.save(f'outputs/visualization/results/guide/unroll/{img_id}.{cid}.good.png')\n",
    "\n",
    "    fig.save(f'outputs/visualization/results/guide/{img_id}.good.png')\n",
    "\n",
    "\n",
    "    fig = origin_fig.copy()\n",
    "    for cid, aid in enumerate(guide_pts['bad']):\n",
    "        if img_id == 1:\n",
    "            if aid in [3]: continue\n",
    "        c1, c2 = (255, 0, 0, 200), (0, 0, 255, 200)\n",
    "        fig = draw_smooth_line_with_dot(fig, guide_pts['bad'][aid][:horizon//2+1], line_color=c1, line_width=4, dot_radius=2, dot_color=c1)\n",
    "        fig = draw_smooth_line_with_dot(fig, guide_pts['bad'][aid][horizon//2:], line_color=c2, line_width=4, dot_radius=4, dot_color=c2)\n",
    "\n",
    "        _fig = origin_fig.copy()\n",
    "        _fig = draw_smooth_line_with_dot(_fig, guide_pts['bad'][aid][:horizon//2+1], line_color=c1, line_width=4, dot_radius=2, dot_color=c1)\n",
    "        _fig = draw_smooth_line_with_dot(_fig, guide_pts['bad'][aid][horizon//2:], line_color=c2, line_width=4, dot_radius=4, dot_color=c2)\n",
    "        _fig.save(f'outputs/visualization/results/guide/unroll/{img_id}.{cid}.bad.png')\n",
    "\n",
    "\n",
    "    fig.save(f'outputs/visualization/results/guide/{img_id}.bad.png')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Likelihood Inference\n",
    "\n",
    "The following code generate the likelihood inference visualization images, and save them to `./outputs/visualization/results/likelihood`. Note for the image `1`, we use the discrete model instead of the continuous one. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.makedirs('outputs/visualization/results/likelihood/unroll', exist_ok=True)\n",
    "\n",
    "for img_id in range(4):\n",
    "    if img_id == 1:\n",
    "        model_v = 'P2'\n",
    "    else:\n",
    "        model_v = 'P1'\n",
    "    horizon = 16\n",
    "    BASELINE = 0.1\n",
    "\n",
    "    origin_fig = Image.open(f'./outputs/visualization/human-data/{str(img_id).zfill(3)}.jpg')\n",
    "    likelihood = {}\n",
    "    state_files = os.listdir(f'./outputs/visualization/likelihood.states/{model_v}/{img_id}/')\n",
    "    for s in state_files:\n",
    "        sid = int(s.split('.')[0])\n",
    "        state = torch.load(f'./outputs/visualization/likelihood.states/{model_v}/{img_id}/{s}')\n",
    "        likelihood[sid] = state\n",
    "\n",
    "    ma = max([v['log_prob'] for v in likelihood.values()])\n",
    "    mi = min([v['log_prob'] for v in likelihood.values()])\n",
    "    for v in likelihood.values():\n",
    "        v['normed_log_prob'] = (v['log_prob'] - mi) / (ma - mi) * 0.9 + 0.1 # some color tweaks\n",
    "\n",
    "    \n",
    "    fig = origin_fig.copy()\n",
    "    for i, s in sorted(likelihood.items()):\n",
    "        if img_id == 3:\n",
    "            if i in [3, 5]: continue\n",
    "        \n",
    "        nlp = round(s['normed_log_prob'], 2)\n",
    "        color = to_color(nlp)\n",
    "        fig = draw_smooth_line_with_dot(fig, s['pts'][0], line_color=color, line_width=4, dot_radius=4, dot_color=color)\n",
    "        _fig = origin_fig.copy()\n",
    "        _fig = draw_smooth_line_with_dot(_fig, s['pts'][0], line_color=color, line_width=4, dot_radius=4, dot_color=color)\n",
    "        _fig.save(f'outputs/visualization/results/likelihood/unroll/{img_id}.{i}.png')\n",
    "\n",
    "    fig.save(f'outputs/visualization/results/likelihood/{img_id}.png')\n",
    "        "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "rvt",
   "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.8.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}