{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The autoreload extension is already loaded. To reload it, use:\n",
      "  %reload_ext autoreload\n"
     ]
    }
   ],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from PIL import Image\n",
    "import torch.nn.functional as F\n",
    "import torch\n",
    "from torch import nn\n",
    "from torchsummary import summary\n",
    "from importlib.util import find_spec\n",
    "if find_spec(\"text_recognizer\") is None:\n",
    "    import sys\n",
    "    sys.path.append('..')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.cuda.is_available()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "from text_recognizer.networks.transformer.layers import Decoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "decoder = Decoder(dim=256, depth=4, num_heads=8, ff_kwargs={}, attn_kwargs={}, cross_attend=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Decoder(\n",
       "  (layers): ModuleList(\n",
       "    (0): ModuleList(\n",
       "      (0): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "      (1): Attention(\n",
       "        (qkv_fn): Sequential(\n",
       "          (0): Linear(in_features=256, out_features=49152, bias=False)\n",
       "          (1): Rearrange('b n (qkv h d) -> qkv b h n d', qkv=3, h=8)\n",
       "        )\n",
       "        (dropout): Dropout(p=0.0, inplace=False)\n",
       "        (fc): Linear(in_features=16384, out_features=256, bias=True)\n",
       "      )\n",
       "      (2): Residual()\n",
       "    )\n",
       "    (1): ModuleList(\n",
       "      (0): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "      (1): Attention(\n",
       "        (qkv_fn): Sequential(\n",
       "          (0): Linear(in_features=256, out_features=49152, bias=False)\n",
       "          (1): Rearrange('b n (qkv h d) -> qkv b h n d', qkv=3, h=8)\n",
       "        )\n",
       "        (dropout): Dropout(p=0.0, inplace=False)\n",
       "        (fc): Linear(in_features=16384, out_features=256, bias=True)\n",
       "      )\n",
       "      (2): Residual()\n",
       "    )\n",
       "    (2): ModuleList(\n",
       "      (0): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "      (1): FeedForward(\n",
       "        (mlp): Sequential(\n",
       "          (0): GEGLU(\n",
       "            (fc): Linear(in_features=256, out_features=2048, bias=True)\n",
       "          )\n",
       "          (1): Dropout(p=0.0, inplace=False)\n",
       "          (2): Linear(in_features=1024, out_features=256, bias=True)\n",
       "        )\n",
       "      )\n",
       "      (2): Residual()\n",
       "    )\n",
       "    (3): ModuleList(\n",
       "      (0): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "      (1): Attention(\n",
       "        (qkv_fn): Sequential(\n",
       "          (0): Linear(in_features=256, out_features=49152, bias=False)\n",
       "          (1): Rearrange('b n (qkv h d) -> qkv b h n d', qkv=3, h=8)\n",
       "        )\n",
       "        (dropout): Dropout(p=0.0, inplace=False)\n",
       "        (fc): Linear(in_features=16384, out_features=256, bias=True)\n",
       "      )\n",
       "      (2): Residual()\n",
       "    )\n",
       "    (4): ModuleList(\n",
       "      (0): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "      (1): Attention(\n",
       "        (qkv_fn): Sequential(\n",
       "          (0): Linear(in_features=256, out_features=49152, bias=False)\n",
       "          (1): Rearrange('b n (qkv h d) -> qkv b h n d', qkv=3, h=8)\n",
       "        )\n",
       "        (dropout): Dropout(p=0.0, inplace=False)\n",
       "        (fc): Linear(in_features=16384, out_features=256, bias=True)\n",
       "      )\n",
       "      (2): Residual()\n",
       "    )\n",
       "    (5): ModuleList(\n",
       "      (0): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "      (1): FeedForward(\n",
       "        (mlp): Sequential(\n",
       "          (0): GEGLU(\n",
       "            (fc): Linear(in_features=256, out_features=2048, bias=True)\n",
       "          )\n",
       "          (1): Dropout(p=0.0, inplace=False)\n",
       "          (2): Linear(in_features=1024, out_features=256, bias=True)\n",
       "        )\n",
       "      )\n",
       "      (2): Residual()\n",
       "    )\n",
       "    (6): ModuleList(\n",
       "      (0): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "      (1): Attention(\n",
       "        (qkv_fn): Sequential(\n",
       "          (0): Linear(in_features=256, out_features=49152, bias=False)\n",
       "          (1): Rearrange('b n (qkv h d) -> qkv b h n d', qkv=3, h=8)\n",
       "        )\n",
       "        (dropout): Dropout(p=0.0, inplace=False)\n",
       "        (fc): Linear(in_features=16384, out_features=256, bias=True)\n",
       "      )\n",
       "      (2): Residual()\n",
       "    )\n",
       "    (7): ModuleList(\n",
       "      (0): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "      (1): Attention(\n",
       "        (qkv_fn): Sequential(\n",
       "          (0): Linear(in_features=256, out_features=49152, bias=False)\n",
       "          (1): Rearrange('b n (qkv h d) -> qkv b h n d', qkv=3, h=8)\n",
       "        )\n",
       "        (dropout): Dropout(p=0.0, inplace=False)\n",
       "        (fc): Linear(in_features=16384, out_features=256, bias=True)\n",
       "      )\n",
       "      (2): Residual()\n",
       "    )\n",
       "    (8): ModuleList(\n",
       "      (0): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "      (1): FeedForward(\n",
       "        (mlp): Sequential(\n",
       "          (0): GEGLU(\n",
       "            (fc): Linear(in_features=256, out_features=2048, bias=True)\n",
       "          )\n",
       "          (1): Dropout(p=0.0, inplace=False)\n",
       "          (2): Linear(in_features=1024, out_features=256, bias=True)\n",
       "        )\n",
       "      )\n",
       "      (2): Residual()\n",
       "    )\n",
       "    (9): ModuleList(\n",
       "      (0): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "      (1): Attention(\n",
       "        (qkv_fn): Sequential(\n",
       "          (0): Linear(in_features=256, out_features=49152, bias=False)\n",
       "          (1): Rearrange('b n (qkv h d) -> qkv b h n d', qkv=3, h=8)\n",
       "        )\n",
       "        (dropout): Dropout(p=0.0, inplace=False)\n",
       "        (fc): Linear(in_features=16384, out_features=256, bias=True)\n",
       "      )\n",
       "      (2): Residual()\n",
       "    )\n",
       "    (10): ModuleList(\n",
       "      (0): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "      (1): Attention(\n",
       "        (qkv_fn): Sequential(\n",
       "          (0): Linear(in_features=256, out_features=49152, bias=False)\n",
       "          (1): Rearrange('b n (qkv h d) -> qkv b h n d', qkv=3, h=8)\n",
       "        )\n",
       "        (dropout): Dropout(p=0.0, inplace=False)\n",
       "        (fc): Linear(in_features=16384, out_features=256, bias=True)\n",
       "      )\n",
       "      (2): Residual()\n",
       "    )\n",
       "    (11): ModuleList(\n",
       "      (0): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
       "      (1): FeedForward(\n",
       "        (mlp): Sequential(\n",
       "          (0): GEGLU(\n",
       "            (fc): Linear(in_features=256, out_features=2048, bias=True)\n",
       "          )\n",
       "          (1): Dropout(p=0.0, inplace=False)\n",
       "          (2): Linear(in_features=1024, out_features=256, bias=True)\n",
       "        )\n",
       "      )\n",
       "      (2): Residual()\n",
       "    )\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "decoder.cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from text_recognizer.networks.transformer.nystromer.nystromer import Nystromer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from functools import partial"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "efficient_transformer = Nystromer(\n",
    "    dim = 128,\n",
    "    depth = 4,\n",
    "    num_heads = 8,\n",
    "    num_landmarks = 128\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from text_recognizer.networks.transformer.vit import ViT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "v = ViT(\n",
    "    dim = 128,\n",
    "    image_size = (576, 640),\n",
    "    patch_size = (32, 32),\n",
    "    transformer = efficient_transformer\n",
    ").cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t = torch.randn(32, 1, 576, 640).cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "v(t).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from text_recognizer.networks.encoders.efficientnet import EfficientNet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "en = EfficientNet()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "(576, 640) // (8, 8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "(576 // 32) ** 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "summary(en, (1, 576, 640))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "type(efficient_transformer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "efficient_transformer = efficient_transformer(num_landmarks=256)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "efficient_transformer()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from omegaconf import OmegaConf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "path = \"../training/configs/vqvae.yaml\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "conf = OmegaConf.load(path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(OmegaConf.to_yaml(conf))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from text_recognizer.networks import VQVAE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vae = VQVAE(**conf.network.args)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vae"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "datum = torch.randn([2, 1, 576, 640])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vae.encoder(datum)[0].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vae(datum)[0].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from text_recognizer.networks.backbones.efficientnet import EfficientNet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "en = EfficientNet()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "datum = torch.randn([2, 1, 576, 640])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trg = torch.randint(0, 1000, [2, 682])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trg.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "datum = torch.randn([2, 1, 224, 224])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "en(datum).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "path = \"../training/configs/cnn_transformer.yaml\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "conf = OmegaConf.load(path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(OmegaConf.to_yaml(conf))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from text_recognizer.networks.cnn_transformer import CNNTransformer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t = CNNTransformer(input_shape=(1, 576, 640), output_shape=(682, 1), **conf.network.args)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t.encode(datum).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trg.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t(datum, trg).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "b, n = 16, 128\n",
    "device = \"cpu\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = lambda: torch.ones((b, n), device=device).bool()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x().shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.ones((b, n), device=device).bool().shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = torch.randn(1, 1, 576, 640)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "576 // 32"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "640 // 32"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "18 * 20"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = torch.randn(1, 1, 144, 160)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from einops import rearrange"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "patch_size=16\n",
    "p = rearrange(x, 'b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1 = patch_size, p2 = patch_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "p.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.9.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}