# Writes a multi page (multi-frame) tiff file

# Inputs:

# Images can be either a list of ndarrays or single ndarray

#

# Output:

# tiff file will be either created ar images added to existing tiff file

#

# !!!! output_filename --> file name must be "tiff"

#

assert output_filename is not None, "Must Provide image file Name"

filename, file_extension = os.path.splitext(output_filename)

assert file_extension.lower() == '.tiff', 'error - file extention must be .tiff'

if type(Images) is np.ndarray:

Images = [Images]

if not type(resize_factors) is list:

resize_factors = [resize_factors]

if len(resize_factors) < len(Images):

resize_factors *= len(Images)

imlist = []

for resize_factor, Frame in zip(resize_factors, Images):

assert type(Frame) is np.ndarray, "imwrite_multi_tiff expects only ndarray within list items "

if Norm_frames:

Frame = ((Frame - Frame.min()) / (Frame.max() - Frame.min()) * 255).astype(np.uint8)

# Frame = (Frame * 65535).astype(np.uint16)

#Frame = (Frame * 255).astype(np.uint8)

if len(Frame.shape) == 3 and Frame.shape[2] == 3:

imlist.append(Image.fromarray(Frame, mode="RGB"))

else:

imlist.append(Image.fromarray(Frame))

if not resize_factor == 1:

imlist[-1] = imlist[-1].resize([int(resize_factor * s) for s in imlist[-1].size])

imlist[0].save(output_filename, save_all=True, append_images=imlist[1:])

pad_x = signal.shape[2] - filter.shape[2]

pad_y = signal.shape[3] - filter.shape[3]

expanded_kernel = F.pad(filter, [0, pad_y, 0, pad_x])

expanded_kernel_np = expanded_kernel.cpu().numpy()

expanded_kernel_shift = np.roll(expanded_kernel_np, -int(filter.shape[2] / 2), axis=2)

expanded_kernel_shift = np.roll(expanded_kernel_shift, -int(filter.shape[2] / 2), axis=3)

expanded_kernel_shift = torch.from_numpy(expanded_kernel_shift).float()

### fft of h*x

expanded_kernel_shift = pad_shift_filter(f, k)

fft_hx = torch.empty([f.shape[0], f.shape[1], f.shape[2], f.shape[3], 2])

for i in range(3):

fft_x = torch.rfft(f[:, i:i + 1, :, :], 2, onesided=False, normalized=True).to(self.device)

fft_kernel = torch.rfft(expanded_kernel_shift, 2, onesided=False, normalized=True).to(self.device)

real = fft_x[:, :, :, :, 0] * fft_kernel[:, :, :, :, 0] - \

fft_x[:, :, :, :, 1] * fft_kernel[:, :, :, :, 1]

im = fft_x[:, :, :, :, 0] * fft_kernel[:, :, :, :, 1] + \

fft_x[:, :, :, :, 1] * fft_kernel[:, :, :, :, 0]

fft_conv = torch.stack([real, im], -1) # (a+bj)*(c+dj) = (ac-bd)+(ad+bc)j

fft_hx[:, i, :, :, :] = fft_conv

L = (input - target)

L_fft = torch.rfft(L, 2, onesided=False, normalized=True).to(self.device)

L_fft = L_fft ** 2

# first part: LS loss

dip_loss_dif = (hx - y)

dip_loss = torch.rfft(dip_loss_dif, 2, onesided=False, normalized=True).to(self.device)

# second part: BP loss

eps = 1e-3

mul_factor = 1e5

eps_ignored = 0.01

sigma = 0

h = torch.from_numpy(get_bicubic(scale))

h = torch.unsqueeze(h, 0).unsqueeze(0)

conv_shape = (h.shape[2] + h.shape[2] - 1, h.shape[3] + h.shape[3] - 1)

H = fft2(h, conv_shape[1], conv_shape[0])

H_flip = fft2(flip(h), conv_shape[1], conv_shape[0])

H_mul_H_flip = mul_complex(H, H_flip)

H_mul_H_flip_ifft = torch.irfft(H_mul_H_flip, signal_ndim=2, normalized=True, onesided=False)

h_downsampled = H_mul_H_flip_ifft[:,:,1::scale, 1::scale]

