Week4_code.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "DGzo73G8O4fT"
   },
   "source": [
    "# WEEK 4 CODE\n",
    "\n",
    "### Content:\n",
    "* TASK 1: Detect keypoints and compute descriptors in Museum and query images\n",
    "* TASK 2: Find tentative matches based on similarity of local appearance and verify matches\n",
    "* TASK 3: Evaluate the system on QSD1-W4, map@k\n",
    "* TASK 4: Evaluate best system from previous week on QSD1-W4\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "EOZEWWHMO4fW"
   },
   "outputs": [],
   "source": [
    "#IMPORTS\n",
    "import sys\n",
    "import cv2\n",
    "import os\n",
    "import os.path\n",
    "import numpy as np\n",
    "import math\n",
    "np.set_printoptions(threshold=np.inf)\n",
    "import matplotlib.patches as patches\n",
    "import ml_metrics as metrics\n",
    "import pickle\n",
    "from matplotlib import pyplot as plt\n",
    "from os.path import isfile, join\n",
    "from os import listdir\n",
    "from sklearn import preprocessing\n",
    "from sklearn.metrics import recall_score,precision_score,f1_score\n",
    "from scipy.spatial import distance\n",
    "from scipy.signal import convolve2d\n",
    "from scipy.fftpack import dct\n",
    "from scipy import ndimage\n",
    "from skimage.feature import local_binary_pattern\n",
    "from skimage.feature import hog\n",
    "from skimage.filters import roberts, sobel, sobel_h, sobel_v\n",
    "from skimage.morphology import skeletonize, dilation, square, opening, closing, erosion\n",
    "from pathlib import Path\n",
    "from cv2 import boundingRect\n",
    "from cv2 import xfeatures2d_SIFT\n",
    "from PIL import Image, ImageDraw\n",
    "import pytesseract"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "SpFrRfRSO4fa"
   },
   "source": [
    "# Pre-processing: Filter noise with linear or non-linear filters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "zWU8LygIO4fb"
   },
   "outputs": [],
   "source": [
    "# DEFINE FUNCTIONS \n",
    "\n",
    "def detectnoise(img):\n",
    "    # Calcula la sigma del soroll. Llavors les imatges que tenen soroll tenen una sigma més gran \n",
    "    #que la resta de les imatges\n",
    "    thr = 8.0\n",
    "    H, W = gray.shape\n",
    "\n",
    "    M = [[1, -2, 1],\n",
    "        [-2, 4, -2],\n",
    "        [1, -2, 1]]\n",
    "    \n",
    "    sigma = np.sum(np.sum(np.absolute(convolve2d(gray, M))))\n",
    "    sigma = sigma * math.sqrt(0.5 * math.pi) / (6 * (W-2) * (H-2))\n",
    "    if sigma>thr:\n",
    "        return True\n",
    "    else:\n",
    "        return False\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "S5R6RI5wO4fe"
   },
   "outputs": [],
   "source": [
    "query2folder = './qsd1_w4/'\n",
    "\n",
    "# CREEM UN DIRECTORI ON GUARDEM LES IMATGES SENSE SOROLL\n",
    "path = 'qsd1_w4_denoise'\n",
    "if not os.path.exists(path):\n",
    "    os.mkdir(path)\n",
    "\n",
    "\n",
    "for filename in sorted(listdir(query2folder)):\n",
    "    if(filename != '.DS_Store' and ('jpg' in filename)):\n",
    "        print (\"\\r Processing image...   {}\".format(query2folder+filename), end=\"\")\n",
    "        query_img = cv2.imread(query2folder + filename)\n",
    "        gray = cv2.cvtColor(query_img, cv2.COLOR_BGR2GRAY)\n",
    "        noise = detectnoise(gray)\n",
    "        if noise:\n",
    "            denoise_img = cv2.bilateralFilter(query_img, 9, 120, 120)\n",
    "            denoise_img_2 = cv2.fastNlMeansDenoisingColored(denoise_img,None,12,10,7,21) \n",
    "        else:\n",
    "            denoise_img_2 = query_img\n",
    "        \n",
    "        cv2.imwrite(os.path.join(path,filename), denoise_img_2)\n",
    "\n",
    "            \n",
    "        "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "P_996i2lO4fg"
   },
   "source": [
    "# Task 1: Detect keypoints and compute descriptors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "xV7bZxEXO4fh"
   },
   "outputs": [],
   "source": [
    "#DEFINING FUNCTIONS FOR TASK 1\n",
    "\n",
    "def findKeypoints(img, descriptor):\n",
    "    \n",
    "    if descriptor is 'SIFT':\n",
    "        # Initiate SIFT detector\n",
    "        sift = xfeatures2d_SIFT.create()\n",
    "\n",
    "        # find the keypoints and descriptors with SIFT\n",
    "        kp, des = sift.detectAndCompute(img, None)\n",
    "    elif descriptor is 'ORB':\n",
    "        # Initiate ORB detector\n",
    "        orb = cv2.ORB_create()\n",
    "        \n",
    "        # find the keypoints and descriptors with ORB\n",
    "        kp = orb.detect(img, None)\n",
    "        kp, des = orb.compute(img, kp)\n",
    "    \n",
    "    elif descriptor is 'SURF':\n",
    "        # Initiate SURF detector\n",
    "        surf = xfeatures2d_SURF.create()\n",
    "\n",
    "        # find the keypoints and descriptors with SIFT\n",
    "        kp, des = surf.detectAndCompute(img, None)\n",
    "\n",
    "    return kp,des\n",
    "\n",
    "\n",
    "# FUNCTIONS BACKGROUND DETECTION \n",
    "\n",
    "def detect_corners(mask):\n",
    "    \"\"\"\n",
    "    Finds four points corresponding to rectangle corners\n",
    "\n",
    "    :param mask: (ndarray) binary image\n",
    "    :return: (int) points from corners\n",
    "    \"\"\"\n",
    "\n",
    "    width = mask.shape[1]\n",
    "    height = mask.shape[0]\n",
    "    coords = np.argwhere(np.ones([height, width]))\n",
    "    coords_x = coords[:, 1]\n",
    "    coords_y = coords[:, 0]\n",
    "\n",
    "    coords_x_filtered = np.extract(mask, coords_x)\n",
    "    coords_y_filtered = np.extract(mask, coords_y)\n",
    "    max_br = np.argmax(coords_x_filtered + coords_y_filtered)\n",
    "    max_tr = np.argmax(coords_x_filtered - coords_y_filtered)\n",
    "    max_tl = np.argmax(-coords_x_filtered - coords_y_filtered)\n",
    "    max_bl = np.argmax(-coords_x_filtered + coords_y_filtered)\n",
    "\n",
    "    tl_x, tl_y = int(coords_x_filtered[max_tl]), int(coords_y_filtered[max_tl])\n",
    "    tr_x, tr_y = int(coords_x_filtered[max_tr]), int(coords_y_filtered[max_tr])\n",
    "    bl_x, bl_y = int(coords_x_filtered[max_bl]), int(coords_y_filtered[max_bl])\n",
    "    br_x, br_y = int(coords_x_filtered[max_br]), int(coords_y_filtered[max_br])\n",
    "\n",
    "    return tl_x, tl_y, bl_x, bl_y, br_x, br_y, tr_x, tr_y\n",
    "\n",
    "def waterShed(query_img, color_space, sat_chan):\n",
    "    \"\"\"\n",
    "    Applies watershed to detect foreground and baskground\n",
    "\n",
    "    :param query_img: (ndarray) query image\n",
    "    :param sat: (ndarray) query image saturation\n",
    "    :param part: (int) one paint(0), two paints left part(1),  two paints right part(2)\n",
    "    :return: mask: predicted mask\n",
    "    \"\"\"\n",
    "    \n",
    "    if sat_chan:\n",
    "        ret, thresh = cv2.threshold(color_space, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)\n",
    "        thresh_pos = np.where(thresh==255)\n",
    "        if np.sum(thresh_pos[0]==0)>.7*thresh.shape[0] or np.sum(thresh_pos[1]==0)>.7*thresh.shape[1]:\n",
    "            thresh = 255 - thresh\n",
    "    else:\n",
    "        thresh = k_means(color_space)\n",
    "        thresh_pos = np.where(thresh==255)\n",
    "        if np.sum(thresh_pos[0]==0)>.7*thresh.shape[0] or np.sum(thresh_pos[1]==0)>.7*thresh.shape[1]:\n",
    "            thresh = 255 - thresh\n",
    "    \n",
    "    # noise removal\n",
    "    kernel = np.ones((3, 3), np.uint8)\n",
    "    morph = cv2.morphologyEx(thresh, cv2.MORPH_CLOSE, kernel, iterations=1)\n",
    "    # fillin holes\n",
    "    kernel = np.ones((5, 5), np.uint8)\n",
    "    opening = cv2.morphologyEx(morph, cv2.MORPH_OPEN, kernel, iterations=2)\n",
    "\n",
    "    # sure background area\n",
    "    sure_bg = cv2.dilate(opening, kernel, iterations=3)  # Fals. El sure\n",
    "\n",
    "    # Finding sure foreground area\n",
    "    dist_transform = cv2.distanceTransform(opening, cv2.DIST_L2, 3)\n",
    "    ret, sure_fg = cv2.threshold(dist_transform, 0.05 * dist_transform.max(), 255, 0)  # 0.15 26.57%\n",
    "\n",
    "    # Finding unknown region\n",
    "    sure_fg = np.uint8(sure_fg)\n",
    "    unknown = cv2.subtract(sure_bg, sure_fg)\n",
    "\n",
    "    # Marker labelling\n",
    "    ret, markers = cv2.connectedComponents(sure_fg)\n",
    "\n",
    "    markers[markers != 0] = 200\n",
    "\n",
    "    border_top = int(sat.shape[0] * .03)\n",
    "    borde_lat = int(sat.shape[1] * .02)\n",
    "\n",
    "    # define known sure background\n",
    "    markers[0:border_top, :] = 100\n",
    "    sure_fg[0:border_top, :] = 100\n",
    "\n",
    "    markers[len(markers) - border_top:len(markers) - 1, :] = 100\n",
    "    sure_fg[len(markers) - border_top:len(markers) - 1, :] = 100\n",
    "\n",
    "    markers[:, 0:borde_lat] = 100\n",
    "    sure_fg[:, 0:borde_lat] = 100\n",
    "    markers[:, len(markers[0, :]) - borde_lat:len(markers[0, :]) - 1] = 100\n",
    "    sure_fg[:, len(markers[0, :]) - borde_lat:len(markers[0, :]) - 1] = 100\n",
    "\n",
    "    # Add one to all labels so that sure background is not 0, but 1\n",
    "    markers = markers + 1\n",
    "\n",
    "    # Now, mark the region of unknown with zero\n",
    "    markers[sure_fg == 0] = 0\n",
    "\n",
    "    markersres = cv2.watershed(query_img, markers)\n",
    "\n",
    "    markersres[markersres == 201] = 255\n",
    "    markersres[markersres == 101] = 0\n",
    "    markersres[markersres == -1] = 255\n",
    "\n",
    "    # Make the mask binary\n",
    "    markersres = markersres / np.max(markersres)\n",
    "            \n",
    "    # Remove imperfections with morphological filers (noise)\n",
    "    markersres = cv2.morphologyEx(markersres, cv2.MORPH_OPEN, np.ones((3, 3), np.uint8), iterations=1)\n",
    "\n",
    "    # Set predicted mask\n",
    "    mask_pred = markersres.astype('uint8')\n",
    "\n",
    "    # Get connected components\n",
    "    ret, labels = cv2.connectedComponents(mask_pred)\n",
    "\n",
    "    # Initialize mask for paintings\n",
    "    mask = np.zeros((mask_pred.shape[0], mask_pred.shape[1], 2), dtype=\"uint8\")\n",
    "\n",
    "    # If exist more than two connected components (background and one painting)\n",
    "    if ret > 2:\n",
    "\n",
    "        # Compute area of each connected component\n",
    "        area = []\n",
    "        for i, lab in enumerate(np.unique(labels)):\n",
    "            area.append(mask_pred[labels == lab].size)\n",
    "\n",
    "        # Sort indexes from length of areas\n",
    "        idx = sorted(range(len(area)), key=lambda k: area[k])\n",
    "\n",
    "        # Check if more than three connected components (background and two paintings)\n",
    "        if ret > 3:\n",
    "            for i, j in enumerate(idx):\n",
    "                # Remove smallest areas\n",
    "                if (len(area)-i) > 3:\n",
    "                    mask_pred[labels == j] = 0\n",
    "\n",
    "        # Remove indexes labels where its pixels are now zero\n",
    "        idx = [n for _, n in enumerate(idx) if np.sum(mask_pred[labels == n]) != 0]\n",
    "\n",
    "        # Positions of non-zero elements of each painting\n",
    "        obj1 = np.where(labels == idx[0])\n",
    "        obj2 = np.where(labels == idx[1])\n",
    "\n",
    "        # Determine if there exist two paintings\n",
    "        if (area[idx[0]] > .28*area[idx[1]]):\n",
    "\n",
    "            # Check that one connected component is not inside the other\n",
    "            if (((np.min(obj1[1][:]) - np.max(obj2[1][:])) < 0) and ((np.max(obj1[1][:]) - np.min(obj2[1][:])) > 0)) \\\n",
    "                    or (((np.min(obj2[1][:]) - np.max(obj1[1][:])) < 0) and\n",
    "                        ((np.max(obj2[1][:]) - np.min(obj1[1][:])) > 0)):\n",
    "                mask = mask_pred\n",
    "            else:\n",
    "                # Set one mask to each painting\n",
    "                mask[labels == idx[0], 0] = 1\n",
    "                mask[labels == idx[1], 1] = 1\n",
    "\n",
    "        else:\n",
    "            # Set mask for the painting\n",
    "            mask = mask_pred\n",
    "\n",
    "    else:\n",
    "        # Set mask for the painting\n",
    "        mask = mask_pred\n",
    "\n",
    "    ## POLYGON - Generate a polygon from mask corners ##\n",
    "\n",
    "    mask_f = np.zeros((mask_pred.shape[0], mask_pred.shape[1], 2), dtype=\"uint8\")\n",
    "\n",
    "    # For two paintings\n",
    "    if len(mask.shape) > 2:\n",
    "        # Get Corners\n",
    "        corners0 = detect_corners(mask[:, :, 0])\n",
    "\n",
    "        # Draw polygon\n",
    "        poly0 = Image.new('RGB', (mask_pred.shape[1], mask_pred.shape[0]), (0))\n",
    "        pdraw0 = ImageDraw.Draw(poly0)\n",
    "        pdraw0.polygon(xy=corners0, fill=(1), outline=None)\n",
    "        mask_pred_f = np.array(poly0)[:, :, 0]\n",
    "        mask_f[:, :, 0] = mask_pred_f\n",
    "\n",
    "        # Get Corners\n",
    "        corners1 = detect_corners(mask[:, :, 1])\n",
    "\n",
    "        # Draw polygon\n",
    "        poly1 = Image.new('RGB', (mask_pred.shape[1], mask_pred.shape[0]), (0))\n",
    "        pdraw1 = ImageDraw.Draw(poly1)\n",
    "        pdraw1.polygon(xy=corners1, fill=(1), outline=None)\n",
    "        mask_f[:, :, 1] = np.array(poly1)[:, :, 0]\n",
    "        mask_pred_f = mask_pred_f + mask_f[:, :, 1]\n",
    "\n",
    "    # Just one painting\n",
    "    else:\n",
    "        # Get Corners\n",
    "        corners = detect_corners(mask)\n",
    "\n",
    "        # Draw polygon\n",
    "        poly = Image.new('RGB', (mask_pred.shape[1], mask_pred.shape[0]), (0))\n",
    "        pdraw = ImageDraw.Draw(poly)\n",
    "        pdraw.polygon(xy=corners, fill=(1), outline=None)\n",
    "        mask_pred_f = np.array(poly)[:, :, 0]\n",
    "        mask_f = mask_pred_f\n",
    "\n",
    "    return mask_pred_f, mask_f\n",
    "\n",
    "\n",
    "def k_means(color_space):\n",
    "    \n",
    "    Z = color_space.reshape((-1))\n",
    "    # convert to np.float32\n",
    "    Z = np.float32(Z)\n",
    "\n",
    "    criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 10, 1.0)\n",
    "    ret, label, center = cv2.kmeans(Z, 5, None, criteria, 10, cv2.KMEANS_RANDOM_CENTERS)\n",
    "\n",
    "    center = np.uint8(center)\n",
    "    res = center[label.flatten()]\n",
    "    res2 = res.reshape((gray.shape))\n",
    "\n",
    "    mask = np.zeros(color_space.shape[:2],np.uint8)\n",
    "\n",
    "    for l in range(len(center)):\n",
    "        lab_pos = np.where(res2==center[l])\n",
    "        if any(lab_pos[0]==0) or any(lab_pos[1]==0) or any(lab_pos[0]==color_space.shape[0]-1) or any(lab_pos[1]==color_space.shape[1]-1):\n",
    "            continue\n",
    "        else:\n",
    "            mask[lab_pos[0],lab_pos[1]] = 255\n",
    "    \n",
    "    if np.sum(mask) == 0:\n",
    "        pos = np.where(res2==res2[int(color_space.shape[0]/2),int(color_space.shape[1]/2)])\n",
    "        mask[pos[0],pos[1]] = 255\n",
    "        \n",
    "    return mask\n",
    "\n",
    "\n",
    "def crop(query_im, mask_crop):\n",
    "    # Coordinates of non-black pixels.\n",
    "    coords = np.argwhere(mask_crop)\n",
    "\n",
    "    # Bounding box of non-black pixels.\n",
    "    x0, y0 = coords.min(axis=0)\n",
    "    x1, y1 = coords.max(axis=0) + 1 # slices are exclusive at the top\n",
    "\n",
    "    # Get the contents of the bounding box.\n",
    "    cropped = query_im[x0:x1, y0:y1]\n",
    "    mask_cropped = mask_crop[x0:x1, y0:y1]\n",
    "\n",
    "    return cropped, mask_cropped, x0, y0\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "n5Uiy81BO4fj"
   },
   "source": [
    "#### Define Parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "LqbCmu9aO4fk"
   },
   "outputs": [],
   "source": [
    "# DEFINING GLOBAL VARIABLES\n",
    "\n",
    "descriptor = 'ORB'\n",
    "\n",
    "if descriptor is 'SIFT':\n",
    "    method = 'FLANN'\n",
    "    resize_factor = 1/4\n",
    "    MIN_MATCHES = 80\n",
    "    \n",
    "elif descriptor is 'ORB':\n",
    "    method = 'BRUTE_FORCE'\n",
    "    resize_factor = 1/4\n",
    "    MIN_MATCHES = 20\n",
    "    \n",
    "elif descriptor is 'SURF':\n",
    "    method = 'FLANN'\n",
    "    resize_factor = 1/4\n",
    "    MIN_MATCHES = 80\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "v7hLQYGJO4fn"
   },
   "source": [
    "### BBDD"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "j9tM7zBNO4fn"
   },
   "outputs": [],
   "source": [
    "folder = './bbdd/'\n",
    "kp_bbdd_array = []\n",
    "des_bbdd_array = []\n",
    "\n",
    "for filename in sorted(listdir(folder)):\n",
    "    if(filename != '.DS_Store' and ('jpg' in filename)):\n",
    "        print (\"\\r Processing image...   {}\".format(folder+filename), end=\"\")\n",
    "        img = cv2.imread(folder + filename,0)#0 is the way of getting B&W image\n",
    "        height, width = img.shape[:2]\n",
    "        resize_size = (int(width*resize_factor), int(height*resize_factor))\n",
    "        img_res = cv2.resize(img, resize_size)\n",
    "        kp_bbdd, des_bbdd = findKeypoints(img_res, descriptor)\n",
    "        kp_bbdd_array.append(kp_bbdd)\n",
    "        des_bbdd_array.append(des_bbdd)\n",
    "        "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "QrfnIpJYO4fq"
   },
   "source": [
    "### QSD1_W4"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "OcOOc_wjO4fr"
   },
   "source": [
    "#### Create masks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "6AsW1v9vO4fr",
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "folder = './qsd1_w4_denoise/'\n",
    "kp_query_array = []\n",
    "des_query_array = []\n",
    "numPaintings = []\n",
    "\n",
    "for filename in sorted(listdir(folder)):\n",
    "    if(filename != '.DS_Store' and ('jpg' in filename)):\n",
    "        print (\"\\r Processing image...   {}\".format(folder+filename), end=\"\")\n",
    "        img = cv2.imread(folder + filename) #0 is the way of getting B&W image\n",
    "        hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)\n",
    "        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)\n",
    "        \n",
    "        # Get saturation and grayscale from query image\n",
    "        sat = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)[:,:,1]\n",
    "        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)\n",
    "        \n",
    "        if ndimage.variance(sat)<ndimage.variance(gray):\n",
    "            color_space = gray\n",
    "            sat_channel=False\n",
    "            \n",
    "        else:\n",
    "            color_space = sat\n",
    "            sat_channel=True\n",
    "            \n",
    "        # Predict mask\n",
    "        mask_pred,mask = waterShed(img, color_space, sat_channel)\n",
    "        \n",
    "        plt.figure()\n",
    "        plt.imshow(mask_pred)\n",
    "        plt.title(filename)\n",
    "        \n",
    "        # Save masks\n",
    "        path = './masks_qd'\n",
    "        if not os.path.exists(path):\n",
    "            os.mkdir(path)\n",
    "\n",
    "        cv2.imwrite('masks_qd/'+filename[:-4]+'.png',mask_pred*255)\n",
    "        if len(mask.shape) > 2:\n",
    "            cv2.imwrite('masks_qd/mask_'+filename[:-4]+'_0.png',mask[:,:,0]*255)\n",
    "            cv2.imwrite('masks_qd/mask_'+filename[:-4]+'_1.png',mask[:,:,1]*255)\n",
    "            numPaintings.append([2])\n",
    "        else:\n",
    "            cv2.imwrite('masks_qd/mask_'+filename[:-4]+'_0.png',mask[:,:]*255)\n",
    "            numPaintings.append([1])\n",
    "        "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "DjkzcqqeO4fu"
   },
   "source": [
    "#### Detect keypoints for cropped images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "YZZYWCFbO4fv",
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "query2masks = './masks_qd/'\n",
    "query2folder = './qsd1_w4_denoise/'\n",
    "\n",
    "kp_query_array = []\n",
    "des_query_array = []\n",
    "\n",
    "#COMPUTE SIFT FOR IMAGES WITHOUT BACKGROUND\n",
    "for filename in sorted(listdir(query2folder)):\n",
    "    if (filename != '.DS_Store' and ('jpg' in filename)):\n",
    "        print (\"\\r Processing image...   {}\".format(query2folder+filename), end=\"\")\n",
    "        query_img = cv2.imread(query2folder + filename)\n",
    "\n",
    "\n",
    "        height, width = query_img.shape[:2]\n",
    "        resize_size = (int(width*resize_factor), int(height*resize_factor))\n",
    "        img_res = cv2.resize(query_img, resize_size)\n",
    "        kp_query, des_query = findKeypoints(img_res, descriptor)\n",
    "        kp_query_array.append(kp_query)\n",
    "        des_query_array.append(des_query)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "bJHEq1dAO4fx"
   },
   "source": [
    "# TASK 2: Find matches based on the descriptors obtained on Task 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "sttQZDxiO4fy"
   },
   "outputs": [],
   "source": [
    "#DEFINING FUNCTIONS FOR TASK 2\n",
    "\n",
    "def findMatches(kp1, kp2, des1, des2, method, MIN_MATCH_COUNT = 10):\n",
    "    \n",
    "    if method is 'BRUTE_FORCE':\n",
    "        bf = cv2.BFMatcher(cv2.NORM_HAMMING, crossCheck=True)\n",
    "        good = bf.match(des1, des2)\n",
    "        \n",
    "    elif method is 'FLANN':\n",
    "        FLANN_INDEX_KDTREE = 0\n",
    "        index_params = dict(algorithm = FLANN_INDEX_KDTREE, trees = 5)\n",
    "        search_params = dict(checks = 50)\n",
    "\n",
    "        # BFMatcher with default params\n",
    "        FLANN_INDEX_KDTREE = 1\n",
    "        index_params= dict(algorithm = FLANN_INDEX_KDTREE,\n",
    "                             trees = 5)\n",
    "        search_param = dict(checks=50)\n",
    "        flann = cv2.FlannBasedMatcher(index_params, search_params)\n",
    "        matches = flann.knnMatch(des1,des2, k=2)\n",
    "\n",
    "        # store all the good matches as per Lowe's ratio test.\n",
    "        good = []\n",
    "        for m,n in matches:\n",
    "            if m.distance < 0.7*n.distance:\n",
    "                good.append(m)\n",
    "    \n",
    "    if len(good)>MIN_MATCH_COUNT:\n",
    "        src_pts = np.float32([ kp1[m.queryIdx].pt for m in good ]).reshape(-1,1,2)\n",
    "        dst_pts = np.float32([ kp2[m.trainIdx].pt for m in good ]).reshape(-1,1,2)\n",
    "\n",
    "        M, mask = cv2.findHomography(src_pts, dst_pts, cv2.RANSAC,5.0)\n",
    "        matchesMask = mask.ravel()\n",
    "\n",
    "    else:\n",
    "        matchesMask = 0\n",
    "    \n",
    "    numMatches = np.sum(matchesMask)\n",
    "    \n",
    "    return numMatches"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "l_5cme_uO4f0",
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "MIN_MATCH_COUNT = 10\n",
    "numMatches_list = []\n",
    "\n",
    "# FINDING MATCHES\n",
    "\n",
    "for i in range(len(kp_query_array)):\n",
    "    numMatches_array = []\n",
    "    print (\"\\nkp_query_array...   {}\".format(i))\n",
    "    for j in range(len(kp_bbdd_array)):\n",
    "        if (des_query_array[i] is not None):\n",
    "            print (\"\\rdes_query_array...   {}\".format(j), end=\"\")\n",
    "            numMatches = findMatches(kp_bbdd_array[j], kp_query_array[i], des_bbdd_array[j], \n",
    "                              des_query_array[i], method, MIN_MATCH_COUNT)\n",
    "            numMatches_array.append(numMatches)\n",
    "        else:\n",
    "            numMatches_array.append(0)\n",
    "        \n",
    "    numMatches_list.append(np.asarray(numMatches_array))\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "gFd5dxPOO4f3"
   },
   "outputs": [],
   "source": [
    "# COMPUTE K-LIST WITH BEST APPROACHES\n",
    "\n",
    "result = []\n",
    "K = 10\n",
    "\n",
    "for i in range(len(numMatches_list)):\n",
    "    maxim = np.max(numMatches_list[i])\n",
    "    \n",
    "    # Check if the are two paintings in one image\n",
    "    if (numPaintings[i] == [2]):\n",
    "        if (maxim != 0):\n",
    "            pos = np.where(numMatches_list[i] == maxim)\n",
    "            non_maxim = np.delete(numMatches_list[i], pos)\n",
    "        else:\n",
    "            non_maxim = numMatches_list[i]\n",
    "        if (maxim > MIN_MATCHES):\n",
    "            imgs_max1 = np.argsort(-numMatches_list[i])[:K]\n",
    "        else:\n",
    "            imgs_max1 = np.argsort(-numMatches_list[i])[:K-1]\n",
    "            imgs_max1 = np.insert(imgs_max1,0,-1)\n",
    "\n",
    "        result.append([imgs_max1.tolist(), imgs_max1.tolist()])\n",
    "    else:\n",
    "        non_maxim = numMatches_list[i]\n",
    "        if (np.max(non_maxim) > MIN_MATCHES):\n",
    "            imgs = np.argsort(-non_maxim)[:K]\n",
    "        else:\n",
    "            imgs = np.argsort(-non_maxim)[:K-1]\n",
    "            imgs = np.insert(imgs,0,-1)\n",
    "            \n",
    "        result.append([imgs.tolist()])\n",
    "\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "pgSXYnS2O4f5"
   },
   "source": [
    "# TASK 3:  Evaluate the system on QSD1-W4, map@k"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "G8VLzDK1O4f6"
   },
   "outputs": [],
   "source": [
    "with open('./qsd1_w4/gt_corresps.pkl', 'rb') as fd:\n",
    "    ll = pickle.load(fd)\n",
    "    \n",
    "\n",
    "gt = [[j] if len(j)==1 else [[j[0]],[j[1]]] for j in ll]  \n",
    "mapk = np.mean([metrics.average_precision.mapk(a,p,10) for a,p in zip(gt, result)])\n",
    "print(mapk*100)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "lo1BHCRiO4f8"
   },
   "source": [
    "# TASK 4: Evaluate best system from previous week on QSD1-W4"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "LzSGsGfxO4f9"
   },
   "source": [
    "## Texture descriptor"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "qNWyW9CKO4f-"
   },
   "source": [
    "### Implement texture descriptors \n",
    "Descriptors:\n",
    "* Color\n",
    "* HoG\n",
    "* Text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "fZuzUeqgO4f-"
   },
   "outputs": [],
   "source": [
    "def colorDescriptor(img,level,numberBins):\n",
    "    img = cv2.cvtColor(img,cv2.COLOR_BGR2Lab)\n",
    "    \n",
    "    # Compute block size\n",
    "    height, width = img.shape[:2]\n",
    "    block_height = int(height/level)\n",
    "    block_width = int(width/level)\n",
    "    hist_img = np.empty([0,0])\n",
    "    colorx = ('x','y','z')\n",
    "    canal = 0\n",
    "    for i,col in enumerate(colorx):\n",
    "        for r in range(level): # For loop Levels\n",
    "            for c in range(level): # For loop Columns\n",
    "                \n",
    "                block = img[r*block_height:r*block_height+block_height, c*block_width:c*block_width+block_width] \n",
    "                hist = cv2.calcHist([block],[i],None,[numberBins],[0,256]) #Calculem histogrames per blocs\n",
    "                hist_t = hist.transpose()\n",
    "                if (i==0) and (r==0) and (c==0):\n",
    "                    hist_img = hist_t\n",
    "                else:\n",
    "                    hist_img = np.concatenate((hist_img, hist_t), axis = 1)\n",
    "\n",
    "    cv2.normalize(hist_img, hist_img, alpha=0, beta=1, norm_type=cv2.NORM_MINMAX)\n",
    "    return hist_img"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "TjC1uWcXO4gB"
   },
   "outputs": [],
   "source": [
    "def HogDescriptor(img):\n",
    "    hog_image = hog(img, orientations=9, pixels_per_cell=(8, 8))\n",
    "    \n",
    "    return hog_image.ravel()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "lphSUriTO4gD"
   },
   "source": [
    "### Test query system using QSD1-W2 using only texture descriptors"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "uohzD0exO4gE"
   },
   "source": [
    "#### 1. Define Parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "PCo0N5vSO4gF"
   },
   "outputs": [],
   "source": [
    "# DEFINING GLOBAL VARIABLES\n",
    "Level = 1\n",
    "numberBins_color = 256"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "ZJfN3q1BO4gH"
   },
   "source": [
    "#### 2. Compute BBDD descriptors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "oOjmQY5aO4gI",
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "folder = './bbdd/'\n",
    "histogram_bbdd_matrix_color = []\n",
    "histogram_bbdd_matrix_texture = []\n",
    "\n",
    "for filename in sorted(listdir(folder)):\n",
    "    if(filename != '.DS_Store' and ('jpg' in filename)):\n",
    "        print (\"\\r Processing image...   {}\".format(folder+filename), end=\"\")\n",
    "        img = cv2.imread(folder + filename)\n",
    "        \n",
    "        equ = np.zeros(img.shape[:], np.uint8)\n",
    "        for c in range(3):\n",
    "            equ[:,:,c] = cv2.equalizeHist(img[:,:,c])\n",
    "        \n",
    "        #COLOR\n",
    "        hist_img_color = colorDescriptor(equ,Level,numberBins_color)\n",
    "        histogram_bbdd_matrix_color.append(hist_img_color)\n",
    "\n",
    "        #TEXTURE\n",
    "        img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)\n",
    "        image_shape = img.shape[:2]\n",
    "        image = cv2.resize(img, (200, 200))\n",
    "        hist_img_texture = HogDescriptor(image)\n",
    "        hist_img_texture = np.array(hist_img_texture).flatten()\n",
    "        histogram_bbdd_matrix_texture.append(hist_img_texture)\n",
    "        \n",
    "histogram_bbdd_matrix_color = np.asarray(histogram_bbdd_matrix_color)\n",
    "histogram_bbdd_matrix_texture = np.asarray(histogram_bbdd_matrix_texture)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "cnQXHi4QO4gM"
   },
   "outputs": [],
   "source": [
    "import codecs\n",
    "# Read txt from bbdd_text\n",
    "textfolder = './bbdd_text/'\n",
    "bbdd_text = []\n",
    "\n",
    "#TEXT\n",
    "for filename in sorted(listdir(textfolder)):\n",
    "    if(filename != '.DS_Store' and ('txt' in filename)):\n",
    "        with codecs.open(textfolder + filename,'r',encoding='latin1') as f:\n",
    "            if  os.stat(textfolder + filename).st_size == 0:\n",
    "                line = ''\n",
    "                bbdd_text.append(line)\n",
    "                #print(filename + line)\n",
    "            else:\n",
    "                for line in f.readlines():\n",
    "                    lim1 = line.find(\"'\",1)\n",
    "                    lim2 = line.find(\"'\",2)\n",
    "                    bbdd_text.append(line[lim1+1:lim2])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "jdftczQPO4gO"
   },
   "source": [
    "#### 3. Compute qsd1_w4 descriptors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "ZFbuU0t3O4gP",
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "histogram_query_matrix_color = []\n",
    "histogram_query_matrix_texture = []\n",
    "query_text = []\n",
    "\n",
    "query2masks = './masks_qd/'\n",
    "query2folder = './qsd1_w4_denoise/'\n",
    "\n",
    "result2 = []\n",
    "\n",
    "with open('qsd1_w4_text.pkl', 'rb') as fd:\n",
    "    ll = pickle.load(fd)\n",
    "\n",
    "for filename in sorted(listdir(query2masks)):\n",
    "    if (filename != '.DS_Store' and ('png' in filename) and ('mask' in filename)):\n",
    "        print (\"\\r Processing image...   {}\".format(query2masks+filename), end=\"\")\n",
    "        \n",
    "        mask = cv2.imread(query2masks + filename)\n",
    "        mask = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)\n",
    "        \n",
    "        img_num = filename[5:10]\n",
    "        num_paint = int(filename[-5])\n",
    "        \n",
    "        if (num_paint == 0) and (len(result2) != int(img_num)-1):\n",
    "            result2.append([])\n",
    "            \n",
    "        result2[int(img_num)].append([])\n",
    "        \n",
    "        img = cv2.imread(query2folder + img_num + '.jpg')\n",
    "        \n",
    "        cropped, mask_cropped, X0, Y0 = crop(img,mask)\n",
    "        \n",
    "        equ = np.zeros(img.shape[:], np.uint8)\n",
    "        for c in range(3):\n",
    "            equ[:,:,c] = cv2.equalizeHist(img[:,:,c])\n",
    "        \n",
    "        #COLOR\n",
    "        hist_img_color = colorDescriptor(equ,Level,numberBins_color)\n",
    "        histogram_query_matrix_color.append(hist_img_color)\n",
    "\n",
    "        #TEXTURE\n",
    "        img = cv2.cvtColor(cropped, cv2.COLOR_BGR2GRAY)\n",
    "        image_shape = img.shape[:2]\n",
    "        image = cv2.resize(img, (200, 200))\n",
    "        hist_img_texture = HogDescriptor(image)\n",
    "        hist_img_texture = np.array(hist_img_texture).flatten()\n",
    "        histogram_query_matrix_texture.append(hist_img_texture)\n",
    "        \n",
    "        #TEXT\n",
    "        img = cropped[ll[int(img_num)][num_paint][1]:ll[int(img_num)][num_paint][3],ll[int(img_num)][num_paint][0]:ll[int(img_num)][num_paint][2]]\n",
    "        if img.shape[0]<4:\n",
    "            query_text.append('')\n",
    "        else:\n",
    "            extractedInformation = pytesseract.image_to_string(img)\n",
    "            query_text.append(extractedInformation)\n",
    "        \n",
    "        with open(img_num+'.txt', 'a') as txtf:\n",
    "            if num_paint==0:\n",
    "                txtf.write(extractedInformation)\n",
    "            else:\n",
    "                txtf.write('\\n'+extractedInformation)\n",
    "            txtf.close()\n",
    "        \n",
    "histogram_query_matrix_color = np.asarray(histogram_query_matrix_color)\n",
    "histogram_query_matrix_texture = np.asarray(histogram_query_matrix_texture)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "c1J5wY_1O4gR"
   },
   "source": [
    "#### 4. Evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "nihlK-6XO4gS"
   },
   "outputs": [],
   "source": [
    "# DEFINE EVALUATION FUNCTIONS\n",
    "\n",
    "def levenshtein(seq1, seq2):\n",
    "    size_x = len(seq1) + 1\n",
    "    size_y = len(seq2) + 1\n",
    "    matrix = np.zeros ((size_x, size_y))\n",
    "    for x in range(size_x):\n",
    "        matrix [x, 0] = x\n",
    "    for y in range(size_y):\n",
    "        matrix [0, y] = y\n",
    "\n",
    "    for x in range(1, size_x):\n",
    "        for y in range(1, size_y):\n",
    "            if seq1[x-1] == seq2[y-1]:\n",
    "                matrix [x,y] = min(\n",
    "                    matrix[x-1, y] + 1,\n",
    "                    matrix[x-1, y-1],\n",
    "                    matrix[x, y-1] + 1\n",
    "                )\n",
    "            else:\n",
    "                matrix [x,y] = min(\n",
    "                    matrix[x-1,y] + 1,\n",
    "                    matrix[x-1,y-1] + 1,\n",
    "                    matrix[x,y-1] + 1\n",
    "                )\n",
    "    \n",
    "    return (matrix[size_x - 1, size_y - 1])\n",
    "\n",
    "def computeMatrixRetrieval_all(histogram_bbdd_color, histogram_query_color, histogram_bbdd_tex, histogram_query_tex, bbdd_text, query_text, K=10):\n",
    "    num_paint = histogram_query_color.shape[0]\n",
    "    \n",
    "    dst_color = np.zeros((num_paint, 279))\n",
    "    dst_tex = np.zeros((num_paint, 279))\n",
    "    dst_text = np.zeros((num_paint, 279))\n",
    "    matrix_retrieval = []\n",
    "    \n",
    "    # Get distances for color and texture descriptors\n",
    "    for query_image_index_c in range(0,len(histogram_query_color)): \n",
    "        for bbdd_image_index_c in range(0, len(histogram_bbdd_color)):\n",
    "            \n",
    "            dst_color[query_image_index_c,bbdd_image_index_c] = distance.braycurtis(histogram_query_color[query_image_index_c,:], histogram_bbdd_color[bbdd_image_index_c,:])\n",
    "            dst_tex[query_image_index_c,bbdd_image_index_c] = distance.braycurtis(histogram_query_tex[query_image_index_c,:], histogram_bbdd_tex[bbdd_image_index_c,:])\n",
    "    \n",
    "    # Get distances for text descriptors\n",
    "    for query_text_index in range(0,len(query_text)):\n",
    "        for bbdd_text_index in range(0,len(bbdd_text)-1):\n",
    "            dst_text[query_text_index,bbdd_text_index] = levenshtein(query_text[query_text_index], bbdd_text[bbdd_text_index])       \n",
    "       \n",
    "    # Get K-list with best approaches\n",
    "    for query_image_index in range(0,len(dst_tex)):\n",
    "        matrix_retrieval_color = np.argsort(dst_color[query_image_index,:])[:K]\n",
    "        matrix_retrieval_tex = np.argsort(dst_tex[query_image_index,:])[:K]\n",
    "        \n",
    "        # Get text retreival when distance is less than 5\n",
    "        dst_text_idx_sorted = np.argsort(dst_text[query_image_index,:])\n",
    "        dst_text_sorted = dst_text[query_image_index, dst_text_idx_sorted]<5\n",
    "        matrix_retrieval_text = dst_text_idx_sorted[:np.sum(dst_text_sorted==True)]\n",
    "        \n",
    "        # Check if any of the text retreivals is equal to one of the 1st pos in texture\n",
    "        most_possible = [int(i) for i in matrix_retrieval_tex.tolist() if i in matrix_retrieval_text.tolist()]\n",
    "        \n",
    "        # To have K retreivals\n",
    "        if len(most_possible)<K:\n",
    "            \n",
    "            dst = 1*(dst_tex[query_image_index,:])+ 0*(dst_color[query_image_index,:]) #only texture\n",
    "            dst_idx_sorted = np.argsort(dst)\n",
    "            \n",
    "            # Get indexes not already in K-list\n",
    "            not_in_most_possible = [int(i) for i in dst_idx_sorted.tolist() if int(i) not in most_possible][:K-len(most_possible)]\n",
    "            \n",
    "            # If the distance is less than 0.25 then is not in bbdd\n",
    "            if (not most_possible) and (dst[not_in_most_possible[0]]>.25):\n",
    "                most_possible = [-1]\n",
    "                \n",
    "            most_possible = most_possible + not_in_most_possible[:K-len(most_possible)]\n",
    "        \n",
    "        matrix_retrieval.append(most_possible[:K])\n",
    "    \n",
    "    return matrix_retrieval"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "xMld_o1lO4gU"
   },
   "outputs": [],
   "source": [
    "K=10\n",
    "matrix_retrieval = computeMatrixRetrieval_all(histogram_bbdd_matrix_color, histogram_query_matrix_color, histogram_bbdd_matrix_texture, histogram_query_matrix_texture, bbdd_text, query_text, K)\n",
    "\n",
    "count = 0\n",
    "for i in range(len(result2)):\n",
    "    for j in range(len(result2[i])):\n",
    "        result2[i][j] = matrix_retrieval[count]\n",
    "        count+=1\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "RAM4BFzJO4gY"
   },
   "outputs": [],
   "source": [
    "with open('./qsd1_w4/gt_corresps.pkl', 'rb') as fd:\n",
    "    ll = pickle.load(fd)\n",
    "    \n",
    "\n",
    "gt = [[j] if len(j)==1 else [[j[0]],[j[1]]] for j in ll]  \n",
    "mapk = np.mean([metrics.average_precision.mapk(a,p,3) for a,p in zip(gt, result2)])\n",
    "print(mapk*100)\n"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "name": "Week4_code_FINAL.ipynb",
   "provenance": []
  },
  "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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}