{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Total variation in-painting\n",
    "\n",
    "## Grayscale Images\n",
    "\n",
    "A grayscale image is represented as an $m \\times n$ matrix of intensities\n",
    "$U^\\mathrm{orig}$ (typically between the values $0$ and $255$).\n",
    "We are given the values $U^\\mathrm{orig}_{ij}$, for $(i,j) \\in \\mathcal K$, where\n",
    "$\\mathcal K \\subset \\{1,\\ldots, m\\} \\times \\{1, \\ldots, n\\}$ is the set of indices\n",
    "corresponding to known pixel values.\n",
    "Our job is to *in-paint* the image by guessing the missing pixel values,\n",
    "*i.e.*, those with indices not in $\\mathcal K$.\n",
    "The reconstructed image will be represented by $U \\in {\\bf R}^{m \\times n}$,\n",
    "where $U$ matches the known pixels, *i.e.*,\n",
    "$U_{ij} = U^\\mathrm{orig}_{ij}$ for $(i,j) \\in \\mathcal K$.\n",
    "\n",
    "The reconstruction $U$ is found by minimizing the total variation of $U$,\n",
    "subject to matching the known pixel values. We will use the $\\ell_2$ total\n",
    "variation, defined as\n",
    "$$\\mathop{\\bf tv}(U) =\n",
    "\\sum_{i=1}^{m-1} \\sum_{j=1}^{n-1}\n",
    "\\left\\| \\left[ \\begin{array}{c}\n",
    " U_{i+1,j}-U_{ij}\\\\ U_{i,j+1}-U_{ij} \\end{array} \\right] \\right\\|_2.$$\n",
    "Note that the norm of the discretized gradient is *not* squared."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We load the original image and the corrupted image and construct the Known matrix. Both images are displayed below. The corrupted image has the missing pixels whited out."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 160332 characters]\n",
      "text/plain": [
       "<Figure size 720x360 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "# Load the images.\n",
    "u_orig = plt.imread(\"data/loki512.png\")\n",
    "u_corr = plt.imread(\"data/loki512_corrupted.png\")\n",
    "rows, cols = u_orig.shape\n",
    "\n",
    "# known is 1 if the pixel is known,\n",
    "# 0 if the pixel was corrupted.\n",
    "known = np.zeros((rows, cols))\n",
    "for i in range(rows):\n",
    "    for j in range(cols):\n",
    "         if u_orig[i, j] == u_corr[i, j]:\n",
    "            known[i, j] = 1\n",
    "\n",
    "%matplotlib inline\n",
    "fig, ax = plt.subplots(1, 2, figsize=(10, 5))\n",
    "ax[0].imshow(u_orig, cmap='gray')\n",
    "ax[0].set_title(\"Original Image\")\n",
    "ax[0].axis('off')\n",
    "ax[1].imshow(u_corr, cmap='gray');\n",
    "ax[1].set_title(\"Corrupted Image\")\n",
    "ax[1].axis('off');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The total variation in-painting problem can be easily expressed in CVXPY. We use the solver SCS, which scales to larger problems than ECOS does."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "----------------------------------------------------------------------------\n",
      "\tSCS v2.0.2 - Splitting Conic Solver\n",
      "\t(c) Brendan O'Donoghue, Stanford University, 2012-2017\n",
      "----------------------------------------------------------------------------\n",
      "Lin-sys: sparse-indirect, nnz in A = 1554199, CG tol ~ 1/iter^(2.00)\n",
      "eps = 1.00e-05, alpha = 1.50, max_iters = 5000, normalize = 1, scale = 1.00\n",
      "acceleration_lookback = 20, rho_x = 1.00e-03\n",
      "Variables n = 523265, constraints m = 1045507\n",
      "Cones:\tprimal zero / dual free vars: 262144\n",
      "\tsoc vars: 783363, soc blks: 261121\n",
      "Setup time: 1.23e-01s\n",
      "----------------------------------------------------------------------------\n",
      " Iter | pri res | dua res | rel gap | pri obj | dua obj | kap/tau | time (s)\n",
      "----------------------------------------------------------------------------\n",
      "     0| 5.19e+00  4.79e+00  1.00e+00 -5.21e+05  1.51e+04  0.00e+00  1.38e+00 \n",
      "   100| 4.46e-03  4.69e-03  3.82e-04  1.10e+04  1.10e+04  3.50e-12  3.95e+01 \n",
      "   200| 3.59e-04  3.83e-04  9.12e-05  1.10e+04  1.10e+04  4.20e-11  7.82e+01 \n",
      "   300| 7.10e-05  6.96e-05  2.77e-05  1.10e+04  1.10e+04  3.75e-11  1.14e+02 \n",
      "   400| 3.30e-05  3.39e-05  2.14e-06  1.10e+04  1.10e+04  6.65e-12  1.47e+02 \n",
      "   500| 2.77e-05  2.85e-05  1.35e-05  1.10e+04  1.10e+04  2.07e-11  1.81e+02 \n",
      "   600| 1.10e-05  1.09e-05  6.45e-06  1.10e+04  1.10e+04  1.48e-11  2.15e+02 \n",
      "   700| 1.00e-05  9.49e-06  1.94e-07  1.10e+04  1.10e+04  2.40e-11  2.48e+02 \n",
      "   720| 9.04e-06  8.24e-06  6.85e-07  1.10e+04  1.10e+04  1.09e-11  2.55e+02 \n",
      "----------------------------------------------------------------------------\n",
      "Status: Solved\n",
      "Timing: Solve time: 2.55e+02s\n",
      "\tLin-sys: avg # CG iterations: 9.58, avg solve time: 1.41e-01s\n",
      "\tCones: avg projection time: 3.42e-03s\n",
      "\tAcceleration: avg step time: 1.71e-01s\n",
      "----------------------------------------------------------------------------\n",
      "Error metrics:\n",
      "dist(s, K) = 2.1720e-04, dist(y, K*) = 3.7180e-04, s'y/|s||y| = -9.9097e-11\n",
      "primal res: |Ax + s - b|_2 / (1 + |b|_2) = 9.0439e-06\n",
      "dual res:   |A'y + c|_2 / (1 + |c|_2) = 8.2388e-06\n",
      "rel gap:    |c'x + b'y| / (1 + |c'x| + |b'y|) = 6.8544e-07\n",
      "----------------------------------------------------------------------------\n",
      "c'x = 11044.2661, -b'y = 11044.2813\n",
      "============================================================================\n",
      "optimal objective value: 11044.28989542425\n"
     ]
    }
   ],
   "source": [
    "# Recover the original image using total variation in-painting.\n",
    "import cvxpy as cp\n",
    "\n",
    "\n",
    "U = cp.Variable(shape=(rows, cols))\n",
    "obj = cp.Minimize(cp.tv(U))\n",
    "constraints = [cp.multiply(known, U) == cp.multiply(known, u_corr)]\n",
    "prob = cp.Problem(obj, constraints)\n",
    "\n",
    "# Use SCS to solve the problem.\n",
    "prob.solve(verbose=True, solver=cp.SCS)\n",
    "print(\"optimal objective value: {}\".format(obj.value))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After solving the problem, the in-painted image is stored in `U.value`. We display the in-painted image and the intensity difference between the original and in-painted images. The intensity difference is magnified by a factor of 10 so it is more visible."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 149748 characters]\n",
      "text/plain": [
       "<Figure size 720x360 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig, ax = plt.subplots(1, 2, figsize=(10, 5))\n",
    "# Display the in-painted image.\n",
    "ax[0].imshow(U.value, cmap='gray');\n",
    "ax[0].set_title(\"In-Painted Image\")\n",
    "ax[0].axis('off')\n",
    "\n",
    "img_diff = 10*np.abs(u_orig - U.value)\n",
    "ax[1].imshow(img_diff, cmap='gray');\n",
    "ax[1].set_title(\"Difference Image\")\n",
    "ax[1].axis('off');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Color Images\n",
    "\n",
    "For color images, the in-painting problem is similar to the grayscale case. A color image is represented as an $m \\times n \\times 3$ matrix of RGB values\n",
    "$U^\\mathrm{orig}$ (typically between the values $0$ and $255$).\n",
    "We are given the pixels $U^\\mathrm{orig}_{ij}$, for $(i,j) \\in \\mathcal K$, where\n",
    "$\\mathcal K \\subset \\{1,\\ldots, m\\} \\times \\{1, \\ldots, n\\}$ is the set of indices\n",
    "corresponding to known pixels. Each pixel $U^\\mathrm{orig}_{ij}$ is a vector in ${\\bf R}^3$ of RGB values.\n",
    "Our job is to *in-paint* the image by guessing the missing pixels,\n",
    "*i.e.*, those with indices not in $\\mathcal K$.\n",
    "The reconstructed image will be represented by $U \\in {\\bf R}^{m \\times n \\times 3}$,\n",
    "where $U$ matches the known pixels, *i.e.*,\n",
    "$U_{ij} = U^\\mathrm{orig}_{ij}$ for $(i,j) \\in \\mathcal K$.\n",
    "\n",
    "The reconstruction $U$ is found by minimizing the total variation of $U$,\n",
    "subject to matching the known pixel values. We will use the $\\ell_2$ total\n",
    "variation, defined as\n",
    "$$\\mathop{\\bf tv}(U) =\n",
    "\\sum_{i=1}^{m-1} \\sum_{j=1}^{n-1}\n",
    "\\left\\| \\left[ \\begin{array}{c}\n",
    " U_{i+1,j}-U_{ij}\\\\ \n",
    " U_{i,j+1}-U_{ij} \n",
    " \\end{array} \\right] \\right\\|_2.$$\n",
    "Note that the norm of the discretized gradient is *not* squared."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We load the original image and construct the Known matrix by randomly selecting 30% of the pixels to keep and discarding the others. The original and corrupted images are displayed below. The corrupted image has the missing pixels blacked out."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 372252 characters]\n",
      "text/plain": [
       "<Figure size 720x360 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "np.random.seed(1)\n",
    "# Load the images.\n",
    "u_orig = plt.imread(\"data/loki512color.png\")\n",
    "rows, cols, colors = u_orig.shape\n",
    "\n",
    "# known is 1 if the pixel is known,\n",
    "# 0 if the pixel was corrupted.\n",
    "# The known matrix is initialized randomly.\n",
    "known = np.zeros((rows, cols, colors))\n",
    "for i in range(rows):\n",
    "    for j in range(cols):\n",
    "        if np.random.random() > 0.7:\n",
    "            for k in range(colors):\n",
    "                known[i, j, k] = 1        \n",
    "u_corr = known * u_orig\n",
    "\n",
    "# Display the images.\n",
    "%matplotlib inline\n",
    "fig, ax = plt.subplots(1, 2, figsize=(10, 5))\n",
    "ax[0].imshow(u_orig, cmap='gray');\n",
    "ax[0].set_title(\"Original Image\")\n",
    "ax[0].axis('off')\n",
    "ax[1].imshow(u_corr);\n",
    "ax[1].set_title(\"Corrupted Image\")\n",
    "ax[1].axis('off');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We express the total variation color in-painting problem in CVXPY using three matrix variables (one for the red values, one for the blue values, and one for the green values). We use the solver SCS; the solvers ECOS and CVXOPT don't scale to this large problem."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARN: A->p (column pointers) not strictly increasing, column 523264 empty\n",
      "WARN: A->p (column pointers) not strictly increasing, column 785408 empty\n",
      "WARN: A->p (column pointers) not strictly increasing, column 1047552 empty\n",
      "----------------------------------------------------------------------------\n",
      "\tSCS v2.0.2 - Splitting Conic Solver\n",
      "\t(c) Brendan O'Donoghue, Stanford University, 2012-2017\n",
      "----------------------------------------------------------------------------\n",
      "Lin-sys: sparse-indirect, nnz in A = 3630814, CG tol ~ 1/iter^(2.00)\n",
      "eps = 1.00e-05, alpha = 1.50, max_iters = 5000, normalize = 1, scale = 1.00\n",
      "acceleration_lookback = 20, rho_x = 1.00e-03\n",
      "Variables n = 1047553, constraints m = 2614279\n",
      "Cones:\tprimal zero / dual free vars: 786432\n",
      "\tsoc vars: 1827847, soc blks: 261121\n",
      "Setup time: 3.00e-01s\n",
      "----------------------------------------------------------------------------\n",
      " Iter | pri res | dua res | rel gap | pri obj | dua obj | kap/tau | time (s)\n",
      "----------------------------------------------------------------------------\n",
      "     0| 1.16e+01  1.18e+01  1.00e+00 -1.02e+06  3.34e+04  1.53e-10  3.81e+00 \n",
      "   100| 2.19e-03  2.32e-03  6.52e-04  1.14e+04  1.15e+04  7.82e-12  1.08e+02 \n",
      "   200| 4.23e-04  3.78e-04  4.97e-05  1.15e+04  1.15e+04  1.34e-11  2.04e+02 \n",
      "   300| 9.58e-05  1.10e-04  5.94e-05  1.15e+04  1.15e+04  1.46e-11  2.96e+02 \n",
      "   400| 4.54e-05  4.57e-05  6.08e-06  1.15e+04  1.15e+04  5.96e-12  3.85e+02 \n",
      "   500| 2.92e-05  3.19e-05  3.42e-06  1.15e+04  1.15e+04  3.37e-11  4.74e+02 \n",
      "   600| 1.77e-05  1.87e-05  1.20e-05  1.15e+04  1.15e+04  3.08e-11  5.60e+02 \n",
      "   700| 1.40e-05  1.43e-05  7.45e-06  1.15e+04  1.15e+04  9.77e-12  6.47e+02 \n",
      "   760| 9.03e-06  9.70e-06  2.43e-06  1.15e+04  1.15e+04  7.02e-12  6.99e+02 \n",
      "----------------------------------------------------------------------------\n",
      "Status: Solved\n",
      "Timing: Solve time: 6.99e+02s\n",
      "\tLin-sys: avg # CG iterations: 11.66, avg solve time: 4.29e-01s\n",
      "\tCones: avg projection time: 4.72e-03s\n",
      "\tAcceleration: avg step time: 3.94e-01s\n",
      "----------------------------------------------------------------------------\n",
      "Error metrics:\n",
      "dist(s, K) = 1.8769e-05, dist(y, K*) = 1.1246e-04, s'y/|s||y| = 6.2851e-11\n",
      "primal res: |Ax + s - b|_2 / (1 + |b|_2) = 9.0269e-06\n",
      "dual res:   |A'y + c|_2 / (1 + |c|_2) = 9.7005e-06\n",
      "rel gap:    |c'x + b'y| / (1 + |c'x| + |b'y|) = 2.4293e-06\n",
      "----------------------------------------------------------------------------\n",
      "c'x = 11465.6528, -b'y = 11465.5971\n",
      "============================================================================\n",
      "optimal objective value: 11465.652787130613\n"
     ]
    }
   ],
   "source": [
    "# Recover the original image using total variation in-painting.\n",
    "import cvxpy as cp\n",
    "\n",
    "\n",
    "variables = []\n",
    "constraints = []\n",
    "for i in range(colors):\n",
    "    U = cp.Variable(shape=(rows, cols))\n",
    "    variables.append(U)\n",
    "    constraints.append(cp.multiply(known[:, :, i], U) == cp.multiply(known[:, :, i], u_corr[:, :, i]))\n",
    "\n",
    "prob = cp.Problem(cp.Minimize(cp.tv(*variables)), constraints)\n",
    "prob.solve(verbose=True, solver=cp.SCS)\n",
    "print(\"optimal objective value: {}\".format(prob.value))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After solving the problem, the RGB values of the in-painted image are stored in the value fields of the three variables. We display the in-painted image and the difference in RGB values at each pixel of the original and in-painted image. Though the in-painted image looks almost identical to the original image, you can see that many of the RGB values differ."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(-0.5, 511.5, 511.5, -0.5)"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 450460 characters]\n",
      "text/plain": [
       "<Figure size 720x360 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.cm as cm\n",
    "%matplotlib inline\n",
    "\n",
    "rec_arr = np.zeros((rows, cols, colors))\n",
    "for i in range(colors):\n",
    "    rec_arr[:, :, i] = variables[i].value\n",
    "rec_arr = np.clip(rec_arr, 0, 1)\n",
    "\n",
    "fig, ax = plt.subplots(1, 2,figsize=(10, 5))\n",
    "ax[0].imshow(rec_arr)\n",
    "ax[0].set_title(\"In-Painted Image\")\n",
    "ax[0].axis('off')\n",
    "\n",
    "img_diff = np.clip(10 * np.abs(u_orig - rec_arr), 0, 1)\n",
    "ax[1].imshow(img_diff)\n",
    "ax[1].set_title(\"Difference Image\")\n",
    "ax[1].axis('off')"
   ]
  },
  {
   "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.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
