notebooks/KNN Grouping to Match Similar Colors in Image.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d65a00ef",
   "metadata": {},
   "source": [
    "# BASALT 2022 KNN Grouping to Match Similar Colors in Image\n",
    "\n",
    "## Background\n",
    "I am working on potential solutions for https://www.aicrowd.com/challenges/neurips-2022-minerl-basalt-competition. A notebook exploration at an attempt to smooth the ColorMap data collected using KNN.\n",
    "\n",
    "## Data Mining\n",
    "In order to get the data for this, I configured a Malmo agent to run around a random Malmo Gym environment while it collected both Image, Depth, and ColorMap data.\n",
    "\n",
    "\n",
    "## Code\n",
    "\n",
    "### Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "df70eb89",
   "metadata": {},
   "outputs": [],
   "source": [
    "import cv2\n",
    "import numpy as np\n",
    "import os\n",
    "import pandas as pd\n",
    "from tqdm.notebook import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "53542658",
   "metadata": {},
   "outputs": [],
   "source": [
    "def getClusters(imgPath):\n",
    "    #Thank you to https://codereview.stackexchange.com/questions/143529/opencv-3-using-k-nearest-neighbors-to-analyse-rgb-image\n",
    "    img_BGR = cv2.imread(imgPath)\n",
    "    if img_BGR is None:\n",
    "        raise FileNotFoundError(\"'{0}' could not be opened!\".format(imgPath))\n",
    "\n",
    "    BGR_COLORS = dict(blue=(255, 0, 0), green=(0, 255, 0), red=(0, 0, 255), black=(0, 0, 0), white=(255, 255, 255))\n",
    "    \n",
    "    LABELS = dict(blue=np.array([0]), green=np.array([1]), red=np.array([2]), black=np.array([3]), white=np.array([4]))\n",
    "    trainData = np.array([BGR_COLORS['blue'], BGR_COLORS['green'], BGR_COLORS['red'], BGR_COLORS['black'], BGR_COLORS['white']], dtype=np.float32)\n",
    "    responses = np.array([[LABELS['blue']], LABELS['green'], LABELS['red'], LABELS['black'], LABELS['white']], dtype=np.float32)\n",
    "\n",
    "    knn = cv2.ml.KNearest_create()\n",
    "    knn.train(trainData, cv2.ml.ROW_SAMPLE, responses)\n",
    "\n",
    "    img_vstacked = np.vstack(img_BGR).astype(np.float32)\n",
    "    ret, results, neighbours, dist = knn.findNearest(img_vstacked, 1)\n",
    "\n",
    "    height, width, depth = img_BGR.shape\n",
    "    return results.reshape(height, width).astype(np.uint8)\n",
    "\n",
    "\n",
    "def colorPixels(image, results, colorName):\n",
    "    image[results[:, :] == LABELS[colorName]] = BGR_COLORS[colorName]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "080b75df",
   "metadata": {},
   "outputs": [],
   "source": [
    "path = \"/Volumes/Extreme SSD/Extra Datasets/video-depth-colormap/\"\n",
    "frames = {\"RGB\": [], \"Depth\": [], \"ColorMap\": []}\n",
    "for run in os.listdir(path):\n",
    "    if not 'csv' in run:\n",
    "        newPath = path + run + \"/\"\n",
    "        for index in range(0, len(os.listdir(newPath + \"video_frames/\"))):\n",
    "            if os.path.exists(newPath + \"video_frames/frame\" + str(index) + \".png\") and os.path.exists(newPath + \"depth_frames/frame\" + str(index) + \".png\") and os.path.exists(newPath + \"colormap_frames/frame\" + str(index) + \".png\"):\n",
    "                frames['RGB'].append(newPath + \"video_frames/frame\" + str(index) + \".png\")\n",
    "                frames['Depth'].append(newPath + \"depth_frames/frame\" + str(index) + \".png\")\n",
    "                frames['ColorMap'].append(newPath + \"colormap_frames/frame\" + str(index) + \".png\")\n",
    "df = pd.DataFrame.from_dict(frames)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "e92c735a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>RGB</th>\n",
       "      <th>Depth</th>\n",
       "      <th>ColorMap</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>/Volumes/Extreme SSD/Extra Datasets/video-dept...</td>\n",
       "      <td>/Volumes/Extreme SSD/Extra Datasets/video-dept...</td>\n",
       "      <td>/Volumes/Extreme SSD/Extra Datasets/video-dept...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>/Volumes/Extreme SSD/Extra Datasets/video-dept...</td>\n",
       "      <td>/Volumes/Extreme SSD/Extra Datasets/video-dept...</td>\n",
       "      <td>/Volumes/Extreme SSD/Extra Datasets/video-dept...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>/Volumes/Extreme SSD/Extra Datasets/video-dept...</td>\n",
       "      <td>/Volumes/Extreme SSD/Extra Datasets/video-dept...</td>\n",
       "      <td>/Volumes/Extreme SSD/Extra Datasets/video-dept...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>/Volumes/Extreme SSD/Extra Datasets/video-dept...</td>\n",
       "      <td>/Volumes/Extreme SSD/Extra Datasets/video-dept...</td>\n",
       "      <td>/Volumes/Extreme SSD/Extra Datasets/video-dept...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>/Volumes/Extreme SSD/Extra Datasets/video-dept...</td>\n",
       "      <td>/Volumes/Extreme SSD/Extra Datasets/video-dept...</td>\n",
       "      <td>/Volumes/Extreme SSD/Extra Datasets/video-dept...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                                 RGB  \\\n",
       "0  /Volumes/Extreme SSD/Extra Datasets/video-dept...   \n",
       "1  /Volumes/Extreme SSD/Extra Datasets/video-dept...   \n",
       "2  /Volumes/Extreme SSD/Extra Datasets/video-dept...   \n",
       "3  /Volumes/Extreme SSD/Extra Datasets/video-dept...   \n",
       "4  /Volumes/Extreme SSD/Extra Datasets/video-dept...   \n",
       "\n",
       "                                               Depth  \\\n",
       "0  /Volumes/Extreme SSD/Extra Datasets/video-dept...   \n",
       "1  /Volumes/Extreme SSD/Extra Datasets/video-dept...   \n",
       "2  /Volumes/Extreme SSD/Extra Datasets/video-dept...   \n",
       "3  /Volumes/Extreme SSD/Extra Datasets/video-dept...   \n",
       "4  /Volumes/Extreme SSD/Extra Datasets/video-dept...   \n",
       "\n",
       "                                            ColorMap  \n",
       "0  /Volumes/Extreme SSD/Extra Datasets/video-dept...  \n",
       "1  /Volumes/Extreme SSD/Extra Datasets/video-dept...  \n",
       "2  /Volumes/Extreme SSD/Extra Datasets/video-dept...  \n",
       "3  /Volumes/Extreme SSD/Extra Datasets/video-dept...  \n",
       "4  /Volumes/Extreme SSD/Extra Datasets/video-dept...  "
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "ad05bfa1",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|                                                                                                | 0/17552 [00:00<?, ?it/s]/var/folders/1v/pkj_pxz51qv0_y_c9mkb0dzh0000gn/T/ipykernel_82893/1830381067.py:11: DeprecationWarning: setting an array element with a sequence. This was supported in some cases where the elements are arrays with a single element. For example `np.array([1, np.array([2])], dtype=int)`. In the future this will raise the same ValueError as `np.array([1, [2]], dtype=int)`.\n",
      "  responses = np.array([[LABELS['blue']], LABELS['green'], LABELS['red'], LABELS['black'], LABELS['white']], dtype=np.float32)\n",
      " 51%|███████████████████████████████████████████▌                                         | 8985/17552 [01:34<01:30, 94.80it/s]\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "Input \u001b[0;32mIn [7]\u001b[0m, in \u001b[0;36m<cell line: 2>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      1\u001b[0m KNNs \u001b[38;5;241m=\u001b[39m []\n\u001b[1;32m      2\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m index, row \u001b[38;5;129;01min\u001b[39;00m tqdm(df\u001b[38;5;241m.\u001b[39miterrows(), total\u001b[38;5;241m=\u001b[39m\u001b[38;5;28mlen\u001b[39m(df)):\n\u001b[0;32m----> 3\u001b[0m     KNN \u001b[38;5;241m=\u001b[39m \u001b[43mgetClusters\u001b[49m\u001b[43m(\u001b[49m\u001b[43mrow\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[38;5;124;43mColorMap\u001b[39;49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m)\u001b[49m\n",
      "Input \u001b[0;32mIn [2]\u001b[0m, in \u001b[0;36mgetClusters\u001b[0;34m(imgPath)\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mgetClusters\u001b[39m(imgPath):\n\u001b[1;32m      2\u001b[0m     \u001b[38;5;66;03m#Thank you to https://codereview.stackexchange.com/questions/143529/opencv-3-using-k-nearest-neighbors-to-analyse-rgb-image\u001b[39;00m\n\u001b[0;32m----> 3\u001b[0m     img_BGR \u001b[38;5;241m=\u001b[39m \u001b[43mcv2\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mimread\u001b[49m\u001b[43m(\u001b[49m\u001b[43mimgPath\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m      4\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m img_BGR \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m      5\u001b[0m         \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mFileNotFoundError\u001b[39;00m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;132;01m{0}\u001b[39;00m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124m could not be opened!\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;241m.\u001b[39mformat(imgPath))\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "KNNs = []\n",
    "for index, row in tqdm(df.iterrows(), total=len(df)):\n",
    "    KNN = getClusters(row['ColorMap'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d31bb1d6",
   "metadata": {},
   "outputs": [],
   "source": [
    "img_clustered = img_BGR.copy()\n",
    "for colorName in BGR_COLORS.keys():\n",
    "    colorPixels(img_clustered, results_int, colorName)\n",
    "\n",
    "cv2.imshow(\"Original vs. Clustered\", np.hstack((img_BGR, img_clustered)))\n",
    "cv2.waitKey(0)\n",
    "cv2.destroyAllWindows()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c17390fc",
   "metadata": {},
   "source": [
    "### Conclusion\n",
    "not be a great choice for this project again, and smoothing was fine, but not great. ColorMap Smoothing.ipynb was better."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9be0c7bf",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "BASALT2022",
   "language": "python",
   "name": "basalt2022"
  },
  "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.9.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}