# Requires PyTorch with MKL when setting to 'cpu'
device = torch.device('cpu')
# Load batch of images [N,1,H,W]
im_batch_numpy = utils.load_image_batch('./assets/lena.jpg',32,600)
img=cv2.imread('./assets/lena.jpg',0)
cv2.imshow('yuantu',img)
im_torch = torch.from_numpy(img).to(device)
im_batch_torch=im_torch.unsqueeze(0).unsqueeze(0).float()
# Initialize Complex Steerbale Pyramid
height = 12
nbands = 4
scale_factor = 2**(1/2)
pyr = SCFpyr_PyTorch(height=height, nbands=nbands, scale_factor=scale_factor, device=device)
pyr_type = 1
# Decompose entire batch of images
coeff = pyr.build(im_batch_torch,pyr_type)
# Reconstruct batch of images again
img_recon = pyr.reconstruct(coeff,pyr_type)
img=im_torch.float()
recon=img_recon.squeeze()
loss=torch.nn.MSELoss()
print('MSE:',loss(img,recon))
cv2.imshow('recon',recon.numpy().astype(np.uint8))
# Visualization
# coeff_single = utils.extract_from_batch(coeff, 0)
# coeff_grid = utils.make_grid_coeff(coeff_single, normalize=True)
# cv2.imshow('Complex Steerable Pyramid', coeff_grid)
cv2.waitKey(0)
for n, Triplets_batch in enumerate(trainloader):
# get images_list [[N,C,H,W],[N,C,H,W],...], usually len(images_list)=batch_size if len(dataset)%batch_size==0
images_list = [torch.stack([Triplets_batch['start'][i],
Triplets_batch['inter'][i],
Triplets_batch['end'][i]]) for i in range(len(Triplets_batch['start']))]
# batch_coeff_list = [pyr.BatchCsp(
# image, channel=channel, type=1) for image in images_list]
batch_coeff_list = [pyr.build(image[:,channel,:,:].unsqueeze(1).to(device), pyr_type=pyr_type)
for image in images_list]
train_coeff, truth_coeff = get_input(batch_coeff_list)
# Forward pass
pre_coeff = model(train_coeff)
truth_img = Triplets_batch['inter'][:, channel, :, :]
pre_img = pyr.reconstruct(output_convert(pre_coeff),pyr_type=pyr_type)
# import pdb;pdb.set_trace()
loss = criterion(truth_coeff, pre_coeff, truth_img, pre_img)
# Backward and optimize
optimizer.zero_grad()
loss.backward()
optimizer.step()
# print(model.state_dict()['alpha'])
# print(model.state_dict()['pred.1.conv.bias'])
total_step+=1
if total_step % 10 == 0:
print('Epoch [{}/{}], Channel [{}], Step [{}], Loss: {:.4f}'
.format(epoch+1, num_epochs, channel, total_step, loss.item()))
# save log
with open(os.path.join(log_dir, log_name), 'at') as f:
coeff = coeffs[0].cpu().numpy().squeeze(0)
all_close = np.allclose(coeff, coeffs_c[0], atol=tolerance)
s = np.sum(coeff - coeffs_c[0])
print('Succesful for subband {}: {}, with tolerance of {}'.format(
0, all_close, tolerance))
print('Sum of difference: {}'.format(s))
for i in range(1, len(coeffs) - 1):
for j in range(len(coeffs[i])):
coeff = coeffs[i][j].cpu().numpy().squeeze(0)
coeff = coeff[..., 0] + 1j * coeff[..., 1]
all_close = np.allclose(coeff, coeffs_c[i][j], atol=tolerance)
s = np.sum(coeff - coeffs_c[i][j])
print('Succesful for subband {} {}: {}'.format(i, j, all_close))
print('Sum of difference: {}'.format(s))
coeff = coeffs[-1].cpu().numpy().squeeze(0)
all_close = np.allclose(coeff, coeffs_c[-1], atol=tolerance)
s = np.sum(coeff - coeffs_c[-1])
print('Succesful for subband {}: {}'.format(len(coeffs), all_close))
print('Sum of difference: {}'.format(s))
# reconstruction
tolerance = 1e-2
rec = pyr_NoSub.reconstruct(coeffs)
rec = rec.cpu().numpy().squeeze(0)
#rec_c = pyr_NoSub_c.reconSCFpyr(coeffs_c)
all_close = np.allclose(rec, img, atol=tolerance)
print(
'Succesful for reconstruction of GPU implementation: {}, with tolerance of {}'
.format(all_close, tolerance))
print('Sum of difference: {}'.format((rec - img).sum()))
reconstruction_numpy = pyr_numpy.reconstruct(coeff_numpy)
reconstruction_numpy = reconstruction_numpy.astype(np.uint8)
print('#' * 60)
################################################################################
# PyTorch
device = torch.device('cuda:0')
im_batch = torch.from_numpy(image[None, None, :, :])
im_batch = im_batch.to(device).float()
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()[0, ]
# Extract first example from the batch and move to CPU
coeff_torch = utils.extract_from_batch(coeff_torch, 0)
################################################################################
# Check correctness
print('#' * 60)
assert len(coeff_numpy) == len(coeff_torch)
for level, _ in enumerate(coeff_numpy):
print('Pyramid Level {level}'.format(level=level))
coeff_level_numpy = coeff_numpy[level]
coeff = coeffs[0].cpu().numpy().squeeze(0)
all_close = np.allclose(coeff, coeffs_c[0], atol=tolerance)
s = np.sum(coeff-coeffs_c[0])
print('Succesful for subband {}: {}, with tolerance of {}'.format(0,all_close, tolerance))
print('Sum of difference: {}'.format(s))
for i in range(1,len(coeffs)-1):
for j in range(len(coeffs[i])):
coeff = coeffs[i][j].cpu().numpy().squeeze(0)
coeff = coeff[...,0] + 1j * coeff[...,1]
all_close = np.allclose(coeff, coeffs_c[i][j], atol=tolerance)
s = np.sum(coeff - coeffs_c[i][j])
print('Succesful for subband {} {}: {}'.format(i,j,all_close))
print('Sum of difference: {}'.format(s))
coeff = coeffs[-1].cpu().numpy().squeeze(0)
all_close = np.allclose(coeff, coeffs_c[-1], atol=tolerance)
s = np.sum(coeff-coeffs_c[-1])
print('Succesful for subband {}: {}'.format(len(coeffs),all_close))
print('Sum of difference: {}'.format(s))
# reconstruction
tolerance = 1e-2
rec = pyr.reconstruct(coeffs)
rec = rec.cpu().numpy().squeeze(0)
#rec_c = pyr_c.reconSCFpyr(coeffs_c)
all_close = np.allclose(rec, img, atol=tolerance)
print('Succesful for reconstruction of GPU implementation: {}, with tolerance of {}'.format(all_close,tolerance))
print('Sum of difference: {}'.format((rec-img).sum()))