def deconv(self,
y,
z0=None,
mode='lbfgsb',
maxfun=100,
alpha=0.,
beta=0.01,
verbose=10,
m=5,
edgetapering=1,
factor=3,
gamma=1e-4):
"""
deconv implements various deconvolution methods. It expects a
blurry image and outputs an estimated blur kernel or a sharp latent
image. Currently, the following algorithms are implemented:
'lbfgsb' uses the lbfgsb optimization code to minimize the following
constrained regularized problem:
|y-Zu|^2 + alpha * |grad(u)|^2 + beta * |u|^2 s.t. u>0
The alpha term promotes smoothness of the solution, while
the beta term is an ordinary Thikhonov regularization
'direct' as above but solves the problem directly, i.e. via
division in Fourier space instead of an iterative
minimization scheme at the cost of the positivity
constraint.
'xdirect' as 'direct' but without corrective term which reduces
artifacts stemming from the windowing
'gdirect' solves the following problem
|grad(y)-grad(Zu)|^2 + alpha * |grad(u)|^2 + beta * |u|^2
This is particularly useful for kernel estimation in the
case of blurred natural images featuring many edges. The
advantage vs. 'direct' is the suppression of noise in the
estimated PSF kernels.
'xdirect' as 'direct' but without corrective term which reduces
artifacts stemming from the windowing
'Fast Image Deconvolution using Hyper-Laplacian Priors'
by Dilip Krishnan and Rob Fergus, NIPS 2009.
It minimizes the following problem
|y-Zu|^2 + gamma * |grad(u)|^(2/3)
via half-quadratic splitting. See paper for details.
----------------------------------------------------------------------
Usage:
Call: Z = OlaGPU(z, sw, mode, winaux)
u = Z.deconv(y)
Input: y blurry image
Ouput: u either image or PSF sized object
"""
from numpy import array
if not all(array(y.shape) == self.sy):
raise IOError('Sizes incompatible. Expected blurred image!')
# Potential data transfer to GPU
if y.__class__ == cua.GPUArray:
y_gpu = 1. * y
else:
y_gpu = cua.to_gpu(y.astype(np.float32))
# --------------------------------------------------------------------
if mode == 'lbfgsb':
from scipy.optimize import fmin_l_bfgs_b
self.res_gpu = cua.empty_like(y_gpu)
if self.__id__ == 'X':
sz = ((int(np.prod(self.winaux.csf)), int(self.sz[0]),
int(self.sz[1])))
elif self.__id__ == 'F':
sz = self.sz
lf = np.prod(sz)
if z0 == None:
z0_gpu = self.cnvtp(y_gpu)
z0 = z0_gpu.get()
z0 = z0.flatten()
#z0 = np.zeros(self.sf) # initialisation with flat kernels
#z0[self.sf[0]/2,self.sf[1]/2] = 1.
#z0 = np.tile(z0, [np.prod(self.csf),1,1])
#z0 = z0.flatten()
else:
z0 = z0.flatten()
lb = 0. # lower bound
ub = np.infty # upper bound
zhat = fmin_l_bfgs_b(func = self.cnvinv_objfun, x0 = z0, \
fprime = self.cnvinv_gradfun,\
args = [sz, y_gpu, alpha, beta],\
factr = 10., pgtol = 10e-15, \
maxfun = maxfun, bounds = [(lb, ub)] * lf,\
m = m, iprint = verbose)
return np.reshape(zhat[0], sz), zhat[1], zhat[2]
# --------------------------------------------------------------------
elif mode == 'gdirect':
# Use this method only for estimating the PSF
if self.__id__ != 'X':
raise Exception('Use direct mode for image estimation!')
# Compute Laplacian
if alpha > 0.:
gx_gpu = gputools.pad_cpu2gpu(np.array([[-1, 1], [-1, 1],
[-1, 1]]),
self.sfft_gpu,
dtype='complex')
gy_gpu = gputools.pad_cpu2gpu(np.array([[-1, -1, -1],
[1, 1, 1]]),
self.sfft_gpu,
dtype='complex')
self.plan.execute(gx_gpu)
self.plan.execute(gy_gpu)
L_gpu = gx_gpu * gx_gpu.conj() + gy_gpu * gy_gpu.conj()
else:
L_gpu = cua.zeros(self.fft_gpu.shape, np.complex64)
if edgetapering == 1:
gputools.edgetaper_gpu(y_gpu, 2 * self.sf, 'barthann')
# Transfer to GPU
if self.x.__class__ == cua.GPUArray:
x_gpu = self.x
else:
x_gpu = cua.to_gpu(self.x)
# Compute gradient images
xx_gpu, xy_gpu = gputools.gradient_gpu(x_gpu)
yx_gpu, yy_gpu = gputools.gradient_gpu(y_gpu)
# Chop and pad business
if self.mode == 'valid':
yx_gpu = gputools.chop_pad_GPU(yx_gpu, self.winaux.csf,
self.winaux.sw,
self.winaux.nhop, self.sfft_gpu,
self.sf - 1, 'complex')
yy_gpu = gputools.chop_pad_GPU(yy_gpu, self.winaux.csf,
self.winaux.sw,
self.winaux.nhop, self.sfft_gpu,
self.sf - 1, 'complex')
elif self.mode == 'same':
yx_gpu = gputools.chop_pad_GPU(yx_gpu, self.winaux.csf,
self.winaux.sw,
self.winaux.nhop, self.sfft_gpu,
np.floor(self.sf / 2),
'complex')
yy_gpu = gputools.chop_pad_GPU(yy_gpu, self.winaux.csf,
self.winaux.sw,
self.winaux.nhop, self.sfft_gpu,
np.floor(self.sf / 2),
'complex')
else:
raise NotImplementedError('Not a valid mode!')
xx_gpu = gputools.chop_pad_GPU(xx_gpu,
self.winaux.csf,
self.winaux.sw,
self.winaux.nhop,
self.sfft_gpu,
dtype='complex')
xy_gpu = gputools.chop_pad_GPU(xy_gpu,
self.winaux.csf,
self.winaux.sw,
self.winaux.nhop,
self.sfft_gpu,
dtype='complex')
# Here each patch should be windowed to reduce ringing artifacts,
# however since we are working in the gradient domain, the effect
# is negligible
# ws_gpu = gputools.pad_stack_GPU(self.winaux.ws_gpu,
# self.sfft_gpu, self.sf-1,
# dtype='complex')
# xx_gpu = ws_gpu * xx_gpu
# xy_gpu = ws_gpu * xy_gpu
# yx_gpu = ws_gpu * yx_gpu
# yy_gpu = ws_gpu * yy_gpu
# Compute Fourier transform
self.fft(yx_gpu, self.fft_gpu.shape[0])
self.fft(yy_gpu, self.fft_gpu.shape[0])
self.fft(xx_gpu, self.fft_gpu.shape[0])
self.fft(xy_gpu, self.fft_gpu.shape[0])
# Do division in Fourier space
z_gpu = cua.zeros(xy_gpu.shape, np.complex64)
z_gpu = gputools.comp_ola_gdeconv(xx_gpu, xy_gpu, yx_gpu, yy_gpu,
L_gpu, alpha, beta)
# Computing the inverse FFT
z_gpu = z_gpu.conj()
self.fft(z_gpu, self.fft_gpu.shape[0])
z_gpu = z_gpu.conj() / np.prod(z_gpu.shape[-2::])
# Crop out the kernels
zc_gpu = gputools.crop_stack_GPU(z_gpu, self.sf)
return zc_gpu
# --------------------------------------------------------------------
elif mode == 'direct':
const_gpu = cua.empty_like(y_gpu)
const_gpu.fill(1.)
# First deconvolution without corrective term
y_gpu = self.deconv(y_gpu,
mode='xdirect',
alpha=alpha,
beta=beta,
edgetapering=edgetapering)
gputools.cliplower_GPU(y_gpu, 0)
# Now same for constant image to get rid of window artifacts
if edgetapering == 1:
gputools.edgetaper_gpu(const_gpu, 2 * self.sf, 'barthann')
const_gpu = self.deconv(const_gpu,
mode='xdirect',
alpha=alpha,
beta=beta,
edgetapering=edgetapering)
gputools.edgetaper_gpu(const_gpu, 2 * self.sf, 'barthann')
gputools.clip_GPU(const_gpu, 0.01, 10.)
# Division of deconvolved latent and constant image to get rid
# of artifacts stemming from windowing
y_gpu = y_gpu / const_gpu
sz = y_gpu.shape
#gputools.clip_GPU(y_gpu, 0., 1.0)
#gputools.edgetaper_gpu(y_gpu, 3*self.sf, 'barthann')
# Do cropping and padding since edges are corrupted by division
y_gpu = gputools.crop_gpu2cpu(y_gpu,
sz - factor * self.sf - 1,
offset=np.floor(
(factor * self.sf - 1) / 2.))
y_gpu = gputools.impad_gpu(y_gpu,
tuple(np.array(sz) - y_gpu.shape))
return y_gpu
# --------------------------------------------------------------------
elif mode == 'xdirect':
# Compute Laplacian
if alpha > 0.:
gx_gpu = gputools.pad_cpu2gpu(np.array([[-1, 1]]),
self.sfft_gpu,
dtype='complex')
gy_gpu = gputools.pad_cpu2gpu(np.array([[-1], [1]]),
self.sfft_gpu,
dtype='complex')
self.plan.execute(gx_gpu)
self.plan.execute(gy_gpu)
L_gpu = gx_gpu * gx_gpu.conj() + gy_gpu * gy_gpu.conj()
else:
L_gpu = cua.zeros(self.fft_gpu.shape, np.complex64)
# Edgetapering of blurry input image
if edgetapering == 1:
gputools.edgetaper_gpu(y_gpu, 3 * self.sf, 'barthann')
if self.mode == 'valid':
#y_gpu = gputools.pad_cpu2gpu(y_gpu, self.sx, self.sf-1, dtype='real')
offset = self.sf - 1
elif self.mode == 'same':
offset = np.floor(self.sf / 2)
else:
raise NotImplementedError('Not a valid mode!')
# Chop and pad business
y_gpu = gputools.chop_pad_GPU(y, self.winaux.csf, self.winaux.sw,
self.winaux.nhop, self.sfft_gpu,
offset, 'complex')
ws_gpu = gputools.pad_stack_GPU(self.winaux.ws_gpu,
self.sfft_gpu,
dtype='complex')
# Windowing
y_gpu = ws_gpu * y_gpu
# Compute FFT
self.fft(y_gpu, self.fft_gpu.shape[0])
# Do division in Fourier space
z_gpu = gputools.comp_ola_deconv(self.fft_gpu, y_gpu, L_gpu, alpha,
beta)
# Computing the inverse FFT
z_gpu = z_gpu.conj()
self.fft(z_gpu, self.fft_gpu.shape[0])
z_gpu = z_gpu.conj() / np.prod(z_gpu.shape[-2::])
# Crop the solution to correct output size
if self.__id__ == 'X':
zc_gpu = gputools.crop_stack_GPU(z_gpu, self.sf)
return zc_gpu
elif self.__id__ == 'F':
zs_gpu = gputools.crop_stack_GPU(z_gpu, self.winaux.sw)
#zs_gpu = self.winaux.ws_gpu * zs_gpu
zc_gpu = gputools.ola_GPU_test(zs_gpu, self.winaux.csf,
self.winaux.sw,
self.winaux.nhop)
zc_gpu = gputools.crop_gpu2cpu(zc_gpu, self.sx)
return zc_gpu
# --------------------------------------------------------------------
elif mode == 'sparse':
# Compute Laplacian
gx_gpu = gputools.pad_cpu2gpu(np.sqrt(2.) / 2. *
np.array([[-1, 1]]),
self.sfft_gpu,
dtype='complex')
gy_gpu = gputools.pad_cpu2gpu(np.sqrt(2.) / 2. *
np.array([[-1], [1]]),
self.sfft_gpu,
dtype='complex')
self.plan.execute(gx_gpu)
self.plan.execute(gy_gpu)
L_gpu = gx_gpu * gx_gpu.conj() + gy_gpu * gy_gpu.conj()
const_gpu = cua.empty_like(y_gpu)
const_gpu.fill(1.)
# Edgetapering
if edgetapering == 1:
gputools.edgetaper_gpu(y_gpu, 2 * self.sf, 'barthann')
gputools.edgetaper_gpu(const_gpu, 2 * self.sf, 'barthann')
# Parameter settings
beta = 1.
beta_rate = 2. * np.sqrt(2.)
beta_max = 2.**8
# Initialisation of x with padded version of y
x_gpu = 1 * y_gpu
if self.mode == 'valid':
offset = self.sf - 1
elif self.mode == 'same':
offset = np.floor(self.sf / 2)
else:
raise NotImplementedError('Not a valid mode!')
# Chop and pad business
y_gpu = gputools.chop_pad_GPU(y_gpu, self.winaux.csf,
self.winaux.sw, self.winaux.nhop,
self.sfft_gpu, offset, 'complex')
const_gpu = gputools.chop_pad_GPU(const_gpu, self.winaux.csf,
self.winaux.sw, self.winaux.nhop,
self.sfft_gpu, offset, 'complex')
ws_gpu = gputools.pad_stack_GPU(self.winaux.ws_gpu,
self.sfft_gpu,
offset,
dtype='complex')
# Windowing
y_gpu = y_gpu * ws_gpu
# Constant image for corrective weighting term
const_gpu = const_gpu * ws_gpu
del ws_gpu
self.fft(const_gpu, self.fft_gpu.shape[0])
const_gpu = gputools.comp_ola_deconv(self.fft_gpu, const_gpu,
L_gpu, alpha, gamma)
const_gpu = const_gpu.conj()
self.fft(const_gpu, self.fft_gpu.shape[0])
const_gpu = const_gpu.conj() / np.prod(const_gpu.shape[-2::])
const_gpu = gputools.crop_stack_GPU(const_gpu, self.winaux.sw)
const_gpu = const_gpu * self.winaux.ws_gpu
const_gpu = gputools.ola_GPU_test(const_gpu, self.winaux.csf,
self.winaux.sw, self.winaux.nhop)
const_gpu = gputools.crop_gpu2cpu(const_gpu, self.sx)
# For debugging purposes
#scipy.misc.imsave('const1.png', const_gpu.get()/const_gpu.get().max())
gputools.cliplower_GPU(const_gpu, 0.01)
const_gpu = 0.01 / const_gpu
# Precompute F'y
self.fft(y_gpu, self.fft_gpu.shape[0])
y_gpu = y_gpu * self.fft_gpu.conj()
while beta < beta_max:
# Compute gradient images of x
xx_gpu, xy_gpu = gputools.gradient_gpu(x_gpu)
del x_gpu
# w sub-problem for alpha 2/3
gputools.modify_sparse23_gpu(xx_gpu, beta)
gputools.modify_sparse23_gpu(xy_gpu, beta)
#gputools.modify_sparse_gpu(xx_gpu, beta, 0.01)
#gputools.modify_sparse_gpu(xy_gpu, beta, 0.01)
# Chop and pad to size of FFT
xx_gpu = gputools.chop_pad_GPU(xx_gpu,
self.winaux.csf,
self.winaux.sw,
self.winaux.nhop,
self.sfft_gpu,
dtype='complex')
xy_gpu = gputools.chop_pad_GPU(xy_gpu,
self.winaux.csf,
self.winaux.sw,
self.winaux.nhop,
self.sfft_gpu,
dtype='complex')
# Compute Fourier transform
self.fft(xx_gpu, self.fft_gpu.shape[0])
self.fft(xy_gpu, self.fft_gpu.shape[0])
# Do division in Fourier space
x_gpu = gputools.comp_ola_sdeconv(gx_gpu, gy_gpu, xx_gpu,
xy_gpu, y_gpu, self.fft_gpu,
L_gpu, alpha, beta, gamma)
del xx_gpu, xy_gpu
# Computing the inverse FFT
x_gpu = x_gpu.conj()
self.fft(x_gpu, self.fft_gpu.shape[0])
x_gpu = x_gpu.conj()
x_gpu /= np.prod(x_gpu.shape[-2::])
# Ola and cropping
x_gpu = gputools.crop_stack_GPU(x_gpu, self.winaux.sw)
x_gpu = x_gpu * self.winaux.ws_gpu
x_gpu = gputools.ola_GPU_test(x_gpu, self.winaux.csf,
self.winaux.sw, self.winaux.nhop)
x_gpu = gputools.crop_gpu2cpu(x_gpu, self.sx)
# Enforce positivity
x_gpu = x_gpu * const_gpu
gputools.cliplower_GPU(x_gpu, 0.)
beta *= beta_rate
return x_gpu
else:
raise NotImplementedError('Not a valid deconv mode!')
def __init__(self, f, sx, mode, winaux):
sf = np.array(f.shape)[-2::]
sx = np.array(sx)
sw = winaux.sw
csf = winaux.csf
nhop = winaux.nhop
# Check what is f and what is x
if (len(f.shape) == 3) and all(sx > sf):
self.f = f
self.x = []
self.__id__ = 'F'
# Safety check
if np.prod(f.shape[0]) != np.prod(csf):
raise IOError('Size missmatch between winaux and PSF size!')
elif (len(f.shape) == 2) and all(sx < sf):
self.f = []
self.x = f
self.__id__ = 'X'
sf = sx
sx = np.array(self.x.shape)
elif any(sf < sx) and any(sf > sx):
raise IOError('Size missmatch')
# Safety check
if any(winaux.sx != sx):
raise IOError('Size missmatch between winaux and image size!')
if mode == 'valid':
sy = sx - sf + 1
elif mode == 'same':
sy = sx
elif mode == 'full':
sy = sx + sf - 1
elif mode == 'circ':
sy = sx
else:
raise NotImplementedError('Not a valid mode!')
# Pad either f or x to be sized a power of 2 and copy it to device
sfft = sw + sf - 1
sfft_gpu = (2**np.ceil(np.log2(sfft)))
sfft_gpu = (int(sfft_gpu[0]), int(sfft_gpu[1]))
if self.__id__ == 'F':
# each kernel of PSF has to be padded
fft_gpu = gputools.pad_stack_GPU(self.f, sfft_gpu, dtype='complex')
self.sz = sx
elif self.__id__ == 'X':
# each patch has to be modulated by window
fft_gpu = gputools.chop_mod_pad_GPU(self.x,
winaux.ws_gpu,
csf,
sw,
nhop,
sz=sfft_gpu,
dtype='complex')
self.sz = sf
# Create FFT plan and compute FFT
plan = cufft.Plan(fft_gpu.shape[-2::])
self.plan = plan
self.fft(fft_gpu, fft_gpu.shape[0])
self.sfft_gpu = sfft_gpu
self.fft_gpu = fft_gpu
self.winaux = winaux
self.csf = csf
self.sfft = sfft
self.sf = sf
self.sx = sx
self.sy = sy
self.mode = mode
def cnv(self, u):
"""
Description:
cnv computes the correlation of the convolution matrix with either
an image or PSF whether the parent class is an instance of F or X,
i.e. Fx or Xf respectively.
----------------------------------------------------------------------
Usage:
Call: Z = OlaGPU(z, sw, mode, winaux)
y = Z.cnv(u)
Input: u either image of PSF
Ouput: y a blurry image
"""
# Pad either f or x and copy it to device
if (len(u.shape) == 3) and (self.__id__ == 'X'):
# Safety check
if np.prod(u.shape[0]) != np.prod(self.winaux.csf):
raise IOError('Size missmatch between winaux and PSF size!')
u_gpu = gputools.pad_stack_GPU(u, self.sfft_gpu, dtype='complex')
elif (len(u.shape) == 2) and (self.__id__ == 'F'):
# Safety check
if sum(u.shape != self.winaux.sx) > 0:
raise IOError('Size missmatch between winaux and image size!')
# Chop input image into overlapping patches, modulate them
# by windows and do appropriate padding
#u_gpu = gputools.chop_mod_pad_GPU(u, self.winaux.ws_gpu,
# self.winaux.csf, self.winaux.sw,
# self.winaux.nhop, sz=self.sfft_gpu,
# dtype='complex')
############
# Something is wrong in above kernel, which should perform
# modulation and padding in one kernel call. For now workaround:
offset = (0, 0)
u_gpu = gputools.chop_pad_GPU(u,
self.winaux.csf,
self.winaux.sw,
self.winaux.nhop,
self.sfft_gpu,
offset,
dtype='complex')
ws_gpu = gputools.pad_stack_GPU(self.winaux.ws_gpu,
self.sfft_gpu,
offset,
dtype='complex')
self.ws = ws_gpu
u_gpu = ws_gpu * u_gpu
# Workauround ends here
############
# Compute FFT of input, do multiplication in Fourier space
# and compute inverse Fourier transform
self.fft(u_gpu, self.fft_gpu.shape[0])
# Strange enough: inverse does not work due to some error in pyfft
# Therefore compute the inverse via conj(F(conj(x)))/length(x)
# see Wikipedia for reference
ys_gpu = (self.fft_gpu * u_gpu).conj()
del u_gpu
self.fft(ys_gpu, self.fft_gpu.shape[0])
ys_gpu = ys_gpu.conj() / np.prod(ys_gpu.shape[-2::])
# Do overlap and add
y_gpu = gputools.ola_GPU_test(ys_gpu, self.winaux.csf,
self.sf - 1 + self.winaux.sw,
self.winaux.nhop)
# Do cropping to correct output size
if self.mode == 'valid':
y = gputools.crop_gpu2cpu(y_gpu, self.sy, self.sf - 1)
elif self.mode == 'same':
y = gputools.crop_gpu2cpu(y_gpu, self.sy, np.floor(self.sf / 2))
elif self.mode == 'full':
y = gputools.crop_gpu2cpu(y_gpu, self.sy)
elif self.mode == 'circ':
if u.__class__ == cua.GPUArray:
raise NotImplementedError('Not a valid mode!')
else:
y = np.real(y_gpu.get())
y = imagetools.circshift(y, floor(self.sf / 2))
else:
raise NotImplementedError('Not a valid mode!')
if u.__class__ == np.ndarray:
return np.array(y.get())
elif u.__class__ == cua.GPUArray:
return y