h_downsampled = pad_shift_filter(y, h_downsampled)

H_downsampled = torch.rfft(h_downsampled, signal_ndim=2, normalized=True, onesided=False)

bp_loss = torch.sqrt(abs2(H_downsampled)[:,:,:,:,0:1])

bp_loss = mul_factor * bp_loss + eps_ignored * (sigma ** 2) + eps

bp_loss = 1 / (torch.sqrt(bp_loss))

bp_loss = torch.repeat_interleave(bp_loss, 2, -1).to(self.device)

loss_mat = bp_loss.to(self.device) * dip_loss

img_name_for_plot = img_name[0:-5]

print(img_name_for_plot)

learning_rate = 1e-3

OPTIMIZER = 'adam'

if factor == 3:

num_iter = 2000

reg_noise_std = 0.03

if factor == 5:

num_iter = 600

reg_noise_std = 0.03

tv_weight = 0

PLOT = True

PLOT_PSNR = True

path_to_image = imgs_dir + img_name

low_res, high_res = load_images_publication.tiff_imgs(path_to_image, factor, use_5_per)

# sanity check

#low_res = create_lr_image.create(high_res)

if use_5_per ==False:

high_res_pil = Image.fromarray(np.uint8(high_res*255), 'L')

low_res_pil = Image.fromarray(np.uint8(low_res*255), 'L')

img_bicubic, img_sharp, img_nearest = get_baselines(low_res_pil, high_res_pil)

high_res = np.expand_dims(high_res, axis=0)

low_res_for_up = torch.from_numpy(low_res)

low_res_for_up = torch.unsqueeze(low_res_for_up, 0).unsqueeze(0)

img_upsampled_torch, img_bicubic1 = upsample_using_h(low_res_for_up, factor, self.device)

img_upsampled = img_upsampled_torch.cpu().numpy()

img_upsampled = np.squeeze(img_upsampled, axis=0)

img_upsampled = (img_upsampled - img_upsampled.min()) / (img_upsampled.max() - img_upsampled.min())

else:

img_bicubic = np.zeros((5, high_res.shape[1], high_res.shape[2]))

img_nearest= np.zeros((5, high_res.shape[1], high_res.shape[2]))

for j in range(5):

img_bicubic[j,:,:] = cv2.resize(low_res[j,:,:], dsize=high_res.shape[1:], interpolation=cv2.INTER_CUBIC)

img_nearest[j,:,:] = cv2.resize(low_res[j,:,:], dsize=high_res.shape[1:], interpolation=cv2.INTER_NEAREST)

high_res_ = high_res

psnr_bicubic = compare_psnr(high_res, img_bicubic, data_range=1)

psnr_nn = compare_psnr(high_res, img_nearest, data_range=1)

psnr_custom = compare_psnr(high_res, img_upsampled, data_range=1)

if use_5_per == False:

perc_sim_bicubic = compute_dists.compute(high_res[0,:,:], img_bicubic[0,:,:])

perc_sim_bicubic = perc_sim_bicubic.cpu().detach().numpy()[0][0][0][0]

perc_sim_nn = compute_dists.compute(high_res[0,:,:], img_nearest[0,:,:])

perc_sim_nn = perc_sim_nn.cpu().detach().numpy()[0][0][0][0]

perc_sim_custom = compute_dists.compute(high_res[0, :, :], img_upsampled[0, :, :])

perc_sim_custom = perc_sim_custom.cpu().detach().numpy()[0][0][0][0]

else:

perc_sim_bicubic = np.zeros(5)

perc_sim_nn = np.zeros(5)

for j in range(5):

perc_sim_bicubic[j] = compute_dists.compute(high_res[j, :, :], img_bicubic[j, :, :])

#perc_sim_bicubic[j] = perc_sim_bicubic.cpu().detach().numpy()[0][0][0][0]

perc_sim_nn[j] = compute_dists.compute(high_res[j, :, :], img_nearest[j, :, :])

#perc_sim_nn[j] = perc_sim_nn.cpu().detach().numpy()[0][0][0][0]

#psnr_basic = compare_psnr(high_res_, low_for_psnr)

if PLOT:

if use_5_per == False:

