tc_clip_inference_demo.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "b82fab3e",
   "metadata": {},
   "source": [
    "# TC-CLIP Inference Demo for Custom Videos"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3e37e892",
   "metadata": {},
   "source": [
    "## Set model path, custom video path, class names."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "b539e8ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "### Change to your settings ###\n",
    "output=\"workspace/inference\"\n",
    "tc_clip_model_path = \"workspace/weights/zero_shot_k400_llm_tc_clip.pth\"   # Your pretrained model saved path\n",
    "class_names = ['swinging baseball bat', 'cutting apple', 'moon walking', 'changing gear in car']  # Class names\n",
    "video_path = \"/131_data/datasets/k600/test/uveYxu3T2Fc_000009_000019.mp4\" # Custom video path"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7e02818a",
   "metadata": {},
   "source": [
    "## No need to change below codes, just run the cells."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "8245b63a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import libraries\n",
    "import os\n",
    "from hydra import initialize, compose\n",
    "from omegaconf import OmegaConf\n",
    "from pathlib import Path\n",
    "import torch\n",
    "\n",
    "from datasets.pipeline import Compose\n",
    "from trainers.build_trainer import returnCLIP\n",
    "from utils.logger import create_logger\n",
    "from utils.print_utils import colorstr\n",
    "from utils.tools import load_checkpoint"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3ccb1a66",
   "metadata": {},
   "source": [
    "### Init configs and logger"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "ec10eb46",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize Hydra configs\n",
    "overrides = [\n",
    "    f\"output={output}\",\n",
    "    \"eval=test\",\n",
    "    \"trainer=tc_clip\",\n",
    "    f\"resume={tc_clip_model_path}\"\n",
    "]\n",
    "\n",
    "# Initialize Hydra with config path\n",
    "with initialize(version_base=None, config_path=\"configs\"):\n",
    "    config = compose(config_name=\"zero_shot.yaml\", overrides=overrides)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "748d8d90",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[32m[09-25 13:51:28 TCCLIP]\u001b[0m\u001b[33m(3979588942.py 9)\u001b[0m: INFO working dir: workspace/inference\n"
     ]
    }
   ],
   "source": [
    "# Init settings\n",
    "OmegaConf.set_struct(config, False)  # Needed to add fields at runtime below\n",
    "\n",
    "# Define working dir\n",
    "Path(config.output).mkdir(parents=True, exist_ok=True)\n",
    "\n",
    "# Logger\n",
    "logger = create_logger(output_dir=config.output, dist_rank=0, name=f\"{config.trainer_name}\")\n",
    "logger.info(f\"working dir: {config.output}\")\n",
    "\n",
    "# Whether to use pytorch or apex amp\n",
    "major, minor = int(torch.__version__.split('.')[0]), int(torch.__version__.split('.')[1])\n",
    "config.use_torch_amp = (major >= 2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e9454e3f",
   "metadata": {},
   "source": [
    "### Build model & load checkpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "a429597c",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[32m[09-25 13:51:32 TCCLIP]\u001b[0m\u001b[33m(build_trainer.py 56)\u001b[0m: INFO Loading CLIP (backbone: ViT-B/16)\n",
      "Using spatial positional embedding\n",
      "Weights not found for some missing keys:  ['visual.transformer.resblocks.1.attn.local_global_bias_table', 'visual.transformer.resblocks.2.attn.local_global_bias_table', 'visual.transformer.resblocks.3.attn.local_global_bias_table', 'visual.transformer.resblocks.4.attn.local_global_bias_table', 'visual.transformer.resblocks.5.attn.local_global_bias_table', 'visual.transformer.resblocks.6.attn.local_global_bias_table', 'visual.transformer.resblocks.7.attn.local_global_bias_table', 'visual.transformer.resblocks.8.attn.local_global_bias_table', 'visual.transformer.resblocks.9.attn.local_global_bias_table', 'visual.transformer.resblocks.10.attn.local_global_bias_table', 'visual.transformer.resblocks.11.attn.local_global_bias_table']\n",
      "\u001b[32m[09-25 13:51:36 TCCLIP]\u001b[0m\u001b[33m(build_trainer.py 59)\u001b[0m: INFO \u001b[34m\u001b[1mBuilding TCCLIP\u001b[0m\n",
      "\u001b[32m[09-25 13:51:36 TCCLIP]\u001b[0m\u001b[33m(tc_clip_prompt_learner.py 54)\u001b[0m: INFO Video-conditional prompt learning\n",
      "\u001b[32m[09-25 13:51:36 TCCLIP]\u001b[0m\u001b[33m(tc_clip_prompt_learner.py 55)\u001b[0m: INFO Initial context: \"a photo of a\"\n",
      "\u001b[32m[09-25 13:51:36 TCCLIP]\u001b[0m\u001b[33m(tc_clip_prompt_learner.py 56)\u001b[0m: INFO Number of learnable text prompt vectors: 4\n",
      "\u001b[32m[09-25 13:51:36 TCCLIP]\u001b[0m\u001b[33m(tc_clip_text_encoder.py 96)\u001b[0m: INFO Copy CLIP transformer 11th layer weights to prompt generation layer\n",
      "\u001b[32m[09-25 13:51:36 TCCLIP]\u001b[0m\u001b[33m(tc_clip_text_encoder.py 106)\u001b[0m: INFO Prompt generation level: [11]\n",
      "\u001b[32m[09-25 13:51:36 TCCLIP]\u001b[0m\u001b[33m(tc_clip_text_encoder.py 107)\u001b[0m: INFO Prompt generation stop grad: True\n",
      "\u001b[32m[09-25 13:51:36 TCCLIP]\u001b[0m\u001b[33m(tc_clip.py 27)\u001b[0m: INFO Using context tokens from vision layer [11]\n",
      "\u001b[32m[09-25 13:51:36 TCCLIP]\u001b[0m\u001b[33m(build_trainer.py 97)\u001b[0m: INFO ----------------------------------------------------\n",
      "\u001b[32m[09-25 13:51:36 TCCLIP]\u001b[0m\u001b[33m(build_trainer.py 98)\u001b[0m: INFO Freezed Parameters\n",
      "\u001b[32m[09-25 13:51:36 TCCLIP]\u001b[0m\u001b[33m(build_trainer.py 102)\u001b[0m: INFO ----------------------------------------------------\n",
      "\u001b[32m[09-25 13:51:36 TCCLIP]\u001b[0m\u001b[33m(build_trainer.py 105)\u001b[0m: INFO Number of Parameters: 127.5M\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "TCCLIP(\n",
       "  (prompt_learner): VPPromptLearner()\n",
       "  (image_encoder): TCVisionTransformer(\n",
       "    (conv1): Conv2d(3, 768, kernel_size=(16, 16), stride=(16, 16), bias=False)\n",
       "    (ln_pre): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "    (transformer): Transformer(\n",
       "      (resblocks): Sequential(\n",
       "        (0): TCAttentionBlock(\n",
       "          (attn): TCAttention(\n",
       "            (q_proj): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (k_proj): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (v_proj): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (out_proj): Linear(in_features=768, out_features=768, bias=True)\n",
       "          )\n",
       "          (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "          (mlp): Sequential(\n",
       "            (c_fc): Linear(in_features=768, out_features=3072, bias=True)\n",
       "            (gelu): QuickGELU()\n",
       "            (c_proj): Linear(in_features=3072, out_features=768, bias=True)\n",
       "          )\n",
       "          (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        )\n",
       "        (1): TCAttentionBlock(\n",
       "          (attn): TCAttention(\n",
       "            (q_proj): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (k_proj): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (v_proj): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (out_proj): Linear(in_features=768, out_features=768, bias=True)\n",
       "          )\n",
       "          (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "          (mlp): Sequential(\n",
       "            (c_fc): Linear(in_features=768, out_features=3072, bias=True)\n",
       "            (gelu): QuickGELU()\n",
       "            (c_proj): Linear(in_features=3072, out_features=768, bias=True)\n",
       "          )\n",
       "          (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        )\n",
       "        (2): TCAttentionBlock(\n",
       "          (attn): TCAttention(\n",
       "            (q_proj): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (k_proj): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (v_proj): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (out_proj): Linear(in_features=768, out_features=768, bias=True)\n",
       "          )\n",
       "          (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "          (mlp): Sequential(\n",
       "            (c_fc): Linear(in_features=768, out_features=3072, bias=True)\n",
       "            (gelu): QuickGELU()\n",
       "            (c_proj): Linear(in_features=3072, out_features=768, bias=True)\n",
       "          )\n",
       "          (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        )\n",
       "        (3): TCAttentionBlock(\n",
       "          (attn): TCAttention(\n",
       "            (q_proj): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (k_proj): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (v_proj): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (out_proj): Linear(in_features=768, out_features=768, bias=True)\n",
       "          )\n",
       "          (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "          (mlp): Sequential(\n",
       "            (c_fc): Linear(in_features=768, out_features=3072, bias=True)\n",
       "            (gelu): QuickGELU()\n",
       "            (c_proj): Linear(in_features=3072, out_features=768, bias=True)\n",
       "          )\n",
       "          (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        )\n",
       "        (4): TCAttentionBlock(\n",
       "          (attn): TCAttention(\n",
       "            (q_proj): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (k_proj): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (v_proj): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (out_proj): Linear(in_features=768, out_features=768, bias=True)\n",
       "          )\n",
       "          (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "          (mlp): Sequential(\n",
       "            (c_fc): Linear(in_features=768, out_features=3072, bias=True)\n",
       "            (gelu): QuickGELU()\n",
       "            (c_proj): Linear(in_features=3072, out_features=768, bias=True)\n",
       "          )\n",
       "          (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        )\n",
       "        (5): TCAttentionBlock(\n",
       "          (attn): TCAttention(\n",
       "            (q_proj): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (k_proj): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (v_proj): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (out_proj): Linear(in_features=768, out_features=768, bias=True)\n",
       "          )\n",
       "          (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "          (mlp): Sequential(\n",
       "            (c_fc): Linear(in_features=768, out_features=3072, bias=True)\n",
       "            (gelu): QuickGELU()\n",
       "            (c_proj): Linear(in_features=3072, out_features=768, bias=True)\n",
       "          )\n",
       "          (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        )\n",
       "        (6): TCAttentionBlock(\n",
       "          (attn): TCAttention(\n",
       "            (q_proj): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (k_proj): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (v_proj): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (out_proj): Linear(in_features=768, out_features=768, bias=True)\n",
       "          )\n",
       "          (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "          (mlp): Sequential(\n",
       "            (c_fc): Linear(in_features=768, out_features=3072, bias=True)\n",
       "            (gelu): QuickGELU()\n",
       "            (c_proj): Linear(in_features=3072, out_features=768, bias=True)\n",
       "          )\n",
       "          (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        )\n",
       "        (7): TCAttentionBlock(\n",
       "          (attn): TCAttention(\n",
       "            (q_proj): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (k_proj): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (v_proj): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (out_proj): Linear(in_features=768, out_features=768, bias=True)\n",
       "          )\n",
       "          (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "          (mlp): Sequential(\n",
       "            (c_fc): Linear(in_features=768, out_features=3072, bias=True)\n",
       "            (gelu): QuickGELU()\n",
       "            (c_proj): Linear(in_features=3072, out_features=768, bias=True)\n",
       "          )\n",
       "          (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        )\n",
       "        (8): TCAttentionBlock(\n",
       "          (attn): TCAttention(\n",
       "            (q_proj): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (k_proj): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (v_proj): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (out_proj): Linear(in_features=768, out_features=768, bias=True)\n",
       "          )\n",
       "          (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "          (mlp): Sequential(\n",
       "            (c_fc): Linear(in_features=768, out_features=3072, bias=True)\n",
       "            (gelu): QuickGELU()\n",
       "            (c_proj): Linear(in_features=3072, out_features=768, bias=True)\n",
       "          )\n",
       "          (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        )\n",
       "        (9): TCAttentionBlock(\n",
       "          (attn): TCAttention(\n",
       "            (q_proj): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (k_proj): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (v_proj): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (out_proj): Linear(in_features=768, out_features=768, bias=True)\n",
       "          )\n",
       "          (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "          (mlp): Sequential(\n",
       "            (c_fc): Linear(in_features=768, out_features=3072, bias=True)\n",
       "            (gelu): QuickGELU()\n",
       "            (c_proj): Linear(in_features=3072, out_features=768, bias=True)\n",
       "          )\n",
       "          (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        )\n",
       "        (10): TCAttentionBlock(\n",
       "          (attn): TCAttention(\n",
       "            (q_proj): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (k_proj): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (v_proj): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (out_proj): Linear(in_features=768, out_features=768, bias=True)\n",
       "          )\n",
       "          (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "          (mlp): Sequential(\n",
       "            (c_fc): Linear(in_features=768, out_features=3072, bias=True)\n",
       "            (gelu): QuickGELU()\n",
       "            (c_proj): Linear(in_features=3072, out_features=768, bias=True)\n",
       "          )\n",
       "          (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        )\n",
       "        (11): TCAttentionBlock(\n",
       "          (attn): TCAttention(\n",
       "            (q_proj): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (k_proj): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (v_proj): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (out_proj): Linear(in_features=768, out_features=768, bias=True)\n",
       "          )\n",
       "          (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "          (mlp): Sequential(\n",
       "            (c_fc): Linear(in_features=768, out_features=3072, bias=True)\n",
       "            (gelu): QuickGELU()\n",
       "            (c_proj): Linear(in_features=3072, out_features=768, bias=True)\n",
       "          )\n",
       "          (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (ln_post): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "  )\n",
       "  (text_encoder): VPTextEncoder(\n",
       "    (transformer): Transformer(\n",
       "      (resblocks): Sequential(\n",
       "        (0): ResidualAttentionBlock(\n",
       "          (attn): MultiheadAttention(\n",
       "            (out_proj): NonDynamicallyQuantizableLinear(in_features=512, out_features=512, bias=True)\n",
       "          )\n",
       "          (ln_1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "          (mlp): Sequential(\n",
       "            (c_fc): Linear(in_features=512, out_features=2048, bias=True)\n",
       "            (gelu): QuickGELU()\n",
       "            (c_proj): Linear(in_features=2048, out_features=512, bias=True)\n",
       "          )\n",
       "          (ln_2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "        )\n",
       "        (1): ResidualAttentionBlock(\n",
       "          (attn): MultiheadAttention(\n",
       "            (out_proj): NonDynamicallyQuantizableLinear(in_features=512, out_features=512, bias=True)\n",
       "          )\n",
       "          (ln_1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "          (mlp): Sequential(\n",
       "            (c_fc): Linear(in_features=512, out_features=2048, bias=True)\n",
       "            (gelu): QuickGELU()\n",
       "            (c_proj): Linear(in_features=2048, out_features=512, bias=True)\n",
       "          )\n",
       "          (ln_2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "        )\n",
       "        (2): ResidualAttentionBlock(\n",
       "          (attn): MultiheadAttention(\n",
       "            (out_proj): NonDynamicallyQuantizableLinear(in_features=512, out_features=512, bias=True)\n",
       "          )\n",
       "          (ln_1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "          (mlp): Sequential(\n",
       "            (c_fc): Linear(in_features=512, out_features=2048, bias=True)\n",
       "            (gelu): QuickGELU()\n",
       "            (c_proj): Linear(in_features=2048, out_features=512, bias=True)\n",
       "          )\n",
       "          (ln_2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "        )\n",
       "        (3): ResidualAttentionBlock(\n",
       "          (attn): MultiheadAttention(\n",
       "            (out_proj): NonDynamicallyQuantizableLinear(in_features=512, out_features=512, bias=True)\n",
       "          )\n",
       "          (ln_1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "          (mlp): Sequential(\n",
       "            (c_fc): Linear(in_features=512, out_features=2048, bias=True)\n",
       "            (gelu): QuickGELU()\n",
       "            (c_proj): Linear(in_features=2048, out_features=512, bias=True)\n",
       "          )\n",
       "          (ln_2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "        )\n",
       "        (4): ResidualAttentionBlock(\n",
       "          (attn): MultiheadAttention(\n",
       "            (out_proj): NonDynamicallyQuantizableLinear(in_features=512, out_features=512, bias=True)\n",
       "          )\n",
       "          (ln_1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "          (mlp): Sequential(\n",
       "            (c_fc): Linear(in_features=512, out_features=2048, bias=True)\n",
       "            (gelu): QuickGELU()\n",
       "            (c_proj): Linear(in_features=2048, out_features=512, bias=True)\n",
       "          )\n",
       "          (ln_2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "        )\n",
       "        (5): ResidualAttentionBlock(\n",
       "          (attn): MultiheadAttention(\n",
       "            (out_proj): NonDynamicallyQuantizableLinear(in_features=512, out_features=512, bias=True)\n",
       "          )\n",
       "          (ln_1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "          (mlp): Sequential(\n",
       "            (c_fc): Linear(in_features=512, out_features=2048, bias=True)\n",
       "            (gelu): QuickGELU()\n",
       "            (c_proj): Linear(in_features=2048, out_features=512, bias=True)\n",
       "          )\n",
       "          (ln_2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "        )\n",
       "        (6): ResidualAttentionBlock(\n",
       "          (attn): MultiheadAttention(\n",
       "            (out_proj): NonDynamicallyQuantizableLinear(in_features=512, out_features=512, bias=True)\n",
       "          )\n",
       "          (ln_1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "          (mlp): Sequential(\n",
       "            (c_fc): Linear(in_features=512, out_features=2048, bias=True)\n",
       "            (gelu): QuickGELU()\n",
       "            (c_proj): Linear(in_features=2048, out_features=512, bias=True)\n",
       "          )\n",
       "          (ln_2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "        )\n",
       "        (7): ResidualAttentionBlock(\n",
       "          (attn): MultiheadAttention(\n",
       "            (out_proj): NonDynamicallyQuantizableLinear(in_features=512, out_features=512, bias=True)\n",
       "          )\n",
       "          (ln_1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "          (mlp): Sequential(\n",
       "            (c_fc): Linear(in_features=512, out_features=2048, bias=True)\n",
       "            (gelu): QuickGELU()\n",
       "            (c_proj): Linear(in_features=2048, out_features=512, bias=True)\n",
       "          )\n",
       "          (ln_2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "        )\n",
       "        (8): ResidualAttentionBlock(\n",
       "          (attn): MultiheadAttention(\n",
       "            (out_proj): NonDynamicallyQuantizableLinear(in_features=512, out_features=512, bias=True)\n",
       "          )\n",
       "          (ln_1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "          (mlp): Sequential(\n",
       "            (c_fc): Linear(in_features=512, out_features=2048, bias=True)\n",
       "            (gelu): QuickGELU()\n",
       "            (c_proj): Linear(in_features=2048, out_features=512, bias=True)\n",
       "          )\n",
       "          (ln_2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "        )\n",
       "        (9): ResidualAttentionBlock(\n",
       "          (attn): MultiheadAttention(\n",
       "            (out_proj): NonDynamicallyQuantizableLinear(in_features=512, out_features=512, bias=True)\n",
       "          )\n",
       "          (ln_1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "          (mlp): Sequential(\n",
       "            (c_fc): Linear(in_features=512, out_features=2048, bias=True)\n",
       "            (gelu): QuickGELU()\n",
       "            (c_proj): Linear(in_features=2048, out_features=512, bias=True)\n",
       "          )\n",
       "          (ln_2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "        )\n",
       "        (10): ResidualAttentionBlock(\n",
       "          (attn): MultiheadAttention(\n",
       "            (out_proj): NonDynamicallyQuantizableLinear(in_features=512, out_features=512, bias=True)\n",
       "          )\n",
       "          (ln_1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "          (mlp): Sequential(\n",
       "            (c_fc): Linear(in_features=512, out_features=2048, bias=True)\n",
       "            (gelu): QuickGELU()\n",
       "            (c_proj): Linear(in_features=2048, out_features=512, bias=True)\n",
       "          )\n",
       "          (ln_2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "        )\n",
       "        (11): ResidualAttentionBlock(\n",
       "          (attn): MultiheadAttention(\n",
       "            (out_proj): NonDynamicallyQuantizableLinear(in_features=512, out_features=512, bias=True)\n",
       "          )\n",
       "          (ln_1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "          (mlp): Sequential(\n",
       "            (c_fc): Linear(in_features=512, out_features=2048, bias=True)\n",
       "            (gelu): QuickGELU()\n",
       "            (c_proj): Linear(in_features=2048, out_features=512, bias=True)\n",
       "          )\n",
       "          (ln_2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (ln_final): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "    (prompt_generation_layer): ModuleList(\n",
       "      (0): PromptGenerationLayer(\n",
       "        (attn): MultiheadAttention(\n",
       "          (out_proj): NonDynamicallyQuantizableLinear(in_features=512, out_features=512, bias=True)\n",
       "        )\n",
       "        (ln_1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "        (ln_1_kv): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "        (mlp): Sequential(\n",
       "          (c_fc): Linear(in_features=512, out_features=2048, bias=True)\n",
       "          (gelu): QuickGELU()\n",
       "          (c_proj): Linear(in_features=2048, out_features=512, bias=True)\n",
       "        )\n",
       "        (ln_2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "      )\n",
       "    )\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Build model\n",
    "model = returnCLIP(config, logger, class_names)\n",
    "model.cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "b1b0a56f",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[32m[09-25 13:51:39 TCCLIP]\u001b[0m\u001b[33m(tools.py 180)\u001b[0m: INFO ==============> Resuming from workspace/weights/zero_shot_k400_llm_tc_clip.pth....................\n",
      "\u001b[32m[09-25 13:51:40 TCCLIP]\u001b[0m\u001b[33m(tools.py 208)\u001b[0m: INFO resume model: _IncompatibleKeys(missing_keys=['prompt_learner.token_prefix', 'prompt_learner.token_suffix'], unexpected_keys=[])\n",
      "\u001b[32m[09-25 13:51:40 TCCLIP]\u001b[0m\u001b[33m(tools.py 218)\u001b[0m: INFO => loaded successfully 'workspace/weights/zero_shot_k400_llm_tc_clip.pth' (epoch 9)\n",
      "\u001b[32m[09-25 13:51:40 TCCLIP]\u001b[0m\u001b[33m(1028797808.py 4)\u001b[0m: INFO Loaded checkpoint at epoch 10 with max accuracy 82.1\n"
     ]
    }
   ],
   "source": [
    "# Load checkpoint\n",
    "if config.resume:\n",
    "    epoch_loaded, max_accuray_loaded = load_checkpoint(config, model, None, None, logger, model_only=True)\n",
    "    logger.info(\n",
    "            f\"Loaded checkpoint at epoch {epoch_loaded} with max accuracy {max_accuray_loaded:.1f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3cc32b34",
   "metadata": {},
   "source": [
    "### Video preprocessing pipeline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "0f1f0416",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Video preprocessing pipeline\n",
    "\n",
    "img_norm_cfg = dict(\n",
    "    mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_bgr=False)\n",
    "\n",
    "scale_resize = int(256 / 224 * config.input_size)\n",
    "collect_keys = ['imgs']\n",
    "\n",
    "val_pipeline = [\n",
    "    dict(type='DecordInit'),\n",
    "    dict(type='SampleFrames', clip_len=1, frame_interval=1, num_clips=config.num_frames, test_mode=True),\n",
    "    dict(type='DecordDecode'),\n",
    "    dict(type='Resize', scale=(-1, scale_resize)),\n",
    "    dict(type='CenterCrop', crop_size=config.input_size),\n",
    "    dict(type='Normalize', **img_norm_cfg),\n",
    "    dict(type='FormatShape', input_format='NCHW'),\n",
    "    dict(type='Collect', keys=collect_keys, meta_keys=[]),\n",
    "    dict(type='ToTensor', keys=['imgs'])\n",
    "]\n",
    "if config.num_crop == 3:\n",
    "    val_pipeline[3] = dict(type='Resize', scale=(-1, config.input_size))\n",
    "    val_pipeline[4] = dict(type='ThreeCrop', crop_size=config.input_size)\n",
    "if config.num_clip > 1:\n",
    "    val_pipeline[1] = dict(type='SampleFrames', clip_len=1, frame_interval=1, num_clips=config.num_frames,\n",
    "                           multiview=config.num_clip)\n",
    "val_pipeline = [p for p in val_pipeline if p is not None]\n",
    "\n",
    "pipeline = Compose(val_pipeline)\n",
    "\n",
    "dict_file = {'filename': video_path, 'tar': False, 'modality': 'RGB', 'start_index': 0}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4845ffa5",
   "metadata": {},
   "source": [
    "### TC-CLIP inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "d30fe299",
   "metadata": {},
   "outputs": [],
   "source": [
    "video = pipeline(dict_file)\n",
    "video_tensor = video['imgs'].unsqueeze(0).cuda().float() # Size: [1, T, 3, H, W]\n",
    "\n",
    "# Inference with TC-CLIP\n",
    "with torch.no_grad():\n",
    "    if config.use_torch_amp:\n",
    "        with torch.cuda.amp.autocast():\n",
    "            output = model(video_tensor)\n",
    "    else:\n",
    "        output = model(video_tensor)\n",
    "    \n",
    "    logits = output['logits']\n",
    "\n",
    "pred_index = logits.argmax(-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "25d0b5a7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Logits: tensor([[30.6250, 23.0625, 23.5000, 24.1094]], device='cuda:0',\n",
      "       dtype=torch.float16)\n",
      "Predicted action category is \"swinging baseball bat\"\n"
     ]
    }
   ],
   "source": [
    "print(f'Logits: {logits}')\n",
    "print(f'Predicted action category is \"{class_names[pred_index]}\"')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2bfc1915",
   "metadata": {},
   "source": [
    "Acknowledgements: [ViFi-CLIP's repository](https://github.com/muzairkhattak/ViFi-CLIP)."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}