CFOCNet_demo.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "id": "qvhoDqESPuqy"
   },
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "from pycocotools.coco import COCO\n",
    "import numpy as np\n",
    "import skimage.io as io\n",
    "import matplotlib.pyplot as plt\n",
    "import pylab\n",
    "import os\n",
    "import cv2\n",
    "# from google.colab.patches import cv2_imshow\n",
    "from skimage.transform import rescale, resize, downscale_local_mean\n",
    "import random\n",
    "from torchvision.models.resnet import resnet50\n",
    "# import Stream_Resnet\n",
    "# import collections\n",
    "# import collections.abc\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.utils.data import DataLoader,Dataset\n",
    "from tqdm import tqdm\n",
    "import os\n",
    "import numpy as np\n",
    "from torchvision import transforms\n",
    "import random\n",
    "import matplotlib.pyplot as plt\n",
    "import cv2\n",
    "import random\n",
    "from pycocotools.coco import COCO\n",
    "import skimage.io as io\n",
    "import PIL\n",
    "# from torchvision.models import resnet50\n",
    "# for the stream resnet\n",
    "# import Stream_Resnet\n",
    "\n",
    "# import Self_Attention\n",
    "# for SSIM LOSS\n",
    "from pytorch_msssim import ssim\n",
    "# Get cpu or gpu device for training.\n",
    "import argparse\n",
    "# import torch.distributed as dist\n",
    "import pandas as pd\n",
    "from matplotlib import cm as CM\n",
    "from model.layers import Self_Attn\n",
    "from torch.nn import MaxPool1d\n",
    "import torch.nn.functional as F"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "id": "zUmBj06wPWlX"
   },
   "outputs": [],
   "source": [
    "def make_four_layers():\n",
    "    net = resnet50()\n",
    "    layer0_name = ['conv1','bn1','relu']\n",
    "    layer1_name = ['maxpool','layer1']\n",
    "    layer2_name = ['layer2']\n",
    "    layer3_name = ['layer3']\n",
    "    \n",
    "    layer0 = nn.Sequential()\n",
    "    layer1 = nn.Sequential()\n",
    "    layer2 = nn.Sequential()\n",
    "    layer3 = nn.Sequential()\n",
    "\n",
    "    for n,c in net.named_children():\n",
    "        if n in layer0_name:\n",
    "            layer0.add_module(n,c)\n",
    "        elif n in layer1_name:\n",
    "            layer1.add_module(n,c)\n",
    "        elif n in layer2_name:\n",
    "            layer2.add_module(n,c)\n",
    "        elif n in layer3_name:\n",
    "            layer3 = c\n",
    "        else:\n",
    "            break\n",
    "\n",
    "    return layer0, layer1,layer2,layer3\n",
    "    \n",
    "\n",
    "class JDimPool(nn.MaxPool1d):\n",
    "    def forward(self, input):\n",
    "        b,n, c, w, h = input.size()\n",
    "        input = input.view(b, n, c* w * h).permute(0, 2, 1)\n",
    "        pooled = nn.functional.max_pool1d(\n",
    "            input,\n",
    "            self.kernel_size,\n",
    "            self.stride,\n",
    "            self.padding,\n",
    "            self.dilation,\n",
    "            self.ceil_mode,\n",
    "            self.return_indices,\n",
    "        )\n",
    "        pooled = pooled.permute(0, 2, 1)\n",
    "        return pooled.view(b,c, w, h)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "id": "2ITr_nRqaaRi"
   },
   "outputs": [],
   "source": [
    "relu = nn.ReLU()\n",
    "\n",
    "res_q0, res_q1,res_q2,res_q3 = make_four_layers()\n",
    "res_r0, res_r1,res_r2,res_r3 = make_four_layers()\n",
    "\n",
    "sa_q1 = Self_Attn(256, nn.ReLU())\n",
    "sa_q2 = Self_Attn(512, nn.ReLU())\n",
    "sa_q3 = Self_Attn(1024, nn.ReLU())\n",
    "j_maxpool = JDimPool(5,1) # stride is useless\n",
    "\n",
    "\n",
    "maxpool_r1 = nn.MaxPool2d(4, stride=4, padding=0)\n",
    "maxpool_r2 = nn.MaxPool2d(2, stride=2, padding=0)\n",
    "# maxpool_r3 = nn.MaxPool2d(2, stride=2, padding=0)\n",
    "\n",
    "match_query_conv1 = nn.Conv2d(in_channels=256, out_channels=256, kernel_size=1)\n",
    "match_query_conv2 = nn.Conv2d(in_channels=512, out_channels=256, kernel_size=1)\n",
    "match_query_conv3 = nn.Conv2d(in_channels=1024, out_channels=256, kernel_size=1)\n",
    "\n",
    "match_reference_conv1 = nn.Conv2d(in_channels=256, out_channels=256, kernel_size=1)\n",
    "match_reference_conv2 = nn.Conv2d(in_channels=512, out_channels=256, kernel_size=1)\n",
    "match_reference_conv3 = nn.Conv2d(in_channels=1024, out_channels=256, kernel_size=1)\n",
    "\n",
    "sum_conv1 = nn.Conv2d(in_channels=1, out_channels=1, kernel_size=1)\n",
    "sum_conv2 = nn.Conv2d(in_channels=1, out_channels=1, kernel_size=1)\n",
    "sum_conv3 = nn.Conv2d(in_channels=1, out_channels=1, kernel_size=1)\n",
    "\n",
    "softmax = nn.Softmax(dim=-1)\n",
    "\n",
    "transpose_convolution = nn.ConvTranspose2d(\n",
    "    in_channels=1,\n",
    "    out_channels=1,\n",
    "    kernel_size=3,\n",
    "    stride=2,\n",
    "    padding=1,\n",
    "    output_padding = 1\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "id": "MEtGBMwMP-x8"
   },
   "outputs": [],
   "source": [
    "q = torch.FloatTensor(5, 3, 256, 256)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "id": "2pThsI1CNxRx"
   },
   "outputs": [],
   "source": [
    "r = torch.FloatTensor(5, 5, 3, 64, 64) # b, k, c, w, h"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "hxwIHsxNbnQi",
    "outputId": "372f6baf-96ac-4df5-c66c-e8998c76158c"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([5, 256, 64, 64]),\n",
       " torch.Size([5, 512, 32, 32]),\n",
       " torch.Size([5, 1024, 16, 16]))"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "q0 = res_q0(q)\n",
    "q1 = res_q1(q0)\n",
    "q2 = res_q2(q1)\n",
    "q3 = res_q3(q2)\n",
    "q1.shape, q2.shape, q3.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "vW2mrdppb-rd",
    "outputId": "e78ec3c6-d332-4865-b44a-1f084717f806"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([5, 5, 3, 64, 64]), torch.Size([25, 3, 64, 64]))"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "org_ref_size = r.size()\n",
    "r = r.view(-1,org_ref_size[-3],org_ref_size[-2],org_ref_size[-1])\n",
    "org_ref_size, r.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "Fketp4N7bnSt",
    "outputId": "7379878b-958f-458f-a3d8-05facbf19603"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([25, 256, 16, 16]),\n",
       " torch.Size([25, 512, 8, 8]),\n",
       " torch.Size([25, 1024, 4, 4]))"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "r0 = res_r0(r)\n",
    "r1 = res_r1(r0)\n",
    "r2 = res_r2(r1)\n",
    "r3 = res_r3(r2)\n",
    "r1.shape, r2.shape, r3.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "w8cyU73HbnUp",
    "outputId": "9ad02fed-86ad-4f7f-f019-2bce577c0daf"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([5, 5, 256, 16, 16]),\n",
       " torch.Size([5, 5, 512, 8, 8]),\n",
       " torch.Size([5, 5, 1024, 4, 4]))"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "r1_size = r1.size()\n",
    "r2_size = r2.size()\n",
    "r3_size = r3.size()\n",
    "r1 = r1.view(org_ref_size[0],org_ref_size[1],r1_size[-3],r1_size[-2],r1_size[-1])\n",
    "r2 = r2.view(org_ref_size[0],org_ref_size[1],r2_size[-3],r2_size[-2],r2_size[-1])\n",
    "r3 = r3.view(org_ref_size[0],org_ref_size[1],r3_size[-3],r3_size[-2],r3_size[-1])\n",
    "r1.shape, r2.shape, r3.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "amxt7yw-cLDs",
    "outputId": "f444ae72-e144-418e-9ede-4d60513de482"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([5, 256, 16, 16]),\n",
       " torch.Size([5, 512, 8, 8]),\n",
       " torch.Size([5, 1024, 4, 4]))"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "r1 = j_maxpool(r1)\n",
    "r2 = j_maxpool(r2)\n",
    "r3 = j_maxpool(r3)\n",
    "r1.shape, r2.shape, r3.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "id": "dayi7XfWcLF4"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([5, 256, 64, 64]),\n",
       " torch.Size([5, 512, 32, 32]),\n",
       " torch.Size([5, 1024, 16, 16]))"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "input1, _ = sa_q1(q1)\n",
    "input2, _ = sa_q2(q2)\n",
    "input3, _ = sa_q3(q3)\n",
    "input1.shape, input2.shape, input3.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "id": "7HV7uGPbc_Hq"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([5, 256, 64, 64]),\n",
       " torch.Size([5, 256, 32, 32]),\n",
       " torch.Size([5, 256, 16, 16]))"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "input1 = match_query_conv1(input1)\n",
    "input2 = match_query_conv2(input2)\n",
    "input3 = match_query_conv3(input3)\n",
    "input1.shape, input2.shape, input3.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "id": "-eYreRm7c_J0"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([5, 256, 4, 4]),\n",
       " torch.Size([5, 512, 4, 4]),\n",
       " torch.Size([5, 1024, 4, 4]))"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "kernel1 = maxpool_r1(r1)\n",
    "kernel2 = maxpool_r2(r2)\n",
    "kernel3 = r3\n",
    "kernel1.shape, kernel2.shape, kernel3.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "id": "rYHmS3PPc_L5"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([5, 256, 4, 4]),\n",
       " torch.Size([5, 256, 4, 4]),\n",
       " torch.Size([5, 256, 4, 4]))"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "kernel1 = match_reference_conv1(kernel1)\n",
    "kernel2 = match_reference_conv2(kernel2)\n",
    "kernel3 = match_reference_conv3(kernel3)\n",
    "kernel1.shape, kernel2.shape, kernel3.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "id": "d7kZs_qLc_N4"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([5, 1, 32, 32]),\n",
       " torch.Size([5, 1, 16, 16]),\n",
       " torch.Size([5, 1, 8, 8]))"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "M1 = []\n",
    "M2 = []\n",
    "M3 = []\n",
    "for i in range(kernel1.size(0)):\n",
    "    tmp_m1 = nn.functional.conv2d(input1[i:i+1], kernel1[i:i+1],padding=1,stride=2)\n",
    "    tmp_m2 = nn.functional.conv2d(input2[i:i+1], kernel2[i:i+1],padding=1,stride=2)\n",
    "    tmp_m3 = nn.functional.conv2d(input3[i:i+1], kernel3[i:i+1],padding=1,stride=2)\n",
    "    M1.append(tmp_m1)\n",
    "    M2.append(tmp_m2)\n",
    "    M3.append(tmp_m3)\n",
    "\n",
    "M1 = torch.cat(M1,0)\n",
    "M2 = torch.cat(M2,0)\n",
    "M3 = torch.cat(M3,0)\n",
    "\n",
    "M1.shape, M2.shape, M3.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "id": "7YY2l2Z2eTlp"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([5, 1, 32, 32]), torch.Size([5, 1, 32, 32]))"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# M1 = nn.functional.max_pool2d(M1,2)\n",
    "\n",
    "M2 = nn.functional.interpolate(M2,scale_factor=2)\n",
    "M3 = nn.functional.interpolate(M3,scale_factor=4)\n",
    "M2.shape, M2.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "id": "mBYNO9P8eTnq"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([5, 1, 32, 32]),\n",
       " torch.Size([5, 1, 32, 32]),\n",
       " torch.Size([5, 1, 32, 32]))"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "S1 = sum_conv1(M1)\n",
    "S2 = sum_conv2(M2)\n",
    "S3 = sum_conv3(M3)\n",
    "S1.shape, S2.shape, S3.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "S1_W = []\n",
    "S2_W = []\n",
    "S3_W = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "for k in range(S1.shape[0]):\n",
    "    S1_W.append(torch.sum(S1[k]))\n",
    "    S2_W.append(torch.sum(S2[k]))\n",
    "    S3_W.append(torch.sum(S3[k]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([5]), torch.Size([5]), torch.Size([5]))"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "S1_W = torch.stack(S1_W)\n",
    "S2_W = torch.stack(S2_W)\n",
    "S3_W = torch.stack(S3_W)\n",
    "S1_W.shape, S2_W.shape, S3_W.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "5"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "S1_W.shape[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "W = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([0., 0., 1.])\n",
      "tensor([0., 0., 1.])\n",
      "tensor([0., 0., 1.])\n",
      "tensor([0., 0., 1.])\n",
      "tensor([0., 0., 1.])\n"
     ]
    }
   ],
   "source": [
    "for k in range(S1_W.shape[0]):\n",
    "    W.append(softmax(torch.tensor([S1_W[k], S2_W[k], S3_W[k]])))\n",
    "    print(W[k])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "FS = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 32, 32])\n",
      "torch.Size([1, 32, 32])\n",
      "torch.Size([1, 32, 32])\n",
      "torch.Size([1, 32, 32])\n",
      "torch.Size([1, 32, 32])\n"
     ]
    }
   ],
   "source": [
    "for k in range(M1.shape[0]):\n",
    "    FS.append(W[k][0] * M1[k] + W[k][1] * M2[k] + W[k][2] * M3[k])\n",
    "    print(FS[k].shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([5, 1, 32, 32])"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "FS = torch.stack(FS)\n",
    "FS.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([5, 1, 64, 64])"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "FS = transpose_convolution(FS)\n",
    "FS.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([5, 1, 256, 256])"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\n",
    "FS = nn.functional.interpolate(FS,\n",
    "                scale_factor=4,\n",
    "                mode='bilinear',\n",
    "                align_corners=True)\n",
    "FS.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "CFOCNet_Three_Blocks.ipynb",
   "provenance": []
  },
  "gpuClass": "standard",
  "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.7.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}