Skip to content

Commit

Permalink
aargh
Browse files Browse the repository at this point in the history
  • Loading branch information
tomrunia committed Dec 12, 2018
1 parent 8d630fb commit 2346521
Show file tree
Hide file tree
Showing 3 changed files with 106 additions and 120 deletions.
38 changes: 21 additions & 17 deletions examples/compare_both.py
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,7 @@
pyr_numpy = SCFpyr_NumPy(pyr_height, pyr_nbands, scale_factor=2)
coeff_numpy = pyr_numpy.build(image)
reconstruction_numpy = pyr_numpy.reconstruct(coeff_numpy)
reconstruction_numpy = reconstruction_numpy.astype(np.uint8)

print('#'*60)

Expand All @@ -60,15 +61,19 @@

pyr_torch = SCFpyr_PyTorch(pyr_height, pyr_nbands, device=device)
coeff_torch = pyr_torch.build(im_batch)
reconstruction_torch = pyr_torch.reconstruct(coeff_torch)
reconstruction_torch = reconstruction_torch.cpu().numpy()
reconstruction_torch_v2 = pyr_torch.reconstruct(coeff_torch)
#reconstruction_torch_v2 = reconstruction_torch_v2.cpu().numpy()[0,]

# Just extract a single example from the batch
# Also moves the example to CPU and NumPy
coeff_torch = utils.extract_from_batch(coeff_torch, 0)

cv2.waitKey(0)
exit()
# NOTE: reconstruction using NumPy implementation
# Gives the same result, so decomposition is correct (!)
# reconstruction_torch = pyr_numpy.reconstruct(coeff_torch)
# reconstruction_torch = reconstruction_torch.astype(np.uint8)

#exit()

################################################################################
# Check correctness
Expand Down Expand Up @@ -125,21 +130,20 @@
################################################################################
# Visualize

coeff_grid_numpy = utils.make_grid_coeff(coeff_numpy, normalize=True)
coeff_grid_torch = utils.make_grid_coeff(coeff_torch, normalize=True)

import cortex.vision

reconstruction_torch = np.ascontiguousarray(reconstruction_torch[0], np.float32)
reconstruction_numpy = np.ascontiguousarray(reconstruction_numpy, np.float32)
coeff_grid_numpy = utils.make_grid_coeff(coeff_numpy, normalize=False)
coeff_grid_torch = utils.make_grid_coeff(coeff_torch, normalize=False)

reconstruction_torch = cortex.vision.normalize_for_display(reconstruction_torch)
reconstruction_numpy = cortex.vision.normalize_for_display(reconstruction_numpy)
# import cortex.vision
# reconstruction_torch = np.ascontiguousarray(reconstruction_torch[0], np.float32)
# reconstruction_numpy = np.ascontiguousarray(reconstruction_numpy, np.float32)
# reconstruction_torch = cortex.vision.normalize_for_display(reconstruction_torch)
# reconstruction_numpy = cortex.vision.normalize_for_display(reconstruction_numpy)

cv2.imshow('image', image)
cv2.imshow('coeff numpy', coeff_grid_numpy)
cv2.imshow('coeff torch', coeff_grid_torch)
#cv2.imshow('image', image)
#cv2.imshow('coeff numpy', coeff_grid_numpy)
#cv2.imshow('coeff torch', coeff_grid_torch)
cv2.imshow('reconstruction numpy', reconstruction_numpy)
cv2.imshow('reconstruction torch', reconstruction_torch)
#cv2.imshow('reconstruction torch', reconstruction_torch)
cv2.imshow('reconstruction torch', reconstruction_torch_v2)

cv2.waitKey(0)
77 changes: 30 additions & 47 deletions steerable/SCFpyr_NumPy.py
Original file line number Diff line number Diff line change
Expand Up @@ -186,51 +186,11 @@ def reconstruct(self, coeff):
tempdft = self._reconstruct_levels(coeff[1:], log_rad, Xrcos, Yrcos, angle)

