# 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()))