k-means.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cv2\n",
    "import matplotlib.pyplot as plt\n",
    "from mpl_toolkits.mplot3d import Axes3D\n",
    "import numpy as np\n",
    "from tqdm import tqdm_notebook as tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "img_1 = cv2.imread('k-means-1.jpeg')\n",
    "img_2 = cv2.imread('k-means-2.png')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(15, 5))\n",
    "plt.subplot(1, 2, 1)\n",
    "plt.imshow(img_1)\n",
    "plt.subplot(1, 2, 2)\n",
    "plt.imshow(img_2)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def choose_random(K, vec):\n",
    "    m = len(vec)\n",
    "    idx = np.random.randint(0, m, K)\n",
    "    return vec[idx]\n",
    "\n",
    "def distortion(mu, c, vec):\n",
    "    return ((mu[c] - vec) ** 2).sum() / vec.shape[0]\n",
    "\n",
    "def distance(x_1, x_2):\n",
    "    return np.matmul((x_1-x_2), (x_1-x_2).transpose())\n",
    "\n",
    "def cluster_assignment(mu, vec):\n",
    "    return ((vec - mu[:, np.newaxis]) ** 2).sum(axis=2).argmin(axis=0)\n",
    "\n",
    "def move_centroid(mu, c, vec):\n",
    "    for i in range(len(mu)):\n",
    "        vec_sub = vec[c==i]\n",
    "        mu[i] = np.mean(vec_sub, axis=0)\n",
    "    return mu\n",
    "\n",
    "def k_means(img, K, plot=True, verbose=False):\n",
    "    l, w, ch = img.shape\n",
    "    vec_img = img.reshape(-1, ch).astype(int)\n",
    "    mu = choose_random(K, vec_img)\n",
    "    c = cluster_assignment(mu, vec_img)\n",
    "    last_dist = distortion(mu, c, vec_img) + 100\n",
    "    curr_dist = last_dist - 100\n",
    "    history = [curr_dist]\n",
    "    while last_dist - curr_dist > 1:\n",
    "        last_dist = curr_dist\n",
    "        c = cluster_assignment(mu, vec_img)\n",
    "        if verbose:\n",
    "            print(curr_dist)\n",
    "        mu = move_centroid(mu, c, vec_img)    \n",
    "        curr_dist = distortion(mu, c, vec_img)\n",
    "        history.append(curr_dist)\n",
    "    if plot:\n",
    "        img_compressed = mu[c].reshape(img.shape)\n",
    "        plt.figure(figsize=(15, 5))\n",
    "        plt.subplot(1, 3, 1)\n",
    "        plt.imshow(img)\n",
    "        plt.title('original image')\n",
    "        plt.axis('off')\n",
    "        plt.subplot(1, 3, 2)\n",
    "        plt.imshow(img_compressed)\n",
    "        plt.title('compressed image')\n",
    "        plt.axis('off')\n",
    "        plt.subplot(1, 3, 3)\n",
    "        plt.plot(range(len(history)), history)\n",
    "        plt.title('distortion over iterations')\n",
    "        plt.xlabel('iterations')\n",
    "        plt.ylabel('distortion')\n",
    "    return mu, c, history[-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mu, c, dist = k_means(img_1, 3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "mu, c, dist = k_means(img_2, 5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def elbow(img):\n",
    "    K_hist = []\n",
    "    dist_hist = []\n",
    "    for K in tqdm(range(1, 10)):\n",
    "        K_hist.append(K)\n",
    "        mu, c, dist = k_means(img, K, plot=False)\n",
    "        dist_hist.append(dist)\n",
    "    plt.plot(K_hist, dist_hist)\n",
    "    plt.xlabel(\"K\")\n",
    "    plt.ylabel(\"final distortion\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(15, 5))\n",
    "plt.subplot(1, 2, 1)\n",
    "plt.title('elbow plot of image 1')\n",
    "elbow(img_1)\n",
    "plt.subplot(1, 2, 2)\n",
    "elbow(img_2)\n",
    "plt.title('elbow plot of image 2')"
   ]
  },
  {
   "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.5.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}