hidft = np.fft.fftshift(np.fft.fft2(coeff[0]))

real, imag = hidft.real, hidft.imag
print(' [numpy] level final. hidft real ({:.3f}, {:.3f}, {:.3f})'.format(
real[0:20,0:20].min().item(), real[0:20,0:20].mean().item(), real[0:20,0:20].max().item()
))
print(' [numpy] level final. hidft imag ({:.3f}, {:.3f}, {:.3f})'.format(
imag[0:20,0:20].min().item(), imag[0:20,0:20].mean().item(), imag[0:20,0:20].max().item()
))

outdft = tempdft * lo0mask + hidft * hi0mask

real, imag = outdft.real, outdft.imag
np.save('./assets/numpy_outdft_real.npy', real)
np.save('./assets/numpy_outdft_imag.npy', imag)

print(' [numpy] level final. outdft real ({:.3f}, {:.3f}, {:.3f})'.format(
real[0:20,0:20].min().item(), real[0:20,0:20].mean().item(), real[0:20,0:20].max().item()
))
print(' [numpy] level final. outdft imag ({:.3f}, {:.3f}, {:.3f})'.format(
imag[0:20,0:20].min().item(), imag[0:20,0:20].mean().item(), imag[0:20,0:20].max().item()
))

import cv2
import cortex.vision

reconstruction = np.fft.ifftshift(outdft)

real, imag = reconstruction.real, reconstruction.imag
print(' [numpy] level final. ifftshift real ({:.3f}, {:.3f}, {:.3f})'.format(
real[0:20,0:20].min().item(), real[0:20,0:20].mean().item(), real[0:20,0:20].max().item()
))
print(' [numpy] level final. ifftshift imag ({:.3f}, {:.3f}, {:.3f})'.format(
imag[0:20,0:20].min().item(), imag[0:20,0:20].mean().item(), imag[0:20,0:20].max().item()
))

real, imag = reconstruction.real, reconstruction.imag
cv2.imshow('numpy real', real)
cv2.imshow('numpy imag', imag)

reconstruction = np.fft.ifft2(reconstruction)
reconstruction = reconstruction.real

print(' [numpy] level final. reconstruction real ({:.3f}, {:.3f}, {:.3f})'.format(
reconstruction.min(), reconstruction.mean(), reconstruction.max()
))
reconstruction = reconstruction.real.astype(np.uint8)

return reconstruction

Expand All @@ -239,6 +199,15 @@ def _reconstruct_levels(self, coeff, log_rad, Xrcos, Yrcos, angle):
if len(coeff) == 1:
dft = np.fft.fft2(coeff[0])
dft = np.fft.fftshift(dft)

real, imag = dft.real, dft.imag
print(' [numpy] levels remaining {}. dft real ({:.3f}, {:.3f}, {:.3f})'.format(
len(coeff), real.mean().item(), real.std().item(), real.sum().item()
))
print(' [numpy] levels remaining {}. dft imag ({:.3f}, {:.3f}, {:.3f})'.format(
len(coeff), imag.mean().item(), imag.std().item(), imag.sum().item()
))

return dft

