annoy_load.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 116,
   "id": "265f7c9f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import re\n",
    "import json\n",
    "import glob\n",
    "import random\n",
    "from PIL import Image, ImageDraw\n",
    "from tqdm import tqdm\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 117,
   "id": "46d0f182",
   "metadata": {},
   "outputs": [],
   "source": [
    "from annoy import AnnoyIndex"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 118,
   "id": "d0c9ff5a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch.utils import data\n",
    "from torchvision import transforms\n",
    "import torchvision.models as models\n",
    "import torch.nn as nn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 119,
   "id": "b6713791",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch.utils import data\n",
    "from torchvision import transforms\n",
    "import torchvision.models as models\n",
    "import torch.nn as nn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 120,
   "id": "0ac7a59b",
   "metadata": {},
   "outputs": [],
   "source": [
    "data_transform = transforms.Compose([\n",
    "        transforms.Resize((224, 224)),\n",
    "        # transforms.CenterCrop(224),\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 121,
   "id": "8fb39e14",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 121,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# loading the trained model and generating embedding based on that\n",
    "base_model = models.resnet18(pretrained=True)\n",
    "for param in base_model.parameters():\n",
    "    param.requires_grad = False\n",
    "num_ftrs = base_model.fc.in_features\n",
    "base_model.fc = nn.Sequential(nn.Linear(num_ftrs, 256), nn.Linear(256, 128))\n",
    "\n",
    "# loading the trained model with trained weights\n",
    "checkpoint = torch.load(\"./weights/model_best.pth\")\n",
    "base_model.load_state_dict(checkpoint['state_dict'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 122,
   "id": "c248d07a",
   "metadata": {},
   "outputs": [],
   "source": [
    "base_model = base_model.cpu()\n",
    "base_model = base_model.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 123,
   "id": "4cca4a75",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 123,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "annoy_index = AnnoyIndex(128, metric=\"euclidean\")\n",
    "annoy_index.load(\"./annoy_index.ann\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 124,
   "id": "8c73b518",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('annoy_index_to_label.json') as f:\n",
    "    annoy_index_to_label = json.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 125,
   "id": "8e62ed1a",
   "metadata": {},
   "outputs": [],
   "source": [
    "query_img_name = './dataset/val_retina/4.png'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 126,
   "id": "9ba4a3bd",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([128])"
      ]
     },
     "execution_count": 126,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "query_img = Image.open(query_img_name)\n",
    "query_img = data_transform(query_img)\n",
    "query_img = query_img.unsqueeze(0)\n",
    "query_img_embedding = base_model(query_img)\n",
    "query_img_embedding = query_img_embedding.squeeze()\n",
    "query_img_embedding.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 127,
   "id": "71ed4728",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([-0.0396, -0.0635,  0.0229,  0.0253,  0.0816, -0.3173, -0.3418, -0.2621,\n",
       "         0.3287,  0.2832,  0.2842,  0.0544,  0.1584, -0.0022,  0.5344, -0.0097,\n",
       "         0.2373, -0.1543,  0.4458, -0.0574,  0.0838,  0.3555, -0.1194, -0.0714,\n",
       "         0.1287, -0.0155,  0.0211,  0.1945,  0.2851, -0.0452,  0.0064,  0.2497,\n",
       "         0.3413,  0.0381, -0.0924,  0.1646,  0.0217, -0.0858, -0.1268,  0.2020,\n",
       "        -0.0409, -0.0884,  0.0089,  0.0843, -0.1467,  0.0396,  0.1318, -0.4844,\n",
       "        -0.0322, -0.4791,  0.2510, -0.3303,  0.4215, -0.0503,  0.2348, -0.3210,\n",
       "         0.1912, -0.0605, -0.2135,  0.1815,  0.2225, -0.0769,  0.0436,  0.0426,\n",
       "         0.0520, -0.1147,  0.1500, -0.1225, -0.0742,  0.1197,  0.1429,  0.1736,\n",
       "         0.1802, -0.2868,  0.1878,  0.2693, -0.0378, -0.0430, -0.1295,  0.0281,\n",
       "        -0.0408,  0.1477, -0.0142,  0.0383,  0.0363,  0.3557,  0.2661,  0.0890,\n",
       "        -0.1227, -0.0353, -0.0621, -0.0224,  0.1996, -0.1658,  0.1955, -0.2141,\n",
       "         0.1688,  0.1162, -0.0088,  0.2131, -0.0277,  0.0167,  0.2778, -0.3077,\n",
       "        -0.4389,  0.2963, -0.0877,  0.1789, -0.1115, -0.0097,  0.2606, -0.0614,\n",
       "         0.0294, -0.1992,  0.2051, -0.0981, -0.0666, -0.2857,  0.1462, -0.1758,\n",
       "        -0.2378, -0.2856, -0.0468,  0.1805, -0.2525, -0.1589, -0.0658,  0.0615],\n",
       "       grad_fn=<SqueezeBackward0>)"
      ]
     },
     "execution_count": 127,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "query_img_embedding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 128,
   "id": "16f2b093",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([-0.0396, -0.0635,  0.0229,  0.0253,  0.0816, -0.3173, -0.3418, -0.2621,\n",
       "         0.3287,  0.2832,  0.2842,  0.0544,  0.1584, -0.0022,  0.5344, -0.0097,\n",
       "         0.2373, -0.1543,  0.4458, -0.0574,  0.0838,  0.3555, -0.1194, -0.0714,\n",
       "         0.1287, -0.0155,  0.0211,  0.1945,  0.2851, -0.0452,  0.0064,  0.2497,\n",
       "         0.3413,  0.0381, -0.0924,  0.1646,  0.0217, -0.0858, -0.1268,  0.2020,\n",
       "        -0.0409, -0.0884,  0.0089,  0.0843, -0.1467,  0.0396,  0.1318, -0.4844,\n",
       "        -0.0322, -0.4791,  0.2510, -0.3303,  0.4215, -0.0503,  0.2348, -0.3210,\n",
       "         0.1912, -0.0605, -0.2135,  0.1815,  0.2225, -0.0769,  0.0436,  0.0426,\n",
       "         0.0520, -0.1147,  0.1500, -0.1225, -0.0742,  0.1197,  0.1429,  0.1736,\n",
       "         0.1802, -0.2868,  0.1878,  0.2693, -0.0378, -0.0430, -0.1295,  0.0281,\n",
       "        -0.0408,  0.1477, -0.0142,  0.0383,  0.0363,  0.3557,  0.2661,  0.0890,\n",
       "        -0.1227, -0.0353, -0.0621, -0.0224,  0.1996, -0.1658,  0.1955, -0.2141,\n",
       "         0.1688,  0.1162, -0.0088,  0.2131, -0.0277,  0.0167,  0.2778, -0.3077,\n",
       "        -0.4389,  0.2963, -0.0877,  0.1789, -0.1115, -0.0097,  0.2606, -0.0614,\n",
       "         0.0294, -0.1992,  0.2051, -0.0981, -0.0666, -0.2857,  0.1462, -0.1758,\n",
       "        -0.2378, -0.2856, -0.0468,  0.1805, -0.2525, -0.1589, -0.0658,  0.0615],\n",
       "       grad_fn=<SqueezeBackward0>)"
      ]
     },
     "execution_count": 128,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "query_img_embedding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 129,
   "id": "c402aa63",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['./dataset/train/salad/5871.png',\n",
       " './dataset/train/salad/6267.png',\n",
       " './dataset/train/salad/1833.png',\n",
       " './dataset/train/salad/476.png',\n",
       " './dataset/train/salad/1010.png',\n",
       " './dataset/train/salad/344.png',\n",
       " './dataset/train/salad/3825.png',\n",
       " './dataset/train/salad/2577.png',\n",
       " './dataset/train/salad/1202.png',\n",
       " './dataset/train/salad/1815.png',\n",
       " './dataset/train/salad/3669.png',\n",
       " './dataset/train/salad/6315.png',\n",
       " './dataset/train/salad/110.png',\n",
       " './dataset/train/salad/1280.png',\n",
       " './dataset/train/salad/6849.png',\n",
       " './dataset/train/salad/122.png',\n",
       " './dataset/train/salad/4671.png',\n",
       " './dataset/train/salad/5145.png',\n",
       " './dataset/train/salad/2469.png',\n",
       " './dataset/train/salad/6747.png']"
      ]
     },
     "execution_count": 129,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "similar_images = annoy_index.get_nns_by_vector(query_img_embedding, 20, include_distances=True)\n",
    "similar_image_labels = [annoy_index_to_label[str(i)] for i in similar_images[0]]\n",
    "similar_image_labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "id": "de06cf39",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "([106,\n",
       "  609,\n",
       "  668,\n",
       "  187,\n",
       "  675,\n",
       "  637,\n",
       "  3,\n",
       "  638,\n",
       "  24,\n",
       "  124,\n",
       "  571,\n",
       "  99,\n",
       "  70,\n",
       "  118,\n",
       "  13,\n",
       "  643,\n",
       "  113,\n",
       "  610,\n",
       "  895,\n",
       "  611],\n",
       " [1.0725980997085571,\n",
       "  1.076981544494629,\n",
       "  1.086986780166626,\n",
       "  1.1012552976608276,\n",
       "  1.1056897640228271,\n",
       "  1.114335298538208,\n",
       "  1.12714684009552,\n",
       "  1.134081482887268,\n",
       "  1.137354850769043,\n",
       "  1.1411970853805542,\n",
       "  1.1431547403335571,\n",
       "  1.1471046209335327,\n",
       "  1.1527646780014038,\n",
       "  1.1545040607452393,\n",
       "  1.159005045890808,\n",
       "  1.1607229709625244,\n",
       "  1.1627691984176636,\n",
       "  1.1657835245132446,\n",
       "  1.1728541851043701,\n",
       "  1.1770964860916138])"
      ]
     },
     "execution_count": 80,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "similar_images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82b52cc3",
   "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.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}