HW2_part1_get_des_default.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## HW2 Solution Part 1\n",
    "\n",
    "**step1**: get keypoints for each images\n",
    "\n",
    "**step2**: get patches for each keypoints\n",
    "\n",
    "**step3**: get descriptions for each keypoints (by forward pass patches to the network)\n",
    "\n",
    "step4: caculate the similarity matrices\n",
    "\n",
    "step5: get topK similar images for each query\n",
    "\n",
    "step6: draw recall vs precision curves\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"1\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load packages\n",
    "from __future__ import division, print_function\n",
    "import glob\n",
    "import os\n",
    "import cv2\n",
    "import PIL\n",
    "import random\n",
    "import numpy as np\n",
    "# import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.image as mpimg\n",
    "import torch\n",
    "import torch.nn.init\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import torch.backends.cudnn as cudnn\n",
    "import torch.nn.functional as F\n",
    "import torchvision.datasets as dset\n",
    "import torchvision.transforms as transforms\n",
    "from tqdm import tqdm\n",
    "from torch.autograd import Variable\n",
    "from copy import deepcopy, copy\n",
    "from Utils import cv2_scale36, cv2_scale, np_reshape, np_reshape64\n",
    "from scipy.optimize import linear_sum_assignment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# parameters setting\n",
    "query_path = \"image_retrieval/query/\"\n",
    "image_path = \"image_retrieval/images/\"\n",
    "query_num = 35\n",
    "image_num = 140\n",
    "kps_num = 30\n",
    "patch_size = 32\n",
    "patches = torch.zeros(query_num+image_num, kps_num,1, 32, 32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initiate SURF detector\n",
    "surf = cv2.xfeatures2d.SURF_create(30)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def getPatches(kps, img, size=32, num=500):\n",
    "    res = torch.zeros(num, 1, size, size)\n",
    "    if type(img) is np.ndarray:\n",
    "        img = torch.from_numpy(img)\n",
    "    h, w = img.shape      # note: for image, the x direction is the verticle, y-direction is the horizontal...\n",
    "    for i in range(num):\n",
    "        cx, cy = kps[i]\n",
    "        cx, cy = int(cx), int(cy)\n",
    "        dd = int(size/2)\n",
    "        xmin, xmax = max(0, cx - dd), min(w, cx + dd ) - 1\n",
    "        ymin, ymax = max(0, cy - dd), min(h, cy + dd ) - 1 \n",
    "        \n",
    "        xmin_res, xmax_res = dd - min(dd,cx), dd + min(dd, w - cx)-1\n",
    "        ymin_res, ymax_res = dd - min(dd,cy), dd + min(dd, h - cy)-1\n",
    "\n",
    "        res[i, 0, ymin_res: ymax_res, xmin_res: xmax_res] = img[ymin: ymax, xmin: xmax]\n",
    "    return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# tensor for query/image patches\n",
    "for idx in range(query_num+image_num):\n",
    "    if idx < query_num:\n",
    "        img_dir = os.path.join(query_path,\"q{}.JPG\".format(idx+1))\n",
    "    else:\n",
    "        img_dir = os.path.join(image_path,\"{}.JPG\".format(idx+1-query_num))\n",
    "    image = cv2.imread(img_dir)\n",
    "    img= cv2.cvtColor(image,cv2.COLOR_BGR2GRAY)\n",
    "    ## find the keypoints and descriptors with SURF\n",
    "    kps, des = surf.detectAndCompute(img, None)\n",
    "    ## find the keypoints and descriptors with SIFT\n",
    "    keypoints_img = [kps[a].pt for a in range(kps_num)] \n",
    "    patches[idx] = getPatches(keypoints_img, img, size=patch_size, num=kps_num)    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import network, , load pretrained weights, and turn on testing mode\n",
    "from descriptor_CNN3 import DesNet\n",
    "model = DesNet()\n",
    "model.cuda()\n",
    "model.eval()\n",
    "trained_weights = torch.load(\"checkpoint.pth\")\n",
    "model.load_state_dict(trained_weights[\"state_dict\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([175, 30, 1, 32, 32])\n"
     ]
    }
   ],
   "source": [
    "print(patches.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get descriptions\n",
    "with torch.no_grad(): # only 3612MiB, without this line, it will be 9132MiB\n",
    "    des1 = model(patches[:query_num].view(-1, 1, 32,32).cuda()).view(-1, 30, 128)\n",
    "    des2 = model(patches[query_num:].view(-1, 1, 32,32).cuda()).view(-1, 30, 128)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([35, 30, 128]) torch.Size([140, 30, 128])\n"
     ]
    }
   ],
   "source": [
    "print(des1.shape, des2.shape)\n",
    "# save des\n",
    "des_dir = \"des.pt\"\n",
    "torch.save([des1, des2], des_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}