def test_conv_halide(self):
"""Test convolution lin op in halide.
"""
np_img = get_test_image(512)
output = np.zeros_like(np_img)
K = get_kernel(11, np_img.ndim)
#Convolve in halide
Halide('A_conv', recompile=True).A_conv(np_img, K, output)
#Convolve in scipy
output_ref = signal.convolve2d(np_img, K, mode='same', boundary='wrap')
# Transpose
output_corr = np.zeros_like(np_img)
Halide('At_conv', recompile=True).At_conv(np_img, K,
output_corr) # Call
output_corr_ref = signal.convolve2d(np_img,
np.flipud(np.fliplr(K)),
mode='same',
boundary='wrap')
self.assertItemsAlmostEqual(output, output_ref)
self.assertItemsAlmostEqual(output_corr, output_corr_ref)
def test_mask_halide(self):
"""Test mask lin op in halide.
"""
if halide_installed():
# Load image
testimg_filename = os.path.join(
os.path.dirname(os.path.realpath(__file__)), 'data',
'angela.jpg')
# opens the file using Pillow - it's not an array yet
img = Image.open(testimg_filename)
np_img = np.asfortranarray(im2nparray(img))
# Test problem
output = np.zeros_like(np_img)
mask = np.asfortranarray(
np.random.randn(*list(np_img.shape)).astype(np.float32))
mask = np.maximum(mask, 0.)
Halide('A_mask.cpp').A_mask(np_img, mask, output) # Call
output_ref = mask * np_img
# Transpose
output_trans = np.zeros_like(np_img)
Halide('At_mask.cpp').At_mask(np_img, mask, output_trans) # Call
self.assertItemsAlmostEqual(output, output_ref)
self.assertItemsAlmostEqual(output_trans, output_ref)
def test_warp_halide(self):
"""Test warp lin op in halide.
"""
if halide_installed():
# Load image
testimg_filename = os.path.join(
os.path.dirname(os.path.realpath(__file__)), 'data',
'angela.jpg')
img = Image.open(testimg_filename)
np_img = np.asfortranarray(im2nparray(img))
# Convert to gray
np_img = np.mean(np_img, axis=2)
# Generate problem
theta_rad = 5.0 * np.pi / 180.0
H = np.array([[np.cos(theta_rad), -np.sin(theta_rad), 0.0001],
[np.sin(theta_rad),
np.cos(theta_rad), 0.0003], [0., 0., 1.]],
dtype=np.float32,
order='F')
# Reference
output_ref = cv2.warpPerspective(np_img,
H.T,
np_img.shape[1::-1],
flags=cv2.INTER_LINEAR,
borderMode=cv2.BORDER_CONSTANT,
borderValue=0.)
# Halide
output = np.zeros_like(np_img)
Hc = np.asfortranarray(
np.linalg.pinv(H)[..., np.newaxis]) # Third axis for halide
Halide('A_warp.cpp').A_warp(np_img, Hc, output) # Call
# Transpose
output_trans = np.zeros_like(np_img)
Hinvc = np.asfortranarray(H[...,
np.newaxis]) # Third axis for halide
Halide('At_warp.cpp').At_warp(output, Hinvc, output_trans) # Call
# Compute reference
output_ref_trans = cv2.warpPerspective(
output_ref,
H.T,
np_img.shape[1::-1],
flags=cv2.INTER_LINEAR + cv2.WARP_INVERSE_MAP,
borderMode=cv2.BORDER_CONSTANT,
borderValue=0.)
# Opencv does inverse warp
self.assertItemsAlmostEqual(output, output_ref, places=1)
# Opencv does inverse warp
self.assertItemsAlmostEqual(output_trans,
output_ref_trans,
places=1)
def test_grad_halide(self):
"""Test gradient lin op in halide.
"""
if halide_installed():
# Load image
testimg_filename = os.path.join(
os.path.dirname(os.path.realpath(__file__)), 'data',
'angela.jpg')
img = Image.open(testimg_filename)
np_img = np.asfortranarray(im2nparray(img))
# Convert to gray
np_img = np.mean(np_img, axis=2)
# Test problem
output = np.zeros(
(np_img.shape[0], np_img.shape[1], np_img.shape[2] if
(len(np_img.shape) > 2) else 1, 2),
dtype=np.float32,
order='FORTRAN')
# Gradient in halide
Halide('A_grad.cpp').A_grad(np_img, output) # Call
# Compute comparison
f = np_img
if len(np_img.shape) == 2:
f = f[..., np.newaxis]
ss = f.shape
fx = f[:, np.r_[1:ss[1], ss[1] - 1], :] - f
fy = f[np.r_[1:ss[0], ss[0] - 1], :, :] - f
Kf = np.asfortranarray(np.stack((fy, fx), axis=-1))
# Transpose
output_trans = np.zeros(f.shape, dtype=np.float32, order='F')
Halide('At_grad.cpp').At_grad(Kf, output_trans) # Call
# Compute comparison (Negative divergence)
Kfy = Kf[:, :, :, 0]
fy = Kfy - Kfy[np.r_[0, 0:ss[0] - 1], :, :]
fy[0, :, :] = Kfy[0, :, :]
fy[-1, :, :] = -Kfy[-2, :, :]
Kfx = Kf[:, :, :, 1]
ss = Kfx.shape
fx = Kfx - Kfx[:, np.r_[0, 0:ss[1] - 1], :]
fx[:, 0, :] = Kfx[:, 0, :]
fx[:, -1, :] = -Kfx[:, -2, :]
def init_kernel(self):
if not self.initialized:
arg = self.input_nodes[0]
# Replicate kernel for multichannel deconv
if len(arg.shape) == 3 and len(self.kernel.shape) == 2:
self.kernel = np.stack((self.kernel,) * arg.shape[2], axis=-1)
# Halide kernel
if self.implementation == Impl['halide'] and \
(len(arg.shape) == 2 or (len(arg.shape) == 3 and arg.dims == 2)):
self.kernel = np.asfortranarray(self.kernel.astype(np.float32))
# Halide FFT (pack into diag)
# TODO: FIX IMREAL LATER
hsize = arg.shape if len(arg.shape) == 3 else arg.shape + (1,)
output_fft_tmp = np.zeros(((hsize[0] + 1) / 2 + 1, hsize[1], hsize[2], 2),
dtype=np.float32, order='F')
Halide('fft2_r2c.cpp').fft2_r2c(self.kernel, self.kernel.shape[1] / 2,
self.kernel.shape[0] / 2, output_fft_tmp)
self.forward_kernel[:] = 0.
if len(arg.shape) == 2:
self.forward_kernel[0:(hsize[0] + 1) / 2 + 1, ...] = 1j * \
output_fft_tmp[..., 0, 1]
self.forward_kernel[0:(hsize[0] + 1) / 2 + 1, ...] += output_fft_tmp[..., 0, 0]
else:
self.forward_kernel[0:(hsize[0] + 1) / 2 + 1, ...] = 1j * output_fft_tmp[..., 1]
self.forward_kernel[0:(hsize[0] + 1) / 2 + 1, ...] += output_fft_tmp[..., 0]
self.tmpout = np.zeros(arg.shape, dtype=np.float32, order='F')
self.initialized = True
def test_norm1_halide(self):
"""Halide Norm 1 test
"""
if halide_installed():
# Load image
testimg_filename = os.path.join(
os.path.dirname(os.path.realpath(__file__)), 'data',
'angela.jpg')
img = Image.open(
testimg_filename
) # opens the file using Pillow - it's not an array yet
np_img = np.asfortranarray(im2nparray(img))
# Convert to gray
np_img = np.mean(np_img, axis=2)
# Test problem
v = np_img
theta = 0.5
# Output
output = np.zeros_like(np_img)
Halide('prox_L1.cpp').prox_L1(v, theta, output) # Call
# Reference
output_ref = np.maximum(0.0, v - theta) - np.maximum(
0.0, -v - theta)
self.assertItemsAlmostEqual(output, output_ref)
def test_poisson_halide(self):
"""Halide Poisson norm test
"""
# Test problem
np_img = get_test_image(512)
v = np_img
theta = 0.5
mask = np.asfortranarray(
np.random.randn(*np_img.shape).astype(np.float32))
mask = np.maximum(mask, 0.)
b = np_img * np_img
# Output
output = np.zeros_like(v)
tic()
Halide('prox_Poisson',
recompile=True).prox_Poisson(v, mask, b, theta, output) # Call
print('Running took: {0:.1f}ms'.format(toc()))
# Reference
output_ref = 0.5 * (v - theta + np.sqrt((v - theta) *
(v - theta) + 4 * theta * b))
output_ref[mask <= 0.5] = v[mask <= 0.5]
self.assertItemsAlmostEqual(output, output_ref)
def test_isonorm1_halide(self):
"""Halide Norm 1 test
"""
# Test problem
theta = 0.5
np_img = get_test_image(512)
f = np_img
if len(np_img.shape) == 2:
f = f[..., np.newaxis]
ss = f.shape
fx = f[:, np.r_[1:ss[1], ss[1] - 1], :] - f
fy = f[np.r_[1:ss[0], ss[0] - 1], :, :] - f
v = np.asfortranarray(np.stack((fx, fy), axis=-1))
# Output
output = np.zeros_like(v)
Halide('prox_IsoL1', recompile=True).prox_IsoL1(v, theta,
output) # Call
# Reference
normv = np.sqrt(
np.multiply(v[:, :, :, 0], v[:, :, :, 0]) +
np.multiply(v[:, :, :, 1], v[:, :, :, 1]))
normv = np.stack((normv, normv), axis=-1)
with np.errstate(divide='ignore'):
output_ref = np.maximum(0.0, 1.0 - theta / normv) * v
self.assertItemsAlmostEqual(output, output_ref)
def test_slicing(self):
"""Test slicing over numpy arrays in halide.
"""
if halide_installed():
# Load image
testimg_filename = os.path.join(
os.path.dirname(os.path.realpath(__file__)), 'data',
'angela.jpg')
np_img = im2nparray(Image.open(testimg_filename))
np_img = np.asfortranarray(
np.tile(np_img[..., np.newaxis], (1, 1, 1, 3)))
# Test problem
output = np.zeros_like(np_img)
mask = np.asfortranarray(
np.random.randn(*list(np_img.shape[0:3])).astype(np.float32))
mask = np.maximum(mask, 0.)
for k in range(np_img.shape[3]):
Halide('A_mask.cpp').A_mask(
np.asfortranarray(np_img[:, :, :, k]), mask,
output[:, :, :, k]) # Call
output_ref = np.zeros_like(np_img)
for k in range(np_img.shape[3]):
output_ref[:, :, :, k] = mask * np_img[:, :, :, k]
self.assertItemsAlmostEqual(output, output_ref)
def test_warp_halide(self):
"""Test warp lin op in halide.
"""
# Load image
WIDTH = 512
np_img = get_test_image(WIDTH)
# Generate problem
theta_rad = 5.0 * np.pi / 180.0
H = np.array(
[[np.cos(theta_rad), -np.sin(theta_rad), WIDTH * 0.5],
[np.sin(theta_rad), np.cos(theta_rad), 0.], [0., 0., 1.]],
dtype=np.float32,
order='F')
# Reference
output_ref = cv2.warpPerspective(np_img,
H,
np_img.shape[1::-1],
flags=cv2.INTER_LINEAR
| cv2.WARP_INVERSE_MAP,
borderMode=cv2.BORDER_CONSTANT,
borderValue=0.)
# Halide
output = np.zeros_like(np_img)
Halide('A_warp', recompile=True).A_warp(np_img, H, output) # Call
self.assertItemsAlmostEqual(output, output_ref, eps=1e-1)
# Transpose
output_trans = np.zeros_like(np_img)
Hinvc = np.asfortranarray(np.linalg.pinv(H))
Halide('At_warp', recompile=True).At_warp(output, Hinvc,
output_trans) # Call
# Compute reference
output_ref_trans = cv2.warpPerspective(output,
H,
np_img.shape[1::-1],
flags=cv2.INTER_LINEAR,
borderMode=cv2.BORDER_CONSTANT,
borderValue=0.)
self.assertItemsAlmostEqual(output_trans, output_ref_trans, eps=1e-1)
def adjoint(self, inputs, outputs):
"""The adjoint operator.
Reads from inputs and writes to outputs.
"""
if self.implementation == Impl['halide'] and \
(len(self.shape) == 3 or len(self.shape) == 4) and self.dims == 2:
# Halide implementation
if len(self.shape) == 3:
tmpin = np.asfortranarray(
np.reshape(inputs[0],
(self.shape[0], self.shape[1], 1, 2)).astype(
np.float32))
else:
tmpin = np.asfortranarray(inputs[0].astype(np.float32))
Halide('At_grad.cpp').At_grad(tmpin, self.tmpadj) # Call
np.copyto(outputs[0], np.reshape(self.tmpadj, self.shape[:-1]))
else:
# Compute comparison (Negative divergence)
f = inputs[0]
outputs[0].fill(0.0)
for j in range(self.dims):
# Get component
fj = f[..., j]
ss = fj.shape
# Add grad for this dimension (same as index)
istart = ()
ic = ()
iend_out = ()
iend_in = ()
for i in range(len(ss)):
if i == j:
istart += np.index_exp[0]
ic += np.index_exp[np.r_[0, 0:ss[j] - 1]]
iend_out += np.index_exp[-1]
iend_in += np.index_exp[-2]
else:
istart += np.index_exp[:]
ic += np.index_exp[:]
iend_in += np.index_exp[:]
iend_out += np.index_exp[:]
# Do the grad operation for dimension j
fd = fj - fj[ic]
fd[istart] = fj[istart]
fd[iend_out] = -fj[iend_in]
outputs[0] += (-fd)
def _prox(self, rho, v, *args, **kwargs):
"""x = v * (1 - (1/rho) * 1/||x||_g )_+
"""
if self.implementation == Impl['halide'] and \
len(self.lin_op.shape) in [3, 4] and self.lin_op.shape[-1] == 2 and \
self.group_dims == [len(self.lin_op.shape) - 1]:
# Halide implementation
if len(self.lin_op.shape) == 3:
tmpin = np.asfortranarray(
np.reshape(v, (self.lin_op.shape[0], self.lin_op.shape[1],
1, 2)).astype(np.float32))
else:
tmpin = np.asfortranarray(v.astype(np.float32))
Halide('prox_IsoL1.cpp').prox_IsoL1(tmpin, 1.0 / rho,
self.tmpout) # Call
np.copyto(v, np.reshape(self.tmpout, self.lin_op.shape))
else:
# Numpy implementation
np.multiply(v, v, self.v_group_norm)
# Sum along dimensions and keep dimensions
orig_s = v.shape
for d in self.group_dims:
self.v_group_norm = np.sum(self.v_group_norm,
axis=d,
keepdims=True)
# Sqrt
np.sqrt(self.v_group_norm, self.v_group_norm)
# Replicate
tiles = ()
for d in range(len(orig_s)):
if d in self.group_dims:
tiles += (orig_s[d], )
else:
tiles += (1, )
self.v_group_norm = np.tile(self.v_group_norm, tiles)
# Thresholded group norm
with np.errstate(divide='ignore'):
np.maximum(0.0, 1.0 - (1.0 / rho) * (1.0 / self.v_group_norm),
self.v_group_norm)
# Mult
v *= self.v_group_norm
return v
def test_mask_halide(self):
"""Test mask lin op in halide.
"""
np_img = get_test_image(512)
# Test problem
output = np.zeros_like(np_img)
mask = np.asfortranarray(
np.random.randn(*np_img.shape).astype(np.float32))
mask = np.maximum(mask, 0.)
Halide('A_mask', recompile=True).A_mask(np_img, mask, output) # Call
output_ref = mask * np_img
# Transpose
output_trans = np.zeros_like(np_img)
Halide('At_mask', recompile=True).At_mask(np_img, mask,
output_trans) # Call
self.assertItemsAlmostEqual(output, output_ref)
self.assertItemsAlmostEqual(output_trans, output_ref)
def test_conv_halide(self):
"""Test convolution lin op in halide.
"""
if halide_installed():
# Load image
testimg_filename = os.path.join(
os.path.dirname(os.path.realpath(__file__)), 'data',
'angela.jpg')
# opens the file using Pillow - it's not an array yet
img = Image.open(testimg_filename)
np_img = np.asfortranarray(im2nparray(img))
# Convert to gray
np_img = np.mean(np_img, axis=2)
# Test problem
output = np.zeros_like(np_img)
K = np.ones((11, 11), dtype=np.float32, order='FORTRAN')
K /= np.prod(K.shape)
# Convolve in halide
Halide('A_conv.cpp').A_conv(np_img, K, output)
# Convolve in scipy
output_ref = signal.convolve2d(np_img,
K,
mode='same',
boundary='wrap')
# Transpose
output_corr = np.zeros_like(np_img)
Halide('At_conv.cpp').At_conv(np_img, K, output_corr) # Call
output_corr_ref = signal.convolve2d(np_img,
np.flipud(np.fliplr(K)),
mode='same',
boundary='wrap')
self.assertItemsAlmostEqual(output, output_ref)
self.assertItemsAlmostEqual(output_corr, output_corr_ref)
def test_grad_halide(self):
"""Test gradient lin op in halide.
"""
# Load image
np_img = get_test_image(512)
# Test problem
output = np.zeros((np_img.shape[0], np_img.shape[1], 1),
dtype=np.float32,
order='F')
#Gradient in halide
Halide('A_grad', recompile=True).A_grad(np_img, output) # Call
# Compute comparison
f = np_img
if len(np_img.shape) == 2:
f = f[..., np.newaxis]
ss = f.shape
fx = f[:, np.r_[1:ss[1], ss[1] - 1], :] - f
fy = f[np.r_[1:ss[0], ss[0] - 1], :, :] - f
Kf = np.asfortranarray(np.stack((fy, fx), axis=-1))
# Transpose
output_trans = np.zeros(f.shape, dtype=np.float32, order='F')
Halide('At_grad', recompile=True).At_grad(Kf, output_trans) # Call
# Compute comparison (Negative divergence)
Kfy = Kf[:, :, :, 0]
fy = Kfy - Kfy[np.r_[0, 0:ss[0] - 1], :, :]
fy[0, :, :] = Kfy[0, :, :]
fy[-1, :, :] = -Kfy[-2, :, :]
Kfx = Kf[:, :, :, 1]
ss = Kfx.shape
fx = Kfx - Kfx[:, np.r_[0, 0:ss[1] - 1], :]
fx[:, 0, :] = Kfx[:, 0, :]
fx[:, -1, :] = -Kfx[:, -2, :]
def forward(self, inputs, outputs):
"""The forward operator.
Reads from inputs and writes to outputs.
"""
if self.implementation == Impl['halide'] and (len(self.shape) in [2, 3]):
# Halide implementation
tmpin = np.asfortranarray(inputs[0].astype(np.float32))
Halide('A_mask.cpp').A_mask(tmpin, self.weight, self.tmpout) # Call
np.copyto(outputs[0], self.tmpout)
else:
# Numpy
np.multiply(inputs[0], self.weight, outputs[0])
def _prox(self, rho, v, *args, **kwargs):
"""x = 1/2* ( v - 1./rho + sqrt( ||v - 1./rho||^2 + 4 * 1./rho * b ) )
"""
if self.implementation == Impl['halide'] and (len(self.lin_op.shape)
in [2, 3, 4]):
# Halide implementation
Halide('prox_Poisson').prox_Poisson(v, self.maskh, self.bph,
np.float32(1. / rho),
self.tmpout)
np.copyto(v, self.tmpout)
else:
v = 0.5 * (v - 1. / rho +
np.sqrt((v - 1. / rho) *
(v - 1. / rho) + 4. * 1. / rho * self.bp))
return v
def adjoint(self, inputs, outputs):
"""The adjoint operator.
Reads from inputs and writes to outputs.
"""
if self.implementation == Impl['halide']:
# Halide implementation
if len(self.H.shape) == 2:
tmpin = np.asfortranarray(inputs[0][..., np.newaxis].astype(
np.float32))
else:
tmpin = np.asfortranarray(inputs[0].astype(np.float32))
Halide('At_warp.cpp').At_warp(tmpin, self.Hf, self.tmpadj) # Call
np.copyto(outputs[0], self.tmpadj)
else:
# CV2 version
inimg = inputs[0]
if len(self.H.shape) == 2:
# + cv2.WARP_INVERSE_MAP
warpedInput = cv2.warpPerspective(
np.asfortranarray(inimg),
self.Hinv.T,
inimg.shape[1::-1],
flags=cv2.INTER_LINEAR,
borderMode=cv2.BORDER_CONSTANT,
borderValue=0.)
np.copyto(outputs[0], warpedInput)
else:
outputs[0][:] = 0.0
for j in range(self.H.shape[2]):
warpedInput = cv2.warpPerspective(
np.asfortranarray(inimg[:, :, :, j]),
self.Hinv[:, :, j].T,
inimg.shape[1::-1],
flags=cv2.INTER_LINEAR,
borderMode=cv2.BORDER_CONSTANT,
borderValue=0.)
# Necessary due to array layout in opencv
outputs[0] += warpedInput
def test_norm1_halide(self):
"""Halide Norm 1 test
"""
# Load image
np_img = get_test_image(512)
# Test problem
v = np_img
theta = 0.5
# Output
output = np.zeros_like(np_img)
Halide('prox_L1', recompile=True).prox_L1(v, theta, output) # Call
# Reference
output_ref = np.maximum(0.0, v - theta) - np.maximum(0.0, -v - theta)
self.assertItemsAlmostEqual(output, output_ref)
def _prox(self, rho, v, *args, **kwargs):
"""x = 1/2* ( v - 1./rho + sqrt( ||v - 1./rho||^2 + 4 * 1./rho * b ) )
"""
if self.implementation == Impl['halide'] and (len(self.lin_op.shape)
in [2, 3, 4]):
# Halide implementation
tmpin = np.asfortranarray(v)
Halide('prox_Poisson.cpp').prox_Poisson(tmpin, self.maskh,
self.bph, 1. / rho,
self.tmpout)
np.copyto(v, self.tmpout)
else:
v = 0.5 * (v - 1. / rho +
np.sqrt((v - 1. / rho) *
(v - 1. / rho) + 4. * 1. / rho * self.bp))
return v
def _prox(self, rho, v, *args, **kwargs):
"""x = sign(v)*(|v| - 1/rho)_+
"""
if self.implementation == Impl['halide'] and (len(self.lin_op.shape)
in [2, 3, 4]):
# Halide implementation
Halide('prox_L1').prox_L1(v, 1. / rho, self.tmpout)
np.copyto(v, self.tmpout)
else:
# Numpy implementation
np.sign(v, self.v_sign)
np.absolute(v, v)
v -= 1. / rho
np.maximum(v, 0, v)
v *= self.v_sign
return v
def forward(self, inputs, outputs):
"""The forward operator.
Reads from inputs and writes to outputs.
"""
self.init_kernel()
if self.implementation == Impl['halide'] and \
(len(self.shape) == 2 or (len(self.shape) == 3 and self.dims == 2)):
# Halide implementation
Halide('A_conv').A_conv(inputs[0], self.kernel,
self.tmpout) # Call
np.copyto(outputs[0], self.tmpout)
else:
# Default numpy using FFT
X = fftd(inputs[0], self.dims)
X *= self.forward_kernel
np.copyto(outputs[0], ifftd(X, self.dims).real)
def adjoint(self, inputs, outputs):
"""The adjoint operator.
Reads from inputs and writes to outputs.
"""
self.init_kernel()
if self.implementation == Impl['halide'] and \
(len(self.shape) == 2 or (len(self.shape) == 3 and self.dims == 2)):
# Halide implementation
tmpin = np.asfortranarray(inputs[0].astype(np.float32))
Halide('At_conv.cpp').At_conv(tmpin, self.kernel, self.tmpout) # Call
np.copyto(outputs[0], self.tmpout)
else:
# Default numpy using FFT
U = fftd(inputs[0], self.dims)
U *= self.adjoint_kernel
np.copyto(outputs[0], ifftd(U, self.dims).real)
def forward(self, inputs, outputs):
"""The forward operator for n-d gradients.
Reads from inputs and writes to outputs.
"""
if self.implementation == Impl['halide'] and \
(len(self.shape) == 3 or len(self.shape) == 4) and self.dims == 2:
# Halide implementation
if len(self.shape) == 3:
tmpin = np.asfortranarray(
(inputs[0][..., np.newaxis]).astype(np.float32))
else:
tmpin = np.asfortranarray((inputs[0]).astype(np.float32))
Halide('A_grad.cpp').A_grad(tmpin, self.tmpfwd) # Call
np.copyto(outputs[0], np.reshape(self.tmpfwd, self.shape))
else:
# Input
f = inputs[0]
# Build up index for shifted array
ss = f.shape
stack_arr = ()
for j in range(self.dims):
# Add grad for this dimension (same as index)
il = ()
for i in range(len(ss)):
if i == j:
il += np.index_exp[np.r_[1:ss[j], ss[j] - 1]]
else:
il += np.index_exp[:]
fgrad_j = f[il] - f
stack_arr += (fgrad_j, )
# Stack all grads as new dimension
np.copyto(outputs[0], np.stack(stack_arr, axis=-1))
def forward(self, inputs, outputs):
"""The forward operator.
Reads from inputs and writes to outputs.
"""
if self.implementation == Impl['halide']:
# Halide implementation
tmpin = np.asfortranarray(inputs[0].astype(np.float32))
Halide('A_warp.cpp').A_warp(tmpin, self.Hinvf, self.tmpfwd) # Call
np.copyto(outputs[0], np.reshape(self.tmpfwd, self.shape))
else:
# CV2 version
inimg = inputs[0]
if len(self.H.shape) == 2:
warpedInput = cv2.warpPerspective(
np.asfortranarray(inimg),
self.H.T,
inimg.shape[1::-1],
flags=cv2.INTER_LINEAR,
borderMode=cv2.BORDER_CONSTANT,
borderValue=0.)
# Necessary due to array layout in opencv
np.copyto(outputs[0], warpedInput)
else:
for j in range(self.H.shape[2]):
warpedInput = cv2.warpPerspective(
np.asfortranarray(inimg),
self.H[:, :, j].T,
inimg.shape[1::-1],
flags=cv2.INTER_LINEAR,
borderMode=cv2.BORDER_CONSTANT,
borderValue=0.)
# Necessary due to array layout in opencv
np.copyto(outputs[0][:, :, :, j], warpedInput)
def test_isonorm1_halide(self):
"""Halide Norm 1 test
"""
if halide_installed():
# Load image
testimg_filename = os.path.join(
os.path.dirname(os.path.realpath(__file__)), 'data',
'angela.jpg')
img = Image.open(testimg_filename)
np_img = np.asfortranarray(im2nparray(img))
# Convert to gray
np_img = np.mean(np_img, axis=2)
# Test problem
theta = 0.5
f = np_img
if len(np_img.shape) == 2:
f = f[..., np.newaxis]
ss = f.shape
fx = f[:, np.r_[1:ss[1], ss[1] - 1], :] - f
fy = f[np.r_[1:ss[0], ss[0] - 1], :, :] - f
v = np.asfortranarray(np.stack((fx, fy), axis=-1))
# Output
output = np.zeros_like(v)
Halide('prox_IsoL1.cpp').prox_IsoL1(v, theta, output) # Call
# Reference
normv = np.sqrt(
np.multiply(v[:, :, :, 0], v[:, :, :, 0]) +
np.multiply(v[:, :, :, 1], v[:, :, :, 1]))
normv = np.stack((normv, normv), axis=-1)
with np.errstate(divide='ignore'):
output_ref = np.maximum(0.0, 1.0 - theta / normv) * v
self.assertItemsAlmostEqual(output, output_ref)
def test_poisson_halide(self):
"""Halide Poisson norm test
"""
if halide_installed():
# Load image
testimg_filename = os.path.join(
os.path.dirname(os.path.realpath(__file__)), 'data',
'angela.jpg')
img = Image.open(testimg_filename)
np_img = np.asfortranarray(im2nparray(img))
# Convert to gray
np_img = np.mean(np_img, axis=2)
# Test problem
v = np_img
theta = 0.5
mask = np.asfortranarray(
np.random.randn(*list(np_img.shape)).astype(np.float32))
mask = np.maximum(mask, 0.)
b = np_img * np_img
# Output
output = np.zeros_like(v)
tic()
Halide('prox_Poisson.cpp').prox_Poisson(v, mask, b, theta,
output) # Call
print('Running took: {0:.1f}ms'.format(toc()))
# Reference
output_ref = 0.5 * (v - theta +
np.sqrt((v - theta) *
(v - theta) + 4 * theta * b))
output_ref[mask <= 0.5] = v[mask <= 0.5]
self.assertItemsAlmostEqual(output, output_ref)
def cg(KtKfun, b, tol, num_iters, verbose, x_init=None, implem=Impl['numpy']):
# Solves KtK x = b with
# KtKfun being a function that computes the matrix vector product KtK x
# TODO: Fix halide later
implem == Impl['numpy']
if implem == Impl['halide']:
output = np.array([0.0], dtype=np.float32)
hl_norm2 = Halide('A_norm_L2.cpp',
generator_name="normL2_1DImg",
func="A_norm_L2_1D").A_norm_L2_1D
hl_dot = Halide('A_dot_prod.cpp',
generator_name="dot_1DImg",
func="A_dot_1D").A_dot_1D
# Temp vars
x = np.zeros(b.shape, dtype=np.float32, order='F')
r = np.zeros(b.shape, dtype=np.float32, order='F')
Ap = np.zeros(b.shape, dtype=np.float32, order='F')
else:
# Temp vars
x = np.zeros(b.shape)
r = np.zeros(b.shape)
Ap = np.zeros(b.shape)
# Initialize x
# Initialize everything to zero.
if x_init is not None:
x = x_init
# Compute residual
# r = b - KtKfun(x)
KtKfun(x, r)
r *= -1.0
r += b
# Do cg iterations
if implem == Impl['halide']:
hl_norm2(b.ravel().astype(np.float32), output)
cg_tol = tol * output[0]
else:
cg_tol = tol * np.linalg.norm(b.ravel(), 2) # Relative tol
# CG iteration
gamma_1 = p = None
cg_iter = np.minimum(num_iters, np.prod(b.shape))
for iter in range(cg_iter):
# Check for convergence
if implem == Impl['halide']:
hl_norm2(r.ravel(), output)
normr = output[0]
else:
normr = np.linalg.norm(r.ravel(), 2)
# Check for convergence
if normr <= cg_tol:
break
# gamma = r'*r;
if implem == Impl['halide']:
hl_norm2(r.ravel(), output)
gamma = output[0]
gamma *= gamma
else:
gamma = np.dot(r.ravel().T, r.ravel())
# direction vector
if iter > 0:
beta = gamma / gamma_1
p = r + beta * p
else:
p = r
# Compute Ap
KtKfun(p, Ap)
# Cg update
q = Ap
# alpha = gamma / (p'*q);
if implem == Impl['halide']:
hl_dot(p.ravel(), q.ravel(), output)
alpha = gamma / output[0]
else:
alpha = gamma / np.dot(p.ravel().T, q.ravel())
x = x + alpha * p # update approximation vector
r = r - alpha * q # compute residual
gamma_1 = gamma
# Iterate
if verbose:
print("CG Iter %03d" % iter)
return x
def _prox(self, rho, v, *args, **kwargs):
"""x = denoise_gaussian_NLM( tonemap(v), sqrt(1/rho))
"""
if self.implementation == Impl['halide'] and \
len(self.lin_op.shape) == 3 and self.lin_op.shape[2] == 3:
# Halide implementation
tmpin = np.asfortranarray(v.astype(np.float32))
Halide('prox_NLM.cpp').prox_NLM(tmpin, 1. / rho, self.paramsh,
self.tmpout)
np.copyto(v, self.tmpout)
else:
# Compute sigma
sigma = np.sqrt(1.0 / rho)
# Fixed sigma if wanted
sigma_estim = sigma
if self.sigma_fixed > 0.0:
sigma_estim = self.sigma_fixed / 30.0 * self.sigma_scale
# Params
print(
("Prox NLM params are: [sigma ={0} prior={1} sigma_scale={2}]".
format(sigma_estim, self.prior, self.sigma_scale)))
# Scale d
v = v.copy()
v_min = np.amin(v)
v_max = np.amax(v)
v_max = np.maximum(v_max, v_min + 0.01)
# Scale and offset parameters d
v -= v_min
v /= (v_max - v_min)
# Denoising params
sigma_luma = sigma_estim
sigma_color = sigma_estim
if self.prior > 0.5:
sigma_luma = sigma_estim * 1.3 # NLM color stronger on luma
sigma_color = sigma_color * 3.0
# Transform
if self.gamma_trans != 1.0:
vuint = np.uint8(
np.clip(v**self.gamma_trans * 255.0, 0,
255)) # Quadratic tranform to account for gamma
else:
# Quadratic tranform to account for gamma
vuint = np.uint8(np.clip(v * 255.0, 0.0, 255.0))
if self.prior <= 0.5:
vdstuint = cv2.fastNlMeansDenoising(vuint, None,
sigma_luma * 255.0,
self.templateWindowSizeNLM,
self.searchWindowSizeNLM)
else:
vdstuint = cv2.fastNlMeansDenoisingColored(
vuint, None, sigma_luma * 255.0, sigma_color * 255.,
self.templateWindowSizeNLM, self.searchWindowSizeNLM)
# Convert to float and inverse scale
if self.gamma_trans != 1.0:
dst = ((vdstuint.astype(np.float32) / 255.0)**
(1.0 / self.gamma_trans)) * (v_max - v_min) + v_min
else:
dst = vdstuint.astype(
np.float32) / 255.0 * (v_max - v_min) + v_min
np.copyto(v, dst)
return v
def test_compile(self):
"""Test compilation and run.
"""
if halide_installed():
# Force recompilation of conv
Halide('A_conv.cpp', recompile=True, verbose=False)
