notebooks/tv-potts.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.core.display import display, HTML\n",
    "display(HTML(\"<style>.container { width:100% !important; }</style>\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib notebook\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import skimage as ski\n",
    "import scipy.sparse\n",
    "import time\n",
    "from sklearn.cluster import KMeans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def forward_differences(M,N):\n",
    "    row = np.arange(0,M*N)\n",
    "    dat = np.ones(M*N)\n",
    "    col = np.arange(0,M*N).reshape(M,N)\n",
    "    col_xp = np.hstack([col[:,1:], col[:,-1:]])\n",
    "    col_yp = np.vstack([col[1:,:], col[-1:,:]])\n",
    "    \n",
    "    FD1 = scipy.sparse.coo_matrix((dat, (row, col_xp.flatten())), shape=(M*N, M*N))- \\\n",
    "          scipy.sparse.coo_matrix((dat, (row, col.flatten())), shape=(M*N, M*N))\n",
    "\n",
    "    FD2 = scipy.sparse.coo_matrix((dat, (row, col_yp.flatten())), shape=(M*N, M*N))- \\\n",
    "          scipy.sparse.coo_matrix((dat, (row, col.flatten())), shape=(M*N, M*N))\n",
    "    \n",
    "    FD = scipy.sparse.vstack([FD1, FD2])\n",
    "    \n",
    "    return FD"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def proj_inf_l2(p, tau):\n",
    "    # size must be (K,N), l2 over K, inf over N\n",
    "    norm_p = np.sqrt(np.sum(p**2, axis=0, keepdims=True))\n",
    "    p /= np.maximum(1, norm_p/tau)\n",
    "    return p\n",
    "\n",
    "def proj_inf(p, tau):\n",
    "    return np.clip(p, -tau, tau)\n",
    "\n",
    "def proj_simplex(x):\n",
    "    # Code adapted from Laurent Condat\n",
    "    # projection onto unit simplex along the first dimension\n",
    "    s = x.shape\n",
    "    K = s[0]\n",
    "    x = x.reshape(K,-1)\n",
    "    k = np.linspace(1, K, K)\n",
    "    x_s = -np.sort(-x, axis=0)\n",
    "    t = (np.cumsum(x_s, axis=0)-1.0)/k[:,None]\n",
    "    mu = np.max(t, axis=0, keepdims=True)\n",
    "    return np.maximum(0.0, x-mu).reshape(s)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def potts_simple_pd(f, maxit=1000, check=100, verbose=0):\n",
    "    \n",
    "    # Simple convex relaxation (Zach)\n",
    "    K,M,N = f.shape\n",
    "    f = f.reshape(K,M*N)\n",
    "    \n",
    "    # primal variables\n",
    "    u = np.zeros((K,M*N))\n",
    "\n",
    "    # dual variable\n",
    "    p = np.zeros((K,2*M*N))\n",
    "    \n",
    "    # make nabla operator\n",
    "    FD = forward_differences(M,N)\n",
    "\n",
    "    # primal and dual step sizes\n",
    "    # Chosen using diagonal preconditioning\n",
    "    fact = 1.0\n",
    "    tau_u = fact*1.0/4.0\n",
    "    sigma_p = 1.0/(2.0)/fact\n",
    "    theta = 1.0\n",
    "    \n",
    "    t0 = time.time()\n",
    "    for it in range(0,maxit):\n",
    "\n",
    "        # remeber old\n",
    "        ui = u.copy()\n",
    "        \n",
    "        # primal update\n",
    "        for i in range(K):\n",
    "            u[i] -= tau_u*(FD.T@(p[i]) + f[i])\n",
    "                \n",
    "        # projection of u onto simplex\n",
    "        u = proj_simplex(u)\n",
    "        \n",
    "        # overrelaxation\n",
    "        ui = u + theta*(u-ui)\n",
    "        \n",
    "        # dual update\n",
    "        for i in range(K):\n",
    "            p[i] += sigma_p*(FD@(ui[i]))\n",
    "            tmp = p[i].reshape(2,M*N)\n",
    "            tmp = proj_inf_l2(tmp, 0.5)\n",
    "            p[i] = tmp.reshape(2*M*N)\n",
    "            \n",
    "        if verbose > 0:\n",
    "            if it%check == check-1:\n",
    "                TV = 0\n",
    "                for i in range(K):\n",
    "                    TV += np.sum((FD@u[i])*p[i])\n",
    "                    \n",
    "                print(\"iter = \", it,\n",
    "                      \", time = \", \"{:.3f}\".format(time.time()-t0),\n",
    "                      \", TV = \", \"{:.3f}\".format(TV),\n",
    "                      end=\"\\r\")\n",
    "    print(\"\\n\")\n",
    "    return u.reshape(K,M,N)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def potts_tight_pd(f, maxit=1000, check=100, verbose=0):\n",
    "    # Tight convex relaxation (CCP 2008)\n",
    "    K,M,N = f.shape\n",
    "    f = f.reshape(K,M*N)\n",
    "    \n",
    "    # number of dual constraints\n",
    "    KK = np.int(K*(K-1)/2)\n",
    "\n",
    "    # primal variables\n",
    "    u = np.zeros((K,M*N))\n",
    "    v = np.zeros((KK,2*M*N))\n",
    "    \n",
    "    # dual variable\n",
    "    p = np.zeros((K,2*M*N))\n",
    "    q = np.zeros((KK,2*M*N))\n",
    "    \n",
    "    # make nabla operator\n",
    "    FD = forward_differences(M,N)\n",
    "\n",
    "    # primal and dual step sizes\n",
    "    # Chosen using diagonal preconditioning\n",
    "    fact = 1.0\n",
    "    tau_u = fact*1.0/4.0\n",
    "    tau_v = fact*1.0/3.0\n",
    "    sigma_p = 1.0/(2.0+K-1)/fact\n",
    "    sigma_q = 1.0/fact\n",
    "    theta = 1.0\n",
    "    \n",
    "    t0 = time.time()\n",
    "    for it in range(0,maxit):\n",
    "\n",
    "        # remeber old\n",
    "        ui = u.copy()\n",
    "        vi = v.copy()\n",
    "        \n",
    "        # primal update\n",
    "        for i in range(K):\n",
    "            u[i] -= tau_u*(FD.T@(p[i]) + f[i])\n",
    "\n",
    "        idx = 0\n",
    "        for i in range(K-1):\n",
    "            for j in range(i+1,K):\n",
    "                v[idx] -= tau_v*(p[i]-p[j]-q[idx])\n",
    "                idx +=1\n",
    "                \n",
    "        # projection of u onto simplex\n",
    "        u = proj_simplex(u)\n",
    "        \n",
    "        # overrelaxation\n",
    "        ui = u + theta*(u-ui)\n",
    "        vi = v + theta*(v-vi)\n",
    "        \n",
    "        # dual update\n",
    "        for i in range(K):\n",
    "            p[i] += sigma_p*(FD@(ui[i]))\n",
    "            \n",
    "        idx = 0\n",
    "        for i in range(K-1):\n",
    "            for j in range(i+1,K):\n",
    "                p[i] += sigma_p*vi[idx]\n",
    "                p[j] -= sigma_p*vi[idx]\n",
    "                q[idx] -= sigma_q*vi[idx]\n",
    "                idx +=1\n",
    "                \n",
    "        # projection\n",
    "        idx = 0\n",
    "        for i in range(K-1):\n",
    "            for j in range(i+1,K):\n",
    "                tmp = q[idx].reshape(2,M*N)\n",
    "                tmp = proj_inf_l2(tmp, 1.0)\n",
    "                q[idx] = tmp.reshape(2*M*N)\n",
    "                idx +=1\n",
    "        \n",
    "        if verbose > 0:\n",
    "            if it%check == check-1:\n",
    "                TV = 0\n",
    "                for i in range(K):\n",
    "                    TV += np.sum((FD@u[i])*p[i])\n",
    "\n",
    "                print(\"iter = \", it,\n",
    "                      \", time = \", \"{:.3f}\".format(time.time()-t0),\n",
    "                      \", TV = \", \"{:.3f}\".format(TV),\n",
    "                      end=\"\\r\")\n",
    "    print(\"\\n\")            \n",
    "    return u.reshape(K,M,N)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_rgb(u,means):\n",
    "    K,M,N = u.shape\n",
    "    u_rgb = np.zeros((M,N,3))\n",
    "    for i in range(K):\n",
    "        mu = means[i]\n",
    "        u_rgb[:,:,0] += u[i]*mu[0]\n",
    "        u_rgb[:,:,1] += u[i]*mu[1]\n",
    "        u_rgb[:,:,2] += u[i]*mu[2]\n",
    "    return u_rgb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_squared_dist(g, means):\n",
    "    M,N,_ = g.shape\n",
    "    K = means.shape[0]\n",
    "    f = np.zeros((K, M,N))\n",
    "    for i in range(K):\n",
    "        mu = means[i]\n",
    "        f[i] = (np.sum(0.5*(pixels-mu)**2, axis=1).reshape(M,N))\n",
    "    return f"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load image\n",
    "g = np.double(ski.io.imread(\"col3.png\"))\n",
    "M,N,_ = g.shape\n",
    "pixels = g.reshape(-1,3)\n",
    "\n",
    "# number of phases\n",
    "K = 3\n",
    "\n",
    "# mean values\n",
    "means_rgb = np.array([[255., 0.,0.],[0.,255,0.],[0.,0.,255.]])\n",
    "\n",
    "# compute field\n",
    "f = compute_squared_dist(g, means_rgb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Simple convex relaxation\")\n",
    "u_simple = potts_simple_pd(f, verbose=1, maxit=5000, check=100)\n",
    "u_rgb = make_rgb(u_simple, means_rgb)\n",
    "\n",
    "plt.figure(figsize=(12,3))\n",
    "\n",
    "plt.subplot(141)\n",
    "plt.imshow(u_rgb/255.0)\n",
    "\n",
    "plt.subplot(142)\n",
    "plt.imshow(u_simple[0], cmap=\"gray\")\n",
    "\n",
    "plt.subplot(143)\n",
    "plt.imshow(u_simple[1], cmap=\"gray\")\n",
    "\n",
    "plt.subplot(144)\n",
    "plt.imshow(u_simple[2], cmap=\"gray\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Tight convex relaxation\")\n",
    "u_tight = potts_tight_pd(f, verbose=1, maxit=5000, check=100)\n",
    "u_rgb = make_rgb(u_tight, means_rgb)\n",
    "\n",
    "plt.figure(figsize=(12,3))\n",
    "\n",
    "plt.subplot(141)\n",
    "plt.imshow(u_rgb/255.0)\n",
    "\n",
    "plt.subplot(142)\n",
    "plt.imshow(u_tight[0], cmap=\"gray\")\n",
    "\n",
    "plt.subplot(143)\n",
    "plt.imshow(u_tight[1], cmap=\"gray\")\n",
    "\n",
    "plt.subplot(144)\n",
    "plt.imshow(u_tight[2], cmap=\"gray\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load image\n",
    "g = ski.io.imread(\"tagpfauenauge.jpg\")/255.0\n",
    "M,N,_ = g.shape\n",
    "pixels = g.reshape(-1,3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Number of phases\n",
    "K = 16\n",
    "\n",
    "# K-means clustering\n",
    "kmeans = KMeans(n_clusters=K, random_state=0).fit(pixels)\n",
    "\n",
    "# compute field\n",
    "f = compute_squared_dist(g, kmeans.cluster_centers_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lamb = 5.0\n",
    "u_simple = potts_simple_pd(f*lamb, maxit=300, check=100, verbose=1)\n",
    "u_rgb = make_rgb(u_simple, kmeans.cluster_centers_)\n",
    "    \n",
    "plt.figure(figsize=(10,3))\n",
    "\n",
    "plt.subplot(121)\n",
    "plt.imshow(g)\n",
    "\n",
    "plt.subplot(122)\n",
    "plt.imshow(u_rgb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lamb = 5.0\n",
    "u_tight = potts_tight_pd(f*lamb, maxit=300, check=100, verbose=1)\n",
    "u_rgb = make_rgb(u_tight, kmeans.cluster_centers_)\n",
    "\n",
    "plt.figure(figsize=(10,3))\n",
    "\n",
    "plt.subplot(121)\n",
    "plt.imshow(g)\n",
    "\n",
    "plt.subplot(122)\n",
    "plt.imshow(u_rgb)"
   ]
  },
  {
   "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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}