{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "430d218d",
   "metadata": {},
   "source": [
    "# BEVFusion Deep Dive - 05: Fuser + BEV Decoder + Detection Head\n",
    "\n",
    "After the two branches produce BEV features on the same grid, BEVFusion:\n",
    "1. **Fuses** them with `ConvFuser` (concat + 3x3 conv + BN + ReLU).\n",
    "2. Runs a **BEV decoder** = SECOND backbone + SECONDFPN neck (multi-scale 2D conv).\n",
    "3. Applies the **TransFusionHead** (object) or a **segmentation head** (map).\n",
    "\n",
    "We reference the real modules and reimplement the fuser + a mini TransFusion-style head to expose all tensor shapes and design choices."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "d7bfcf56",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-08-02T03:38:55.525785Z",
     "iopub.status.busy": "2026-08-02T03:38:55.525464Z",
     "iopub.status.idle": "2026-08-02T03:38:56.744291Z",
     "shell.execute_reply": "2026-08-02T03:38:56.743771Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "IN_CONTAINER   : True\n",
      "BEVFUSION_SRC  : /workspace/perception/bevfusion exists= True\n",
      "NUSC_ROOT      : /data/nuscenes exists= True\n",
      "CKPT_DIR       : /weights exists= True\n",
      "INFOS_DIR      : /data/nuscenes exists= True\n",
      "HAS_TORCH      : True DEVICE= cuda\n",
      "GPU            : NVIDIA RTX 6000 Ada Generation | torch 1.10.1+cu113\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.8/dist-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": [
    "# --- Common bootstrap (safe to run first in any BEVFusion notebook) ---\n",
    "#\n",
    "# This cell auto-detects whether the notebook runs INSIDE the reproducible\n",
    "# `bevfusion:cu113` Docker container (recommended; see\n",
    "# perception/docker/run_notebooks.sh) or on the bare host, and picks paths\n",
    "# accordingly. Everything is overridable via environment variables.\n",
    "#\n",
    "#   In-container layout (set up by run_notebooks.sh):\n",
    "#     BEVFUSION_SRC = /workspace/perception/bevfusion   (mmcv/mmdet/mmdet3d + compiled CUDA ops)\n",
    "#     NUSC_ROOT     = /data/nuscenes                    (writable mini_root: samples/sweeps/maps/v1.0-mini + *_infos_*.pkl)\n",
    "#     CKPT_DIR      = /weights                          (bevfusion-det.pth)\n",
    "#     INFOS_DIR     = /data/nuscenes                    (nuscenes_infos_{train,val}.pkl live next to the data)\n",
    "#\n",
    "#   Host fallback (bare-metal, needs perception/setup.sh built):\n",
    "#     BEVFUSION_SRC = <repo>/perception/bevfusion\n",
    "#     NUSC_ROOT     = /mnt/mydata/nuscenes/mini_root  (or /mnt/mydata/nuscenes)\n",
    "#     CKPT_DIR      = /mnt/mydata/huggingface/bevfusion\n",
    "import os, sys, math, json, textwrap\n",
    "import numpy as np\n",
    "\n",
    "def _first_existing(paths, default):\n",
    "    for p in paths:\n",
    "        if p and os.path.isdir(p):\n",
    "            return p\n",
    "    return default\n",
    "\n",
    "IN_CONTAINER = os.path.isdir(\"/workspace/perception/bevfusion\")\n",
    "\n",
    "BEVFUSION_SRC = os.environ.get(\"BEVFUSION_SRC\") or _first_existing(\n",
    "    [\"/workspace/perception/bevfusion\"], \"/home/nadurthi/repos/deeplearning/perception/bevfusion\")\n",
    "NUSC_ROOT = os.environ.get(\"NUSC_ROOT\") or _first_existing(\n",
    "    [\"/data/nuscenes\", \"/mnt/mydata/nuscenes/mini_root\"], \"/mnt/mydata/nuscenes\")\n",
    "CKPT_DIR = os.environ.get(\"BEVFUSION_CKPT_DIR\") or _first_existing(\n",
    "    [\"/weights\"], \"/mnt/mydata/huggingface/bevfusion\")\n",
    "# Where the generated nuscenes_infos_*.pkl live (defaults to NUSC_ROOT).\n",
    "INFOS_DIR = os.environ.get(\"BEVFUSION_INFOS_DIR\", NUSC_ROOT)\n",
    "\n",
    "# Make mmcv/mmdet/mmdet3d importable from the vendored source tree.\n",
    "if BEVFUSION_SRC and BEVFUSION_SRC not in sys.path:\n",
    "    sys.path.insert(0, BEVFUSION_SRC)\n",
    "\n",
    "# Make matplotlib headless-safe.\n",
    "import matplotlib\n",
    "if os.environ.get(\"MPLBACKEND\") is None:\n",
    "    try:\n",
    "        matplotlib.use(\"Agg\")\n",
    "    except Exception:\n",
    "        pass\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# Torch is expected; but guard so markdown-only runs still import.\n",
    "try:\n",
    "    import torch\n",
    "    import torch.nn as nn\n",
    "    import torch.nn.functional as F\n",
    "    HAS_TORCH = True\n",
    "    DEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "except Exception as e:\n",
    "    HAS_TORCH = False\n",
    "    DEVICE = \"cpu\"\n",
    "    print(\"torch not available:\", e)\n",
    "\n",
    "print(\"IN_CONTAINER   :\", IN_CONTAINER)\n",
    "print(\"BEVFUSION_SRC  :\", BEVFUSION_SRC, \"exists=\", os.path.isdir(BEVFUSION_SRC))\n",
    "print(\"NUSC_ROOT      :\", NUSC_ROOT, \"exists=\", os.path.isdir(NUSC_ROOT))\n",
    "print(\"CKPT_DIR       :\", CKPT_DIR, \"exists=\", os.path.isdir(CKPT_DIR))\n",
    "print(\"INFOS_DIR      :\", INFOS_DIR, \"exists=\", os.path.isdir(INFOS_DIR))\n",
    "print(\"HAS_TORCH      :\", HAS_TORCH, \"DEVICE=\", DEVICE)\n",
    "if HAS_TORCH and torch.cuda.is_available():\n",
    "    print(\"GPU            :\", torch.cuda.get_device_name(0),\n",
    "          \"| torch\", torch.__version__)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aac26ebb",
   "metadata": {},
   "source": [
    "## 1. `ConvFuser` - why concatenation is enough\n",
    "Exactly [`ConvFuser`](../../perception/bevfusion/mmdet3d/models/fusers/conv.py:11): concatenate `[cam_bev(80ch), lidar_bev(256ch)]` on the channel axis and mix with a single 3x3 conv. Because both live on the *same* BEV grid, no attention or resampling is needed - a key efficiency win of BEVFusion."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "48c08cee",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-08-02T03:38:56.746981Z",
     "iopub.status.busy": "2026-08-02T03:38:56.746638Z",
     "iopub.status.idle": "2026-08-02T03:38:56.906037Z",
     "shell.execute_reply": "2026-08-02T03:38:56.905467Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cam BEV  : (1, 80, 180, 180)\n",
      "lidar BEV: (1, 256, 180, 180)\n",
      "fused BEV: (1, 256, 180, 180) (B,256,180,180)\n"
     ]
    }
   ],
   "source": [
    "if HAS_TORCH:\n",
    "    class ConvFuser(nn.Sequential):\n",
    "        def __init__(self, in_channels, out_channels):\n",
    "            super().__init__(nn.Conv2d(sum(in_channels), out_channels, 3, padding=1, bias=False),\n",
    "                             nn.BatchNorm2d(out_channels), nn.ReLU(True))\n",
    "        def forward(self, inputs): return super().forward(torch.cat(inputs, dim=1))\n",
    "    fuser = ConvFuser([80, 256], 256).eval()\n",
    "    cam_bev   = torch.randn(1, 80, 180, 180)\n",
    "    lidar_bev = torch.randn(1, 256, 180, 180)\n",
    "    with torch.no_grad():\n",
    "        fused = fuser([cam_bev, lidar_bev])\n",
    "    print('cam BEV  :', tuple(cam_bev.shape))\n",
    "    print('lidar BEV:', tuple(lidar_bev.shape))\n",
    "    print('fused BEV:', tuple(fused.shape), '(B,256,180,180)')\n",
    "else:\n",
    "    print('skipped')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5d2cc658",
   "metadata": {},
   "source": [
    "## 2. BEV decoder = SECOND + SECONDFPN\n",
    "The decoder ([`decoder.backbone`](../../perception/bevfusion/mmdet3d/models/backbones/second.py) + [`SECONDFPN`](../../perception/bevfusion/mmdet3d/models/necks/second.py)) is a lightweight 2D CNN: two downsampling stages produce multi-scale BEV features that are upsampled and concatenated. This enlarges the receptive field before detection."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "fec96f54",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-08-02T03:38:56.908161Z",
     "iopub.status.busy": "2026-08-02T03:38:56.907889Z",
     "iopub.status.idle": "2026-08-02T03:38:56.998807Z",
     "shell.execute_reply": "2026-08-02T03:38:56.998130Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "decoder multi-scale BEV out: (1, 256, 180, 180) (B,256,180,180)\n"
     ]
    }
   ],
   "source": [
    "if HAS_TORCH:\n",
    "    class MiniSECOND(nn.Module):\n",
    "        def __init__(self, c=256):\n",
    "            super().__init__()\n",
    "            def blk(i,o,s): return nn.Sequential(nn.Conv2d(i,o,3,s,1,bias=False),\n",
    "                                                 nn.BatchNorm2d(o), nn.ReLU(True))\n",
    "            self.s1 = blk(c,128,1); self.s2 = blk(128,256,2)\n",
    "            self.up1 = nn.ConvTranspose2d(128,128,1,1)\n",
    "            self.up2 = nn.ConvTranspose2d(256,128,2,2)\n",
    "        def forward(self,x):\n",
    "            a=self.s1(x); b=self.s2(a)\n",
    "            return torch.cat([self.up1(a), self.up2(b)],1)\n",
    "    dec = MiniSECOND().eval()\n",
    "    with torch.no_grad(): y = dec(fused)\n",
    "    print('decoder multi-scale BEV out:', tuple(y.shape), '(B,256,180,180)')\n",
    "else:\n",
    "    print('skipped')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0d602e70",
   "metadata": {},
   "source": [
    "## 3. TransFusionHead - transformer set prediction on BEV\n",
    "The real head is [`TransFusionHead`](../../perception/bevfusion/mmdet3d/models/heads/bbox/transfusion.py). Its pipeline:\n",
    "* A **shared conv** + **heatmap head** predict class centers over the BEV grid   `(B, num_classes, X, Y)`.\n",
    "* The **top-K** heatmap peaks initialize object **queries** (learned content + positional   embedding from the peak location).\n",
    "* A **transformer decoder** cross-attends queries to the BEV feature map and refines them.\n",
    "* **Regression heads** output per-query `center(2), height(1), dim(3), rot(2=sin/cos),   vel(2)` + class logits -> decoded to `(x,y,z,w,l,h,yaw,vx,vy)`.\n",
    "* **Training** uses `mmdet`'s Hungarian matcher (`build_match_cost`) + focal/L1/GIoU losses   (set prediction, no NMS).\n",
    "\n",
    "Below is a compact stand-in showing the heatmap -> top-K query -> box-regression shapes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "ff5989ea",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-08-02T03:38:57.001071Z",
     "iopub.status.busy": "2026-08-02T03:38:57.000779Z",
     "iopub.status.idle": "2026-08-02T03:38:57.161544Z",
     "shell.execute_reply": "2026-08-02T03:38:57.160945Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "heatmap: (1, 10, 180, 180) (B,num_classes,H,W)\n",
      "  query center : (1, 200, 2)\n",
      "  query height : (1, 200, 1)\n",
      "  query dim    : (1, 200, 3)\n",
      "  query rot    : (1, 200, 2)\n",
      "  query vel    : (1, 200, 2)\n",
      "  query cls    : (1, 200, 10)\n",
      "decoded box vector per query = center(2)+height(1)+dim(3)+rot(2)+vel(2) = 10, +cls\n"
     ]
    }
   ],
   "source": [
    "if HAS_TORCH:\n",
    "    class MiniTransFusionHead(nn.Module):\n",
    "        def __init__(self, c=256, num_classes=10, num_queries=200):\n",
    "            super().__init__(); self.nc=num_classes; self.nq=num_queries\n",
    "            self.shared = nn.Sequential(nn.Conv2d(c,c,3,1,1), nn.BatchNorm2d(c), nn.ReLU(True))\n",
    "            self.heatmap = nn.Conv2d(c, num_classes, 1)\n",
    "            self.query_proj = nn.Linear(c, c)\n",
    "            self.attn = nn.MultiheadAttention(c, 8, batch_first=True)\n",
    "            self.reg = nn.ModuleDict({\n",
    "                'center': nn.Linear(c,2), 'height': nn.Linear(c,1), 'dim': nn.Linear(c,3),\n",
    "                'rot': nn.Linear(c,2), 'vel': nn.Linear(c,2), 'cls': nn.Linear(c,num_classes)})\n",
    "        def forward(self, bev):\n",
    "            B,C,H,W = bev.shape\n",
    "            f = self.shared(bev)\n",
    "            hm = self.heatmap(f).sigmoid()               # (B,nc,H,W)\n",
    "            flat = f.flatten(2).transpose(1,2)            # (B,HW,C)\n",
    "            score = hm.amax(1).flatten(1)                 # (B,HW)\n",
    "            topk = score.topk(self.nq, dim=1).indices     # (B,nq)\n",
    "            q = torch.gather(flat, 1, topk.unsqueeze(-1).expand(-1,-1,C))\n",
    "            q = self.query_proj(q)\n",
    "            q,_ = self.attn(q, flat, flat)                # cross-attend to BEV\n",
    "            out = {k: layer(q) for k,layer in self.reg.items()}\n",
    "            return hm, out\n",
    "    head = MiniTransFusionHead().eval()\n",
    "    with torch.no_grad(): hm, out = head(y)\n",
    "    print('heatmap:', tuple(hm.shape), '(B,num_classes,H,W)')\n",
    "    for k,v in out.items(): print(f'  query {k:7s}:', tuple(v.shape))\n",
    "    print('decoded box vector per query = center(2)+height(1)+dim(3)+rot(2)+vel(2) = 10, +cls')\n",
    "else:\n",
    "    print('skipped')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4c7e1a10",
   "metadata": {},
   "source": [
    "## 4. CenterHead alternative\n",
    "The camera-only detection config uses [`CenterHead`](../../perception/bevfusion/mmdet3d/models/heads/bbox/centerpoint.py) instead: dense per-pixel heatmaps + regression maps (`reg`, `height`, `dim`, `rot`, `vel`), decoded by top-K peak picking and circular/rotated NMS ([`box3d_nms`](../../perception/bevfusion/mmdet3d/core/post_processing/box3d_nms.py)). TransFusion (query/set-prediction) generally scores higher on the fusion model; CenterHead is simpler and anchor-free."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e57f2c63",
   "metadata": {},
   "source": [
    "---\n",
    "**Next:** [`06_end_to_end_inference.ipynb`](06_end_to_end_inference.ipynb)"
   ]
  }
 ],
 "metadata": {
  "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}