Total variation in-painting¶
Grayscale Images¶
A grayscale image is represented as an $m \times n$ matrix of intensities $U^\mathrm{orig}$ (typically between the values $0$ and $255$). We are given the values $U^\mathrm{orig}_{ij}$, for $(i,j) \in \mathcal K$, where $\mathcal K \subset \{1,\ldots, m\} \times \{1, \ldots, n\}$ is the set of indices corresponding to known pixel values. Our job is to in-paint the image by guessing the missing pixel values, i.e., those with indices not in $\mathcal K$. The reconstructed image will be represented by $U \in {\bf R}^{m \times n}$, where $U$ matches the known pixels, i.e., $U_{ij} = U^\mathrm{orig}_{ij}$ for $(i,j) \in \mathcal K$.
The reconstruction $U$ is found by minimizing the total variation of $U$, subject to matching the known pixel values. We will use the $\ell_2$ total variation, defined as $$\mathop{\bf tv}(U) = \sum_{i=1}^{m-1} \sum_{j=1}^{n-1} \left\| \left[ \begin{array}{c} U_{i+1,j}-U_{ij}\\ U_{i,j+1}-U_{ij} \end{array} \right] \right\|_2.$$ Note that the norm of the discretized gradient is not squared.
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.
from PIL import Image
import matplotlib.pyplot as plt
import numpy as np
# Load the images.
orig_img = Image.open("data/lena512.png")
corr_img = Image.open("data/lena512_corrupted.png")
# Convert to arrays.
Uorig = np.array(orig_img)
Ucorr = np.array(corr_img)
rows, cols = Uorig.shape
# Known is 1 if the pixel is known,
# 0 if the pixel was corrupted.
Known = np.zeros((rows, cols))
for i in xrange(rows):
for j in xrange(cols):
if Uorig[i, j] == Ucorr[i, j]:
Known[i, j] = 1
%matplotlib inline
fig, ax = plt.subplots(1, 2,figsize=(10, 5))
ax[0].imshow(orig_img);
ax[0].set_title("Original Image")
ax[0].axis('off')
ax[1].imshow(corr_img);
ax[1].set_title("Corrupted Image")
ax[1].axis('off');
The total variation in-painting problem can be easily expressed in CVXPY. We use the solver SCS, which finds the optimal value in a few seconds. The solvers ECOS and CVXOPT take much longer to solve this large problem.
# Recover the original image using total variation in-painting.
from cvxpy import *
U = Variable(rows, cols)
obj = Minimize(tv(U))
constraints = [mul_elemwise(Known, U) == mul_elemwise(Known, Ucorr)]
prob = Problem(obj, constraints)
# Use SCS to solve the problem.
prob.solve(verbose=True, solver=SCS)
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.
fig, ax = plt.subplots(1, 2,figsize=(10, 5))
# Display the in-painted image.
img_rec = Image.fromarray(U.value)
ax[0].imshow(img_rec);
ax[0].set_title("In-Painted Image")
ax[0].axis('off')
img_diff = Image.fromarray(10*np.abs(Uorig - U.value))
ax[1].imshow(img_diff);
ax[1].set_title("Difference Image")
ax[1].axis('off');
Color Images¶
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 $U^\mathrm{orig}$ (typically between the values $0$ and $255$). We are given the pixels $U^\mathrm{orig}_{ij}$, for $(i,j) \in \mathcal K$, where $\mathcal K \subset \{1,\ldots, m\} \times \{1, \ldots, n\}$ is the set of indices corresponding to known pixels. Each pixel $U^\mathrm{orig}_{ij}$ is a vector in ${\bf R}^3$ of RGB values. Our job is to in-paint the image by guessing the missing pixels, i.e., those with indices not in $\mathcal K$. The reconstructed image will be represented by $U \in {\bf R}^{m \times n \times 3}$, where $U$ matches the known pixels, i.e., $U_{ij} = U^\mathrm{orig}_{ij}$ for $(i,j) \in \mathcal K$.
The reconstruction $U$ is found by minimizing the total variation of $U$, subject to matching the known pixel values. We will use the $\ell_2$ total variation, defined as $$\mathop{\bf tv}(U) = \sum_{i=1}^{m-1} \sum_{j=1}^{n-1} \left\| \left[ \begin{array}{c} U_{i+1,j}-U_{ij}\\ U_{i,j+1}-U_{ij} \end{array} \right] \right\|_2.$$ Note that the norm of the discretized gradient is not squared.
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.
from PIL import Image
import matplotlib.pyplot as plt
import numpy as np
np.random.seed(1)
# Load the images.
orig_img = Image.open("data/lena512color.tiff")
# Convert to arrays.
Uorig = np.array(orig_img)
rows, cols, colors = Uorig.shape
# Known is 1 if the pixel is known,
# 0 if the pixel was corrupted.
# The Known matrix is initialized randomly.
Known = np.zeros((rows, cols, colors))
for i in xrange(rows):
for j in xrange(cols):
if np.random.random() > 0.7:
for k in xrange(colors):
Known[i, j, k] = 1
Ucorr = Known*Uorig
corr_img = Image.fromarray(np.uint8(Ucorr))
# Display the images.
%matplotlib inline
fig, ax = plt.subplots(1, 2,figsize=(10, 5))
ax[0].imshow(orig_img);
ax[0].set_title("Original Image")
ax[0].axis('off')
ax[1].imshow(corr_img);
ax[1].set_title("Corrupted Image")
ax[1].axis('off');
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, which finds the optimal value in 25 seconds. The solvers ECOS and CVXOPT don't scale to this large problem.
# Recover the original image using total variation in-painting.
from cvxpy import *
variables = []
constraints = []
for i in xrange(colors):
U = Variable(rows, cols)
variables.append(U)
constraints.append(mul_elemwise(Known[:, :, i], U) == mul_elemwise(Known[:, :, i], Ucorr[:, :, i]))
prob = Problem(Minimize(tv(*variables)), constraints)
prob.solve(verbose=True, solver=SCS)
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.
import matplotlib.pyplot as plt
import matplotlib.cm as cm
%matplotlib inline
# Load variable values into a single array.
rec_arr = np.zeros((rows, cols, colors), dtype=np.uint8)
for i in xrange(colors):
rec_arr[:, :, i] = variables[i].value
fig, ax = plt.subplots(1, 2,figsize=(10, 5))
# Display the in-painted image.
img_rec = Image.fromarray(rec_arr)
ax[0].imshow(img_rec);
ax[0].set_title("In-Painted Image")
ax[0].axis('off')
img_diff = Image.fromarray(np.abs(Uorig - rec_arr))
ax[1].imshow(img_diff);
ax[1].set_title("Difference Image")
ax[1].axis('off');