Xrcos = Xrcos - np.log2(self.scale_factor)
Expand Down Expand Up @@ -267,10 +236,10 @@ def _reconstruct_levels(self, coeff, log_rad, Xrcos, Yrcos, angle):
real = orientdft.real
imag = orientdft.imag
print(' [numpy] levels remaining {}. orientdft real ({:.3f}, {:.3f}, {:.3f})'.format(
len(coeff), real.min().item(), real.mean().item(), real.max().item()
len(coeff), real.mean().item(), real.std().item(), real.sum().item()
))
print(' [numpy] levels remaining {}. orientdft real ({:.3f}, {:.3f}, {:.3f})'.format(
len(coeff), imag.min().item(), imag.mean().item(), imag.max().item()
print(' [numpy] levels remaining {}. orientdft imag ({:.3f}, {:.3f}, {:.3f})'.format(
len(coeff), imag.mean().item(), imag.std().item(), imag.sum().item()
))

####################################################################
Expand All @@ -279,20 +248,34 @@ def _reconstruct_levels(self, coeff, log_rad, Xrcos, Yrcos, angle):

dims = np.array(coeff[0][0].shape)

lostart = (np.ceil((dims+0.5)/2) -
np.ceil((np.ceil((dims-0.5)/2)+0.5)/2)).astype(np.int32)
lostart = (np.ceil((dims+0.5)/2) - np.ceil((np.ceil((dims-0.5)/2)+0.5)/2)).astype(np.int32)
loend = lostart + np.ceil((dims-0.5)/2).astype(np.int32)

nlog_rad = log_rad[lostart[0]:loend[0], lostart[1]:loend[1]]
nangle = angle[lostart[0]:loend[0], lostart[1]:loend[1]]
YIrcos = np.sqrt(np.abs(1 - Yrcos * Yrcos))
YIrcos = np.sqrt(np.abs(1 - Yrcos**2))
lomask = pointOp(nlog_rad, YIrcos, Xrcos)

print(' [numpy] levels remaining {}. nlog_rad = {:.3f}, nangle = {:.3f}, YIrcos = {:.3f}, lomask = {:.3f}'.format(
len(coeff), nlog_rad.sum(), nangle.sum(), YIrcos.sum(), lomask.sum()
))

################################################################################

# Recursive call for image reconstruction
nresdft = self._reconstruct_levels(coeff[1:], nlog_rad, Xrcos, Yrcos, nangle)

resdft = np.zeros(dims, 'complex')
resdft[lostart[0]:loend[0], lostart[1]:loend[1]] = nresdft * lomask

real, imag = nresdft.real, nresdft.imag
print(' [numpy] levels remaining {}. nresdft real ({:.3f}, {:.3f}, {:.3f})'.format(
len(coeff), real.mean().item(), real.std().item(), real.sum().item()
))
print(' [numpy] levels remaining {}. nresdft imag ({:.3f}, {:.3f}, {:.3f})'.format(
len(coeff), imag.mean().item(), imag.std().item(), imag.sum().item()
))

return resdft + orientdft


Expand Down
111 changes: 55 additions & 56 deletions steerable/SCFpyr_PyTorch.py
Original file line number Diff line number Diff line change
Expand Up @@ -58,7 +58,7 @@ def __init__(self, height=5, nbands=4, scale_factor=2, device=None):
self.lutsize = 1024
self.Xcosn = np.pi * np.array(range(-(2*self.lutsize+1), (self.lutsize+2)))/self.lutsize
self.alpha = (self.Xcosn + np.pi) % (2*np.pi) - np.pi
self.complex_fact_construct = np.power(np.complex(0, -1), self.nbands-1)
self.complex_fact_construct = np.power(np.complex(0, -1), self.nbands-1)
self.complex_fact_reconstruct = np.power(np.complex(0, 1), self.nbands-1)

################################################################################
Expand Down Expand Up @@ -229,70 +229,50 @@ def reconstruct(self, coeff):
# Start recursive reconstruction
tempdft = self._reconstruct_levels(coeff[1:], log_rad, Xrcos, Yrcos, angle)

hidft = torch.rfft(coeff[0], signal_ndim=2, onesided=False)
hidft = math_utils.batch_fftshift2d(hidft)

real, imag = torch.unbind(hidft, -1)
print(' [torch] level final. hidft real ({:.3f}, {:.3f}, {:.3f})'.format(
real[0,0:20,0:20].min().item(), real[0,0:20,0:20].mean().item(), real[0,0:20,0:20].max().item()
))
print(' [torch] level final. hidft imag ({:.3f}, {:.3f}, {:.3f})'.format(
imag[0,0:20,0:20].min().item(), imag[0,0:20,0:20].mean().item(), imag[0,0:20,0:20].max().item()
))

outdft = tempdft * lo0mask + hidft * hi0mask

real, imag = torch.unbind(outdft, -1)
np.save('./assets/torch_outdft_real.npy', real[0].cpu().numpy())
np.save('./assets/torch_outdft_imag.npy', imag[0].cpu().numpy())

print(' [torch] level final. outdft real ({:.3f}, {:.3f}, {:.3f})'.format(
real[0,0:20,0:20].min().item(), real[0,0:20,0:20].mean().item(), real[0,0:20,0:20].max().item()
))
print(' [torch] level final. outdft imag ({:.3f}, {:.3f}, {:.3f})'.format(
imag[0,0:20,0:20].min().item(), imag[0,0:20,0:20].mean().item(), imag[0,0:20,0:20].max().item()
))
################################################################################
# This NumPy reconstruction still gives the wrong result, so mistake is somewhere before

# TODO: check these ops (!)
# hidft = np.fft.fftshift(np.fft.fft2(coeff[0]))
# outdft = tempdft * lo0mask + hidft * hi0mask

import cv2
import cortex.vision
tempdft = tempdft[0].cpu().numpy()
tempdft = tempdft[:,:,0] + 1j*tempdft[:,:,1]

lo0mask = pointOp(log_rad, YIrcos, Xrcos)
hi0mask = pointOp(log_rad, Yrcos, Xrcos)

coeff_0 = coeff[0].cpu().numpy()[0,]

# reconstruction = np.fft.ifft2(np.fft.ifftshift(outdft))
# reconstruction = reconstruction.real.astype(int)
reconstruction = math_utils.batch_ifftshift2d(outdft)
hidft = np.fft.fftshift(np.fft.fft2(coeff_0))
outdft = tempdft * lo0mask + hidft * hi0mask
reconstruction = np.fft.ifftshift(outdft)
reconstruction = np.fft.ifft2(reconstruction)
reconstruction = reconstruction.real.astype(np.uint8)

real, imag = torch.unbind(reconstruction, -1)
print(' [torch] level final. ifftshift real ({:.3f}, {:.3f}, {:.3f})'.format(
real[0,0:20,0:20].min().item(), real[0,0:20,0:20].mean().item(), real[0,0:20,0:20].max().item()
))
print(' [torch] level final. ifftshift imag ({:.3f}, {:.3f}, {:.3f})'.format(
imag[0,0:20,0:20].min().item(), imag[0,0:20,0:20].mean().item(), imag[0,0:20,0:20].max().item()
))
return reconstruction

#real, imag = reconstruction, reconstruction #qtorch.unbind(reconstruction, -1)
cv2.imshow('real', real.cpu().numpy()[0])
cv2.imshow('imag', imag.cpu().numpy()[0])

# this fucks it up ...
reconstruction = torch.ifft(reconstruction, signal_ndim=2) #, onesided=False)
reconstruction_real = torch.unbind(reconstruction, -1)[0]
################################################################################

print(' [torch] level final. reconstruction real ({:.3f}, {:.3f}, {:.3f})'.format(
reconstruction_real.min().item(), reconstruction_real.mean().item(), reconstruction_real.max().item()
))
# hidft = torch.rfft(coeff[0], signal_ndim=2, onesided=False)
# hidft = math_utils.batch_fftshift2d(hidft)
# outdft = tempdft * lo0mask + hidft * hi0mask

return reconstruction_real
# reconstruction = torch.ifft(outdft, signal_ndim=2) #, onesided=False)
# reconstruction_real = torch.unbind(reconstruction, -1)[0]
# return reconstruction_real

def _reconstruct_levels(self, coeff, log_rad, Xrcos, Yrcos, angle):

if len(coeff) == 1:
dft = torch.rfft(coeff[0], signal_ndim=2, onesided=False)
dft = math_utils.batch_fftshift2d(dft)

real, imag = torch.unbind(dft, -1)
real = real[0,].cpu().numpy()
imag = imag[0,].cpu().numpy()
print(' [torch] levels remaining {}. dft real ({:.3f}, {:.3f}, {:.3f})'.format(
len(coeff), real.mean().item(), real.std().item(), real.sum().item()
))
print(' [torch] levels remaining {}. dft imag ({:.3f}, {:.3f}, {:.3f})'.format(
len(coeff), imag.mean().item(), imag.std().item(), imag.sum().item()
))
return dft

Xrcos = Xrcos - np.log2(self.scale_factor)
Expand Down Expand Up @@ -323,18 +303,22 @@ def _reconstruct_levels(self, coeff, log_rad, Xrcos, Yrcos, angle):
# Now multiply with complex number
# (x+yi)(u+vi) = (xu-yv) + (xv+yu)i
banddft = torch.unbind(banddft, -1)
# banddft_real = self.complex_fact_construct.real*banddft[0] - self.complex_fact_construct.imag*banddft[1]
# banddft_imag = self.complex_fact_construct.real*banddft[1] + self.complex_fact_construct.imag*banddft[0]
banddft_real = self.complex_fact_reconstruct.real*banddft[0] - self.complex_fact_reconstruct.imag*banddft[1]
banddft_imag = self.complex_fact_reconstruct.real*banddft[1] + self.complex_fact_reconstruct.imag*banddft[0]
banddft = torch.stack((banddft_real, banddft_imag), -1)
banddft = banddft * anglemask * himask
orientdft += banddft

real, imag = torch.unbind(orientdft, -1)
real = real[0,].cpu().numpy()
imag = imag[0,].cpu().numpy()
print(' [torch] levels remaining {}. orientdft real ({:.3f}, {:.3f}, {:.3f})'.format(
len(coeff), real.min().item(), real.mean().item(), real.max().item()
len(coeff), real.mean().item(), real.std().item(), real.sum().item()
))
print(' [torch] levels remaining {}. orientdft real ({:.3f}, {:.3f}, {:.3f})'.format(
len(coeff), imag.min().item(), imag.mean().item(), imag.max().item()
print(' [torch] levels remaining {}. orientdft imag ({:.3f}, {:.3f}, {:.3f})'.format(
len(coeff), imag.mean().item(), imag.std().item(), imag.sum().item()
))

####################################################################
Expand All @@ -348,7 +332,12 @@ def _reconstruct_levels(self, coeff, log_rad, Xrcos, Yrcos, angle):

nlog_rad = log_rad[lostart[0]:loend[0], lostart[1]:loend[1]]
nangle = angle[lostart[0]:loend[0], lostart[1]:loend[1]]
YIrcos = np.sqrt(np.abs(1 - Yrcos * Yrcos))
YIrcos = np.sqrt(np.abs(1 - Yrcos**2))
lomask = pointOp(nlog_rad, YIrcos, Xrcos)

print(' [torch] levels remaining {}. nlog_rad = {:.3f}, nangle = {:.3f}, YIrcos = {:.3f}, lomask = {:.3f}'.format(
len(coeff), nlog_rad.sum(), nangle.sum(), YIrcos.sum(), lomask.sum()
))

# Filtering
lomask = pointOp(nlog_rad, YIrcos, Xrcos)
Expand All @@ -363,6 +352,16 @@ def _reconstruct_levels(self, coeff, log_rad, Xrcos, Yrcos, angle):
resdft = torch.zeros_like(coeff[0][0]).to(self.device)
resdft[:,lostart[0]:loend[0], lostart[1]:loend[1],:] = nresdft * lomask

real, imag = torch.unbind(nresdft, -1)
real = real[0,].cpu().numpy()
imag = imag[0,].cpu().numpy()
print(' [torch] levels remaining {}. nresdft real ({:.3f}, {:.3f}, {:.3f})'.format(
len(coeff), real.mean().item(), real.std().item(), real.sum().item()
))
print(' [torch] levels remaining {}. nresdft imag ({:.3f}, {:.3f}, {:.3f})'.format(
len(coeff), imag.mean().item(), imag.std().item(), imag.sum().item()
))

return resdft + orientdft

################################################################################
Expand Down

0 comments on commit 2346521

Please sign in to comment.