plot_image_grid([high_res, img_bicubic, img_nearest], directory, 'basic_compare_img_'+img_name_for_plot, 3, 12)

else:

for j in range(5):

plot_image_grid([np.expand_dims(high_res[j,:,:], axis=0), np.expand_dims(img_bicubic[j,:,:], axis=0),

np.expand_dims(img_nearest[j,:,:], axis=0)], directory,

'basic_compare_img_' + img_name_for_plot + 'frame_' + str(j), 3, 12)

print ('PSNR bicubic: %.4f PSNR nearest: %.4f' % (psnr_bicubic, psnr_nn))

print ('per_sim bicubic: %.4f per_sim nearest: %.4f' % (perc_sim_bicubic.mean(), perc_sim_nn.mean()))

input_depth = 32

INPUT = 'noise'

pad = 'reflection'

OPT_OVER = 'net'

#KERNEL_TYPE = 'lanczos2'

net_input = get_noise(input_depth, INPUT, (high_res.shape[-1], high_res.shape[-2])).type(dtype).detach()

NET_TYPE = 'skip' # UNet, ResNet

net = get_net(input_depth, 'skip', pad,

n_channels=high_res.shape[0],

skip_n33d=128,

skip_n33u=128,

skip_n11=4,

num_scales=5,

upsample_mode='bilinear').type(dtype)

# Losses

img_LR_var = np_to_torch(low_res).type(dtype)

def closure():

global i, net_input, psnr_history, psnr_history_short_HR, psnr_history_short_LR, orig_img_HR, per_sim_history

if reg_noise_std > 0:

net_input = net_input_saved + (noise.normal_() * reg_noise_std)

out_HR = net(net_input)

out_LR = custom_downsample(out_HR, factor, use_5_per)

if loss_type == 'dip':

total_loss = mse_fft(out_LR, img_LR_var)

#total_loss = mse(out_LR, img_LR_var)

elif loss_type == 'bp':

total_loss = mse_bp(out_LR, img_LR_var.unsqueeze(0), factor)

#print(total_loss)

#print(tv_loss_try_abs(out_HR))

if tv_weight > 0:

mul_factor = 1

total_loss = total_loss + tv_weight * tv_loss(out_HR, mul_factor).to(self.device)

total_loss.backward()

#high_res_ = np.squeeze(high_res, axis=0)

#orig_img_HR = high_res_

orig_img_HR = np.squeeze(high_res_, axis=0)

orig_img_LR = low_res

out_HR = out_HR.squeeze()

out_LR = out_LR.squeeze()

# History

if i % 100 == 0:

psnr_LR = compare_psnr(orig_img_LR, torch_to_np(out_LR), data_range=1)

#psnr_HR = compare_psnr(orig_img_HR[:,pix_ignore:-pix_ignore, pix_ignore:-pix_ignore], torch_to_np(out_HR.unsqueeze(0))[:,pix_ignore:-pix_ignore, pix_ignore:-pix_ignore], data_range=1)

psnr_HR = compare_psnr(orig_img_HR[pix_ignore:-pix_ignore, pix_ignore:-pix_ignore], torch_to_np(out_HR)[pix_ignore:-pix_ignore, pix_ignore:-pix_ignore], data_range=1)

if use_5_per == False:

perc_sim_HR = compute_dists.compute(orig_img_HR[pix_ignore:-pix_ignore, pix_ignore:-pix_ignore],

torch_to_np(out_HR)[pix_ignore:-pix_ignore, pix_ignore:-pix_ignore])

else:

perc_sim_HR = np.zeros(5)

for j in range(5):

perc_sim_HR[j] = compute_dists.compute(orig_img_HR[j, pix_ignore:-pix_ignore, pix_ignore:-pix_ignore],

torch_to_np(out_HR)[j, pix_ignore:-pix_ignore, pix_ignore:-pix_ignore])

psnr_history_short_HR.append([psnr_HR])

psnr_history_short_LR.append([psnr_LR])

per_sim_history.append([perc_sim_HR.mean()])

if PLOT and i % 100 == 0:

out_HR_np = torch_to_np(out_HR)

if use_5_per == False:

# plot_image_grid([np.expand_dims(high_res_, axis=0), img_nearest, img_bicubic, np.expand_dims(np.clip(out_HR_np, 0, 1), axis=0)],

