{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step A"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings\n",
    "warnings.filterwarnings('ignore')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import cv2\n",
    "from matplotlib import pyplot as plt\n",
    "import collections\n",
    "import os"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The ```compute_all_keypoints``` function calculates all keypoints of all query and train images and stores them in a dictionary, in order to easily access them later."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_all_keypoints(query_imgs, train_imgs, sift):\n",
    "\n",
    "    img_dict = {}\n",
    "\n",
    "    for img in query_imgs:\n",
    "        file = 'models/' + img + '.jpg'\n",
    "        query = cv2.imread(file, 0)\n",
    "        kp, des = sift.detectAndCompute(query, None)\n",
    "        img_dict[img] = {'kp': kp, 'des': des, 'shape': query.shape}\n",
    "\n",
    "    for img in train_imgs:\n",
    "        file = 'scenes/' + img + '.png'\n",
    "        train = cv2.imread(file, 0)\n",
    "        kp, des = sift.detectAndCompute(train, None)\n",
    "        img_dict[img] = {'kp': kp, 'des': des, 'shape': train.shape}\n",
    "\n",
    "    return img_dict"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The ```apply_ratio_test``` function takes all the matches found between the query and the train image, it chooses the good ones with the usual ratio test and it stores them in a dictionary using the indexes of the query keypoints as keys and the indexes of the train keypoints as values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def apply_ratio_test(all_matches):\n",
    "\n",
    "    # map of matches kp_query_idx -> kp_train_idx\n",
    "    good_matches = {}\n",
    "  \n",
    "    for m, n in all_matches:\n",
    "        if m.distance < LOWE_COEFF * n.distance:\n",
    "            good_matches[m.queryIdx] = m.trainIdx\n",
    "\n",
    "    return good_matches"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The ```check_matches``` function orders the good matches in decreasing number of keypoints and it runs a series of tests on them, checking the geometric arrangement and the color consistency."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def check_matches(global_matches, train_img, img_dict):\n",
    "    \n",
    "    sorted_global_matches = collections.OrderedDict(sorted(global_matches.items(), key=lambda item: item[1][0], reverse=True))\n",
    "    \n",
    "    recognised = {}\n",
    "    \n",
    "    train_file = 'scenes/' + train_img + '.png'\n",
    "    train_bgr = cv2.imread(train_file)\n",
    "\n",
    "    for k, v in sorted_global_matches.items():\n",
    "\n",
    "        if v[0] > MIN_MATCH_COUNT:\n",
    "            \n",
    "            query_file = 'models/' + k + '.jpg'\n",
    "            query_bgr = cv2.imread(query_file)\n",
    "            \n",
    "            src_pts = np.float32([img_dict[k]['kp'][p].pt for p in v[1].keys()]).reshape(-1, 1, 2)\n",
    "            dst_pts = np.float32([img_dict[train_img]['kp'][p].pt for p in v[1].values()]).reshape(-1, 1, 2)\n",
    "            M, _ = cv2.findHomography(src_pts, dst_pts, cv2.RANSAC, 5.0)\n",
    "            h, w, d = query_bgr.shape\n",
    "            pts = np.float32([[0, 0], [0, h - 1], [w - 1, h - 1], [w - 1, 0]]).reshape(-1, 1, 2)\n",
    "            dst = cv2.perspectiveTransform(pts, M)\n",
    "\n",
    "            center = tuple((dst[0, 0, i] + dst[1, 0, i] + dst[2, 0, i] + dst[3, 0, i]) / 4 for i in (0, 1))\n",
    "\n",
    "            x_min = int(max((dst[0, 0, 0] + dst[1, 0, 0]) / 2, 0))\n",
    "            y_min = int(max((dst[0, 0, 1] + dst[3, 0, 1]) / 2, 0))\n",
    "            x_max = int(min((dst[2, 0, 0] + dst[3, 0, 0]) / 2, img_dict[train_img]['shape'][1]))\n",
    "            y_max = int(min((dst[1, 0, 1] + dst[2, 0, 1]) / 2, img_dict[train_img]['shape'][0]))\n",
    "\n",
    "            query_color = query_bgr.mean(axis=0).mean(axis=0)\n",
    "            train_crop = train_bgr[y_min:y_max,x_min:x_max]\n",
    "            train_color = train_crop.mean(axis=0).mean(axis=0)   \n",
    "            color_diff = np.sqrt(np.sum([value ** 2 for value in abs(query_color - train_color)]))\n",
    "\n",
    "            temp = True \n",
    "            if color_diff < COLOR_T :\n",
    "                for r, corners in recognised.items():\n",
    "                    r_center = tuple((corners[0, 0, i] + corners[1, 0, i] + corners[2, 0, i] + corners[3, 0, i]) / 4 for i in (0, 1))\n",
    "                    if (center[0] > min(corners[0, 0, 0], corners[1, 0, 0]) and center[0] < max(corners[2, 0, 0], corners[3, 0, 0])\\\n",
    "                        and center[1] > min(corners[0, 0, 1], corners[3, 0, 1]) and center[1] < max(corners[1, 0, 1], corners[2, 0, 1]))\\\n",
    "                        or (r_center[0] > x_min and r_center[0] < x_max\\\n",
    "                        and r_center[1] > y_min and r_center[1] < y_max):\n",
    "                        temp = False\n",
    "                        break\n",
    "                if temp:\n",
    "                    recognised[k] = dst\n",
    "                    \n",
    "    return recognised"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The ```print_matches``` function takes all the recognised images and prints their details, i.e. their position, width, and height."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def print_matches(train_img, query_imgs, recognised, true_imgs, verbose):\n",
    "    \n",
    "    print('Scene: ' + train_img + '\\n')\n",
    "\n",
    "    for query_img in query_imgs:\n",
    "                     \n",
    "        total = int(query_img in recognised.keys())    \n",
    "        true_total = int(query_img in true_imgs[train_img])\n",
    "        \n",
    "        if total != true_total:\n",
    "            print('\\033[1m' + 'Product ' + query_img + ' – ' + str(total) + '/' + str(true_total) + ' instances found' + '\\033[0m')\n",
    "        elif total > 0 or verbose == True:\n",
    "            print('Product ' + query_img + ' – ' + str(total) + '/' + str(true_total) + ' instances found')\n",
    "        \n",
    "        if total == 1:\n",
    "            dst = recognised[query_img]\n",
    "            center = tuple(int((dst[0, 0, i] + dst[1, 0, i] + dst[2, 0, i] + dst[3, 0, i]) / 4) for i in (0, 1))\n",
    "            w = int(((dst[3, 0, 0] - dst[0, 0, 0]) + (dst[2, 0, 0] - dst[1, 0, 0])) /2)\n",
    "            h = int(((dst[1, 0, 1] - dst[0, 0, 1]) + (dst[2, 0, 1] - dst[3, 0, 1])) /2)\n",
    "            print('\\t' + 'Position: ' + str(center)\\\n",
    "                 + '\\t' + 'Width: ' + str(w)\\\n",
    "                 + '\\t' + 'Height: ' + str(h))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The ```draw_matches``` function draws on the train image the boxes' homographies and the numbers corresponding to the query images."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def draw_matches(recognised, train_img, color):\n",
    "    \n",
    "    train_file = 'scenes/' + train_img + '.png'\n",
    "    if color == True:\n",
    "        train_bgr = cv2.imread(train_file)\n",
    "        train_temp = cv2.cvtColor(train_bgr, cv2.COLOR_BGR2RGB)\n",
    "        train_rgb = np.zeros(train_bgr.shape, train_bgr.dtype)\n",
    "        for y in range(train_temp.shape[0]):\n",
    "            for x in range(train_temp.shape[1]):\n",
    "                for c in range(train_temp.shape[2]):\n",
    "                    train_rgb[y, x, c] = np.clip(0.5 * train_temp[y, x, c], 0, 255)\n",
    "    else:\n",
    "        train_gray = cv2.imread(train_file, 0)\n",
    "        train_rgb = cv2.cvtColor(train_gray // 2, cv2.COLOR_GRAY2RGB)\n",
    "\n",
    "    for k, v in recognised.items():\n",
    "\n",
    "        train_rgb = cv2.polylines(train_rgb, [np.int32(v)], True, (0, 255, 0), 3, cv2.LINE_AA)\n",
    "        font = cv2.FONT_HERSHEY_SIMPLEX\n",
    "        cv2.putText(train_rgb, k,\\\n",
    "                        (int((v[3, 0, 0] - v[0, 0, 0]) * 0.25 + v[0, 0, 0]), int((v[1, 0, 1] - v[0, 0, 1]) * 0.67 + v[0, 0, 1])),\\\n",
    "                        font, 5, (0, 255, 0), 10, cv2.LINE_AA)\n",
    "    \n",
    "    plt.imshow(train_rgb),plt.show();\n",
    "    \n",
    "    if color == True:\n",
    "        \n",
    "        if not os.path.exists('output/step_A/'):\n",
    "            os.mkdir('output/step_A/')\n",
    "\n",
    "        cv2.imwrite('output/step_A/' + train_img + '.png', cv2.cvtColor(train_rgb, cv2.COLOR_RGB2BGR))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The ```step_A``` function takes the lists of query and train images and performs the product recognition."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "def step_A(query_imgs, train_imgs, true_imgs, verbose, color):\n",
    "\n",
    "    sift = cv2.xfeatures2d.SIFT_create()\n",
    "\n",
    "    bf = cv2.BFMatcher()\n",
    "\n",
    "    img_dict = compute_all_keypoints(query_imgs, train_imgs, sift)\n",
    "\n",
    "    for train_img in train_imgs:\n",
    "\n",
    "        kp_train, des_train = img_dict[train_img]['kp'], img_dict[train_img]['des']\n",
    "\n",
    "        global_matches = {}\n",
    "\n",
    "        for query_img in query_imgs:\n",
    "\n",
    "            kp_query, des_query = img_dict[query_img]['kp'], img_dict[query_img]['des']\n",
    "            \n",
    "            all_matches = bf.knnMatch(des_query, des_train, k=2)\n",
    "            good_matches = apply_ratio_test(all_matches)\n",
    "            global_matches[query_img] = (len(good_matches), good_matches)\n",
    "\n",
    "        recognised = check_matches(global_matches, train_img, img_dict)\n",
    "\n",
    "        print_matches(train_img, query_imgs, recognised, true_imgs, verbose)\n",
    "        \n",
    "        draw_matches(recognised, train_img, color)\n",
    "\n",
    "        print('\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Parameters:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "LOWE_COEFF = 0.5\n",
    "MIN_MATCH_COUNT = 30\n",
    "COLOR_T = 50"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "query_imgs = ['0', '1', '11', '19', '24', '25', '26']\n",
    "train_imgs = ['e1', 'e2', 'e3', 'e4', 'e5']\n",
    "true_imgs = {\n",
    "    'e1': {'0', '11'},\n",
    "    'e2': {'24', '25', '26'},\n",
    "    'e3': {'0', '1', '11'},\n",
    "    'e4': {'0', '11', '25', '26'},\n",
    "    'e5': {'19', '25'},\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# verbose=False does not print the true negative instances\n",
    "# color=True outputs all the scenes in color instead of grayscale and saves them, but the process is quite slow\n",
    "step_A(query_imgs, train_imgs, true_imgs, verbose=False, color=False)"
   ]
  }
 ],
 "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.5.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
