Approach C.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# load libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "#load libraries\n",
    "import os\n",
    "import numpy as np\n",
    "import cv2\n",
    "import json\n",
    "import skimage.io as io\n",
    "import shutil\n",
    "import os\n",
    "from matplotlib import pyplot as plt\n",
    "import seaborn as sns\n",
    "import csv\n",
    "\n",
    "\n",
    "import torch\n",
    "import torchvision # torch package for vision related things\n",
    "import torch.nn.functional as F  # Parameterless functions, like (some) activation functions\n",
    "import torchvision.datasets as datasets  # Standard datasets\n",
    "import torchvision.transforms as transforms  # Transformations we can perform on our dataset for augmentation\n",
    "from torch import optim  # For optimizers like SGD, Adam, etc.\n",
    "from torch import nn  # All neural network modules, nn.Linear, nn.Conv2d, BatchNorm, Loss functions\n",
    "from torch.utils.data import DataLoader  # Gives easier dataset managment by creating mini batches etc.\n",
    "from tqdm import tqdm  # For nice progress bar!\n",
    "from random import sample\n",
    "import copy"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load Image Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import pandas as pd \n",
    "import torch\n",
    "from torch.utils.data import Dataset\n",
    "from PIL import Image\n",
    "import torchvision.transforms as transforms\n",
    "\n",
    "class trainDataset(Dataset):\n",
    "    '''\n",
    "    csv_file contain the image name and the label (x,y) of each image\n",
    "        first column: name of the image\n",
    "        second column: label (the x,y coordinations) of the image\n",
    "    img_dir: the path to where the train or test images are stored\n",
    "    flag_resize (bool): a flag, if True, then do the resize. Resize the image into suqare.\n",
    "    \n",
    "    '''\n",
    "    def __init__(self, csv_file, img_dir, flag_resize=True, img_size=224, transform=None, cmap='RGB'):\n",
    "        self.annotations = pd.read_csv(csv_file) #read csv as dataframe\n",
    "        self.img_dir = img_dir\n",
    "        self.img_names = self.annotations['id'].to_numpy() #array of object, size: (7500,)\n",
    "        self.flag_resize = flag_resize\n",
    "        self.img_size = img_size\n",
    "        self.transform = transform\n",
    "        self.cmap = cmap\n",
    "    \n",
    "    \n",
    "    def __len__(self):\n",
    "        return len(self.img_names)\n",
    "        \n",
    "\n",
    "    '''\n",
    "    must have __getitem__ function\n",
    "    to return specific image and the target/label of that image\n",
    "    '''\n",
    "    def __getitem__(self, index):\n",
    "        img_filename = self.img_names[index] + \".jpg\" #output will look like this ('IMG2760_5.jpg',)\n",
    "        img_path = os.path.join(self.img_dir, img_filename)\n",
    "        image = Image.open(img_path).convert(self.cmap) # w, h, channel #image raed by PIL is channel last. shape(490, 680, 3)\n",
    "        #image = io.imread(img_path)\n",
    "        geo_label_x = torch.tensor(self.annotations.iloc[index,1])\n",
    "        geo_label_y = torch.tensor(self.annotations.iloc[index,2])\n",
    "        \n",
    "        \n",
    "        #Resize PIL image to square\n",
    "        if self.flag_resize:\n",
    "            ori_img_width = image.size[0]\n",
    "            ori_img_height = image.size[1]\n",
    "            image = image.resize((self.img_size,self.img_size),resample=Image.LANCZOS) #LANCZOS is the best\n",
    "      \n",
    "        \n",
    "        \n",
    "        if self.transform: #transforms.ToTensor() convert PIL image into a tensor with shape (channel, height, width)\n",
    "            image = self.transform(image) #after convert to tensor it is channel first (shape: 3,490,680)\n",
    "\n",
    "\n",
    "\n",
    "        return (img_filename, image, (geo_label_x, geo_label_y))\n",
    "    \n",
    "    '''\n",
    "    This function is used to get the id, x, and y of spdcific image in the dataset by its image file name\n",
    "    '''\n",
    "    def get_item_by_imgName(self, imgName):\n",
    "        dataframe = self.annotations\n",
    "        imgID = imgName.strip('.jpg')\n",
    "        out = dataframe[dataframe['id'] == imgID] #out is a dataframe that contain 3 columns: 'id', 'x', 'y'\n",
    "        ID = out.iloc[0,0]\n",
    "        assert ID==imgID\n",
    "        x = out.iloc[0,1]\n",
    "        y = out.iloc[0,2]\n",
    "        return (ID, x, y)\n",
    "        \n",
    "        \n",
    "        \n",
    "'''\n",
    "This dataset is used for test. Because during the test time, we don't have the label (x,y)\n",
    "'''\n",
    "class testDataset(Dataset):\n",
    "    '''\n",
    "    csv_file contain the image name and the label (x,y) of each image\n",
    "        first column: name of the image\n",
    "        second column: label (the x,y coordinations) of the image\n",
    "    img_dir: the path to where the train or test images are stored\n",
    "    flag_resize (bool): a flag, if True, then do the resize. Resize the image into suqare.\n",
    "    \n",
    "    '''\n",
    "    def __init__(self, csv_file, img_dir, flag_resize=True, img_size=224, transform=None, cmap = 'RGB'):\n",
    "        self.annotations = pd.read_csv(csv_file)\n",
    "        self.img_dir = img_dir\n",
    "        self.img_names = self.annotations['id'].to_numpy() #array of object, size: (7500,)\n",
    "        self.flag_resize = flag_resize\n",
    "        self.img_size = img_size\n",
    "        self.transform = transform\n",
    "        self.cmap = cmap\n",
    "    \n",
    "    '''\n",
    "    must have __len__ function\n",
    "    '''\n",
    "    def __len__(self):\n",
    "        return len(self.img_names)\n",
    "        \n",
    "\n",
    "    '''\n",
    "    must have __getitem__ function\n",
    "    to return specific image and the target/label of that image\n",
    "    '''\n",
    "    def __getitem__(self, index):\n",
    "        img_filename = self.img_names[index] + \".jpg\" #output will look like this ('IMG2760_5.jpg',)\n",
    "        img_path = os.path.join(self.img_dir, img_filename)\n",
    "        image = Image.open(img_path).convert(self.cmap) # w, h, channel #image raed by PIL is channel last. shape(490, 680, 3)\n",
    "        #image = io.imread(img_path)\n",
    "        #geo_label_x = torch.tensor(self.annotations.iloc[index,1])\n",
    "        #geo_label_y = torch.tensor(self.annotations.iloc[index,2])\n",
    "        \n",
    "        \n",
    "        #Resize PIL image to square\n",
    "        if self.flag_resize:\n",
    "            ori_img_width = image.size[0]\n",
    "            ori_img_height = image.size[1]\n",
    "            image = image.resize((self.img_size,self.img_size),resample=Image.LANCZOS) #LANCZOS is the best\n",
    "        \n",
    "        if self.transform: #transforms.ToTensor() convert PIL image into a tensor with shape (channel, height, width)\n",
    "            image = self.transform(image) #after convert to tensor it is channel first (shape: 3,490,680)\n",
    "\n",
    "        return (img_filename, image)\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Feature Extraction from ViT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load Image Data\n",
    "import os\n",
    "import pandas as pd \n",
    "import torch\n",
    "from torch.utils.data import Dataset\n",
    "#from skimage import io\n",
    "from PIL import Image\n",
    "import torchvision.transforms as transforms\n",
    "\n",
    "class myDataset(Dataset):\n",
    "    def __init__(self, csv_file, img_dir, flag_resize=True, img_size=224, transform=None):\n",
    "        self.annotations = pd.read_csv(csv_file)\n",
    "        self.img_dir = img_dir\n",
    "        self.img_names = self.annotations['id'].to_numpy() #array of object, size: (7500,)\n",
    "        self.flag_resize = flag_resize\n",
    "        self.img_size = img_size\n",
    "        self.transform = transform\n",
    "    \n",
    "    def __len__(self):\n",
    "        return len(self.img_names)\n",
    "        \n",
    "    def __getitem__(self, index):\n",
    "        img_filename = self.img_names[index] + \".jpg\"\n",
    "        img_path = os.path.join(self.img_dir, img_filename)\n",
    "        image = Image.open(img_path).convert('RGB') # w, h, channel #image raed by PIL is channel last. shape(490, 680, 3)\n",
    "        #image = io.imread(img_path)\n",
    "        \n",
    "#         geo_label_x = torch.tensor(self.annotations.iloc[index,1])\n",
    "#         geo_label_y = torch.tensor(self.annotations.iloc[index,2])\n",
    "        \n",
    "        #Resize PIL image to square\n",
    "        if self.flag_resize:\n",
    "            ori_img_width = image.size[0]\n",
    "            ori_img_height = image.size[1]\n",
    "            image = image.resize((self.img_size,self.img_size),resample=Image.LANCZOS) #LANCZOS is the best\n",
    "        \n",
    "        if self.transform: #transforms.ToTensor() convert PIL image into a tensor with shape (channel, height, width)\n",
    "            image = self.transform(image) #after convert to tensor it is channel first (shape: 3,490,680)\n",
    "\n",
    "        return (img_filename, image) # (geo_label_x, geo_label_y)) \n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "'''\n",
    "This file is the main module of the vision transformer\n",
    "The code is extracted from \n",
    "https://github.com/jankrepl/mildlyoverfitted/blob/master/github_adventures/vision_transformer/custom.py\n",
    "\n",
    "'''\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "class PatchEmbed(nn.Module):\n",
    "    def __init__(self, img_size, patch_size, in_chans=3, embed_dim=1024):\n",
    "        super().__init__()\n",
    "        self.img_size = img_size #assume square image, image_size = the length of the image (int)\n",
    "        self.patch_size = patch_size #assume square, patch_size = the length of the patch (int)\n",
    "        self.n_patches = (img_size // patch_size) ** 2 # total number of patches (flatten)\n",
    "\n",
    "        #kernel will exactly fall into the patches and never overlap.\n",
    "        self.proj = nn.Conv2d(\n",
    "                in_chans,\n",
    "                embed_dim,\n",
    "                kernel_size=patch_size,\n",
    "                stride=patch_size,\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        # feed x to proj layer\n",
    "        x = self.proj(\n",
    "                x\n",
    "            )  # (n_samples, embed_dim, n_patches ** 0.5, n_patches ** 0.5) #64 patch -> 8x8\n",
    "        x = x.flatten(2)  # (n_samples, embed_dim, n_patches)\n",
    "        x = x.transpose(1, 2)  # (n_samples, n_patches, embed_dim) #swap the axis\n",
    "\n",
    "        return x\n",
    "\n",
    "# Below code are just copied from NLP\n",
    "\n",
    "class Attention(nn.Module):\n",
    "    def __init__(self, dim, n_heads=12, qkv_bias=True, attn_p=0., proj_p=0.):\n",
    "        super().__init__()\n",
    "        self.n_heads = n_heads\n",
    "        self.dim = dim\n",
    "        self.head_dim = dim // n_heads\n",
    "        self.scale = self.head_dim ** -0.5 #come from \"attention is all you need paper\" to prevent too large head to dominant\n",
    "\n",
    "        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) #linear mapping, take the embedding token and mapping/split to query, key, value\n",
    "        #input size = (?,?,dim) output size = (?,?,3*dim) #linear mapping applied to the last dimension\n",
    "\n",
    "        self.attn_drop = nn.Dropout(attn_p)\n",
    "        self.proj = nn.Linear(dim, dim) #take the concatenate heads #input size = dim #output size = dim\n",
    "\n",
    "        self.proj_drop = nn.Dropout(proj_p)\n",
    "\n",
    "    def forward(self, x):\n",
    "        n_samples, n_tokens, dim = x.shape\n",
    "\n",
    "        if dim != self.dim:\n",
    "            raise ValueError\n",
    "\n",
    "        qkv = self.qkv(x)  # (n_samples, n_patches + 1, 3 * dim)\n",
    "        \n",
    "        # create extra dimension for heads and q,k,v\n",
    "        qkv = qkv.reshape(\n",
    "                n_samples, n_tokens, 3, self.n_heads, self.head_dim\n",
    "        )  # (n_smaples, n_patches + 1, 3, n_heads, head_dim)\n",
    "        \n",
    "        # change the order\n",
    "        qkv = qkv.permute(\n",
    "                2, 0, 3, 1, 4\n",
    "        )  # (3, n_samples, n_heads, n_patches + 1, head_dim) \n",
    "\n",
    "        q, k, v = qkv[0], qkv[1], qkv[2]\n",
    "        k_t = k.transpose(-2, -1)  # (n_samples, n_heads, head_dim, n_patches + 1) \n",
    "        dp = (\n",
    "           q @ k_t #matrix multiplication\n",
    "        ) * self.scale # (n_samples, n_heads, n_patches + 1, n_patches + 1)\n",
    "        \n",
    "        #apply softmax to get discrete probabilities that some up to 1. This can be use as the weight.\n",
    "        attn = dp.softmax(dim=-1)  # (n_samples, n_heads, n_patches + 1, n_patches + 1) \n",
    "        attn = self.attn_drop(attn) #dropout\n",
    "\n",
    "        weighted_avg = attn @ v  # (n_samples, n_heads, n_patches +1, head_dim) #compute the weighted average\n",
    "        weighted_avg = weighted_avg.transpose(\n",
    "                1, 2\n",
    "        )  # (n_samples, n_patches + 1, n_heads, head_dim)\n",
    "        weighted_avg = weighted_avg.flatten(2)  # (n_samples, n_patches + 1, dim) #flatten(2), start flatten from axis 2\n",
    "        # dim = n_heards x head_dim\n",
    "\n",
    "        x = self.proj(weighted_avg)  # (n_samples, n_patches + 1, dim)\n",
    "        x = self.proj_drop(x)  # (n_samples, n_patches + 1, dim) #dropout\n",
    "\n",
    "        return x\n",
    "\n",
    "\n",
    "class MLP(nn.Module):\n",
    "    def __init__(self, in_features, hidden_features, out_features, p=0.):\n",
    "        super().__init__()\n",
    "        self.fc1 = nn.Linear(in_features, hidden_features)\n",
    "        self.act = nn.GELU()\n",
    "        self.fc2 = nn.Linear(hidden_features, out_features)\n",
    "        self.drop = nn.Dropout(p)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.fc1(x) # (n_samples, n_patches + 1, hidden_features)\n",
    "        x = self.act(x)  # (n_samples, n_patches + 1, hidden_features)\n",
    "        x = self.drop(x)  # (n_samples, n_patches + 1, hidden_features)\n",
    "        x = self.fc2(x)  # (n_samples, n_patches + 1, hidden_features)\n",
    "        x = self.drop(x)  # (n_samples, n_patches + 1, hidden_features) \n",
    "\n",
    "        return x\n",
    "\n",
    "\n",
    "class Block(nn.Module):\n",
    "\n",
    "    def __init__(self, dim, n_heads, mlp_ratio=4.0, qkv_bias=True, p=0., attn_p=0.):\n",
    "        super().__init__()\n",
    "        self.norm1 = nn.LayerNorm(dim, eps=1e-6) #Applies Layer Normalization over a mini-batch of inputs\n",
    "\n",
    "        self.attn = Attention(\n",
    "                dim,\n",
    "                n_heads=n_heads,\n",
    "                qkv_bias=qkv_bias,\n",
    "                attn_p=attn_p,\n",
    "                proj_p=p\n",
    "        )\n",
    "        self.norm2 = nn.LayerNorm(dim, eps=1e-6)\n",
    "        hidden_features = int(dim * mlp_ratio)\n",
    "        #remove the mlp layer\n",
    "        self.mlp = MLP(\n",
    "                in_features=dim,\n",
    "                hidden_features=hidden_features,\n",
    "                out_features=dim,\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "\n",
    "        x = x + self.attn(self.norm1(x)) #residule (add x to output)\n",
    "        x = x + self.mlp(self.norm2(x)) #residule (add x to output)\n",
    "\n",
    "        return x\n",
    "\n",
    "\n",
    "class VisionTransformer(nn.Module):\n",
    "    def __init__(\n",
    "            self,\n",
    "            img_size=384,\n",
    "            patch_size=16,\n",
    "            in_chans=3,\n",
    "            n_classes=1, #output dimension\n",
    "            embed_dim=1024,\n",
    "            depth=12,\n",
    "            n_heads=12,\n",
    "            mlp_ratio=4.,\n",
    "            qkv_bias=True,\n",
    "            p=0.,\n",
    "            attn_p=0.,\n",
    "    ):\n",
    "        super().__init__()\n",
    "\n",
    "        self.patch_embed = PatchEmbed(\n",
    "                img_size=img_size,\n",
    "                patch_size=patch_size,\n",
    "                in_chans=in_chans,\n",
    "                embed_dim=embed_dim,\n",
    "        )\n",
    "        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) #3D\n",
    "        #nn.Parameter is a kind of Tensor that is to be considered a module parameter.\n",
    "\n",
    "        # determine where is exactly the location of the given patch in the image\n",
    "        self.pos_embed = nn.Parameter(\n",
    "                torch.zeros(1, 1 + self.patch_embed.n_patches, embed_dim) #1+ for CLS token\n",
    "        )\n",
    "        self.pos_drop = nn.Dropout(p=p)\n",
    "\n",
    "        self.blocks = nn.ModuleList(\n",
    "            [\n",
    "                Block(\n",
    "                    dim=embed_dim,\n",
    "                    n_heads=n_heads,\n",
    "                    mlp_ratio=mlp_ratio,\n",
    "                    qkv_bias=qkv_bias,\n",
    "                    p=p,\n",
    "                    attn_p=attn_p,\n",
    "                )\n",
    "                for _ in range(depth)\n",
    "            ]\n",
    "        )\n",
    "\n",
    "        self.norm = nn.LayerNorm(embed_dim, eps=1e-6) #normalize over the lasr axis (which size is embed_dim)\n",
    "        self.projEmb = nn.Linear(embed_dim, n_classes) #self.head\n",
    "        #self.sig = nn.Sigmoid() #convert to [0,1]\n",
    "        \n",
    "\n",
    "    def forward(self, x):\n",
    "        \"\"\"Run the forward pass.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        x : torch.Tensor\n",
    "            Shape `(n_samples, in_chans, img_size, img_size)`.\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        logits : torch.Tensor\n",
    "            Logits over all the classes - `(n_samples, n_classes)`.\n",
    "        \"\"\"\n",
    "        n_samples = x.shape[0]\n",
    "        x = self.patch_embed(x) #turn input image x into patch embedding\n",
    "\n",
    "\n",
    "        # learnable cks token\n",
    "        cls_token = self.cls_token.expand(\n",
    "                n_samples, -1, -1\n",
    "        )  # (n_samples, 1, embed_dim)\n",
    "\n",
    "        # prepend to the patch embedding\n",
    "        x = torch.cat((cls_token, x), dim=1)  # (n_samples, 1 + n_patches, embed_dim)\n",
    "        # add the learnable location embedding \n",
    "        x = x + self.pos_embed  # (n_samples, 1 + n_patches, embed_dim)\n",
    "        x = self.pos_drop(x)\n",
    "\n",
    "        for block in self.blocks:\n",
    "            x = block(x) # (n_samples, 1 + n_patches, embed_dim)?\n",
    "\n",
    "        x = self.norm(x) # (n_samples, 1 + n_patches, embed_dim)?\n",
    "\n",
    "        \n",
    "        # for classification task\n",
    "        cls_token_final = x[:, 0] #only take the CLS token to represent the whole image \n",
    "        \n",
    "        x = cls_token_final #shape [batch_size, 768(embed_dim)]\n",
    "        \n",
    "        return x\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "total train image is 7500\n",
      "total test image is 1200\n"
     ]
    }
   ],
   "source": [
    "# load ground true locations\n",
    "train_dic = {}\n",
    "with open('train.csv', 'r') as file:\n",
    "    reader = csv.reader(file)\n",
    "    first = True\n",
    "    for row in reader:\n",
    "        if first:\n",
    "            first = False\n",
    "            continue\n",
    "        train_dic[row[0]+\".jpg\"] = {'x':row[1], 'y':row[2]}\n",
    "print(f'total train image is {len(train_dic)}')\n",
    "        \n",
    "test_dic = {}\n",
    "with open('test.csv', 'r') as file:\n",
    "    reader = csv.reader(file)\n",
    "    first = True\n",
    "    for row in reader:\n",
    "        if first:\n",
    "            first = False\n",
    "            continue\n",
    "        test_dic[row[0]+'.jpg'] = [] \n",
    "print(f'total test image is {len(test_dic)}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set device\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "# load pre-trained model\n",
    "\n",
    "import numpy as np\n",
    "import timm\n",
    "import torch\n",
    "\n",
    "model_name = \"vit_base_patch32_224\" \n",
    "model_official = timm.create_model(model_name, pretrained=True).to(device)\n",
    "model_official.eval() #switch to eval mode\n",
    "print(type(model_official))\n",
    "\n",
    "custom_config = {\n",
    "        \"img_size\": 224, #384\n",
    "        \"in_chans\": 3,\n",
    "        \"patch_size\": 32,\n",
    "        \"embed_dim\": 768,#1024 for large #768 for base #192 for tiny\n",
    "        \"depth\": 12, #24 for large #12 for base #12 for tiny\n",
    "        \"n_heads\": 12, #16 for large #12 for base #3 for tiny, #6 for small\n",
    "        \"qkv_bias\": True,\n",
    "        \"mlp_ratio\": 4,\n",
    "        \"n_classes\": 1\n",
    "}\n",
    "\n",
    "model_custom = VisionTransformer(**custom_config).to(device)\n",
    "model_custom.eval() #set to evaluation mode"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "collapsed": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cls_token\n",
      "pos_embed\n",
      "patch_embed.proj.weight\n",
      "patch_embed.proj.bias\n",
      "blocks.0.norm1.weight\n",
      "blocks.0.norm1.bias\n",
      "blocks.0.attn.qkv.weight\n",
      "blocks.0.attn.qkv.bias\n",
      "blocks.0.attn.proj.weight\n",
      "blocks.0.attn.proj.bias\n",
      "blocks.0.norm2.weight\n",
      "blocks.0.norm2.bias\n",
      "blocks.0.mlp.fc1.weight\n",
      "blocks.0.mlp.fc1.bias\n",
      "blocks.0.mlp.fc2.weight\n",
      "blocks.0.mlp.fc2.bias\n",
      "blocks.1.norm1.weight\n",
      "blocks.1.norm1.bias\n",
      "blocks.1.attn.qkv.weight\n",
      "blocks.1.attn.qkv.bias\n",
      "blocks.1.attn.proj.weight\n",
      "blocks.1.attn.proj.bias\n",
      "blocks.1.norm2.weight\n",
      "blocks.1.norm2.bias\n",
      "blocks.1.mlp.fc1.weight\n",
      "blocks.1.mlp.fc1.bias\n",
      "blocks.1.mlp.fc2.weight\n",
      "blocks.1.mlp.fc2.bias\n",
      "blocks.2.norm1.weight\n",
      "blocks.2.norm1.bias\n",
      "blocks.2.attn.qkv.weight\n",
      "blocks.2.attn.qkv.bias\n",
      "blocks.2.attn.proj.weight\n",
      "blocks.2.attn.proj.bias\n",
      "blocks.2.norm2.weight\n",
      "blocks.2.norm2.bias\n",
      "blocks.2.mlp.fc1.weight\n",
      "blocks.2.mlp.fc1.bias\n",
      "blocks.2.mlp.fc2.weight\n",
      "blocks.2.mlp.fc2.bias\n",
      "blocks.3.norm1.weight\n",
      "blocks.3.norm1.bias\n",
      "blocks.3.attn.qkv.weight\n",
      "blocks.3.attn.qkv.bias\n",
      "blocks.3.attn.proj.weight\n",
      "blocks.3.attn.proj.bias\n",
      "blocks.3.norm2.weight\n",
      "blocks.3.norm2.bias\n",
      "blocks.3.mlp.fc1.weight\n",
      "blocks.3.mlp.fc1.bias\n",
      "blocks.3.mlp.fc2.weight\n",
      "blocks.3.mlp.fc2.bias\n",
      "blocks.4.norm1.weight\n",
      "blocks.4.norm1.bias\n",
      "blocks.4.attn.qkv.weight\n",
      "blocks.4.attn.qkv.bias\n",
      "blocks.4.attn.proj.weight\n",
      "blocks.4.attn.proj.bias\n",
      "blocks.4.norm2.weight\n",
      "blocks.4.norm2.bias\n",
      "blocks.4.mlp.fc1.weight\n",
      "blocks.4.mlp.fc1.bias\n",
      "blocks.4.mlp.fc2.weight\n",
      "blocks.4.mlp.fc2.bias\n",
      "blocks.5.norm1.weight\n",
      "blocks.5.norm1.bias\n",
      "blocks.5.attn.qkv.weight\n",
      "blocks.5.attn.qkv.bias\n",
      "blocks.5.attn.proj.weight\n",
      "blocks.5.attn.proj.bias\n",
      "blocks.5.norm2.weight\n",
      "blocks.5.norm2.bias\n",
      "blocks.5.mlp.fc1.weight\n",
      "blocks.5.mlp.fc1.bias\n",
      "blocks.5.mlp.fc2.weight\n",
      "blocks.5.mlp.fc2.bias\n",
      "blocks.6.norm1.weight\n",
      "blocks.6.norm1.bias\n",
      "blocks.6.attn.qkv.weight\n",
      "blocks.6.attn.qkv.bias\n",
      "blocks.6.attn.proj.weight\n",
      "blocks.6.attn.proj.bias\n",
      "blocks.6.norm2.weight\n",
      "blocks.6.norm2.bias\n",
      "blocks.6.mlp.fc1.weight\n",
      "blocks.6.mlp.fc1.bias\n",
      "blocks.6.mlp.fc2.weight\n",
      "blocks.6.mlp.fc2.bias\n",
      "blocks.7.norm1.weight\n",
      "blocks.7.norm1.bias\n",
      "blocks.7.attn.qkv.weight\n",
      "blocks.7.attn.qkv.bias\n",
      "blocks.7.attn.proj.weight\n",
      "blocks.7.attn.proj.bias\n",
      "blocks.7.norm2.weight\n",
      "blocks.7.norm2.bias\n",
      "blocks.7.mlp.fc1.weight\n",
      "blocks.7.mlp.fc1.bias\n",
      "blocks.7.mlp.fc2.weight\n",
      "blocks.7.mlp.fc2.bias\n",
      "blocks.8.norm1.weight\n",
      "blocks.8.norm1.bias\n",
      "blocks.8.attn.qkv.weight\n",
      "blocks.8.attn.qkv.bias\n",
      "blocks.8.attn.proj.weight\n",
      "blocks.8.attn.proj.bias\n",
      "blocks.8.norm2.weight\n",
      "blocks.8.norm2.bias\n",
      "blocks.8.mlp.fc1.weight\n",
      "blocks.8.mlp.fc1.bias\n",
      "blocks.8.mlp.fc2.weight\n",
      "blocks.8.mlp.fc2.bias\n",
      "blocks.9.norm1.weight\n",
      "blocks.9.norm1.bias\n",
      "blocks.9.attn.qkv.weight\n",
      "blocks.9.attn.qkv.bias\n",
      "blocks.9.attn.proj.weight\n",
      "blocks.9.attn.proj.bias\n",
      "blocks.9.norm2.weight\n",
      "blocks.9.norm2.bias\n",
      "blocks.9.mlp.fc1.weight\n",
      "blocks.9.mlp.fc1.bias\n",
      "blocks.9.mlp.fc2.weight\n",
      "blocks.9.mlp.fc2.bias\n",
      "blocks.10.norm1.weight\n",
      "blocks.10.norm1.bias\n",
      "blocks.10.attn.qkv.weight\n",
      "blocks.10.attn.qkv.bias\n",
      "blocks.10.attn.proj.weight\n",
      "blocks.10.attn.proj.bias\n",
      "blocks.10.norm2.weight\n",
      "blocks.10.norm2.bias\n",
      "blocks.10.mlp.fc1.weight\n",
      "blocks.10.mlp.fc1.bias\n",
      "blocks.10.mlp.fc2.weight\n",
      "blocks.10.mlp.fc2.bias\n",
      "blocks.11.norm1.weight\n",
      "blocks.11.norm1.bias\n",
      "blocks.11.attn.qkv.weight\n",
      "blocks.11.attn.qkv.bias\n",
      "blocks.11.attn.proj.weight\n",
      "blocks.11.attn.proj.bias\n",
      "blocks.11.norm2.weight\n",
      "blocks.11.norm2.bias\n",
      "blocks.11.mlp.fc1.weight\n",
      "blocks.11.mlp.fc1.bias\n",
      "blocks.11.mlp.fc2.weight\n",
      "blocks.11.mlp.fc2.bias\n",
      "norm.weight\n",
      "norm.bias\n"
     ]
    }
   ],
   "source": [
    "# check if the pre-trained model is correctly loaded\n",
    "#this function count the number of learnable parameters\n",
    "def get_n_params(module):\n",
    "    return sum(p.numel() for p in module.parameters() if p.requires_grad)\n",
    "\n",
    "# compare whethere 2 tensor are equal\n",
    "def assert_tensors_equal(t1, t2):\n",
    "    a1, a2 = t1.detach().numpy(), t2.detach().numpy()\n",
    "\n",
    "    np.testing.assert_allclose(a1, a2)\n",
    "#==============================================================================\n",
    "\n",
    "# copy the pretrained weighted into our customer model\n",
    "for n_o, p_o in model_official.named_parameters():\n",
    "    for n_c, p_c in model_custom.named_parameters():\n",
    "        if n_o == n_c:\n",
    "            print(n_o)\n",
    "            p_c.data[:] = p_o.data\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# save the model\n",
    "torch.save(model_custom, \"model_custom_with_pretrained.pth\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "img_transform = transforms.Compose([\n",
    "        transforms.ToTensor(), #transform image from [0,1] PIL Image into a tensor\n",
    "        transforms.Normalize(mean=torch.tensor([0.5000, 0.5000, 0.5000]), std=torch.tensor([0.5000, 0.5000, 0.5000])) #transform image into [-1,1] #(three value for there channel)\n",
    "])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset = myDataset(csv_file = 'train.csv', img_dir = 'train_img', flag_resize=True, img_size = 224, transform = img_transform)\n",
    "batch_size = 1\n",
    "train_loader = DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=False)\n",
    "\n",
    "test_dataset = myDataset(csv_file = 'test.csv', img_dir = 'test_img', flag_resize=True, img_size = 224, transform = img_transform)\n",
    "batch_size = 1\n",
    "test_loader = DataLoader(dataset=test_dataset, batch_size=batch_size, shuffle=False)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# train_dataset is the whole 7500 image in \"COMP90086_2021_Project_train\"\n",
    "# del_set and valid_set is for develop feature matching algorithm purpose!\n",
    "# cmap=\"L\" meas read image as grayscale\n",
    "train_dataset_original_size = trainDataset(csv_file = 'train.csv', img_dir = 'train_img', flag_resize=False, \n",
    "                                           img_size = None, transform = transforms.ToTensor(), cmap='L') #return image is channel first, value between [0,1]\n",
    "# divide training dataset into validation and develop dataset (10:90)\n",
    "dev_set, valid_set = torch.utils.data.random_split(train_dataset_original_size, [6750, 750]) #10% for validation\n",
    "\n",
    "test_dataset = testDataset(csv_file = 'test.csv', img_dir = 'test_img', flag_resize=False, img_size = None, transform = transforms.ToTensor(), cmap='L')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 750/750 [00:15<00:00, 47.19it/s]\n"
     ]
    }
   ],
   "source": [
    "# save the images\n",
    "all_image_names_in_valid_set = []\n",
    "for img_name, _, _ in tqdm(valid_set):\n",
    "    all_image_names_in_valid_set.append(img_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 6750/6750 [02:30<00:00, 44.93it/s]\n"
     ]
    }
   ],
   "source": [
    "# save the images\n",
    "all_image_names_in_dev_set = []\n",
    "for img_name, _, _ in tqdm(dev_set):\n",
    "    all_image_names_in_dev_set.append(img_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Cosine Similarity of image features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "collapsed": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 88%|████████▊ | 6635/7500 [06:21<00:49, 17.40it/s]\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-15-4b3319602602>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[0mpred_all\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mdict\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 2\u001b[0;31m \u001b[0;32mfor\u001b[0m \u001b[0mbatch_idx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mfilename\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mimg\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32min\u001b[0m \u001b[0menumerate\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtqdm\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtrain_loader\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      3\u001b[0m     \u001b[0;32mwith\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mno_grad\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      4\u001b[0m         \u001b[0mimg\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mimg\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdevice\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mdevice\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      5\u001b[0m         \u001b[0mout\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmodel_custom\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mimg\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.7/site-packages/tqdm/std.py\u001b[0m in \u001b[0;36m__iter__\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m   1183\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1184\u001b[0m         \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1185\u001b[0;31m             \u001b[0;32mfor\u001b[0m \u001b[0mobj\u001b[0m \u001b[0;32min\u001b[0m \u001b[0miterable\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1186\u001b[0m                 \u001b[0;32myield\u001b[0m \u001b[0mobj\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1187\u001b[0m                 \u001b[0;31m# Update and possibly print the progressbar.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.7/site-packages/torch/utils/data/dataloader.py\u001b[0m in \u001b[0;36m__next__\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    519\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_sampler_iter\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    520\u001b[0m                 \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_reset\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 521\u001b[0;31m             \u001b[0mdata\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_next_data\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    522\u001b[0m             \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_num_yielded\u001b[0m \u001b[0;34m+=\u001b[0m \u001b[0;36m1\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    523\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_dataset_kind\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0m_DatasetKind\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mIterable\u001b[0m \u001b[0;32mand\u001b[0m\u001b[0;31m \u001b[0m\u001b[0;31m\\\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.7/site-packages/torch/utils/data/dataloader.py\u001b[0m in \u001b[0;36m_next_data\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    559\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m_next_data\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    560\u001b[0m         \u001b[0mindex\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_next_index\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# may raise StopIteration\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 561\u001b[0;31m         \u001b[0mdata\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_dataset_fetcher\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfetch\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mindex\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# may raise StopIteration\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    562\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_pin_memory\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    563\u001b[0m             \u001b[0mdata\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0m_utils\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpin_memory\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpin_memory\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdata\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.7/site-packages/torch/utils/data/_utils/fetch.py\u001b[0m in \u001b[0;36mfetch\u001b[0;34m(self, possibly_batched_index)\u001b[0m\n\u001b[1;32m     42\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mfetch\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mpossibly_batched_index\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     43\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mauto_collation\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 44\u001b[0;31m             \u001b[0mdata\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdataset\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0midx\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0midx\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mpossibly_batched_index\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     45\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     46\u001b[0m             \u001b[0mdata\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdataset\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mpossibly_batched_index\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.7/site-packages/torch/utils/data/_utils/fetch.py\u001b[0m in \u001b[0;36m<listcomp>\u001b[0;34m(.0)\u001b[0m\n\u001b[1;32m     42\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mfetch\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mpossibly_batched_index\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     43\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mauto_collation\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 44\u001b[0;31m             \u001b[0mdata\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdataset\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0midx\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0midx\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mpossibly_batched_index\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     45\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     46\u001b[0m             \u001b[0mdata\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdataset\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mpossibly_batched_index\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m<ipython-input-3-610630d10359>\u001b[0m in \u001b[0;36m__getitem__\u001b[0;34m(self, index)\u001b[0m\n\u001b[1;32m     36\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     37\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtransform\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0;31m#transforms.ToTensor() convert PIL image into a tensor with shape (channel, height, width)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 38\u001b[0;31m             \u001b[0mimage\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtransform\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mimage\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;31m#after convert to tensor it is channel first (shape: 3,490,680)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     39\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     40\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mimg_filename\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mimage\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;31m# (geo_label_x, geo_label_y))\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.7/site-packages/torchvision/transforms/transforms.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, img)\u001b[0m\n\u001b[1;32m     58\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m__call__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mimg\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     59\u001b[0m         \u001b[0;32mfor\u001b[0m \u001b[0mt\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtransforms\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 60\u001b[0;31m             \u001b[0mimg\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mt\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mimg\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     61\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mimg\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     62\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.7/site-packages/torchvision/transforms/transforms.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, pic)\u001b[0m\n\u001b[1;32m     95\u001b[0m             \u001b[0mTensor\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mConverted\u001b[0m \u001b[0mimage\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     96\u001b[0m         \"\"\"\n\u001b[0;32m---> 97\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0mF\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto_tensor\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpic\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     98\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     99\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m__repr__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.7/site-packages/torchvision/transforms/functional.py\u001b[0m in \u001b[0;36mto_tensor\u001b[0;34m(pic)\u001b[0m\n\u001b[1;32m    134\u001b[0m     \u001b[0mimg\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mimg\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mview\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpic\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msize\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mpic\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msize\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mlen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpic\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mgetbands\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    135\u001b[0m     \u001b[0;31m# put it from HWC to CHW format\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 136\u001b[0;31m     \u001b[0mimg\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mimg\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpermute\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m2\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m0\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcontiguous\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    137\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0misinstance\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mimg\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mByteTensor\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    138\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mimg\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdtype\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mdefault_float_dtype\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdiv\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m255\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "\n",
    "pred_all = dict()\n",
    "for batch_idx, (filename, img) in enumerate(tqdm(train_loader)):\n",
    "    with torch.no_grad():\n",
    "        img = img.to(device=device)\n",
    "        out = model_custom(img)\n",
    "        filename = filename[0]\n",
    "        pred_all[filename] = out\n",
    "        \n",
    "pred_test = dict()\n",
    "for batch_idx, (filename, img) in enumerate(tqdm(test_loader)):\n",
    "    with torch.no_grad():\n",
    "        img = img.to(device=device)\n",
    "        out = model_custom(img)\n",
    "        filename = filename[0]\n",
    "        pred_test[filename] = out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "# save the similarity into dictionary format\n",
    "from tqdm import tqdm \n",
    "\n",
    "cos = nn.CosineSimilarity(dim=1, eps=1e-6)\n",
    "sim_dic = dict()\n",
    "\n",
    "for filename1, output1 in tqdm(pred_all.items()):\n",
    "    filename1 = filename1[0]\n",
    "    sim_dic[filename1] = dict()\n",
    "    for filename2, output2 in pred_all.items():\n",
    "        filename2 = filename2[0]\n",
    "        if filename1 != filename2:\n",
    "            sim = cos(output1,output2)\n",
    "            sim_dic[filename1][filename2] = sim.item()\n",
    "            \n",
    "for key in sim_dic.keys():\n",
    "    sim_dic[key] = dict(sorted(sim_dic[key].items(), key=lambda item: item[1], reverse=True))\n",
    "\n",
    "\n",
    "\n",
    "cos = nn.CosineSimilarity(dim=1, eps=1e-6)\n",
    "test_sim_dic = dict()\n",
    "\n",
    "for filename1, output1 in tqdm(pred_test.items()):\n",
    "    filename1 = filename1\n",
    "    test_sim_dic[filename1] = dict()\n",
    "    for filename2, output2 in pred_all.items():\n",
    "        filename2 = filename2\n",
    "        if filename1 != filename2:\n",
    "            sim = cos(output1,output2)\n",
    "            test_sim_dic[filename1][filename2] = sim.item()\n",
    "            \n",
    "for key in test_sim_dic.keys():\n",
    "    test_sim_dic[key] = dict(sorted(test_sim_dic[key].items(), key=lambda item: item[1], reverse=True))\n",
    "  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export the similarity result into json file\n",
    "with open('train_sim_dic.json', 'w') as file:\n",
    "    file.write(json.dumps(sim_dic2))\n",
    "\n",
    "with open('test_sim_dic.json', 'w') as file:\n",
    "    file.write(json.dumps(test_sim_dic))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load the similarity result from json file\n",
    "f = open('train_sim_dic.json',)\n",
    "sim_dic = json.load(f)\n",
    "\n",
    "f = open('test_sim_dic.json',)\n",
    "test_sim_dic = json.load(f)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Build Method 1\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Feature (Affine Matrix) extraction for Linear Regression"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "# install kornia if not downloaded\n",
    "# pip install kornia"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 84,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/easybuild-2019/easybuild/software/mpi/gcc-cuda/8.3.0-10.1.243/openmpi/3.1.4/jupyter/1.0.0-python-3.7.4/lib/python3.7/site-packages/ipykernel/ipkernel.py:287: DeprecationWarning: `should_run_async` will not call `transform_cell` automatically in the future. Please pass the result to `transformed_cell` argument and any exception that happen during thetransform in `preprocessing_exc_tuple` in IPython 7.17 and above.\n",
      "  and should_run_async(code)\n"
     ]
    }
   ],
   "source": [
    "'''\n",
    "Input:\n",
    "    img0: the source image. Must be a grayscale image, channel first\n",
    "    img1: the target image to match. Must be a grayscale image, channel first\n",
    "        The input image with the shape (1,h,w)\n",
    "    threshold (float): a value in [0,1]. It is the threshold to filter bad matching points\n",
    "'''\n",
    "# load the pre-trained LoFTR for point matching and setup\n",
    "from kornia.feature import LoFTR\n",
    "loftr = LoFTR('indoor')\n",
    "\n",
    "# get matching points by LoFTR model\n",
    "def get_matched_points_by_loftr(img0, img1, threshold =0.5):\n",
    "    # Add batch axis\n",
    "    img0 = img0.unsqueeze(0)\n",
    "    img1 = img1.unsqueeze(0)\n",
    "    \n",
    "    input = {\"image0\": img0, \"image1\": img1}\n",
    "    out = loftr(input)\n",
    "    # out is a dict with keys: dict_keys(['keypoints0', 'keypoints1', 'confidence', 'batch_indexes'])\n",
    "    pts_img0 = out['keypoints0']\n",
    "    pts_img1 = out['keypoints1']\n",
    "    confidence = out['confidence']\n",
    "    filter = (confidence>threshold)\n",
    "    good_pts_img0 = pts_img0[filter]\n",
    "    good_pts_img1 = pts_img1[filter]\n",
    "    good_confidence = confidence[filter]\n",
    "    return good_pts_img0, good_pts_img1, good_confidence\n",
    "\n",
    "# get the top N similar image from the input dictionary\n",
    "def find_top_similar_imgs(img, n_top, dic):\n",
    "    cnt = 0\n",
    "    top_imgs = []\n",
    "    for img2 in dic[img].keys():\n",
    "        top_imgs.append(img2)\n",
    "        cnt += 1\n",
    "        if cnt == n_top:\n",
    "            break\n",
    "    return top_imgs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 7500/7500 [01:28<00:00, 84.55it/s]\n",
      "100%|██████████| 1200/1200 [00:14<00:00, 85.11it/s]\n"
     ]
    }
   ],
   "source": [
    "# load the image information into dictionary format\n",
    "dev_set_dic = dict()\n",
    "for dev in tqdm(dev_set):\n",
    "    img1_name, img1, labels = dev\n",
    "    dev_set_dic[img1_name] = {'value': img1, 'label': labels}\n",
    "    \n",
    "val_set_dic = dict()\n",
    "for val in tqdm(valid_set):\n",
    "    img1_name, img1, labels = val\n",
    "    val_set_dic[img1_name] = {'value': img1, 'label': labels}\n",
    "\n",
    "all_set_dic = dict()\n",
    "for data in tqdm(train_dataset_original_size):\n",
    "    img1_name, img1, labels = data\n",
    "    all_set_dic[img1_name] = {'value': img1, 'label': labels}\n",
    "    \n",
    "test_set_dic = dict()\n",
    "for data in tqdm(test_dataset):\n",
    "    img1_name, img1 = data\n",
    "    test_set_dic[img1_name] = {'value': img1}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 198,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/easybuild-2019/easybuild/software/mpi/gcc-cuda/8.3.0-10.1.243/openmpi/3.1.4/jupyter/1.0.0-python-3.7.4/lib/python3.7/site-packages/ipykernel/ipkernel.py:287: DeprecationWarning: `should_run_async` will not call `transform_cell` automatically in the future. Please pass the result to `transformed_cell` argument and any exception that happen during thetransform in `preprocessing_exc_tuple` in IPython 7.17 and above.\n",
      "  and should_run_async(code)\n",
      " 32%|███▏      | 2434/7500 [10:04:04<41:01:05, 29.15s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "here 1\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 7500/7500 [31:02:34<00:00, 14.90s/it]   \n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'Hailey'"
      ]
     },
     "execution_count": 198,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from tqdm import tqdm\n",
    "from numba import jit, cuda\n",
    "import scipy.stats as ss\n",
    "import tensorflow as tf\n",
    "import time\n",
    "\n",
    "# get affine matrix, and others from training images\n",
    "for img1_name in tqdm(all_set_dic.keys()):\n",
    "    max_num_pts = 0\n",
    "    best_label, best_pts1, best_pts2, best_confidence = None, [], [], None\n",
    "    Affine_M, Essential_M, R_est, t_est = None, None, None, None\n",
    "    img1, label1 = all_set_dic[img1_name]['value'], all_set_dic[img1_name]['label']\n",
    "\n",
    "    selected_imgs_10 = find_top_similar_imgs(img1_name, n_top = 10, dic = sim_dic)\n",
    "    loftr_train_dic[img1_name] = dict()\n",
    "    \n",
    "    t = 0.7\n",
    "    s, e = 0, 10 # search the top 10 matching images\n",
    "    while True:\n",
    "        for img2_name in selected_imgs_10[s:e]:   # params to avoid tqdm mutiple printing\n",
    "            img2 = all_set_dic[img2_name]['value']\n",
    "            pts1, pts2, confidence = get_matched_points_by_loftr(img1, img2, threshold = t)\n",
    "            \n",
    "            if len(pts1) > max_num_pts:\n",
    "                max_num_pts = len(pts1)\n",
    "                best_img_name = img2_name\n",
    "                best_label = all_set_dic[img2_name]['label']\n",
    "                best_pts1 = pts1\n",
    "                best_pts2 = pts2\n",
    "                best_confidence = confidence\n",
    "\n",
    "        if (max_num_pts <=5 and t == 0.7):\n",
    "            t = 0\n",
    "#             s, e = 3, 10 \n",
    "        else:\n",
    "            break\n",
    "    \n",
    "    # if there is no matching point, record the image and save empty values\n",
    "    if max_num_pts == 0 :\n",
    "        train_skipped.append(img1_name)\n",
    "        print('here', len(train_skipped))\n",
    "        loftr_train_dic[img1_name][best_img_name] = {'Affine_M': Affine_M, 'Essential_M': Essential_M, \n",
    "                                                     'R_est': R_est, 't_est': t_est,\n",
    "                                                     'pts1':best_pts1, 'pts2': best_pts2, \n",
    "                                                     'confidence': best_confidence}\n",
    "    else:\n",
    "        if len(best_pts1) >= 6:\n",
    "            e_src = np.float32(best_pts1)\n",
    "            e_dst = np.float32(best_pts2)\n",
    "        \n",
    "            Essential_M, mask = cv2.findEssentialMat(e_src, e_dst)  # occurs error with only 5 points!!!\n",
    "            points, R_est, t_est, mask_pose = cv2.recoverPose(Essential_M, e_src, e_dst)  # R_est: Rotation & t_est: transformation\n",
    "\n",
    "        if len(best_pts1) >= 3:\n",
    "            best3_pts1, best3_pts2 = zip(*random.sample(list(zip(best_pts1, best_pts2)), 3))\n",
    "            best3_pts1, best3_pts2 = torch.stack(best3_pts1), torch.stack(best3_pts2)\n",
    "\n",
    "            a_src = np.float32(best3_pts1)\n",
    "            a_dst = np.float32(best3_pts2)\n",
    "            Affine_M = cv2.getAffineTransform(a_src, a_dst)\n",
    "            \n",
    "        loftr_train_dic[img1_name][best_img_name] = {'Affine_M': Affine_M, 'Essential_M': Essential_M, \n",
    "                                                     'R_est': R_est, 't_est': t_est,\n",
    "                                                     'pts1':best_pts1, 'pts2': best_pts2, \n",
    "                                                     'confidence': best_confidence}\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 247,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/easybuild-2019/easybuild/software/mpi/gcc-cuda/8.3.0-10.1.243/openmpi/3.1.4/jupyter/1.0.0-python-3.7.4/lib/python3.7/site-packages/ipykernel/ipkernel.py:287: DeprecationWarning: `should_run_async` will not call `transform_cell` automatically in the future. Please pass the result to `transformed_cell` argument and any exception that happen during thetransform in `preprocessing_exc_tuple` in IPython 7.17 and above.\n",
      "  and should_run_async(code)\n",
      "100%|██████████| 7500/7500 [00:12<00:00, 585.20it/s] \n"
     ]
    }
   ],
   "source": [
    "# save the result into a dictionary and export it\n",
    "for k1 in tqdm(loftr_train_dic.keys()):\n",
    "    for k2 in loftr_train_dic[k1].keys():\n",
    "        for k3 in loftr_train_dic[k1][k2].keys():\n",
    "            a = loftr_train_dic[k1][k2][k3]\n",
    "            if type(a) is not list and a is not None:\n",
    "                loftr_train_dic[k1][k2][k3] = loftr_train_dic[k1][k2][k3].tolist()\n",
    "                \n",
    "with open('loftr_train_dic.json', 'w') as file:\n",
    "    file.write(json.dumps(loftr_train_dic))\n",
    "    \n",
    "# load saved file if needed    \n",
    "f = open('loftr_train_dic.json',)\n",
    "loftr_train_dic = json.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get the features and labels (difference of x, y coordinates) for linear regression model \n",
    "X, Y = [] , [] \n",
    "real_Y, match_Y = [], []\n",
    "valid_keys = ['Affine_M','R_est',  't_est']\n",
    "Affine_M_empty = [0, 0, 0, 0, 0, 0]\n",
    "R_est_empty = [0, 0, 0,0, 0, 0,0, 0, 0]\n",
    "t_est_empty = [0, 0, 0]\n",
    "for k1 in loftr_train_dic.keys():\n",
    "    for k2 in loftr_train_dic[k1].keys():\n",
    "        x = []\n",
    "        for k3 in valid_keys:\n",
    "            a = loftr_train_dic[k1][k2][k3]\n",
    "            if a != None:\n",
    "                x.extend(np.array(a).reshape(-1))\n",
    "            else:\n",
    "                if k3 == 'Affine_M':\n",
    "                    x.extend(Affine_M_empty)\n",
    "                elif k3 == 'R_est':\n",
    "                    x.extend(R_est_empty)\n",
    "                elif k3 == 't_est':\n",
    "                    x.extend(t_est_empty)\n",
    "                else:\n",
    "                    print(k3)\n",
    "                    print(a)\n",
    "        \n",
    "    diff_x = all_set_dic[k2]['label'][0].item() - all_set_dic[k1]['label'][0].item()\n",
    "    diff_y = all_set_dic[k2]['label'][1].item() - all_set_dic[k1]['label'][1].item()\n",
    "    diff_xy = [diff_x, diff_y]\n",
    "    X.append(x)\n",
    "    Y.append(diff_xy)\n",
    "    real_Y.append([all_set_dic[k1]['label'][0].item(), all_set_dic[k1]['label'][1].item()])\n",
    "    match_Y.append([all_set_dic[k2]['label'][0].item(), all_set_dic[k2]['label'][1].item()])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 250,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/easybuild-2019/easybuild/software/mpi/gcc-cuda/8.3.0-10.1.243/openmpi/3.1.4/jupyter/1.0.0-python-3.7.4/lib/python3.7/site-packages/ipykernel/ipkernel.py:287: DeprecationWarning: `should_run_async` will not call `transform_cell` automatically in the future. Please pass the result to `transformed_cell` argument and any exception that happen during thetransform in `preprocessing_exc_tuple` in IPython 7.17 and above.\n",
      "  and should_run_async(code)\n",
      "100%|██████████| 1200/1200 [5:09:25<00:00, 15.47s/it]  \n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'Hailey'"
      ]
     },
     "execution_count": 250,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# get affine matrix, and others from test images\n",
    "\n",
    "for img1_name in tqdm(test_set_dic.keys()):\n",
    "    max_num_pts = 0\n",
    "    best_label, best_pts1, best_pts2, best_confidence = None, [], [], None\n",
    "    Affine_M, Essential_M, R_est, t_est = None, None, None, None\n",
    "    img1 = test_set_dic[img1_name]['value'] \n",
    "\n",
    "    selected_imgs_10 = find_top_similar_imgs(img1_name, n_top = 10, dic = test_sim_dic)\n",
    "    loftr_test_dic[img1_name] = dict()\n",
    "    \n",
    "    t = 0.7\n",
    "    s, e = 0, 10 # search the top 10 matching images\n",
    "    while True:\n",
    "        for img2_name in selected_imgs_10[s:e]:   # params to avoid tqdm mutiple printing\n",
    "            img2 = all_set_dic[img2_name]['value']\n",
    "            pts1, pts2, confidence = get_matched_points_by_loftr(img1, img2, threshold = t)\n",
    "            \n",
    "            if len(pts1) > max_num_pts:\n",
    "                max_num_pts = len(pts1)\n",
    "                best_img_name = img2_name\n",
    "                best_label = all_set_dic[img2_name]['label']\n",
    "                best_pts1 = pts1\n",
    "                best_pts2 = pts2\n",
    "                best_confidence = confidence\n",
    "\n",
    "        if (max_num_pts <=5 and t == 0.7):\n",
    "            t = 0\n",
    "#             s, e = 3, 10\n",
    "        else:\n",
    "            break\n",
    "    \n",
    "\n",
    "    if max_num_pts == 0 :\n",
    "        test_skipped.append(img1_name)\n",
    "        print('here', len(test_skipped))\n",
    "        loftr_test_dic[img1_name][best_img_name] = {'Affine_M': Affine_M, 'Essential_M': Essential_M, \n",
    "                                                     'R_est': R_est, 't_est': t_est,\n",
    "                                                     'pts1':best_pts1, 'pts2': best_pts2, \n",
    "                                                     'confidence': best_confidence}\n",
    "    else:\n",
    "        if len(best_pts1) >= 6:\n",
    "            e_src = np.float32(best_pts1)\n",
    "            e_dst = np.float32(best_pts2)\n",
    "        \n",
    "            Essential_M, mask = cv2.findEssentialMat(e_src, e_dst)  # occurs error with only 5 points!!!\n",
    "            points, R_est, t_est, mask_pose = cv2.recoverPose(Essential_M, e_src, e_dst)  # R_est: Rotation & t_est: transformation\n",
    "\n",
    "        if len(best_pts1) >= 3:\n",
    "            best3_pts1, best3_pts2 = zip(*random.sample(list(zip(best_pts1, best_pts2)), 3))\n",
    "            best3_pts1, best3_pts2 = torch.stack(best3_pts1), torch.stack(best3_pts2)\n",
    "\n",
    "            a_src = np.float32(best3_pts1)\n",
    "            a_dst = np.float32(best3_pts2)\n",
    "            Affine_M = cv2.getAffineTransform(a_src, a_dst)\n",
    "            \n",
    "        loftr_test_dic[img1_name][best_img_name] = {'Affine_M': Affine_M, 'Essential_M': Essential_M, \n",
    "                                                     'R_est': R_est, 't_est': t_est,\n",
    "                                                     'pts1':best_pts1, 'pts2': best_pts2, \n",
    "                                                     'confidence': best_confidence}\n",
    "'''Hailey'''\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 251,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/easybuild-2019/easybuild/software/mpi/gcc-cuda/8.3.0-10.1.243/openmpi/3.1.4/jupyter/1.0.0-python-3.7.4/lib/python3.7/site-packages/ipykernel/ipkernel.py:287: DeprecationWarning: `should_run_async` will not call `transform_cell` automatically in the future. Please pass the result to `transformed_cell` argument and any exception that happen during thetransform in `preprocessing_exc_tuple` in IPython 7.17 and above.\n",
      "  and should_run_async(code)\n",
      "100%|██████████| 1200/1200 [00:00<00:00, 1449.18it/s]\n"
     ]
    }
   ],
   "source": [
    "# save the result into a dictionary and export it\n",
    "for k1 in tqdm(loftr_test_dic.keys()):\n",
    "    for k2 in loftr_test_dic[k1].keys():\n",
    "        for k3 in loftr_test_dic[k1][k2].keys():\n",
    "            a = loftr_test_dic[k1][k2][k3]\n",
    "            if type(a) is not list and a is not None:\n",
    "                loftr_test_dic[k1][k2][k3] = loftr_test_dic[k1][k2][k3].tolist()\n",
    "\n",
    "with open('loftr_test_dic.json', 'w') as file:\n",
    "    file.write(json.dumps(loftr_test_dic))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Build LR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mean Absolute Error:  -5.55179654706102\n",
      "Standard Deviation:  0.45471223183850656\n"
     ]
    }
   ],
   "source": [
    "from sklearn.linear_model import LinearRegression\n",
    "from sklearn.model_selection import cross_val_score\n",
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "x_train, x_test, y_train, y_test = train_test_split(X, Y, test_size = 0.1, random_state = 10)\n",
    "# x_train, y_train = X, real_Y\n",
    "lr = LinearRegression()\n",
    "lr.fit(x_train,y_train)\n",
    "\n",
    "scoring = \"neg_mean_absolute_error\"\n",
    "results = cross_val_score(lr, x_train, y_train, scoring=scoring, n_jobs = 4)  # 2 n_jobs\n",
    "print(\"Mean Absolute Error: \", results.mean()); print(\"Standard Deviation: \", results.std())\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MSE in train : 11.284431491847418\n"
     ]
    }
   ],
   "source": [
    "# MAE from all training data\n",
    "train_pred = lr.predict(X)\n",
    "train_result = []\n",
    "for i in range(len(train_pred)):\n",
    "    result_x = match_Y[i][0] + train_pred[i][0]\n",
    "    result_y = match_Y[i][1] + train_pred[i][1]\n",
    "    train_result.append([result_x, result_y])\n",
    "    \n",
    "total = 0\n",
    "for i in range(len(train_result)):\n",
    "    total += abs(train_result[i][0] - real_Y[i][0]) + abs(train_result[i][1] - real_Y[i][1])\n",
    "print('MAE in train :',total/len(train_result) )    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Real prediction on test images for Kaggle"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 317,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/easybuild-2019/easybuild/software/mpi/gcc-cuda/8.3.0-10.1.243/openmpi/3.1.4/jupyter/1.0.0-python-3.7.4/lib/python3.7/site-packages/ipykernel/ipkernel.py:287: DeprecationWarning: `should_run_async` will not call `transform_cell` automatically in the future. Please pass the result to `transformed_cell` argument and any exception that happen during thetransform in `preprocessing_exc_tuple` in IPython 7.17 and above.\n",
      "  and should_run_async(code)\n"
     ]
    }
   ],
   "source": [
    "test_X= []  \n",
    "test_result = []\n",
    "valid_keys = ['Affine_M','R_est',  't_est']\n",
    "\n",
    "for k1 in loftr_test_dic.keys():\n",
    "    for k2 in loftr_test_dic[k1].keys():\n",
    "        x = []\n",
    "        for k3 in valid_keys:\n",
    "            a = loftr_test_dic[k1][k2][k3]\n",
    "            if a != None:\n",
    "                x.extend(np.array(a).reshape(-1))\n",
    "            else:\n",
    "                if k3 == 'Affine_M':\n",
    "                    x.extend(Affine_M_empty)\n",
    "                elif k3 == 'R_est':\n",
    "                    x.extend(R_est_empty)\n",
    "                elif k3 == 't_est':\n",
    "                    x.extend(t_est_empty)\n",
    "                else:\n",
    "                    print(k3)\n",
    "                    print(a)\n",
    "    test_X.append(x)\n",
    "    test_pred = np.array(lr.predict([x])).reshape(-1)\n",
    "    result_x = all_set_dic[k2]['label'][0].item() + test_pred[0]\n",
    "    result_y = all_set_dic[k2]['label'][1].item() + test_pred[1]\n",
    "    test_result.append([result_x, result_y])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 318,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/easybuild-2019/easybuild/software/mpi/gcc-cuda/8.3.0-10.1.243/openmpi/3.1.4/jupyter/1.0.0-python-3.7.4/lib/python3.7/site-packages/ipykernel/ipkernel.py:287: DeprecationWarning: `should_run_async` will not call `transform_cell` automatically in the future. Please pass the result to `transformed_cell` argument and any exception that happen during thetransform in `preprocessing_exc_tuple` in IPython 7.17 and above.\n",
      "  and should_run_async(code)\n"
     ]
    }
   ],
   "source": [
    "# export the results\n",
    "header = ['id','x','y']\n",
    "predicts = test_result\n",
    "\n",
    "with open('new_result2.csv', 'w') as file:\n",
    "    writer = csv.writer(file)\n",
    "    writer.writerow(header)\n",
    "    for imgfile, predict in zip(test_dic.keys(), predicts):\n",
    "        imgfile = [imgfile.strip(\".jpg\")]\n",
    "        imgfile.extend([predict[0], predict[1]])\n",
    "        writer.writerow(imgfile)  "
   ]
  }
 ],
 "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.7.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}