# directory, 'compare_img_'+img_name_for_plot, factor=13, nrow=4)

plot_image_grid([high_res_, img_nearest, img_bicubic, np.expand_dims(np.clip(out_HR_np, 0, 1), axis=0)],

directory, 'compare_img_'+img_name_for_plot, factor=13, nrow=4)

plt.imsave(directory + 'orig_'+img_name_for_plot + '.png', np.squeeze(high_res_), cmap='gray')

plt.imsave(directory + 'nn_' + img_name_for_plot + '.png', np.squeeze(img_nearest), cmap='gray')

plt.imsave(directory + 'bicubic_' + img_name_for_plot + '.png', np.squeeze(img_bicubic), cmap='gray')

plt.imsave(directory + 'dip_' + img_name_for_plot + '.png', np.clip(out_HR_np, 0, 1), cmap='gray')

else:

for j in range(5):

plot_image_grid([np.expand_dims(high_res_[j,:,:], axis=0), np.expand_dims(img_bicubic[j,:,:], axis=0),

np.expand_dims(np.clip(out_HR_np[j,:,:], 0, 1),

axis=0)],directory, 'compare_img_' + img_name_for_plot + 'frame' + str(j), factor=13, nrow=3)

if PLOT_PSNR and i % 100 == 0:

print('Iteration %04d PSNR %.3f perc_sim %.3f' % (i, psnr_HR, perc_sim_HR.mean()), '\r')

if i == num_iter - 1:

out_HR_np = torch_to_np(out_HR)

if use_5_per == False:

# imwrite_multi_tiff([high_res_, np.squeeze(img_bicubic), np.clip(out_HR_np, 0, 1)],

# directory + 'final_comparison_' + img_name_for_plot + '.tiff')

imwrite_multi_tiff([np.squeeze(high_res_), np.clip(out_HR_np, 0, 1), np.squeeze(img_bicubic)],

directory + 'final_comparison_' + img_name_for_plot + '.tiff')

else:

for j in range(5):

imwrite_multi_tiff([high_res_[j,:,:], np.clip(out_HR_np[j,:,:], 0, 1), np.squeeze(img_bicubic[j,:,:])],

directory + 'final_comparison_' + img_name_for_plot + 'frame' + str(j) +'.tiff')

#imwrite_multi_tiff([out_LR.detach().cpu().numpy(), img_LR_var.detach().cpu().numpy().squeeze()], directory + '1.tiff')

#imwrite_multi_tiff([out_LR.detach().cpu().numpy(), img_LR_var.detach().cpu().numpy().squeeze()], directory + '1.tiff')

# ### heat map ###

# img_diff = (high_res_ - np.clip(out_HR_np, 0, 1)) ** 2

# heat_map = ndimage.filters.gaussian_filter(img_diff, sigma=16)

# max_val = np.max(heat_map)

# min_val = np.min(heat_map)

# norm_heat_map = (heat_map - min_val) / (max_val - min_val)

# plt.imshow(high_res_)

# plt.imshow(255 * norm_heat_map, alpha=0.5, cmap='viridis')

# plt.axis('on')

# plt.savefig(directory + 'heat_map_' + img_name_for_plot + '.jpg', bbox_inches='tight', pad_inches=0)

i += 1

return total_loss

global psnr_history, psnr_history_short_HR, psnr_history_short_LR, orig_img_HR, per_sim_history

psnr_history = []

psnr_history_short_HR = []

psnr_history_short_LR = []

per_sim_history = []

net_input_saved = net_input.detach().clone()

noise = net_input.detach().clone()

global i

i = 0

p = get_params(OPT_OVER, net, net_input)

optimize(OPTIMIZER, p, closure, learning_rate, num_iter)

out_HR_np = np.clip(torch_to_np(net(net_input)), 0, 1)

#result_deep_prior = put_in_center(out_HR_np, imgs['orig_np'].shape[1:])

return out_HR_np, orig_img_HR, net, net_input, psnr_history_short_HR, psnr_history_short_LR, \

per_sim_history, psnr_bicubic, psnr_nn, psnr_custom, perc_sim_bicubic, perc_sim_nn, perc_sim_custom