├── README.rst ├── .gitignore ├── tv_denoise ├── __init__.py ├── utils.py ├── cli.py ├── chambolle.py └── gradient_descent.py └── setup.py /README.rst: -------------------------------------------------------------------------------- 1 | tv-denoise 2 | ========== 3 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | __pycache__ 2 | build 3 | dist 4 | *.egg-info 5 | venv* 6 | .ipynb_checkpoints 7 | out* 8 | -------------------------------------------------------------------------------- /tv_denoise/__init__.py: -------------------------------------------------------------------------------- 1 | from .chambolle import tv_denoise_chambolle 2 | from .gradient_descent import tv_denoise_gradient_descent 3 | from .utils import to_float32, to_uint8 4 | -------------------------------------------------------------------------------- /tv_denoise/utils.py: -------------------------------------------------------------------------------- 1 | """Utility functions.""" 2 | 3 | import numpy as np 4 | 5 | 6 | def to_float32(image): 7 | """Converts a uint8 image, in numpy or Pillow format, to float32.""" 8 | return np.float32(image) / 255 9 | 10 | 11 | def to_uint8(image): 12 | """Converts a float32 image to a numpy uint8 image.""" 13 | return np.uint8(np.round(np.clip(image, 0, 1) * 255)) 14 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup 2 | 3 | setup( 4 | name='tv-denoise', 5 | version='0.1', 6 | description='Total variation denoising for images.', 7 | long_description=open('README.rst').read(), 8 | url='https://www.example.com/', 9 | author='Katherine Crowson', 10 | author_email='crowsonkb@gmail.com', 11 | # license='MIT', 12 | packages=['tv_denoise'], 13 | install_requires=['dataclasses>=0.6;python_version<"3.7"', 14 | 'numpy>=1.14.3', 15 | 'pillow>=5.1.0'], 16 | entry_points={ 17 | 'console_scripts': ['tv_denoise=tv_denoise.cli:main'], 18 | }, 19 | ) 20 | -------------------------------------------------------------------------------- /tv_denoise/cli.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | 3 | """Total variation denoising for images.""" 4 | 5 | import argparse 6 | 7 | from PIL import Image 8 | 9 | from tv_denoise import to_float32, to_uint8, tv_denoise_chambolle, tv_denoise_gradient_descent 10 | 11 | 12 | def main(): 13 | """The main function.""" 14 | ap = argparse.ArgumentParser(description=__doc__, 15 | formatter_class=argparse.ArgumentDefaultsHelpFormatter) 16 | ap.add_argument('input', help='the input image') 17 | ap.add_argument('output', help='the output image') 18 | ap.add_argument('strength', type=float, nargs='?', default=0.1, 19 | help='the luma denoising strength') 20 | ap.add_argument('strength_chroma', type=float, nargs='?', 21 | help='the chroma denoising strength (gradient descent method only, twice ' 22 | 'luma strength if not specified)') 23 | ap.add_argument('--method', default='gradient', choices=['gradient', 'chambolle'], 24 | help='the denoising method to use') 25 | args = ap.parse_args() 26 | 27 | if args.strength_chroma is None: 28 | args.strength_chroma = args.strength * 2 29 | 30 | def grad_callback(status): 31 | print(f'step: {status.i}, loss: {status.loss:g}') 32 | 33 | def chambolle_callback(status): 34 | print(f'step: {status.i}, max diff: {status.diff:g}') 35 | 36 | image = to_float32(Image.open(args.input).convert('RGB')) 37 | 38 | if args.method == 'gradient': 39 | out_arr = tv_denoise_gradient_descent(image, 40 | args.strength, 41 | args.strength_chroma, 42 | callback=grad_callback) 43 | elif args.method == 'chambolle': 44 | out_arr = tv_denoise_chambolle(image, args.strength, callback=chambolle_callback) 45 | else: 46 | raise ValueError('Invalid method') 47 | 48 | out = Image.fromarray(to_uint8(out_arr)) 49 | out.save(args.output) 50 | 51 | 52 | if __name__ == '__main__': 53 | main() 54 | -------------------------------------------------------------------------------- /tv_denoise/chambolle.py: -------------------------------------------------------------------------------- 1 | """Implements Chambolle's projection algorithm for total variation image denoising. See 2 | https://www.ipol.im/pub/art/2013/61/article.pdf.""" 3 | 4 | from dataclasses import dataclass 5 | 6 | import numpy as np 7 | 8 | 9 | def grad(arr): 10 | """Computes the discrete gradient of an image.""" 11 | out = np.zeros((2,) + arr.shape, arr.dtype) 12 | out[0, :-1, :, ...] = arr[1:, :, ...] - arr[:-1, :, ...] 13 | out[1, :, :-1, ...] = arr[:, 1:, ...] - arr[:, :-1, ...] 14 | return out 15 | 16 | 17 | def div(arr): 18 | """Computes the discrete divergence of a vector array.""" 19 | out = np.zeros_like(arr) 20 | out[0, 0, :, ...] = arr[0, 0, :, ...] 21 | out[0, -1, :, ...] = -arr[0, -2, :, ...] 22 | out[0, 1:-1, :, ...] = arr[0, 1:-1, :, ...] - arr[0, :-2, :, ...] 23 | out[1, :, 0, ...] = arr[1, :, 0, ...] 24 | out[1, :, -1, ...] = -arr[1, :, -2, ...] 25 | out[1, :, 1:-1, ...] = arr[1, :, 1:-1, ...] - arr[1, :, :-2, ...] 26 | return np.sum(out, axis=0) 27 | 28 | 29 | def magnitude(arr, axis=0, keepdims=False): 30 | """Computes the element-wise magnitude of a vector array.""" 31 | return np.sqrt(np.sum(arr**2, axis=axis, keepdims=keepdims)) 32 | 33 | 34 | @dataclass 35 | class ChambolleDenoiseStatus: 36 | """A status object supplied to the callback specified in tv_denoise_chambolle().""" 37 | i: int 38 | diff: float 39 | 40 | 41 | def tv_denoise_chambolle(image, strength, step_size=0.25, tol=3.2e-3, callback=None): 42 | """Total variation image denoising with Chambolle's projection algorithm.""" 43 | image = np.atleast_3d(image) 44 | p = np.zeros((2,) + image.shape, image.dtype) 45 | image_over_strength = image / strength 46 | diff = np.inf 47 | i = 0 48 | while diff > tol: 49 | i += 1 50 | grad_div_p_i = grad(div(p) - image_over_strength) 51 | mag_gdpi = magnitude(grad_div_p_i, axis=(0, -1), keepdims=True) 52 | new_p = (p + step_size * grad_div_p_i) / (1 + step_size * mag_gdpi) 53 | diff = np.max(magnitude(new_p - p)) 54 | if callback is not None: 55 | callback(ChambolleDenoiseStatus(i, float(diff))) 56 | p[:] = new_p 57 | 58 | return np.squeeze(image - strength * div(p)) 59 | -------------------------------------------------------------------------------- /tv_denoise/gradient_descent.py: -------------------------------------------------------------------------------- 1 | """Total variation denoising.""" 2 | 3 | from dataclasses import dataclass 4 | 5 | import numpy as np 6 | 7 | 8 | RGB_TO_YUV = np.float32([ 9 | [0.2126, 0.7152, 0.0722], 10 | [-0.09991, -0.33609, 0.436], 11 | [0.615, -0.55861, -0.05639], 12 | ]) 13 | 14 | YUV_TO_RGB = np.linalg.inv(RGB_TO_YUV) 15 | 16 | 17 | def tv_norm(image, eps=1e-8): 18 | """Computes the isotropic total variation norm and its gradient. Modified from 19 | https://github.com/jcjohnson/cnn-vis.""" 20 | x_diff = image[:-1, :-1, ...] - image[:-1, 1:, ...] 21 | y_diff = image[:-1, :-1, ...] - image[1:, :-1, ...] 22 | grad_mag = np.sqrt(x_diff**2 + y_diff**2 + eps) 23 | loss = np.sum(grad_mag) 24 | dx_diff = x_diff / grad_mag 25 | dy_diff = y_diff / grad_mag 26 | grad = np.zeros_like(image) 27 | grad[:-1, :-1, ...] = dx_diff + dy_diff 28 | grad[:-1, 1:, ...] -= dx_diff 29 | grad[1:, :-1, ...] -= dy_diff 30 | return loss, grad 31 | 32 | 33 | def l2_norm(image, orig_image): 34 | """Computes 1/2 the square of the L2-norm of the difference between the image and 35 | the original image and its gradient.""" 36 | grad = image - orig_image 37 | loss = np.sum(grad**2) / 2 38 | return loss, grad 39 | 40 | 41 | def eval_loss_and_grad(image, orig_image, strength_luma, strength_chroma): 42 | """Computes the loss function for TV denoising and its gradient.""" 43 | tv_loss_y, tv_grad_y = tv_norm(image[:, :, 0]) 44 | tv_loss_uv, tv_grad_uv = tv_norm(image[:, :, 1:]) 45 | tv_grad = np.zeros_like(image) 46 | tv_grad[..., 0] = tv_grad_y * strength_luma 47 | tv_grad[..., 1:] = tv_grad_uv * strength_chroma 48 | l2_loss, l2_grad = l2_norm(image, orig_image) 49 | loss = tv_loss_y * strength_luma + tv_loss_uv * strength_chroma + l2_loss 50 | grad = tv_grad + l2_grad 51 | return loss, grad 52 | 53 | 54 | @dataclass 55 | class GradientDescentDenoiseStatus: 56 | """A status object supplied to the callback specified in tv_denoise_gradient_descent().""" 57 | i: int 58 | loss: float 59 | 60 | 61 | # pylint: disable=too-many-arguments, too-many-locals 62 | def tv_denoise_gradient_descent(image, 63 | strength_luma, 64 | strength_chroma, 65 | callback=None, 66 | step_size=1e-2, 67 | tol=3.2e-3): 68 | """Total variation image denoising with gradient descent.""" 69 | image = image @ RGB_TO_YUV.T 70 | orig_image = image.copy() 71 | momentum = np.zeros_like(image) 72 | momentum_beta = 0.9 73 | loss_smoothed = 0 74 | loss_smoothing_beta = 0.9 75 | i = 0 76 | while True: 77 | i += 1 78 | 79 | loss, grad = eval_loss_and_grad(image, orig_image, strength_luma, strength_chroma) 80 | 81 | if callback is not None: 82 | callback(GradientDescentDenoiseStatus(i, loss)) 83 | 84 | # Stop iterating if the loss has not been decreasing recently 85 | loss_smoothed = loss_smoothed * loss_smoothing_beta + loss * (1 - loss_smoothing_beta) 86 | loss_smoothed_debiased = loss_smoothed / (1 - loss_smoothing_beta**i) 87 | if i > 1 and loss_smoothed_debiased / loss < tol + 1: 88 | break 89 | 90 | # Calculate the step size per channel 91 | step_size_luma = step_size / (strength_luma + 1) 92 | step_size_chroma = step_size / (strength_chroma + 1) 93 | step_size_arr = np.float32([[[step_size_luma, step_size_chroma, step_size_chroma]]]) 94 | 95 | # Gradient descent step 96 | momentum *= momentum_beta 97 | momentum += grad * (1 - momentum_beta) 98 | image -= step_size_arr / (1 - momentum_beta**i) * momentum 99 | 100 | return image @ YUV_TO_RGB.T 101 | --------------------------------------------------------------------------------