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