def admm_denoise(y, Phi_sum, A, At, _lambda=1, gamma=0.0,
denoiser='tv', iter_max=50, noise_estimate=False, sigma=None,
tv_weight=0.1, tv_iter_max=5, multichannel=True, x0=None, model=None, X_orig=None, show_iqa=True, tvm='tv_chambolle'):
'''
Alternating direction method of multipliers (ADMM)[1]-based denoising
regularization for snapshot compressive imaging (SCI).
Parameters
----------
y : two-dimensional (2D) ndarray of ints, uints or floats
Input single measurement of the snapshot compressive imager (SCI).
Phi : three-dimensional (3D) ndarray of ints, uints or floats, omitted
Input sensing matrix of SCI with the third dimension as the
time-variant, spectral-variant, volume-variant, or angular-variant
masks, where each mask has the same pixel resolution as the snapshot
measurement.
Phi_sum : 2D ndarray
Sum of the sensing matrix `Phi` along the third dimension.
A : function
Forward model of SCI, where multiple encoded frames are collapsed into
a single measurement.
At : function
Transpose of the forward model.
proj_meth : {'admm' or 'gap'}, optional
Projection method of the data term. Alternating direction method of
multipliers (ADMM)[1] and generalizedv alternating projection (GAP)[2]
are used, where ADMM for noisy data, especially real data and GAP for
noise-free data.
gamma : float, optional
Parameter in the ADMM projection, where more noisy measurements require
greater gamma.
denoiser : string, optional
Denoiser used as the regularization imposing on the prior term of the
reconstruction.
_lambda : float, optional
Regularization factor balancing the data term and the prior term,
where larger `_lambda` imposing more constrains on the prior term.
iter_max : int or uint, optional
Maximum number of iterations.
accelerate : boolean, optional
Enable acceleration in GAP.
noise_estimate : boolean, optional
Enable noise estimation in the denoiser.
sigma : one-dimensional (1D) ndarray of ints, uints or floats
Input noise standard deviation for the denoiser if and only if noise
estimation is disabled(i.e., noise_estimate==False). The scale of sigma
is [0, 255] regardless of the the scale of the input measurement and
masks.
tv_weight : float, optional
weight in total variation (TV) denoising.
x0 : 3D ndarray
Start point (initialized value) for the iteration process of the
reconstruction.
Returns
-------
x : 3D ndarray
Reconstructed 3D scene captured by the SCI system.
References
----------
.. [1] S. Boyd, N. Parikh, E. Chu, B. Peleato, and J. Eckstein,
"Distributed Optimization and Statistical Learning via the
Alternating Direction Method of Multipliers," Foundations and
Trends® in Machine Learning, vol. 3, no. 1, pp. 1-122, 2011.
.. [2] X. Yuan, "Generalized alternating projection based total variation
minimization for compressive sensing," in IEEE International
Conference on Image Processing (ICIP), 2016, pp. 2539-2543.
.. [3] Y. Liu, X. Yuan, J. Suo, D. Brady, and Q. Dai, "Rank Minimization
for Snapshot Compressive Imaging," IEEE Transactions on Pattern
Analysis and Machine Intelligence, doi:10.1109/TPAMI.2018.2873587,
2018.
Code credit
-----------
Xin Yuan, Bell Labs, [email protected], created Aug 7, 2018.
Yang Liu, Tsinghua University, [email protected],
updated Jan 22, 2019.
See Also
--------
gap_denoise
'''
# [0] initialization
if x0 is None:
x0 = At(y) # default start point (initialized value)
if not isinstance(sigma, list):
sigma = [sigma]
if not isinstance(iter_max, list):
iter_max = [iter_max] * len(sigma)
# [1] start iteration for reconstruction
x = x0 # initialization
theta = x0
b = np.zeros_like(x0)
psnr_all = []
k = 0
for idx, nsig in enumerate(sigma): # iterate all noise levels
for it in range(iter_max[idx]):
# Euclidean projection
yb = A(theta+b)
x = (theta+b) + _lambda*(At((y-yb)/(Phi_sum+gamma))) # ADMM
# switch denoiser
if denoiser.lower() == 'tv': # total variation (TV) denoising
try:
if tvm == 'tv_chambolle':
theta = denoise_tv_chambolle(x-b, tv_weight, n_iter_max=tv_iter_max, multichannel=multichannel)
elif tvm == 'ITV3D_FGP':
theta = denoise_tv_chambolle(x-b, tv_weight, n_iter_max=tv_iter_max, multichannel=multichannel)
elif tvm == 'ITV2D_cham':
theta = denoise_tv_chambolle(x-b, tv_weight, n_iter_max=tv_iter_max, multichannel=multichannel)
else:
raise TypeError("no such tv denoiser")
except TypeError as e:
print("Exception: ",repr(e))
elif denoiser.lower() == 'wavelet': # wavelet denoising
if noise_estimate or nsig is None: # noise estimation enabled
theta = denoise_wavelet(x-b, multichannel=multichannel)
else:
theta = denoise_wavelet(x-b, sigma=nsig, multichannel=multichannel)
# elif denoiser.lower() == 'vnlnet': # Video Non-local net denoising
# theta = vnlnet(np.expand_dims((x-b).transpose(2,0,1),3), nsig)
# theta = np.transpose(theta.squeeze(3),(1,2,0))
elif denoiser.lower() == 'ffdnet': # FFDNet frame-wise video denoising
# x = ffdnet_vdenoiser(x, nsig, model) # [zzh] original
theta = ffdnet_vdenoiser(x-b, nsig, model) # [zzh] new code from xinyuan(1/3)
elif denoiser.lower() == 'fastdvdnet': # FastDVDnet video denoising
# x = fastdvdnet_denoiser(x, nsig, model, gray=True) # [zzh] original
theta = fastdvdnet_denoiser(x-b, nsig, model, gray=True) # [zzh] new code from xinyuan(2/3)
else:
raise ValueError('Unsupported denoiser {}!'.format(denoiser))
theta = np.clip(theta,0,1) # [zzh] new code from xinyuan(3/3)
b = b - (x-theta) # update residual
# [optional] calculate image quality assessment, i.e., PSNR for
# every five iterations
if show_iqa and X_orig is not None:
psnr_all.append(psnr(X_orig, x))
if (k+1)%5 == 0:
if not noise_estimate and nsig is not None:
if nsig < 1:
print(' ADMM-{0} iteration {1: 3d}, sigma {2: 3g}/255, '
'PSNR {3:2.2f} dB.'.format(denoiser.upper(),
k+1, nsig*255, psnr_all[k]))
else:
print(' ADMM-{0} iteration {1: 3d}, sigma {2: 3g}, '
'PSNR {3:2.2f} dB.'.format(denoiser.upper(),
k+1, nsig, psnr_all[k]))
else:
print(' ADMM-{0} iteration {1: 3d}, '
'PSNR {2: 2.2f} dB.'.format(denoiser.upper(),
k+1, psnr_all[k]))
k = k+1
psnr_ = []
ssim_ = []
nmask = x.shape[2]
if X_orig is not None:
for imask in range(nmask):
psnr_.append(compare_psnr(X_orig[:,:,imask], x[:,:,imask], data_range=1.))
ssim_.append(compare_ssim(X_orig[:,:,imask], x[:,:,imask], data_range=1.))
return x, psnr_, ssim_, psnr_all
def gap_denoise(y, Phi_sum, A, At, _lambda=1, accelerate=True,
denoiser='tv', iter_max=50, noise_estimate=False, sigma=None,
tv_weight=0.1, tv_iter_max=5, multichannel=True, x0=None,
X_orig=None, model=None, show_iqa=True, tvm='tv_chambolle'):
'''
Alternating direction method of multipliers (ADMM)[1]-based denoising
regularization for snapshot compressive imaging (SCI).
Parameters
----------
y : two-dimensional (2D) ndarray of ints, uints or floats
Input single measurement of the snapshot compressive imager (SCI).
Phi : three-dimensional (3D) ndarray of ints, uints or floats, omitted
Input sensing matrix of SCI with the third dimension as the
time-variant, spectral-variant, volume-variant, or angular-variant
masks, where each mask has the same pixel resolution as the snapshot
measurement.
Phi_sum : 2D ndarray,
Sum of the sensing matrix `Phi` along the third dimension.
A : function
Forward model of SCI, where multiple encoded frames are collapsed into
a single measurement.
At : function
Transpose of the forward model.
proj_meth : {'admm' or 'gap'}, optional
Projection method of the data term. Alternating direction method of
multipliers (ADMM)[1] and generalizedv alternating projection (GAP)[2]
are used, where ADMM for noisy data, especially real data and GAP for
noise-free data.
gamma : float, optional
Parameter in the ADMM projection, where more noisy measurements require
greater gamma.
denoiser : string, optional
Denoiser used as the regularization imposing on the prior term of the
reconstruction.
_lambda : float, optional
Regularization factor balancing the data term and the prior term,
where larger `_lambda` imposing more constrains on the prior term.
iter_max : int or uint, optional
Maximum number of iterations.
accelerate : boolean, optional
Enable acceleration in GAP.
noise_estimate : boolean, optional
Enable noise estimation in the denoiser.
sigma : one-dimensional (1D) ndarray of ints, uints or floats
Input noise standard deviation for the denoiser if and only if noise
estimation is disabled(i.e., noise_estimate==False). The scale of sigma
is [0, 255] regardless of the the scale of the input measurement and
masks.
tv_weight : float, optional
weight in total variation (TV) denoising.
x0 : 3D ndarray
Start point (initialized value) for the iteration process of the
reconstruction.
model : pretrained model for image/video denoising.
tvm : string, optional, {'tv_chambolle', 'ATV_ClipA', 'ATV_ClipB','ATV_cham','ATV_FGP',
'ITV2D_cham','ITV2D_FGP','ITV3D_cham','ITV3D_FGP'}
tv denoiser type, default value = 'tv_chambolle' (zzh)
Returns
-------
x : 3D ndarray
Reconstructed 3D scene captured by the SCI system.
References
----------
.. [1] X. Liao, H. Li, and L. Carin, "Generalized Alternating Projection
for Weighted-$\ell_{2,1}$ Minimization with Applications to
Model-Based Compressive Sensing," SIAM Journal on Imaging Sciences,
vol. 7, no. 2, pp. 797-823, 2014.
.. [2] X. Yuan, "Generalized alternating projection based total variation
minimization for compressive sensing," in IEEE International
Conference on Image Processing (ICIP), 2016, pp. 2539-2543.
.. [3] Y. Liu, X. Yuan, J. Suo, D. Brady, and Q. Dai, "Rank Minimization
for Snapshot Compressive Imaging," IEEE Transactions on Pattern
Analysis and Machine Intelligence, doi:10.1109/TPAMI.2018.2873587,
2018.
Code credit
-----------
Xin Yuan, Bell Labs, [email protected], created Aug 7, 2018.
Yang Liu, Tsinghua University, [email protected],
updated Jan 22, 2019.
See Also
--------
admm_denoise
'''
# [0] initialization
if x0 is None:
# x0 = At(y, Phi) # default start point (initialized value)
x0 = At(y) # default start point (initialized value)
if not isinstance(sigma, list):
sigma = [sigma]
if not isinstance(iter_max, list):
iter_max = [iter_max] * len(sigma)
# y1 = np.zeros(y.shape)
y1 = np.zeros_like(y)
# [1] start iteration for reconstruction
x = x0 # initialization
psnr_all = []
k = 0
time_start = time.time() # timing
print('---> gap_denoise')
for idx, nsig in enumerate(sigma): # iterate all noise levels
for it in range(iter_max[idx]):
yb = A(x)
if accelerate: # accelerated version of GAP
y1 = y1 + (y-yb)
x = x + _lambda*(At((y1-yb)/Phi_sum)) # GAP_acc
else:
x = x + _lambda*(At((y-yb)/Phi_sum)) # GAP
# switch denoiser
if denoiser.lower() == 'tv': # total variation (TV) denoising
try:
if tvm == 'tv_chambolle':
x = denoise_tv_chambolle(x, tv_weight, n_iter_max=tv_iter_max, multichannel=multichannel)
elif tvm == 'ITV3D_FGP':
x = denoise_tv_FGP_ITV3D(x, tv_weight, n_iter_max=tv_iter_max)
elif tvm == 'ITV2D_cham':
x = denoise_tv_cham_ITV2D(x, tv_weight, n_iter_max=tv_iter_max)
else:
raise TypeError("no such tv denoiser")
except TypeError as e:
print("Exception: ",repr(e))
elif denoiser.lower() == 'wavelet': # wavelet denoising
if noise_estimate or nsig is None: # noise estimation enabled
x = denoise_wavelet(x, multichannel=multichannel)
else:
x = denoise_wavelet(x, sigma=nsig, multichannel=multichannel)
# elif denoiser.lower() == 'vnlnet': # Video Non-local net denoising
# x = vnlnet(np.expand_dims(x.transpose(2,0,1),3), nsig)
# x = np.transpose(x.squeeze(3),(1,2,0))
elif denoiser.lower() == 'ffdnet': # FFDNet frame-wise video denoising
x = ffdnet_vdenoiser(x, nsig, model)
elif denoiser.lower() == 'fastdvdnet': # FastDVDnet video denoising
x = fastdvdnet_denoiser(x, nsig, model, gray=True)
else:
raise ValueError('Unsupported denoiser {}!'.format(denoiser))
# [optional] calculate image quality assessment, i.e., PSNR for
# every five iterations
if show_iqa and X_orig is not None:
psnr_all.append(psnr(X_orig, x))
if (k+1)%5 == 0:
if not noise_estimate and nsig is not None:
if nsig < 1:
print(' GAP-{0} iteration {1: 3d}, sigma {2: 3g}/255, '
'PSNR {3:2.2f} dB.'.format(denoiser.upper(),
k+1, nsig*255, psnr_all[k]))
else:
print(' GAP-{0} iteration {1: 3d}, sigma {2: 3g}, '
'PSNR {3:2.2f} dB.'.format(denoiser.upper(),
k+1, nsig, psnr_all[k]))
else:
print(' GAP-{0} iteration {1: 3d}, '
'PSNR {2:2.2f} dB.'.format(denoiser.upper(),
k+1, psnr_all[k]))
k = k+1
time_now = time.time()
print('----> finish {}/{} time cost {:.2f} min'.format(idx+1, len(sigma),(time_now-time_start)/60))
psnr_ = []
ssim_ = []
nmask = x.shape[2]
if X_orig is not None:
for imask in range(nmask):
psnr_.append(compare_psnr(X_orig[:,:,imask], x[:,:,imask], data_range=1.))
ssim_.append(compare_ssim(X_orig[:,:,imask], x[:,:,imask], data_range=1.))
return x, psnr_, ssim_, psnr_all
def admm_denoise_bayer(y_bayer, Phi_bayer, _lambda=1, gamma=0.01,
denoiser='tv', iter_max=50, noise_estimate=True, sigma=None,
tv_weight=0.1, tv_iter_max=5, multichannel=True, x0_bayer=None,
X_orig=None, model=None, show_iqa=True):
'''
Generalized alternating projection (GAP)[1]-based denoising regularization
for snapshot compressive imaging (SCI).
Parameters
----------
y_bayer : two-dimensional (2D) ndarray of ints, uints or floats
Input single measurement of the snapshot compressive imager (SCI).
Phi_bayer : three-dimensional (3D) ndarray of ints, uints or floats, omitted
Input sensing matrix of SCI with the third dimension as the
time-variant, spectral-variant, volume-variant, or angular-variant
masks, where each mask has the same pixel resolution as the snapshot
measurement.
Phi : 3D ndarray,
Sensing matrix `Phi`.
proj_meth : {'admm' or 'gap'}, optional
Projection method of the data term. Alternating direction method of
multipliers (ADMM)[1] and generalizedv alternating projection (GAP)[2]
are used, where ADMM for noisy data, especially real data and GAP for
noise-free data.
gamma : float, optional
Parameter in the ADMM projection, where more noisy measurements require
greater gamma.
denoiser : string, optional
Denoiser used as the regularization imposing on the prior term of the
reconstruction.
_lambda : float, optional
Regularization factor balancing the data term and the prior term,
where larger `_lambda` imposing more constrains on the prior term.
iter_max : int or uint, optional
Maximum number of iterations.
accelerate : boolean, optional
Enable acceleration in GAP.
noise_estimate : boolean, optional
Enable noise estimation in the denoiser.
sigma : one-dimensional (1D) ndarray of ints, uints or floats
Input noise standard deviation for the denoiser if and only if noise
estimation is disabled(i.e., noise_estimate==False). The scale of sigma
is [0, 255] regardless of the the scale of the input measurement and
masks.
tv_weight : float, optional
weight in total variation (TV) denoising.
x0_bayer : 3D ndarray
Start point (initialized value) for the iteration process of the
reconstruction.
model : pretrained model for image/video denoising.
Returns
-------
x : 3D ndarray
Reconstructed 3D scene captured by the SCI system.
References
----------
.. [1] X. Liao, H. Li, and L. Carin, "Generalized Alternating Projection
for Weighted-$\ell_{2,1}$ Minimization with Applications to
Model-Based Compressive Sensing," SIAM Journal on Imaging Sciences,
vol. 7, no. 2, pp. 797-823, 2014.
.. [2] X. Yuan, "Generalized alternating projection based total variation
minimization for compressive sensing," in IEEE International
Conference on Image Processing (ICIP), 2016, pp. 2539-2543.
.. [3] Y. Liu, X. Yuan, J. Suo, D. Brady, and Q. Dai, "Rank Minimization
for Snapshot Compressive Imaging," IEEE Transactions on Pattern
Analysis and Machine Intelligence, vol. 41, no. 12, pp. 2990-3006,
2019.
Code credit
-----------
Xin Yuan, Bell Labs, [email protected], basic version created Aug 7, 2018.
Yang Liu, MIT CSAIL, [email protected], updated Dec 5, 2019.
See Also
--------
admm_denoise
'''
bayer = [[0,0], [0,1], [1,0], [1,1]] # `BGGR` Bayer pattern
if not isinstance(sigma, list):
sigma = [sigma]
if not isinstance(iter_max, list):
iter_max = [iter_max] * len(sigma)
# stack the bayer channels at the last dimension [consistent to image color channels]
(nrow, ncol, nmask) = Phi_bayer.shape
yall = np.zeros([nrow//2, ncol//2, 4], dtype=np.float32)
Phiall = np.zeros([nrow//2, ncol//2, nmask, 4], dtype=np.float32)
Phi_sumall = np.zeros([nrow//2, ncol//2, 4], dtype=np.float32)
x0all = np.zeros([nrow//2, ncol//2, nmask, 4], dtype=np.float32)
# iterative solve for each Bayer channel
for ib in range(len(bayer)):
b = bayer[ib]
yall[...,ib] = y_bayer[b[0]::2, b[1]::2]
Phiall[...,ib] = Phi_bayer[b[0]::2, b[1]::2]
# y = y_bayer[b[0]::2][b[1]::2]
# Phi = Phi_bayer[b[0]::2][b[1]::2]
# A = lambda x : A_(x, Phi) # forward model function handle
# At = lambda y : At_(y, Phi) # transpose of forward model
Phib = Phiall[...,ib]
Phib_sum = np.sum(Phib, axis=2)
Phib_sum[Phib_sum==0] = 1
Phi_sumall[...,ib] = Phib_sum
# [0] initialization
if x0_bayer is None:
# x0 = At(y, Phi) # default start point (initialized value)
x0all[...,ib] = At_(yall[...,ib], Phiall[...,ib]) # default start point (initialized value)
else:
x0all[...,ib] = x0_bayer[b[0]::2,b[1]::2]
# y1 = np.zeros(y.shape)
y1all = np.zeros_like(yall)
# [1] start iteration for reconstruction
xall = x0all # initialization
thetaall = x0all
x_bayer = np.zeros_like(Phi_bayer)
b = np.zeros_like(x0all)
psnr_all = []
k = 0
for idx, nsig in enumerate(sigma): # iterate all noise levels
for it in range(iter_max[idx]):
start_time = time.time()
for ib in range(len(bayer)): # iterate all bayer channels
yb = A_(thetaall[...,ib]+ball[...,ib], Phiall[...,ib])
xall[...,ib] = thetaall[...,ib]+ball[...,ib] + _lambda*(At_((yall[...,ib]-yb)/(Phi_sumall[...,ib]+gamma), Phiall[...,ib])) # GAP
end_time = time.time()
# print(' Euclidean projection eclipsed in {:.3f}s.'.format(end_time-start_time))
# joint Bayer multi-channel denoising
# switch denoiser
if denoiser.lower() == 'tv': # total variation (TV) denoising
thetaall_vch = (xall-ball).reshape([nrow//2, ncol//2, nmask*4])
thetaall_vch = denoise_tv_chambolle(thetaall_vch, tv_weight, n_iter_max=tv_iter_max,
multichannel=multichannel)
thetaall = thetaall_vch.reshape([nrow//2, ncol//2, nmask, 4])
# xall = xall.clip(0., 1.) # [0,1]
elif denoiser.lower() == 'wavelet': # wavelet denoising
thetaall_vch = (xall-ball).reshape([nrow//2, ncol//2, nmask*4])
if noise_estimate or nsig is None: # noise estimation enabled
thetaall_vch = denoise_wavelet(thetaall_vch, multichannel=multichannel)
else:
thetaall_vch = denoise_wavelet(thetaall_vch, sigma=nsig, multichannel=multichannel)
thetaall = thetaall_vch.reshape([nrow//2, ncol//2, nmask, 4])
# elif denoiser.lower() == 'vnlnet': # Video Non-local net denoising
# x = vnlnet(np.expand_dims(x.transpose(2,0,1),3), nsig)
# x = np.transpose(x.squeeze(3),(1,2,0))
elif denoiser.lower() == 'ffdnet': # FFDNet frame-wise video denoising
xall_vch = xall.reshape([nrow//2, ncol//2, nmask*4])
xall_vch = ffdnet_vdenoiser(xall_vch, nsig, model)
xall = xall_vch.reshape([nrow//2, ncol//2, nmask, 4])
elif denoiser.lower() == 'fastdvdnet': # FastDVDnet video denoising
# # option 1 - run denoising twice
# xrgb1 = xall[..., [0,1,3]] # R-G1-B (H x W x F x C)
# xrgb2 = xall[..., [0,2,3]] # R-G2-B (H x W x F x C)
# xrgb1 = fastdvdnet_denoiser(xrgb1, nsig, model)
# xrgb2 = fastdvdnet_denoiser(xrgb2, nsig, model)
# xall[...,0] = (xrgb1[...,0] + xrgb2[...,0])/2 # R channel (average over two)
# xall[...,1] = xrgb1[...,1] # G1 channel (average over two)
# xall[...,2] = xrgb2[...,1] # G2 channel (average over two)
# xall[...,3] = (xrgb1[...,2] + xrgb2[...,2])/2 # B channel (average over two)
# option 2 - run deniosing once
thetargb1 = (xall-ball)[..., [3,1,0]] # R-G1-B (H x W x F x C)
thetargb1 = fastdvdnet_denoiser(thetargb1, nsig, model)
thetaall[...,3] = thetargb1[...,0] # R channel (average over two)
thetaall[...,2] = thetargb1[...,1] # G1=G2 channel (average over two)
thetaall[...,1] = thetargb1[...,1] # G2=G1 channel (average over two)
thetaall[...,0] = thetargb1[...,2] # B channel (average over two)
else:
raise ValueError('Unsupported denoiser {}!'.format(denoiser))
ball = ball - (xall-thetaall) # update residual
# [optional] calculate image quality assessment, i.e., PSNR for
# every five iterations
if show_iqa and X_orig is not None:
for ib in range(len(bayer)):
b = bayer[ib]
x_bayer[b[0]::2, b[1]::2] = xall[...,ib]
psnr_all.append(psnr(X_orig, x_bayer))
if (k+1)%5 == 0:
if not noise_estimate and nsig is not None:
if nsig < 1:
print(' GAP-{0} iteration {1: 3d}, sigma {2: 3g}/255, '
'PSNR {3:2.2f} dB.'.format(denoiser.upper(),
k+1, nsig*255, psnr_all[k]))
else:
print(' GAP-{0} iteration {1: 3d}, sigma {2: 3g}, '
'PSNR {3:2.2f} dB.'.format(denoiser.upper(),
k+1, nsig, psnr_all[k]))
else:
print(' GAP-{0} iteration {1: 3d}, '
'PSNR {2:2.2f} dB.'.format(denoiser.upper(),
k+1, psnr_all[k]))
k = k+1
for ib in range(len(bayer)):
b = bayer[ib]
x_bayer[b[0]::2, b[1]::2] = xall[...,ib]
return x_bayer, psnr_all