{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0d498eb4-9486-4086-b62c-65ad3a6d5833",
   "metadata": {},
   "source": [
    "# Sparse4D Tutorial\n",
    "\n",
    "- [GitHub](https://github.com/linxuewu/Sparse4D), [Sparse4D v1](https://arxiv.org/abs/2211.10581), [Sparse4D v2](https://arxiv.org/abs/2305.14018), [Sparse4D v3](https://arxiv.org/abs/2311.11722)\n",
    " \n",
    "  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "8ec48941-cce4-4c50-81cd-07efb5649f10",
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "import os\n",
    "cwd = os.getcwd()\n",
    "if cwd.endswith(\"tutorial\"):\n",
    "    os.chdir(\"../\")\n",
    "\n",
    "from pprint import pprint\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import torch\n",
    "from mmcv import Config\n",
    "from mmdet.datasets import build_dataset\n",
    "from mmdet.models import build_detector\n",
    "from mmcv.parallel import scatter\n",
    "from mmcv.cnn.utils.flops_counter import add_flops_counting_methods\n",
    "\n",
    "from projects.mmdet3d_plugin.datasets.builder import build_dataloader\n",
    "from projects.mmdet3d_plugin.datasets.utils import draw_lidar_bbox3d"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "997076c6-93ae-441d-a0fd-63002d2ed579",
   "metadata": {},
   "source": [
    "### Base Config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "222af1e8-f615-4c11-a65b-d245088ec3bc",
   "metadata": {},
   "outputs": [],
   "source": [
    "gpu_id = 0\n",
    "config = \"sparse4dv3_temporal_r50_1x8_bs6_256x704\"\n",
    "checkpoint = \"ckpt/sparse4dv3_temporal_r50_1x8_bs6_256x704.pth\"\n",
    "cfg = Config.fromfile(f\"projects/configs/{config}.py\")\n",
    "# cfg.model[\"use_deformable_func\"] = False\n",
    "# cfg.model[\"head\"][\"deformable_model\"][\"use_deformable_func\"] = False\n",
    "img_norm_mean = np.array(cfg.img_norm_cfg[\"mean\"])\n",
    "img_norm_std = np.array(cfg.img_norm_cfg[\"std\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4eb2437c-13fd-40d5-b028-9c41dad5b2dc",
   "metadata": {},
   "source": [
    "### Build Dataloader and Get the Test Data Sample"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "fb7105cb-d8d9-4cfa-ac84-c7188183113c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'version': 'v1.0-trainval'}\n",
      "WARNING!!!!, Only can be used for obtain inference speed!!!!\n"
     ]
    }
   ],
   "source": [
    "dataset = build_dataset(cfg.data.val)\n",
    "dataloader = build_dataloader(\n",
    "    dataset,\n",
    "    samples_per_gpu=1,\n",
    "    workers_per_gpu=0,\n",
    "    dist=False,\n",
    "    shuffle=False,\n",
    ")\n",
    "data_iter = dataloader.__iter__()\n",
    "data = next(data_iter)\n",
    "data = scatter(data, [gpu_id])[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f3344d61-c900-415f-ad49-56217e975a46",
   "metadata": {},
   "source": [
    "### Build Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "32de9c6e-a588-490d-8518-96bf02d920fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = build_detector(cfg.model)\n",
    "model = model.cuda(gpu_id)\n",
    "_ = model.load_state_dict(torch.load(checkpoint)[\"state_dict\"], strict=False)\n",
    "model = model.eval()\n",
    "assert model.use_deformable_func, \"Please compile deformable aggregation first !!!\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5fd79769-0ee1-4f3c-87bd-f8a0c2d6153a",
   "metadata": {},
   "source": [
    "### Anchor Visualization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "e8648b05-952e-4150-95d8-048d9f7aa45e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.image.AxesImage at 0x7fa8987d2070>"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 1048780 characters]",
      "text/plain": [
       "<Figure size 2000x1000 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "raw_imgs = data[\"img\"][0].permute(0, 2, 3, 1).cpu().numpy()\n",
    "raw_imgs = raw_imgs * img_norm_std + img_norm_mean\n",
    "\n",
    "plt.figure(figsize=(20, 10))\n",
    "anchor = model.head.decoder.decode_box(model.head.instance_bank.anchor)\n",
    "plt.imshow(\n",
    "    draw_lidar_bbox3d(anchor[::5], raw_imgs, data[\"projection_mat\"][0])\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3df137b9-a140-42ab-82c6-ab9580cf3b27",
   "metadata": {},
   "source": [
    "### Forward Image Encoder (Backbone and Neck)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "35e7bd45-d804-488b-b12d-8e255f9def20",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Feature map format for deformable aggregation: (flatten feature maps, feature size, start inidces).\n",
      "Deformable aggregation can accept feature maps from multiple views with different resolutions as input !!!\n",
      "N_camera = 6\n",
      "N_scale= 4\n",
      "\n",
      "Multi-camera multi-scale feature maps, flatten to shape [batch_size, N_feature, channel]: torch.Size([1, 89760, 256])\n",
      "\n",
      "Deformable aggregation\n",
      "Size of all feature maps, [N_camera, N_scale, 2]: \n",
      "[[[64, 176], [32, 88], [16, 44], [8, 22]],\n",
      " [[64, 176], [32, 88], [16, 44], [8, 22]],\n",
      " [[64, 176], [32, 88], [16, 44], [8, 22]],\n",
      " [[64, 176], [32, 88], [16, 44], [8, 22]],\n",
      " [[64, 176], [32, 88], [16, 44], [8, 22]],\n",
      " [[64, 176], [32, 88], [16, 44], [8, 22]]]\n",
      "\n",
      "Start indices of all feature maps, [N_camera, N_scale]:\n",
      "[[0, 11264, 14080, 14784],\n",
      " [14960, 26224, 29040, 29744],\n",
      " [29920, 41184, 44000, 44704],\n",
      " [44880, 56144, 58960, 59664],\n",
      " [59840, 71104, 73920, 74624],\n",
      " [74800, 86064, 88880, 89584]]\n"
     ]
    }
   ],
   "source": [
    "feature_maps = model.extract_feat(data[\"img\"], metas=data)\n",
    "if model.use_deformable_func:\n",
    "    print(\"Feature map format for deformable aggregation: (flatten feature maps, feature size, start inidces).\")\n",
    "    print(\"Deformable aggregation can accept feature maps from multiple views with different resolutions as input !!!\")\n",
    "    \n",
    "    print(f\"N_camera = {feature_maps[1].shape[0]}\")\n",
    "    print(f\"N_scale= {feature_maps[1].shape[1]}\\n\")\n",
    "    print(\n",
    "        \"Multi-camera multi-scale feature maps, \"\n",
    "        \"flatten to shape [batch_size, N_feature, channel]: \"\n",
    "        f\"{feature_maps[0].shape}\\n\"\n",
    "    )\n",
    "    print(\"Deformable aggregation\")\n",
    "    print(\"Size of all feature maps, [N_camera, N_scale, 2]: \")\n",
    "    pprint(feature_maps[1].tolist())\n",
    "    \n",
    "    print(\"\\nStart indices of all feature maps, [N_camera, N_scale]:\")\n",
    "    pprint(feature_maps[2].tolist())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "62ff5dc9-7c5c-4aa2-a659-e8c00c8f0998",
   "metadata": {},
   "source": [
    "### Forward Sparse4D Decoder Layers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "0f3165c1-d5d3-475a-9231-ab57bac4b70e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output Keys: dict_keys(['classification', 'prediction', 'quality', 'instance_id'])\n",
      "Number of transformer layers: len(model_outs['classification']).\n",
      "Shape of classification is [batch_size, N_instance, N_class] = torch.Size([1, 900, 10])\n",
      "Shape of prediction is [batch_size, N_instance, box_dim] = torch.Size([1, 900, 11]).\n",
      "Shape of quality is [batch_size, N_instance, (centerness, yawness)] = torch.Size([1, 900, 2]).\n",
      "Shape of instance_id is [batch_size, N_instance] = torch.Size([1, 900]).\n"
     ]
    }
   ],
   "source": [
    "model.head.instance_bank.reset()\n",
    "model_outs = model.head(feature_maps, data)\n",
    "print(\"Output Keys:\", model_outs.keys())\n",
    "print(\"Number of transformer layers: len(model_outs['classification']).\")\n",
    "print(\n",
    "    f\"Shape of classification is [batch_size, N_instance, N_class] = \"\n",
    "    f\"{model_outs['classification'][-1].shape}\"\n",
    ")\n",
    "print(\n",
    "    f\"Shape of prediction is [batch_size, N_instance, box_dim] = \"\n",
    "    f\"{model_outs['prediction'][-1].shape}.\"\n",
    ")\n",
    "if \"quality\" in model_outs:\n",
    "    print(\n",
    "        f\"Shape of quality is [batch_size, N_instance, (centerness, yawness)] = \"\n",
    "        f\"{model_outs['quality'][-1].shape}.\"\n",
    "    )\n",
    "if \"instance_id\" in model_outs:\n",
    "    print(\n",
    "        f\"Shape of instance_id is [batch_size, N_instance] = \"\n",
    "        f\"{model_outs['instance_id'].shape}.\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "369b2dee-3cb3-484e-b540-314cf51a8d9a",
   "metadata": {},
   "source": [
    "### Detection Results Visualization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "15f65538-8d9b-40d5-9b56-797d87006b0f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.image.AxesImage at 0x7fa8988f03d0>"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 930264 characters]",
      "text/plain": [
       "<Figure size 2000x1000 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.figure(figsize=(20, 10))\n",
    "pred_bbox3d = model.head.decoder.decode_box(model_outs[\"prediction\"][-1][0])\n",
    "confidence = model_outs[\"classification\"][-1].max(dim=-1)[0][0].sigmoid()\n",
    "mask = confidence > 0.35\n",
    "num_det = mask.sum()\n",
    "img = draw_lidar_bbox3d(\n",
    "    torch.cat([pred_bbox3d[mask], anchor[mask]]),\n",
    "    raw_imgs, data[\"projection_mat\"][0],\n",
    "    color=[(0, 255, 0)] * num_det + [(255, 0, 0)] * num_det\n",
    ")\n",
    "plt.imshow(img)  # The green boxes denotes model detections, and red ones are the corresponding anchors."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb2283d4-9820-4c8a-9e98-6d3de696ca7c",
   "metadata": {},
   "source": [
    "### Temporal Instance Visualization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "e5668fff-ad7d-4657-87f3-66879e1d9261",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "time interval: 0.4999 s.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<matplotlib.image.AxesImage at 0x7fa8989cb550>"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 923464 characters]",
      "text/plain": [
       "<Figure size 2000x1000 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "data_next = next(data_iter)\n",
    "data_next = scatter(data_next, [gpu_id])[0]\n",
    "raw_imgs_next = data_next[\"img\"][0].permute(0, 2, 3, 1).cpu().numpy()\n",
    "raw_imgs_next = raw_imgs_next * img_norm_std + img_norm_mean\n",
    "\n",
    "plt.figure(figsize=(20, 10))\n",
    "temporal_anchor = model.head.decoder.decode_box(model.head.instance_bank.cached_anchor[0])\n",
    "print(f\"time interval: {(data_next['timestamp'] - data['timestamp']).tolist()[0]:.4f} s.\")\n",
    "mask = model.head.instance_bank.confidence[0] > 0.35\n",
    "plt.imshow(\n",
    "    draw_lidar_bbox3d(temporal_anchor[mask], raw_imgs_next, data[\"projection_mat\"][0])\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f3c0b02f-e98a-43b3-b625-7d3af89e84eb",
   "metadata": {},
   "source": [
    "### Model Complexity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "934f9a6a-f223-4df7-9d8c-197cb0979208",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "total         complexity: FLOPs=168.1413 G / 100.00%, Params= 47.9257 M / 100.00%.\n",
      "img_backbone  complexity: FLOPs= 88.7854 G /  52.80%, Params= 23.5080 M /  49.05%.\n",
      "img_neck      complexity: FLOPs= 61.2933 G /  36.45%, Params=  3.3444 M /   6.98%.\n",
      "head          complexity: FLOPs= 18.0626 G /  10.74%, Params= 21.0725 M /  43.97%.\n"
     ]
    }
   ],
   "source": [
    "num_key_pts = (\n",
    "    cfg.model[\"head\"][\"deformable_model\"][\"kps_generator\"][\"num_learnable_pts\"]\n",
    "    + len(cfg.model[\"head\"][\"deformable_model\"][\"kps_generator\"][\"fix_scale\"])\n",
    ")\n",
    "bilinear_flops = 11\n",
    "deformable_agg_flops = (\n",
    "    cfg.num_decoder\n",
    "    * cfg.embed_dims\n",
    "    * cfg.num_levels\n",
    "    * cfg.model[\"head\"][\"instance_bank\"][\"num_anchor\"]\n",
    "    * cfg.model[\"head\"][\"deformable_model\"][\"num_cams\"]\n",
    "    * num_key_pts\n",
    "    * bilinear_flops\n",
    ")\n",
    "\n",
    "for module in [\"total\", \"img_backbone\", \"img_neck\", \"head\"]:\n",
    "    if module != \"total\":\n",
    "        flops_model = add_flops_counting_methods(getattr(model, module))\n",
    "    else:\n",
    "        flops_model = add_flops_counting_methods(model)\n",
    "    flops_model.eval()\n",
    "    flops_model.start_flops_count()\n",
    "    \n",
    "    if module == \"img_backbone\":\n",
    "        flops_model(data[\"img\"].flatten(0, 1))\n",
    "    elif module == \"img_neck\":\n",
    "        flops_model(model.img_backbone(data[\"img\"].flatten(0, 1)))\n",
    "    elif module == \"head\":\n",
    "        flops_model(model.extract_feat(data[\"img\"], metas=data), data)\n",
    "    else:\n",
    "        flops_model(**data)\n",
    "    flops_count, params_count = flops_model.compute_average_flops_cost()\n",
    "    flops_count *= flops_model.__batch_counter__\n",
    "    flops_model.stop_flops_count()\n",
    "    if module == \"head\" or module == \"total\":\n",
    "        flops_count += deformable_agg_flops\n",
    "    if module == \"total\":\n",
    "        total_flops = flops_count\n",
    "        total_params = params_count\n",
    "    print(\n",
    "        f\"{module:<13} complexity: \"\n",
    "        f\"FLOPs={flops_count/ 10.**9:>8.4f} G / {flops_count/total_flops*100:>6.2f}%, \"\n",
    "        f\"Params={params_count/10**6:>8.4f} M / {params_count/total_params*100:>6.2f}%.\"\n",
    "    )"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "mm_py38",
   "language": "python",
   "name": "mm_py38"
  },
  "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
