From 13d1a3d111aa6cf04e42b5c20b824fa6d20119c4 Mon Sep 17 00:00:00 2001 From: Emmanuelle Gouillart Date: Sat, 28 May 2011 14:55:01 +0200 Subject: [PATCH] New module for total variation denoising, for 2D and 3D arrays. --- scikits/image/filter/__init__.py | 1 + scikits/image/filter/tv_denoise.py | 313 +++++++++++++++++++++++++++++ 2 files changed, 314 insertions(+) create mode 100644 scikits/image/filter/tv_denoise.py diff --git a/scikits/image/filter/__init__.py b/scikits/image/filter/__init__.py index ef33468d..bec0034d 100644 --- a/scikits/image/filter/__init__.py +++ b/scikits/image/filter/__init__.py @@ -2,3 +2,4 @@ from lpi_filter import * from ctmf import median_filter from canny import canny from edges import sobel, hsobel, vsobel, hprewitt, vprewitt, prewitt +from tv_denoise import tv_denoise diff --git a/scikits/image/filter/tv_denoise.py b/scikits/image/filter/tv_denoise.py new file mode 100644 index 00000000..da5b0903 --- /dev/null +++ b/scikits/image/filter/tv_denoise.py @@ -0,0 +1,313 @@ +import numpy as np + +def _tv_denoise_3d(im, eps=2.e-4, weight=100, keep_type=False, n_iter_max=200): + """ + Perform total-variation denoising on 3-D arrays + + Parameters + ---------- + im: ndarray + 3-D input data to be denoised + + eps: float, optional + relative difference of the value of the cost function that determines + the stop criterion. The algorithm stops when + (E_(n-1) - E_n) < eps * E_0 + + weight: float, optional + denoising weight. The greater ``weight``, the more denoising (at + the expense of fidelity to ``input``) + + keep_type: bool, optional (False) + whether the output has the same dtype as the input array. + keep_type is False by default, and the dtype of the output + is np.float + + n_iter_max: int, optional + maximal number of iterations used for the optimization. + + Returns + ------- + out: ndarray + denoised array + + Notes + ----- + Rudin, Osher and Fatemi algorithm + + Examples + --------- + First build synthetic noisy data + >>> x, y, z = np.ogrid[0:40, 0:40, 0:40] + >>> mask = (x -22)**2 + (y - 20)**2 + (z - 17)**2 < 8**2 + >>> mask = mask.astype(np.float) + >>> mask += 0.2*np.random.randn(*mask.shape) + >>> res = tv_denoise_3d(mask, weight=100) + """ + im_type = im.dtype + if im_type is not np.float: + im = im.astype(np.float) + px = np.zeros_like(im) + py = np.zeros_like(im) + pz = np.zeros_like(im) + gx = np.zeros_like(im) + gy = np.zeros_like(im) + gz = np.zeros_like(im) + d = np.zeros_like(im) + i = 0 + while i < n_iter_max: + d = - px - py - pz + d[1:] += px[:-1] + d[:, 1:] += py[:, :-1] + d[:, :, 1:] += pz[:, :, :-1] + + out = im + d + E = (d**2).sum() + + gx[:-1] = np.diff(out, axis=0) + gy[:, :-1] = np.diff(out, axis=1) + gz[:, :, :-1] = np.diff(out, axis=2) + norm = np.sqrt(gx**2 + gy**2 + gz**2) + E += weight * norm.sum() + norm *= 0.5 / weight + norm += 1. + px -= 1./6.*gx + px /= norm + py -= 1./6.*gy + py /= norm + pz -= 1/6.*gz + pz /= norm + E /= float(im.size) + print E + if i == 0: + E_init = E + E_previous = E + else: + if np.abs(E_previous - E) < eps * E_init: + print E_previous, E + break + else: + E_previous = E + i += 1 + if keep_type: + return out.astype(im_type) + else: + return out + +def _tv_denoise_2d(im, weight=50, eps=2.e-4, keep_type=False, n_iter_max=200): + """ + Perform total-variation denoising + + Parameters + ---------- + im: ndarray + input data to be denoised + + eps: float, optional + relative difference of the value of the cost function that determines + the stop criterion. The algorithm stops when + (E_(n-1) - E_n) < eps * E_0 + + weight: float, optional + denoising weight. The greater ``weight``, the more denoising (at + the expense of fidelity to ``input``) + + keep_type: bool, optional (False) + whether the output has the same dtype as the input array. + keep_type is False by default, and the dtype of the output + is np.float + + n_iter_max: int, optional + maximal number of iterations used for the optimization. + + Returns + ------- + out: ndarray + denoised array + + Notes + ----- + The principle of total variation denoising is explained in + http://en.wikipedia.org/wiki/Total_variation_denoising + + This code is an implementation of the algorithm of Rudin, Fatemi and Osher + that was proposed by Chambolle in [1]_. + + References + ---------- + + .. [1] A. Chambolle, An algorithm for total variation minimization and + applications, Journal of Mathematical Imaging and Vision, + Springer, 2004, 20, 89-97. + + Examples + --------- + >>> import scipy + >>> lena = scipy.lena() + >>> import scipy + >>> lena = scipy.lena().astype(np.float) + >>> lena += 0.5 * lena.std()*np.random.randn(*lena.shape) + >>> denoised_lena = tv_denoise(lena, weight=60.0) + """ + im_type = im.dtype + if im_type is not np.float: + im = im.astype(np.float) + px = np.zeros_like(im) + py = np.zeros_like(im) + gx = np.zeros_like(im) + gy = np.zeros_like(im) + d = np.zeros_like(im) + i = 0 + while i < n_iter_max: + d = -px -py + d[1:] += px[:-1] + d[:, 1:] += py[:, :-1] + + out = im + d + E = (d**2).sum() + gx[:-1] = np.diff(out, axis=0) + gy[:, :-1] = np.diff(out, axis=1) + norm = np.sqrt(gx**2 + gy**2) + E += weight * norm.sum() + norm *= 0.5 / weight + norm += 1 + px -= 0.25*gx + px /= norm + py -= 0.25*gy + py /= norm + E /= float(im.size) + print E + if i == 0: + E_init = E + E_previous = E + else: + if np.abs(E_previous - E) < eps * E_init: + break + else: + E_previous = E + i += 1 + print i + if keep_type: + return out.astype(im_type) + else: + return out + +def tv_denoise(im, eps=2.e-4, weight=50, keep_type=False, n_iter_max=200): + """ + Perform total-variation denoising on 2-d and 3-d images + + Parameters + ---------- + im: ndarray (2d or 3d) + input data to be denoised + + eps: float, optional + relative difference of the value of the cost function that determines + the stop criterion. The algorithm stops when + (E_(n-1) - E_n) < eps * E_0 + + weight: float, optional + denoising weight. The greater ``weight``, the more denoising (at + the expense of fidelity to ``input``) + + keep_type: bool, optional (False) + whether the output has the same dtype as the input array. + keep_type is False by default, and the dtype of the output + is np.float + + n_iter_max: int, optional + maximal number of iterations used for the optimization. + + Returns + ------- + out: ndarray + denoised array + + + Notes + ----- + The principle of total variation denoising is explained in + http://en.wikipedia.org/wiki/Total_variation_denoising + + The principle of total variation denoising is to minimize the + total variation of the image, which can be roughly described as + the integral of the norm of the image gradient. Total variation + denoising tends to produce "cartoon-like" images, that is, + piecewise-constant images. + + This code is an implementation of the algorithm of Rudin, Fatemi and Osher + that was proposed by Chambolle in [1]_. + + References + ---------- + + .. [1] A. Chambolle, An algorithm for total variation minimization and + applications, Journal of Mathematical Imaging and Vision, + Springer, 2004, 20, 89-97. + + Examples + --------- + >>> import scipy + >>> # 2D example using lena + >>> lena = scipy.lena() + >>> import scipy + >>> lena = scipy.lena().astype(np.float) + >>> lena += 0.5 * lena.std()*np.random.randn(*lena.shape) + >>> denoised_lena = tv_denoise(lena, weight=60) + >>> # 3D example on synthetic data + >>> x, y, z = np.ogrid[0:40, 0:40, 0:40] + >>> mask = (x -22)**2 + (y - 20)**2 + (z - 17)**2 < 8**2 + >>> mask = mask.astype(np.float) + >>> mask += 0.2*np.random.randn(*mask.shape) + >>> res = tv_denoise_3d(mask, weight=100) + """ + + if im.ndim == 2: + return _tv_denoise_2d(im, eps, weight, keep_type, n_iter_max) + elif im.ndim == 3: + return _tv_denoise_3d(im, eps, weight, keep_type, n_iter_max) + else: + raise ValueError('only 2-d and 3-d images may be denoised with this function') + +def test_tv_denoise(): + """ + Apply the TV denoising algorithm on the lena image provided + by scipy + """ + import scipy + lena = scipy.lena().astype(np.float) + lena += 0.5 * lena.std()*np.random.randn(*lena.shape) + denoised_lena = tv_denoise(lena, weight=60.0) + assert denoised_lena.dtype in [np.float, np.float32, np.float64] + from scipy import ndimage + grad = ndimage.morphological_gradient(lena, size=((3,3))) + grad_denoised = ndimage.morphological_gradient(denoised_lena, size=((3,3))) + assert np.sqrt((grad_denoised**2).sum()) < np.sqrt((grad**2).sum()) + denoised_lena_int = tv_denoise(lena.astype(np.int32), \ + weight=60.0, keep_type=True) + assert denoised_lena_int.dtype is np.dtype('int32') + +def test_tv_denoise_3d(): + """ + Apply the TV denoising algorithm on a 3D image representing + a sphere. + """ + x, y, z = np.ogrid[0:40, 0:40, 0:40] + mask = (x -22)**2 + (y - 20)**2 + (z - 17)**2 < 8**2 + mask = 100 * mask.astype(np.float) + mask += 60 + mask += 20*np.random.randn(*mask.shape) + mask[mask < 0] = 0 + mask[mask > 255] = 255 + res = tv_denoise(mask.astype(np.uint8), weight=100, keep_type=True) + assert res.std() < mask.std() + assert res.dtype is np.dtype('uint8') + res = tv_denoise(mask.astype(np.uint8), weight=100) + assert res.std() < mask.std() + assert res.dtype is not np.dtype('uint8') + # test wrong number of dimensions + a = np.random.random((8, 8, 8, 8)) + try: + res = tv_denoise(a) + except ValueError: + pass