def preprocess_content_image(opt, reals,scale): real = functions.read_image(opt) functions.adjust_scales2image(real, opt) ref = functions.read_image_dir('%s/%s' % (opt.ref_dir, opt.ref_name), opt) if ref.shape[3] != real.shape[3]: ref = imresize_to_shape(ref, [real.shape[2], real.shape[3]], opt) ref = ref[:, :, :real.shape[2], :real.shape[3]] N = len(reals) - 1 n = scale in_s = imresize(ref, pow(opt.scale_factor, (N - n + 1)), opt) in_s = in_s[:, :, :reals[n - 1].shape[2], :reals[n - 1].shape[3]] in_s = imresize(in_s, 1 / opt.scale_factor, opt) in_s = in_s[:, :, :reals[n].shape[2], :reals[n].shape[3]] return in_s
im_mask = cv2.imread( '%s/%s_mask%s' % (opt.ref_dir, opt.ref_name[:-4], opt.ref_name[-4:])) img = cv2.imread('%s/%s' % (opt.input_dir, opt.ref_name)) im_mask = im_mask / 255 for k in range(3): img[:, :, k][im_mask[:, :, k] == 1] = img[:, :, k][im_mask[:, :, k] == 0].mean() cv2.imwrite( '%s/%s_global_mean%s' % (opt.input_dir, opt.ref_name[:-4], opt.ref_name[-4:]), img) ref = functions.read_image_dir( '%s/%s_global_mean%s' % (opt.input_dir, opt.ref_name[:-4], opt.ref_name[-4:]), opt) mask = functions.read_image_dir( '%s/%s_mask%s' % (opt.ref_dir, opt.ref_name[:-4], opt.ref_name[-4:]), opt) if ref.shape[3] != real.shape[3]: mask = imresize_to_shape(mask, [real.shape[2], real.shape[3]], opt) mask = mask[:, :, :real.shape[2], :real.shape[3]] ref = imresize_to_shape(ref, [real.shape[2], real.shape[3]], opt) ref = ref[:, :, :real.shape[2], :real.shape[3]] mask = functions.dilate_mask(mask, opt) N = len(reals) - 1 n = opt.inpainting_start_scale
print('task does not exist') #elif (os.path.exists(dir2save)): # print("output already exist") else: try: os.makedirs(dir2save) except OSError: pass real = functions.read_image(opt) real = functions.adjust_scales2image(real, opt) Gs, Zs, reals, NoiseAmp = functions.load_trained_pyramid(opt) if (opt.editing_start_scale < 1) | (opt.editing_start_scale > (len(Gs) - 1)): print("injection scale should be between 1 and %d" % (len(Gs) - 1)) else: ref = functions.read_image_dir( '%s/%s' % (opt.ref_dir, opt.ref_name), opt) mask = functions.read_image_dir( '%s/%s_mask%s' % (opt.ref_dir, opt.ref_name[:-4], opt.ref_name[-4:]), opt) if ref.shape[3] != real.shape[3]: mask = imresize(mask, real.shape[3] / ref.shape[3], opt) mask = mask[:, :, :real.shape[2], :real.shape[3]] ref = imresize(ref, real.shape[3] / ref.shape[3], opt) ref = ref[:, :, :real.shape[2], :real.shape[3]] mask = functions.dilate_mask(mask, opt) N = len(reals) - 1 n = opt.editing_start_scale in_s = imresize(ref, pow(opt.scale_factor, (N - n + 1)), opt) in_s = in_s[:, :, :reals[n - 1].shape[2], :reals[n - 1].shape[3]] in_s = imresize(in_s, 1 / opt.scale_factor, opt)
if dir2save is None: print('task does not exist') #elif (os.path.exists(dir2save)): # print("output already exist") else: try: os.makedirs(dir2save) except OSError: pass real = functions.read_image(opt) real = functions.adjust_scales2image(real, opt) Gs, Zs, reals, NoiseAmp = functions.load_trained_pyramid(opt) if (opt.gen_start_scale < 1) | (opt.gen_start_scale > (len(Gs) - 1)): print("injection scale should be between 1 and %d" % (len(Gs) - 1)) else: ref = functions.read_image_dir( '%s/%s' % (opt.ref_dir, opt.ref_name), opt) #mask = functions.read_image_dir('%s/%s_mask%s' % (opt.ref_dir,opt.ref_name[:-4],opt.ref_name[-4:]), opt) if ref.shape[3] != real.shape[3]: #mask = imresize_to_shape(mask, [real.shape[2], real.shape[3]], opt) #mask = mask[:, :, :real.shape[2], :real.shape[3]] ref = imresize_to_shape(ref, [real.shape[2], real.shape[3]], opt) ref = ref[:, :, :real.shape[2], :real.shape[3]] #mask = functions.dilate_mask(mask, opt) N = len(reals) - 1 n = opt.gen_start_scale in_s = imresize(ref, pow(opt.scale_factor, (N - n + 1)), opt) in_s = in_s[:, :, :reals[n - 1].shape[2], :reals[n - 1].shape[3]] in_s = imresize(in_s, 1 / opt.scale_factor, opt) in_s = in_s[:, :, :reals[n].shape[2], :reals[n].shape[3]]
print('task does not exist') # elif (os.path.exists(dir2save)): # print("output already exist") else: try: os.makedirs(dir2save) except OSError: pass real = functions.read_image(opt) real = functions.adjust_scales2image(real, opt) Gs, Zs, reals, NoiseAmp = functions.load_trained_pyramid(opt) if (opt.editing_start_scale < 1) | (opt.editing_start_scale > (len(Gs) - 1)): print("injection scale should be between 1 and %d" % (len(Gs) - 1)) else: ref = functions.read_image_dir( '%s/%s' % (opt.ref_dir, opt.ref_name), opt) mask = functions.read_image_dir( 'Input/Editing/3-VAM9owl-o/start_scale=2,input=0.5/mask.png', opt) # mask = functions.read_image_dir('%s/%s_mask%s' % (opt.ref_dir,opt.ref_name[:-4],opt.ref_name[-4:]), opt) if ref.shape[3] != real.shape[3]: ''' mask = imresize(mask, real.shape[3]/ref.shape[3], opt) mask = mask[:, :, :real.shape[2], :real.shape[3]] ref = imresize(ref, real.shape[3] / ref.shape[3], opt) ref = ref[:, :, :real.shape[2], :real.shape[3]] ''' mask = imresize_to_shape(mask, [real.shape[2], real.shape[3]], opt) mask = mask[:, :, :real.shape[2], :real.shape[3]] ref = imresize_to_shape(ref, [real.shape[2], real.shape[3]],
def train(opt, Gs, Zs, reals, NoiseAmp): real_ = functions.read_image(opt) in_s = 0 scale_num = 0 real = imresize(real_, opt.scale1, opt) reals = functions.creat_reals_pyramid(real, reals, opt) nfc_prev = 0 #creating a pyramid of masks the same way we did for the img and thus to train on only the correct pixels #at all scales if opt.inpainting: m = functions.read_image_dir( '%s/%s_mask%s' % (opt.ref_dir, opt.input_name[:-4], opt.input_name[-4:]), opt) m = imresize(m, opt.scale1, opt) m_s = [] #pyramid of masks opt.m_s = functions.creat_reals_pyramid(m, m_s, opt) while scale_num < opt.stop_scale + 1: opt.nfc = min(opt.nfc_init * pow(2, math.floor(scale_num / 4)), 128) opt.min_nfc = min(opt.min_nfc_init * pow(2, math.floor(scale_num / 4)), 128) opt.out_ = functions.generate_dir2save(opt) opt.outf = '%s/%d' % (opt.out_, scale_num) try: os.makedirs(opt.outf) except OSError: pass #plt.imsave('%s/in.png' % (opt.out_), functions.convert_image_np(real), vmin=0, vmax=1) #plt.imsave('%s/original.png' % (opt.out_), functions.convert_image_np(real_), vmin=0, vmax=1) plt.imsave('%s/real_scale.png' % (opt.outf), functions.convert_image_np(reals[scale_num]), vmin=0, vmax=1) D_curr, G_curr = init_models(opt) if (nfc_prev == opt.nfc): G_curr.load_state_dict( torch.load('%s/%d/netG.pth' % (opt.out_, scale_num - 1))) D_curr.load_state_dict( torch.load('%s/%d/netD.pth' % (opt.out_, scale_num - 1))) z_curr, in_s, G_curr = train_single_scale(D_curr, G_curr, reals, Gs, Zs, in_s, NoiseAmp, opt) G_curr = functions.reset_grads(G_curr, False) G_curr.eval() D_curr = functions.reset_grads(D_curr, False) D_curr.eval() Gs.append(G_curr) Zs.append(z_curr) NoiseAmp.append(opt.noise_amp) torch.save(Zs, '%s/Zs.pth' % (opt.out_)) torch.save(Gs, '%s/Gs.pth' % (opt.out_)) torch.save(reals, '%s/reals.pth' % (opt.out_)) torch.save(NoiseAmp, '%s/NoiseAmp.pth' % (opt.out_)) scale_num += 1 nfc_prev = opt.nfc del D_curr, G_curr return
#writing coloured img if opt.on_drive != None: cv2.imwrite( "%s/%s_coloured_%s.jpg" % (dir2save, opt.mask_name[:-4], opt.fill_method), coloured_image) else: cv2.imwrite( "%s/%s_coloured_%s.jpg" % (dir2save, opt.mask_name[:-4], opt.fill_method), coloured_image) #Reading in coloured image if opt.on_drive != None: ref = functions.read_image_dir( '%s/%s_coloured_%s.jpg' % (dir2save, opt.mask_name[:-4], opt.fill_method), opt) mask = functions.read_image_dir( '%s/%s/%s' % (opt.on_drive, opt.input_dir, opt.mask_name), opt) else: ref = functions.read_image_dir( '%s/%s_coloured_%s.jpg' % (dir2save, opt.mask_name[:-4], opt.fill_method), opt) mask = functions.read_image_dir( '%s/%s' % (opt.input_dir, opt.mask_name), opt) #ref=cv2.imread('%s/%s' % (opt.input_dir, opt.input_name)) if ref.shape[3] != real.shape[3]: mask = imresize_to_shape(mask, [real.shape[2], real.shape[3]],
#plt.figure() #plt.imshow(ref_img) #plt.show() #noisy_img = clean_img + np.random.normal(0, float(opt.noise), clean_img.shape) #noisy_img = functions.np2torch(noisy_img, opt) #noisy_img = noisy_img[:,0:3,:,:] #clean_img = clean_img[:,:,0:3] / 255.0 N = len(reals) - 1 n = opt.denoise_scale noisy_bm3d = [] noisy_img = img.imread( '%s/%s/%s/noisy.png' % (opt.noise_dir, opt.input_name[:-4], str(opt.noise)))[:, :, 0:3] noisy_img_torch = functions.read_image_dir( '%s/%s/%s/noisy.png' % (opt.noise_dir, opt.input_name[:-4], str(opt.noise)), opt) noisy_bm3d.append(noisy_img_torch) bm3d_img = img.imread( '%s/%s/%s/denoised.png' % (opt.noise_dir, opt.input_name[:-4], str(opt.noise)))[:, :, 0:3] bm3d_img_torch = functions.read_image_dir( '%s/%s/%s/denoised.png' % (opt.noise_dir, opt.input_name[:-4], str(opt.noise)), opt) noisy_bm3d.append(bm3d_img_torch) print('BM3D result:', compare_psnr(bm3d_img, clean_img)) for index, ref in enumerate(noisy_bm3d): temp = functions.convert_image_np(ref) * 255 temp = temp.astype(np.uint8) plt.figure()