{
    "cells": [
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "# GroundingDINO: Intermediate Layer Visualization\n",
                "\n",
                "## Notebook 11 of 14\n",
                "\n",
                "This notebook hooks into the actual GroundingDINO model to extract and visualize intermediate activations:\n",
                "- Swin Transformer feature maps at each stage\n",
                "- BERT token embeddings evolution\n",
                "- Feature Enhancer outputs\n",
                "- Decoder query evolution\n",
                "- Box refinement at each decoder layer\n",
                "\n",
                "---"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "import os, sys\n",
                "import torch\n",
                "import torch.nn as nn\n",
                "import numpy as np\n",
                "import matplotlib.pyplot as plt\n",
                "import matplotlib.patches as patches\n",
                "from PIL import Image\n",
                "import warnings\n",
                "warnings.filterwarnings('ignore')\n",
                "\n",
                "MODEL_DIR = \"/mnt/mydata/huggingface\"\n",
                "GDINO_DIR = os.path.join(MODEL_DIR, \"GroundingDINO\")\n",
                "if GDINO_DIR not in sys.path:\n",
                "    sys.path.insert(0, GDINO_DIR)\n",
                "\n",
                "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
                "print(f\"Device: {device}\")"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## 1. Hook-Based Feature Extraction"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "class FeatureExtractor:\n",
                "    \"\"\"Extract intermediate features using PyTorch hooks.\"\"\"\n",
                "    \n",
                "    def __init__(self, model):\n",
                "        self.model = model\n",
                "        self.features = {}\n",
                "        self.hooks = []\n",
                "    \n",
                "    def register_hook(self, module, name):\n",
                "        \"\"\"Register a forward hook on a module.\"\"\"\n",
                "        def hook_fn(module, input, output):\n",
                "            if isinstance(output, torch.Tensor):\n",
                "                self.features[name] = output.detach().cpu()\n",
                "            elif isinstance(output, (list, tuple)):\n",
                "                self.features[name] = [o.detach().cpu() if isinstance(o, torch.Tensor) else o \n",
                "                                        for o in output]\n",
                "        hook = module.register_forward_hook(hook_fn)\n",
                "        self.hooks.append(hook)\n",
                "        return hook\n",
                "    \n",
                "    def remove_hooks(self):\n",
                "        for hook in self.hooks:\n",
                "            hook.remove()\n",
                "        self.hooks = []\n",
                "    \n",
                "    def clear(self):\n",
                "        self.features = {}\n",
                "\n",
                "# Load model if available\n",
                "try:\n",
                "    from groundingdino.util.inference import load_model, load_image, predict\n",
                "    CONFIG_PATH = os.path.join(GDINO_DIR, \"groundingdino/config/GroundingDINO_SwinT_OGC.py\")\n",
                "    WEIGHTS_PATH = os.path.join(MODEL_DIR, \"groundingdino_swint_ogc.pth\")\n",
                "    model = load_model(CONFIG_PATH, WEIGHTS_PATH, device=str(device))\n",
                "    model.eval()\n",
                "    extractor = FeatureExtractor(model)\n",
                "    GDINO_AVAILABLE = True\n",
                "    print(\"✓ Model loaded\")\n",
                "    \n",
                "    # Print model structure to find hook points\n",
                "    print(\"\\nModel structure (top-level):\")\n",
                "    for name, module in model.named_children():\n",
                "        print(f\"  {name}: {type(module).__name__}\")\n",
                "        \n",
                "except Exception as e:\n",
                "    print(f\"Model not available: {e}\")\n",
                "    GDINO_AVAILABLE = False"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## 2. Register Hooks on Key Layers"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "if GDINO_AVAILABLE:\n",
                "    extractor.clear()\n",
                "    extractor.remove_hooks()\n",
                "    \n",
                "    # Register hooks on key components\n",
                "    hooks_registered = []\n",
                "    \n",
                "    # Try to hook into backbone stages\n",
                "    try:\n",
                "        if hasattr(model, 'backbone'):\n",
                "            backbone = model.backbone\n",
                "            # Hook Swin stages\n",
                "            for i, layer in enumerate(backbone.layers if hasattr(backbone, 'layers') else []):\n",
                "                extractor.register_hook(layer, f'swin_stage_{i}')\n",
                "                hooks_registered.append(f'swin_stage_{i}')\n",
                "    except Exception as e:\n",
                "        print(f\"Could not hook backbone: {e}\")\n",
                "    \n",
                "    # Hook BERT\n",
                "    try:\n",
                "        if hasattr(model, 'bert'):\n",
                "            extractor.register_hook(model.bert, 'bert_output')\n",
                "            hooks_registered.append('bert_output')\n",
                "    except Exception as e:\n",
                "        print(f\"Could not hook BERT: {e}\")\n",
                "    \n",
                "    print(f\"Registered {len(hooks_registered)} hooks: {hooks_registered}\")\n",
                "else:\n",
                "    print(\"Model not available - will use synthetic features for visualization\")"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## 3. Run Inference and Extract Features"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "if GDINO_AVAILABLE:\n",
                "    img_path = os.path.join(MODEL_DIR, \"sample_cats.jpg\")\n",
                "    if not os.path.exists(img_path):\n",
                "        os.system(f\"wget -q http://images.cocodataset.org/val2017/000000039769.jpg -O {img_path}\")\n",
                "    \n",
                "    if os.path.exists(img_path):\n",
                "        image_source, image = load_image(img_path)\n",
                "        \n",
                "        with torch.no_grad():\n",
                "            boxes, logits, phrases = predict(\n",
                "                model=model,\n",
                "                image=image,\n",
                "                caption=\"cat . remote control\",\n",
                "                box_threshold=0.3,\n",
                "                text_threshold=0.25,\n",
                "                device=str(device)\n",
                "            )\n",
                "        \n",
                "        print(f\"Inference complete. Extracted features:\")\n",
                "        for name, feat in extractor.features.items():\n",
                "            if isinstance(feat, torch.Tensor):\n",
                "                print(f\"  {name}: {feat.shape}\")\n",
                "            elif isinstance(feat, list):\n",
                "                print(f\"  {name}: list of {len(feat)} tensors\")\n",
                "        \n",
                "        print(f\"\\nDetections: {len(boxes)}\")\n",
                "        for box, logit, phrase in zip(boxes, logits, phrases):\n",
                "            print(f\"  '{phrase}' score={logit:.3f}\")\n",
                "else:\n",
                "    print(\"Using synthetic features for visualization\")"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## 4. Visualize Swin Feature Maps"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "# Visualize Swin feature maps (real or synthetic)\n",
                "import timm\n",
                "\n",
                "# Load Swin-T from timm for feature visualization\n",
                "swin = timm.create_model('swin_tiny_patch4_window7_224', pretrained=True, features_only=True)\n",
                "swin.eval()\n",
                "\n",
                "# Load and preprocess image\n",
                "import torchvision.transforms as TF\n",
                "\n",
                "img_path = os.path.join(MODEL_DIR, \"sample_cats.jpg\")\n",
                "if os.path.exists(img_path):\n",
                "    img = Image.open(img_path).convert('RGB')\n",
                "else:\n",
                "    img = Image.fromarray(np.random.randint(0, 255, (480, 640, 3), dtype=np.uint8))\n",
                "\n",
                "transform = TF.Compose([\n",
                "    TF.Resize((224, 224)),\n",
                "    TF.ToTensor(),\n",
                "    TF.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n",
                "])\n",
                "img_t = transform(img).unsqueeze(0)\n",
                "\n",
                "with torch.no_grad():\n",
                "    features = swin(img_t)\n",
                "\n",
                "print(\"Swin-T Feature Maps:\")\n",
                "for i, feat in enumerate(features):\n",
                "    print(f\"  Stage {i+1}: {feat.shape}\")\n",
                "\n",
                "# Visualize\n",
                "fig, axes = plt.subplots(3, 5, figsize=(20, 12))\n",
                "\n",
                "# Row 0: Original image and stage mean activations\n",
                "axes[0, 0].imshow(img.resize((224, 224)))\n",
                "axes[0, 0].set_title('Input Image\\n224×224', fontsize=10, fontweight='bold')\n",
                "axes[0, 0].axis('off')\n",
                "\n",
                "stage_names = ['Stage 1\\n56×56×96', 'Stage 2\\n28×28×192', 'Stage 3\\n14×14×384', 'Stage 4\\n7×7×768']\n",
                "for i, (feat, name) in enumerate(zip(features, stage_names)):\n",
                "    feat_np = feat[0].mean(0).numpy()  # H, W\n",
                "    im = axes[0, i+1].imshow(feat_np, cmap='viridis')\n",
                "    axes[0, i+1].set_title(f'{name}\\n(mean channels)', fontsize=9, fontweight='bold')\n",
                "    axes[0, i+1].axis('off')\n",
                "    plt.colorbar(im, ax=axes[0, i+1], fraction=0.046)\n",
                "\n",
                "# Row 1: Individual channels from Stage 2\n",
                "feat2 = features[1][0].numpy()  # C, H, W\n",
                "axes[1, 0].set_title('Stage 2 Channels', fontsize=10, fontweight='bold')\n",
                "axes[1, 0].axis('off')\n",
                "for i in range(4):\n",
                "    ch_idx = i * 48\n",
                "    im = axes[1, i+1].imshow(feat2[ch_idx], cmap='hot')\n",
                "    axes[1, i+1].set_title(f'Stage 2, Ch {ch_idx}', fontsize=9)\n",
                "    axes[1, i+1].axis('off')\n",
                "\n",
                "# Row 2: Individual channels from Stage 3\n",
                "feat3 = features[2][0].numpy()  # C, H, W\n",
                "axes[2, 0].set_title('Stage 3 Channels', fontsize=10, fontweight='bold')\n",
                "axes[2, 0].axis('off')\n",
                "for i in range(4):\n",
                "    ch_idx = i * 96\n",
                "    im = axes[2, i+1].imshow(feat3[ch_idx], cmap='plasma')\n",
                "    axes[2, i+1].set_title(f'Stage 3, Ch {ch_idx}', fontsize=9)\n",
                "    axes[2, i+1].axis('off')\n",
                "\n",
                "plt.suptitle('Swin Transformer Intermediate Feature Maps', fontsize=14, fontweight='bold')\n",
                "plt.tight_layout()\n",
                "plt.savefig('swin_intermediate_features.png', dpi=150, bbox_inches='tight')\n",
                "plt.show()"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## 5. Visualize BERT Token Embeddings"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "from transformers import BertTokenizer, BertModel\n",
                "\n",
                "tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')\n",
                "bert = BertModel.from_pretrained('bert-base-uncased', output_hidden_states=True, output_attentions=True)\n",
                "bert.eval()\n",
                "\n",
                "# Test with GroundingDINO-style query\n",
                "query = \"cat . remote control . couch\"\n",
                "inputs = tokenizer(query, return_tensors='pt')\n",
                "tokens = tokenizer.convert_ids_to_tokens(inputs['input_ids'][0])\n",
                "\n",
                "with torch.no_grad():\n",
                "    outputs = bert(**inputs)\n",
                "\n",
                "hidden_states = outputs.hidden_states  # 13 tensors\n",
                "attentions = outputs.attentions  # 12 tensors\n",
                "\n",
                "print(f\"Query: '{query}'\")\n",
                "print(f\"Tokens: {tokens}\")\n",
                "print(f\"Sequence length: {len(tokens)}\")\n",
                "\n",
                "# Visualize token feature evolution\n",
                "fig, axes = plt.subplots(2, 3, figsize=(18, 10))\n",
                "\n",
                "# Token features at different layers\n",
                "for ax, layer_idx in zip(axes[0], [0, 6, 12]):\n",
                "    feat = hidden_states[layer_idx][0].numpy()  # L, 768\n",
                "    im = ax.imshow(feat[:, :64], cmap='RdBu', aspect='auto', vmin=-3, vmax=3)\n",
                "    ax.set_yticks(range(len(tokens)))\n",
                "    ax.set_yticklabels(tokens, fontsize=9)\n",
                "    ax.set_xlabel('Feature dim (first 64)')\n",
                "    ax.set_title(f'BERT Layer {layer_idx} Features', fontsize=11, fontweight='bold')\n",
                "    plt.colorbar(im, ax=ax)\n",
                "\n",
                "# Token similarity matrices\n",
                "for ax, layer_idx in zip(axes[1], [0, 6, 12]):\n",
                "    feat = hidden_states[layer_idx][0].numpy()\n",
                "    feat_norm = feat / (np.linalg.norm(feat, axis=-1, keepdims=True) + 1e-8)\n",
                "    sim = feat_norm @ feat_norm.T\n",
                "    im = ax.imshow(sim, cmap='viridis', aspect='equal', vmin=-0.5, vmax=1)\n",
                "    ax.set_xticks(range(len(tokens)))\n",
                "    ax.set_yticks(range(len(tokens)))\n",
                "    ax.set_xticklabels(tokens, rotation=45, ha='right', fontsize=8)\n",
                "    ax.set_yticklabels(tokens, fontsize=8)\n",
                "    ax.set_title(f'Token Similarity (Layer {layer_idx})', fontsize=11, fontweight='bold')\n",
                "    plt.colorbar(im, ax=ax)\n",
                "\n",
                "plt.suptitle(f'BERT Token Embeddings Evolution\\nQuery: \"{query}\"', fontsize=13, fontweight='bold')\n",
                "plt.tight_layout()\n",
                "plt.savefig('bert_token_evolution.png', dpi=150, bbox_inches='tight')\n",
                "plt.show()"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## 6. Visualize Decoder Query Evolution"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "# Simulate decoder query evolution across 6 layers\n",
                "np.random.seed(42)\n",
                "K = 50  # queries\n",
                "D = 256\n",
                "\n",
                "# Simulate query features evolving\n",
                "query_features_per_layer = []\n",
                "query_boxes_per_layer = []\n",
                "\n",
                "# Initial queries\n",
                "q = np.random.randn(K, D)\n",
                "boxes = np.random.rand(K, 4)\n",
                "boxes[:, 2:] = np.abs(boxes[:, 2:]) * 0.3 + 0.05  # reasonable box sizes\n",
                "\n",
                "# GT boxes (simulated)\n",
                "gt_boxes = np.array([[0.3, 0.4, 0.2, 0.25], [0.7, 0.6, 0.15, 0.2]])\n",
                "\n",
                "for layer in range(7):  # initial + 6 layers\n",
                "    query_features_per_layer.append(q.copy())\n",
                "    query_boxes_per_layer.append(boxes.copy())\n",
                "    \n",
                "    if layer < 6:\n",
                "        # Simulate refinement\n",
                "        q = q + np.random.randn(K, D) * 0.1\n",
                "        # Boxes converge toward GT\n",
                "        for i in range(K):\n",
                "            # Find nearest GT\n",
                "            dists = np.linalg.norm(boxes[i, :2] - gt_boxes[:, :2], axis=-1)\n",
                "            nearest_gt = gt_boxes[np.argmin(dists)]\n",
                "            boxes[i] = boxes[i] + (nearest_gt - boxes[i]) * 0.3 + np.random.randn(4) * 0.02\n",
                "            boxes[i] = np.clip(boxes[i], 0, 1)\n",
                "\n",
                "# Visualize\n",
                "fig, axes = plt.subplots(2, 4, figsize=(20, 10))\n",
                "\n",
                "# Feature norm evolution\n",
                "ax = axes[0, 0]\n",
                "norms_per_layer = [np.linalg.norm(q, axis=-1).mean() for q in query_features_per_layer]\n",
                "ax.plot(range(7), norms_per_layer, 'b-o', linewidth=2, markersize=8)\n",
                "ax.set_xlabel('Layer')\n",
                "ax.set_ylabel('Mean Feature Norm')\n",
                "ax.set_title('Query Feature Norm\\nAcross Decoder Layers', fontsize=10, fontweight='bold')\n",
                "ax.set_xticks(range(7))\n",
                "ax.set_xticklabels(['Init'] + [f'L{i+1}' for i in range(6)])\n",
                "ax.grid(True, alpha=0.3)\n",
                "\n",
                "# Box center evolution\n",
                "ax = axes[0, 1]\n",
                "for q_idx in range(5):  # show 5 queries\n",
                "    centers_x = [b[q_idx, 0] for b in query_boxes_per_layer]\n",
                "    centers_y = [b[q_idx, 1] for b in query_boxes_per_layer]\n",
                "    ax.plot(centers_x, centers_y, '-o', markersize=5, alpha=0.7)\n",
                "# Mark GT\n",
                "for gt in gt_boxes:\n",
                "    ax.plot(gt[0], gt[1], 'r*', markersize=15, zorder=5)\n",
                "ax.set_xlim(0, 1)\n",
                "ax.set_ylim(0, 1)\n",
                "ax.set_title('Box Center Trajectories\\n(Red stars = GT)', fontsize=10, fontweight='bold')\n",
                "ax.set_xlabel('X center')\n",
                "ax.set_ylabel('Y center')\n",
                "ax.invert_yaxis()\n",
                "ax.grid(True, alpha=0.3)\n",
                "\n",
                "# Box predictions at each layer\n",
                "for layer_idx, ax in zip([0, 2, 4, 6], axes[0, 2:].tolist() + axes[1, :2].tolist()):\n",
                "    ax.set_xlim(0, 1)\n",
                "    ax.set_ylim(0, 1)\n",
                "    ax.set_facecolor('#f8f8f8')\n",
                "    ax.invert_yaxis()\n",
                "    \n",
                "    # Draw GT boxes\n",
                "    for gt in gt_boxes:\n",
                "        xc, yc, w, h = gt\n",
                "        rect = patches.Rectangle((xc-w/2, yc-h/2), w, h, linewidth=3,\n",
                "                                   edgecolor='green', facecolor='lightgreen', alpha=0.4)\n",
                "        ax.add_patch(rect)\n",
                "    \n",
                "    # Draw predicted boxes\n",
                "    boxes_layer = query_boxes_per_layer[layer_idx]\n",
                "    for i in range(min(20, K)):\n",
                "        xc, yc, w, h = boxes_layer[i]\n",
                "        rect = patches.Rectangle((xc-w/2, yc-h/2), w, h, linewidth=1,\n",
                "                                   edgecolor='blue', facecolor='none', alpha=0.5)\n",
                "        ax.add_patch(rect)\n",
                "    \n",
                "    label = 'Initial' if layer_idx == 0 else f'Layer {layer_idx}'\n",
                "    ax.set_title(f'{label} Predictions\\n(Green=GT, Blue=Pred)', fontsize=10, fontweight='bold')\n",
                "    ax.grid(True, alpha=0.3)\n",
                "\n",
                "# Feature PCA\n",
                "from sklearn.decomposition import PCA\n",
                "ax = axes[1, 2]\n",
                "pca = PCA(n_components=2)\n",
                "q_init_2d = pca.fit_transform(query_features_per_layer[0])\n",
                "q_final_2d = pca.transform(query_features_per_layer[-1])\n",
                "ax.scatter(q_init_2d[:, 0], q_init_2d[:, 1], c='blue', alpha=0.6, s=30, label='Initial')\n",
                "ax.scatter(q_final_2d[:, 0], q_final_2d[:, 1], c='red', alpha=0.6, s=30, label='Final')\n",
                "ax.set_title('Query Features PCA\\n(Initial vs Final)', fontsize=10, fontweight='bold')\n",
                "ax.legend()\n",
                "ax.grid(True, alpha=0.3)\n",
                "\n",
                "# IoU improvement\n",
                "ax = axes[1, 3]\n",
                "def compute_max_iou(pred_boxes, gt_boxes):\n",
                "    ious = []\n",
                "    for pred in pred_boxes:\n",
                "        max_iou = 0\n",
                "        for gt in gt_boxes:\n",
                "            px1, py1 = pred[0]-pred[2]/2, pred[1]-pred[3]/2\n",
                "            px2, py2 = pred[0]+pred[2]/2, pred[1]+pred[3]/2\n",
                "            gx1, gy1 = gt[0]-gt[2]/2, gt[1]-gt[3]/2\n",
                "            gx2, gy2 = gt[0]+gt[2]/2, gt[1]+gt[3]/2\n",
                "            ix = max(0, min(px2, gx2) - max(px1, gx1))\n",
                "            iy = max(0, min(py2, gy2) - max(py1, gy1))\n",
                "            inter = ix * iy\n",
                "            union = pred[2]*pred[3] + gt[2]*gt[3] - inter\n",
                "            iou = inter / (union + 1e-6)\n",
                "            max_iou = max(max_iou, iou)\n",
                "        ious.append(max_iou)\n",
                "    return np.mean(ious)\n",
                "\n",
                "mean_ious = [compute_max_iou(b, gt_boxes) for b in query_boxes_per_layer]\n",
                "ax.plot(range(7), mean_ious, 'g-o', linewidth=2, markersize=8)\n",
                "ax.fill_between(range(7), mean_ious, alpha=0.3, color='green')\n",
                "ax.set_xlabel('Layer')\n",
                "ax.set_ylabel('Mean Max IoU')\n",
                "ax.set_title('IoU Improvement\\nAcross Decoder Layers', fontsize=10, fontweight='bold')\n",
                "ax.set_xticks(range(7))\n",
                "ax.set_xticklabels(['Init'] + [f'L{i+1}' for i in range(6)])\n",
                "ax.grid(True, alpha=0.3)\n",
                "\n",
                "plt.suptitle('Decoder Query Evolution Across 6 Layers', fontsize=14, fontweight='bold')\n",
                "plt.tight_layout()\n",
                "plt.savefig('decoder_query_evolution.png', dpi=150, bbox_inches='tight')\n",
                "plt.show()"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## Summary\n",
                "\n",
                "Key observations from intermediate layer visualization:\n",
                "1. **Swin features**: Progressively more abstract, lower resolution, higher channels\n",
                "2. **BERT features**: Token representations become more contextual in deeper layers\n",
                "3. **Decoder queries**: Boxes converge toward GT objects across 6 layers\n",
                "4. **Feature norms**: Stabilize after initial layers\n",
                "\n",
                "**Next**: [Notebook 12 - Attention Visualization](12_attention_visualization.ipynb)"
            ]
        }
    ],
    "metadata": {
        "kernelspec": {
            "display_name": "Python 3",
            "language": "python",
            "name": "python3"
        },
        "language_info": {
            "name": "python",
            "version": "3.10.0"
        }
    },
    "nbformat": 4,
    "nbformat_minor": 4
}