{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 代码实现"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "我们来编写$k$均值算法，并利用RGB值对图像进行分割。先导入一张图像。\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.cluster import KMeans\n",
    "from matplotlib.image import imread\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from PIL import Image\n",
    "\n",
    "# 输入一张图像，图像来源于参考文献[5]\n",
    "image = imread('segmentation.jpeg')[:,:,:3]\n",
    "# 将RGB值统一到0-1内\n",
    "if np.max(image)>1:\n",
    "    image = image / 255\n",
    "\n",
    "X = image.reshape(-1, image.shape[2])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0 1 2 3]\n"
     ]
    }
   ],
   "source": [
    "# 利用k-means算法进行聚类\n",
    "segmented_imgs = []\n",
    "\n",
    "# 设定聚类中心个数\n",
    "n_cluster= 4\n",
    "kmeans = KMeans(n_clusters=n_cluster, random_state=42).fit(X)\n",
    "print(np.unique(kmeans.labels_))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-10-12T07:10:20.112994Z",
     "iopub.status.busy": "2022-10-12T07:10:20.111763Z",
     "iopub.status.idle": "2022-10-12T07:10:21.589547Z",
     "shell.execute_reply": "2022-10-12T07:10:21.588090Z",
     "shell.execute_reply.started": "2022-10-12T07:10:20.112945Z"
    }
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 159024 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 43168 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# 为每一个类别赋予一个对应的颜色，用于展示\n",
    "def decode_segmap(label_mask, plot=False):\n",
    "    label_colours = np.asarray([[79, 103, 67], [143, 146, 126], \n",
    "                                [129, 94, 64], [52, 53, 55],\n",
    "                                [96, 84, 70], [164, 149, 129]])\n",
    "    r = label_mask.copy()\n",
    "    g = label_mask.copy()\n",
    "    b = label_mask.copy()\n",
    "    # 为每个类别赋予对应的R、G、B值\n",
    "    for ll in range(0, 6):\n",
    "        r[label_mask == ll] = label_colours[ll, 0]\n",
    "        g[label_mask == ll] = label_colours[ll, 1]\n",
    "        b[label_mask == ll] = label_colours[ll, 2]\n",
    "\n",
    "    rgb = np.zeros((label_mask.shape[0], label_mask.shape[1], 3))\n",
    "    rgb[:, :, 0] = r\n",
    "    rgb[:, :, 1] = g\n",
    "    rgb[:, :, 2] = b\n",
    "\n",
    "    return rgb\n",
    "\n",
    "\n",
    "# 获得预测的标签\n",
    "segmented_img = kmeans.cluster_centers_[kmeans.labels_]\n",
    "segmented_imgs = decode_segmap(kmeans.labels_.reshape(\n",
    "                        image.shape[0],image.shape[1]))\n",
    "\n",
    "# 展示结果\n",
    "plt.imshow(image[:,:,:3])\n",
    "plt.title('Original image')\n",
    "plt.axis('off')\n",
    "plt.show()\n",
    "\n",
    "plt.imshow(segmented_imgs.astype(np.uint8))\n",
    "plt.title('{} center'.format(n_cluster))\n",
    "plt.axis('off')\n",
    "plt.show()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "再次动手尝试图像分割，并比较设置不同个数的聚类中心对图像分割的影响。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-10-12T07:10:02.205430Z",
     "iopub.status.busy": "2022-10-12T07:10:02.204848Z",
     "iopub.status.idle": "2022-10-12T07:10:08.992311Z",
     "shell.execute_reply": "2022-10-12T07:10:08.991733Z",
     "shell.execute_reply.started": "2022-10-12T07:10:02.205401Z"
    }
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 242992 characters]",
      "text/plain": [
       "<Figure size 1200x800 with 6 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "image = imread('segmentation.jpeg')[:,:,:3]\n",
    "# 将RGB值统一到0-255内\n",
    "if np.max(image)>1:\n",
    "    image = image / 255\n",
    "    \n",
    "\n",
    "sp = image.shape\n",
    "# 增加xy坐标的信息\n",
    "# 设定一个权重，对坐标信息加权\n",
    "weight = 2\n",
    "y = weight * np.array([[i for i in range(sp[1])] \n",
    "                       for j in range(sp[0])]) / sp[0] / sp[1]\n",
    "x = weight * np.array([[j for i in range(sp[1])] \n",
    "                       for j in range(sp[0])])/ sp[0] / sp[1]\n",
    "image = np.append(image, x.reshape(sp[0], sp[1], 1), axis=2)\n",
    "image = np.append(image, y.reshape(sp[0], sp[1], 1), axis=2)\n",
    "\n",
    "X = image.reshape(-1, image.shape[2])\n",
    "segmented_imgs = []\n",
    "\n",
    "# 将 K 分别设置为6、5、4、3、2\n",
    "n_colors = (6, 5, 4, 3, 2)\n",
    "for n_cluster in n_colors:\n",
    "    kmeans = KMeans(n_clusters=n_cluster, random_state=42).fit(X)\n",
    "    segmented_img = kmeans.cluster_centers_[kmeans.labels_]\n",
    "    segmented_imgs.append(decode_segmap(\n",
    "        kmeans.labels_.reshape(image.shape[0],\n",
    "                               image.shape[1])).astype(np.uint8))\n",
    "\n",
    "# 展示结果\n",
    "plt.figure(figsize=(12,8))\n",
    "plt.subplot(231)\n",
    "plt.imshow(image[:,:,:3])\n",
    "plt.title('Original image')\n",
    "plt.axis('off')\n",
    "\n",
    "for idx,n_clusters in enumerate(n_colors):\n",
    "    plt.subplot(232+idx)\n",
    "    plt.imshow(segmented_imgs[idx])\n",
    "    plt.title('{} center'.format(n_clusters))\n",
    "    plt.axis('off')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "接下来，我们将动手编程实现归一化图切割算法。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-11-20T06:59:29.560525Z",
     "iopub.status.busy": "2022-11-20T06:59:29.559668Z",
     "iopub.status.idle": "2022-11-20T06:59:30.809302Z",
     "shell.execute_reply": "2022-11-20T06:59:30.808285Z",
     "shell.execute_reply.started": "2022-11-20T06:59:29.560394Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.image.AxesImage at 0x162442580>"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 98704 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import numpy as np\n",
    "import cv2\n",
    "from PIL import Image\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "sample_img = cv2.imread('segmentation.jpeg', cv2.IMREAD_GRAYSCALE).astype(\n",
    "                                                np.float32) / 255\n",
    "plt.imshow(sample_img, cmap='gray')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "我们接着编写权重矩阵。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-11-20T06:59:30.812754Z",
     "iopub.status.busy": "2022-11-20T06:59:30.811390Z",
     "iopub.status.idle": "2022-11-20T07:24:25.257332Z",
     "shell.execute_reply": "2022-11-20T07:24:25.256304Z",
     "shell.execute_reply.started": "2022-11-20T06:59:30.812708Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[1.        , 0.81873075, 0.44863849, ..., 0.        , 0.        ,\n",
       "        0.        ],\n",
       "       [0.81873075, 1.        , 0.81747264, ..., 0.        , 0.        ,\n",
       "        0.        ],\n",
       "       [0.44863849, 0.81747264, 1.        , ..., 0.        , 0.        ,\n",
       "        0.        ],\n",
       "       ...,\n",
       "       [0.        , 0.        , 0.        , ..., 1.        , 0.81873075,\n",
       "        0.44863849],\n",
       "       [0.        , 0.        , 0.        , ..., 0.81873075, 1.        ,\n",
       "        0.81747264],\n",
       "       [0.        , 0.        , 0.        , ..., 0.44863849, 0.81747264,\n",
       "        1.        ]])"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def cal_dist_weighted_matrix(size, r, sig_X):\n",
    "    h, w = size\n",
    "    X_matrix = np.zeros((h*w, h*w))\n",
    "    for i in range(h*w):\n",
    "        for j in range(h*w):\n",
    "            i_row, i_col = i // w, i % w\n",
    "            j_row, j_col = j // w, j % w\n",
    "            dist = np.power(i_row - j_row, 2) + \\\n",
    "                np.power(i_col - j_col, 2)\n",
    "            \n",
    "            if np.sqrt(dist) < r:    \n",
    "                X_matrix[i, j] = np.exp(-dist / sig_X)\n",
    "    return X_matrix\n",
    "    \n",
    "    \n",
    "def set_weighted_matrix(img, sig_I=0.01, sig_X=5, r=10):\n",
    "    vec_img = img.flatten()\n",
    "    F_matrix = np.power(vec_img[None, :] - vec_img[:, None], 2)\n",
    "    F_matrix /= sig_I\n",
    "    F_matrix = np.exp(-F_matrix)\n",
    "    \n",
    "    X_matrix = cal_dist_weighted_matrix(img.shape, r, sig_X)\n",
    "    return F_matrix * X_matrix\n",
    "\n",
    "\n",
    "weighted_matrix = set_weighted_matrix(sample_img)\n",
    "\n",
    "weighted_matrix"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "接着求解每一个特征值对应的特征向量。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-11-20T07:29:40.250671Z",
     "iopub.status.busy": "2022-11-20T07:29:40.250227Z",
     "iopub.status.idle": "2022-11-20T07:29:40.259628Z",
     "shell.execute_reply": "2022-11-20T07:29:40.258310Z",
     "shell.execute_reply.started": "2022-11-20T07:29:40.250634Z"
    }
   },
   "outputs": [],
   "source": [
    "from scipy import sparse\n",
    "from scipy.sparse.linalg import eigs, svds\n",
    "import skimage\n",
    "\n",
    "def n_cuts(W, image):\n",
    "    n = W.shape[0]\n",
    "    s_D = sparse.csr_matrix((n, n))\n",
    "    d_i = W.sum(axis=0)\n",
    "    for i in range(n):\n",
    "        s_D[i, i] = d_i[i]\n",
    "    s_W = sparse.csr_matrix(W)\n",
    "    # print(s_W.shape)\n",
    "    s_D_nhalf = np.sqrt(s_D).power(-1)\n",
    "    # print(s_D.shape)\n",
    "\n",
    "    L = s_D_nhalf @ (s_D - s_W) @ s_D_nhalf\n",
    "    # print(L.shape)\n",
    "\n",
    "    _, eigenvalues, eigenvectors = svds(L, which='SM')\n",
    "    print(eigenvectors.shape)\n",
    "\n",
    "    for i in range(1, 5):\n",
    "        print(eigenvectors[i].shape)\n",
    "        Partition = eigenvectors[i] > \\\n",
    "            np.sum(eigenvectors[i])/len(eigenvectors[i])\n",
    "        print(Partition.shape)\n",
    "        skimage.io.imshow(Partition.reshape(image.shape))\n",
    "        plt.title('Ncut')\n",
    "        plt.show()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "展示相应的分割结果。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_cuts(weighted_matrix, sample_img)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<center>\n",
    "    <img style=\"border-radius: 0.3125em;\" \n",
    "    src=\"https://pic1.zhimg.com/80/v2-31642a71d6d95f791284a142399d1800_1440w.webp\" width=400>\n",
    "    <br>\n",
    "    <div style=\"color:orange; \n",
    "    display: inline-block;\n",
    "    color: #999;\n",
    "    padding: 2px;\"></div>\n",
    "</center>"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
