{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "a2b1851d",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/anaconda3/envs/scit/lib/python3.8/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "from torch.utils.data import Dataset\n",
    "from torch.utils.data import DataLoader\n",
    "import pickle\n",
    "import glob\n",
    "# from utils import *\n",
    "from metrics import *\n",
    "from model import SocialImplicit\n",
    "from CFG import CFG"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "0a60ad04",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "SocialImplicit(\n",
       "  (implicit_cells): ModuleList(\n",
       "    (0): SocialCellGlobal(\n",
       "      (feat): Conv2d(2, 2, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      (feat_act): ReLU()\n",
       "      (highway_input): Conv2d(2, 2, kernel_size=(1, 1), stride=(1, 1))\n",
       "      (highway): Conv2d(8, 12, kernel_size=(1, 1), stride=(1, 1))\n",
       "      (tpcnn): Conv2d(8, 12, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      (ped): SocialCellLocal(\n",
       "        (feat): Conv1d(2, 2, kernel_size=(3,), stride=(1,), padding=(1,))\n",
       "        (feat_act): ReLU()\n",
       "        (highway_input): Conv1d(2, 2, kernel_size=(1,), stride=(1,))\n",
       "        (highway): Conv1d(8, 12, kernel_size=(1,), stride=(1,))\n",
       "        (tpcnn): Conv1d(8, 12, kernel_size=(3,), stride=(1,), padding=(1,))\n",
       "      )\n",
       "    )\n",
       "    (1): SocialCellGlobal(\n",
       "      (feat): Conv2d(2, 2, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      (feat_act): ReLU()\n",
       "      (highway_input): Conv2d(2, 2, kernel_size=(1, 1), stride=(1, 1))\n",
       "      (highway): Conv2d(8, 12, kernel_size=(1, 1), stride=(1, 1))\n",
       "      (tpcnn): Conv2d(8, 12, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      (ped): SocialCellLocal(\n",
       "        (feat): Conv1d(2, 2, kernel_size=(3,), stride=(1,), padding=(1,))\n",
       "        (feat_act): ReLU()\n",
       "        (highway_input): Conv1d(2, 2, kernel_size=(1,), stride=(1,))\n",
       "        (highway): Conv1d(8, 12, kernel_size=(1,), stride=(1,))\n",
       "        (tpcnn): Conv1d(8, 12, kernel_size=(3,), stride=(1,), padding=(1,))\n",
       "      )\n",
       "    )\n",
       "    (2): SocialCellGlobal(\n",
       "      (feat): Conv2d(2, 2, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      (feat_act): ReLU()\n",
       "      (highway_input): Conv2d(2, 2, kernel_size=(1, 1), stride=(1, 1))\n",
       "      (highway): Conv2d(8, 12, kernel_size=(1, 1), stride=(1, 1))\n",
       "      (tpcnn): Conv2d(8, 12, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      (ped): SocialCellLocal(\n",
       "        (feat): Conv1d(2, 2, kernel_size=(3,), stride=(1,), padding=(1,))\n",
       "        (feat_act): ReLU()\n",
       "        (highway_input): Conv1d(2, 2, kernel_size=(1,), stride=(1,))\n",
       "        (highway): Conv1d(8, 12, kernel_size=(1,), stride=(1,))\n",
       "        (tpcnn): Conv1d(8, 12, kernel_size=(3,), stride=(1,), padding=(1,))\n",
       "      )\n",
       "    )\n",
       "    (3): SocialCellGlobal(\n",
       "      (feat): Conv2d(2, 2, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      (feat_act): ReLU()\n",
       "      (highway_input): Conv2d(2, 2, kernel_size=(1, 1), stride=(1, 1))\n",
       "      (highway): Conv2d(8, 12, kernel_size=(1, 1), stride=(1, 1))\n",
       "      (tpcnn): Conv2d(8, 12, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      (ped): SocialCellLocal(\n",
       "        (feat): Conv1d(2, 2, kernel_size=(3,), stride=(1,), padding=(1,))\n",
       "        (feat_act): ReLU()\n",
       "        (highway_input): Conv1d(2, 2, kernel_size=(1,), stride=(1,))\n",
       "        (highway): Conv1d(8, 12, kernel_size=(1,), stride=(1,))\n",
       "        (tpcnn): Conv1d(8, 12, kernel_size=(3,), stride=(1,), padding=(1,))\n",
       "      )\n",
       "    )\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\n",
    "model_path = './checkpoint/social-implicit-eth/val_best.pth'\n",
    "\n",
    "KSTEPS= 20\n",
    "obs_seq_len = 8\n",
    "pred_seq_len = 12\n",
    "\n",
    "model = SocialImplicit(spatial_input=CFG[\"spatial_input\"],\n",
    "                       spatial_output=CFG[\"spatial_output\"],\n",
    "                       temporal_input=CFG[\"temporal_input\"],\n",
    "                       temporal_output=CFG[\"temporal_output\"],\n",
    "                       bins=CFG[\"bins\"],\n",
    "                       noise_weight=CFG[\"noise_weight\"]).cuda()\n",
    "model.load_state_dict(torch.load(model_path))\n",
    "model.cuda().double()\n",
    "model.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "8f8da27f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def anorm(p1, p2):\n",
    "    NORM = math.sqrt((p1[0] - p2[0])**2 + (p1[1] - p2[1])**2)\n",
    "    if NORM == 0:\n",
    "        return 0\n",
    "    return 1 / (NORM)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "9cf207a7",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[[[-5.1900, -4.3000, -3.3600, -2.4200, -1.3900, -0.5600,  0.2600,\n",
      "            1.2700],\n",
      "          [ 6.3800,  6.5000,  6.5600,  6.5100,  6.3900,  6.3800,  6.1800,\n",
      "            6.0200]],\n",
      "\n",
      "         [[-5.4500, -4.5600, -3.6600, -2.7200, -1.5900, -0.5900,  0.3600,\n",
      "            1.3200],\n",
      "          [ 7.3100,  7.4200,  7.4800,  7.4200,  7.3100,  7.2200,  7.1400,\n",
      "            7.0500]]]])\n",
      "tensor([[[[ 0.0000,  0.0000],\n",
      "          [ 0.0000,  0.0000]],\n",
      "\n",
      "         [[ 0.8900,  0.1200],\n",
      "          [ 0.8900,  0.1100]],\n",
      "\n",
      "         [[ 0.9400,  0.0600],\n",
      "          [ 0.9000,  0.0600]],\n",
      "\n",
      "         [[ 0.9400, -0.0500],\n",
      "          [ 0.9400, -0.0600]],\n",
      "\n",
      "         [[ 1.0300, -0.1200],\n",
      "          [ 1.1300, -0.1100]],\n",
      "\n",
      "         [[ 0.8300, -0.0100],\n",
      "          [ 1.0000, -0.0900]],\n",
      "\n",
      "         [[ 0.8200, -0.2000],\n",
      "          [ 0.9500, -0.0800]],\n",
      "\n",
      "         [[ 1.0100, -0.1600],\n",
      "          [ 0.9600, -0.0900]]]])\n",
      "torch.Size([1, 2, 2, 8])\n",
      "torch.Size([1, 8, 2, 2])\n",
      "<class 'torch.Tensor'>\n",
      "<class 'torch.Tensor'>\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_1124962/3389512904.py:101: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).\n",
      "  obs_traj = torch.tensor(obs_traj, device='cuda:0', dtype=torch.float64)\n",
      "/tmp/ipykernel_1124962/3389512904.py:104: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).\n",
      "  V_obs = torch.tensor(V_obs, device='cuda:0', dtype=torch.float64)\n"
     ]
    }
   ],
   "source": [
    "obs_traj = torch.tensor([[[[-5.1900, -4.3000, -3.3600, -2.4200, -1.3900, -0.5600,  0.2600,\n",
    "                            1.2700],\n",
    "                          [ 6.3800,  6.5000,  6.5600,  6.5100,  6.3900,  6.3800,  6.1800,\n",
    "                            6.0200]],\n",
    "                         [[-5.4500, -4.5600, -3.6600, -2.7200, -1.5900, -0.5900,  0.3600,\n",
    "                            1.3200],\n",
    "                          [ 7.3100,  7.4200,  7.4800,  7.4200,  7.3100,  7.2200,  7.1400,\n",
    "                            7.0500]]]])\n",
    "\n",
    "# obs_traj = torch.tensor([[[[815.6164, 815.9497, 826.1590, 841.0939, 861.3014, 881.9042, 913.6781,\n",
    "#                            930.4001],\n",
    "#                           [627.8403, 601.6028, 563.9097, 525.0204, 508.3439, 468.5327, 448.9330,\n",
    "#                            422.7056]],\n",
    "#                          [[683.9027, 686.3411, 704.4994, 725.0726, 747.9325, 766.6326, 781.7536,\n",
    "#                            823.3508],\n",
    "#                           [370.3823, 370.7501, 375.2039, 403.3822, 413.6887, 449.8790, 470.3914,\n",
    "#                            478.2605]]]])\n",
    "\n",
    "# obs_traj = torch.tensor([[[[1192.4408, 1224.5033, 1242.3634, 1252.1580, 1263.6909, 1268.7311,\n",
    "#                            1291.1096, 1306.3733],\n",
    "#                           [ 274.1519,  283.8012,  310.8191,  314.6419,  320.4248,  322.6592,\n",
    "#                             332.1011,  338.6567]]]])\n",
    "\n",
    "# pred_traj_gt = torch.tensor([[[[ 4.8700,  4.5100,  4.2000,  3.9500,  3.4700,  2.8200,  2.0100,\n",
    "#                                 1.2800,  0.5400, -0.1800, -0.8300, -1.5200],\n",
    "#                               [ 7.1600,  7.5800,  7.3000,  7.7100,  7.8600,  8.0000,  8.0000,\n",
    "#                                 7.8200,  7.4000,  7.0600,  6.4300,  6.0500]],\n",
    "#                              [[ 6.2900,  5.6200,  5.0600,  4.6900,  4.3500,  3.7600,  3.1900,\n",
    "#                                 2.6200,  1.7800,  1.0100,  0.0700, -0.7200],\n",
    "#                               [ 7.0000,  7.1000,  7.0400,  7.0000,  7.0100,  6.9900,  6.8900,\n",
    "#                                 7.1300,  7.1500,  6.9600,  6.9100,  6.6600]]]])\n",
    "\n",
    "# pred_traj_gt = torch.tensor([[[[ 4.8700,  4.5100,  4.2000,  3.9500,  3.4700,  2.8200,  2.0100,\n",
    "#                                 1.2800,  0.5400, -0.1800, -0.8300, -1.5200],\n",
    "#                               [ 7.1600,  7.5800,  7.3000,  7.7100,  7.8600,  8.0000,  8.0000,\n",
    "#                                 7.8200,  7.4000,  7.0600,  6.4300,  6.0500]]]])\n",
    "\n",
    "V_obs = torch.tensor([[[[ 0.0000,  0.0000],\n",
    "                      [ 0.0000,  0.0000]],\n",
    "                     [[ 0.8900,  0.1200],\n",
    "                      [ 0.8900,  0.1100]],\n",
    "                     [[ 0.9400,  0.0600],\n",
    "                      [ 0.9000,  0.0600]],\n",
    "                     [[ 0.9400, -0.0500],\n",
    "                      [ 0.9400, -0.0600]],\n",
    "                     [[ 1.0300, -0.1200],\n",
    "                      [ 1.1300, -0.1100]],\n",
    "                     [[ 0.8300, -0.0100],\n",
    "                      [ 1.0000, -0.0900]],\n",
    "                     [[ 0.8200, -0.2000],\n",
    "                      [ 0.9500, -0.0800]],\n",
    "                     [[ 1.0100, -0.1600],\n",
    "                      [ 0.9600, -0.0900]]]])\n",
    "\n",
    "# V_obs = torch.tensor([[[[ 0.0000,  0.0000]],\n",
    "#                      [[32.0625,  9.6493]],\n",
    "#                      [[17.8601, 27.0179]],\n",
    "#                      [[ 9.7946,  3.8229]],\n",
    "#                      [[11.5330,  5.7829]],\n",
    "#                      [[ 5.0402,  2.2344]],\n",
    "#                      [[22.3786,  9.4419]],\n",
    "#                      [[15.2637,  6.5555]]]])\n",
    "\n",
    "\n",
    "\n",
    "V_tr = torch.tensor([[[[-0.3700,  0.1800],\n",
    "                      [-0.6700,  0.1600]],\n",
    "                     [[-0.3600,  0.4200],\n",
    "                      [-0.6700,  0.1000]],\n",
    "                     [[-0.3100, -0.2800],\n",
    "                      [-0.5600, -0.0600]],\n",
    "                     [[-0.2500,  0.4100],\n",
    "                      [-0.3700, -0.0400]],\n",
    "                     [[-0.4800,  0.1500],\n",
    "                      [-0.3400,  0.0100]],\n",
    "                     [[-0.6500,  0.1400],\n",
    "                      [-0.5900, -0.0200]],\n",
    "                     [[-0.8100,  0.0000],\n",
    "                      [-0.5700, -0.1000]],\n",
    "                     [[-0.7300, -0.1800],\n",
    "                      [-0.5700,  0.2400]],\n",
    "                     [[-0.7400, -0.4200],\n",
    "                      [-0.8400,  0.0200]],\n",
    "                     [[-0.7200, -0.3400],\n",
    "                      [-0.7700, -0.1900]],\n",
    "                     [[-0.6500, -0.6300],\n",
    "                      [-0.9400, -0.0500]],\n",
    "                     [[-0.6900, -0.3800],\n",
    "                      [-0.7900, -0.2500]]]])\n",
    "\n",
    "\n",
    "print(obs_traj)\n",
    "print(V_obs)\n",
    "\n",
    "print(obs_traj.shape)\n",
    "print(V_obs.shape)\n",
    "\n",
    "print(type(obs_traj))\n",
    "print(type(V_obs))\n",
    "\n",
    "obs_traj = torch.tensor(obs_traj, device='cuda:0', dtype=torch.float64)\n",
    "# obs_traj2 = torch.tensor(obs_traj2, device='cuda:0', dtype=torch.float64)\n",
    "# pred_traj_gt = torch.tensor(pred_traj_gt, device='cuda:0', dtype=torch.float64)\n",
    "V_obs = torch.tensor(V_obs, device='cuda:0', dtype=torch.float64)\n",
    "# V_obs2 = torch.tensor(V_obs2, device='cuda:0', dtype=torch.float64)\n",
    "# V_tr = torch.tensor(V_tr, device='cuda:0', dtype=torch.float64)\n",
    "\n",
    "\n",
    "# V_tr = V_tr.squeeze()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "d67e73b7-3294-42fc-a4c4-375a012fc2dd",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "V_obs_tmp = V_obs.permute(0, 3, 1, 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68655f27-69a5-476c-a02d-fc8ea9a99793",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "15128f90-b085-493e-9110-f1c201313915",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# def seq_to_nodes(seq_):\n",
    "#     seq_ = seq_.squeeze().permute(2, 1, 0)\n",
    "#     return seq_\n",
    "\n",
    "def seq_to_nodes(seq_):\n",
    "    seq_ = seq_[0].permute(2, 1, 0)\n",
    "    return seq_\n",
    "\n",
    "def nodes_rel_to_nodes_abs(nodes, init_node):\n",
    "    nodes_ = torch.cumsum(nodes, dim=0) + init_node\n",
    "    return nodes_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "86e6f47f-f275-49c4-ad93-d97164cd4791",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "V_predx = model(V_obs_tmp, obs_traj, KSTEPS=KSTEPS)\n",
    "\n",
    "V_x = seq_to_nodes(obs_traj)\n",
    "num_of_objs = V_obs.shape[-2]\n",
    "pred_result = []\n",
    "\n",
    "for k in range(KSTEPS):\n",
    "    V_pred = V_predx[k:k + 1, ...]\n",
    "    V_pred = V_pred.permute(0, 2, 3, 1)\n",
    "    \n",
    "    V_pred = V_pred.squeeze()\n",
    "\n",
    "    V_pred_rel_to_abs = nodes_rel_to_nodes_abs(\n",
    "        V_pred, V_x[-1, :, :].clone())\n",
    "\n",
    "    for n in range(num_of_objs):\n",
    "        # target = []\n",
    "        # obsrvs = []\n",
    "        # number_of = []\n",
    "        pred_result.append(V_pred_rel_to_abs[:, n:n + 1, :])\n",
    "        \n",
    "pred_result_final = []\n",
    "for n in range(num_of_objs):\n",
    "    pred_result_final.append(pred_result[n])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "b62dade0-3d32-45a5-adaa-ac4e3b1dc6f9",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[ 2.1789,  1.1712]],\n",
       "\n",
       "        [[ 3.0693,  1.0040]],\n",
       "\n",
       "        [[ 3.9621,  0.8512]],\n",
       "\n",
       "        [[ 4.8520,  0.6834]],\n",
       "\n",
       "        [[ 5.7727,  0.5236]],\n",
       "\n",
       "        [[ 6.6680,  0.3720]],\n",
       "\n",
       "        [[ 7.5314,  0.2315]],\n",
       "\n",
       "        [[ 8.4358,  0.0950]],\n",
       "\n",
       "        [[ 9.3548, -0.0513]],\n",
       "\n",
       "        [[10.2703, -0.2050]],\n",
       "\n",
       "        [[11.1680, -0.3551]],\n",
       "\n",
       "        [[12.0495, -0.4947]]], device='cuda:0', dtype=torch.float64,\n",
       "       grad_fn=<SliceBackward0>)"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pred_result_final[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b0b4f98-e0a0-4989-a2d0-75ad0ab5510c",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a8ed513a-fcff-4082-b021-e31c0aac3c87",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "scit",
   "language": "python",
   "name": "scit"
  },
  "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.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}