def get_mask(mask_file_name, real_img, opt):
mask = functions.read_image_dir(mask_file_name, opt)
if mask.shape[3] != real_img.shape[3]:
mask = imresize_to_shape(mask, [real_img.shape[2], real_img.shape[3]],
opt)
mask = functions.dilate_mask(mask, opt)
return mask
def generate_samples(netG,
reals_shapes,
noise_amp,
scale_w=1.0,
scale_h=1.0,
reconstruct=False,
n=50):
if reconstruct:
reconstruction = netG(fixed_noise, reals_shapes, noise_amp)
if opt.train_mode == "generation" or opt.train_mode == "retarget":
functions.save_image('{}/reconstruction.jpg'.format(dir2save),
reconstruction.detach())
functions.save_image('{}/real_image.jpg'.format(dir2save),
reals[-1].detach())
elif opt.train_mode == "harmonization" or opt.train_mode == "editing":
functions.save_image('{}/{}_wo_mask.jpg'.format(dir2save, _name),
reconstruction.detach())
functions.save_image(
'{}/real_image.jpg'.format(dir2save),
imresize_to_shape(real, reals_shapes[-1][2:], opt).detach())
return reconstruction
if scale_w == 1. and scale_h == 1.:
dir2save_parent = os.path.join(dir2save, "random_samples")
else:
reals_shapes = [[
r_shape[0], r_shape[1],
int(r_shape[2] * scale_h),
int(r_shape[3] * scale_w)
] for r_shape in reals_shapes]
dir2save_parent = os.path.join(
dir2save,
"random_samples_scale_h_{}_scale_w_{}".format(scale_h, scale_w))
make_dir(dir2save_parent)
for idx in range(n):
noise = functions.sample_random_noise(opt.train_stages - 1,
reals_shapes, opt)
sample = netG(noise, reals_shapes, noise_amp)
functions.save_image(
'{}/gen_sample_{}.jpg'.format(dir2save_parent, idx),
sample.detach())
noise_amp,
scale_w=1,
scale_h=2,
n=opt.num_samples)
generate_samples(netG,
reals_shapes,
noise_amp,
scale_w=2,
scale_h=2,
n=opt.num_samples)
elif opt.train_mode == "harmonization" or opt.train_mode == "editing":
opt.noise_scaling = 0.1
_name = "harmonized" if opt.train_mode == "harmonization" else "edited"
real = functions.read_image_dir(opt.naive_img, opt)
real = imresize_to_shape(real, reals_shapes[0][2:], opt)
fixed_noise[0] = real
if opt.train_mode == "editing":
fixed_noise[0] = fixed_noise[0] + opt.noise_scaling * \
functions.generate_noise([opt.nc_im, fixed_noise[0].shape[2],
fixed_noise[0].shape[3]],
device=opt.device)
out = generate_samples(netG, reals_shapes, noise_amp, reconstruct=True)
mask_file_name = '{}_mask{}'.format(opt.naive_img[:-4],
opt.naive_img[-4:])
if os.path.exists(mask_file_name):
mask = functions.read_image_dir(mask_file_name, opt)
if mask.shape[3] != out.shape[3]:
mask = imresize_to_shape(mask, [out.shape[2], out.shape[3]],
def generate_samples(netG,
img_to_augment,
naive_img,
naive_img_large,
aug,
opt,
depth,
noise_amp,
writer,
reals,
iter,
n=16):
opt.out_ = functions.generate_dir2save(opt)
dir2save = '{}/harmonized_samples_stage_{}'.format(opt.out_, depth)
reals_shapes = [r.shape for r in reals]
_name = "harmonized" if opt.train_mode == "harmonization" else "edited"
images = []
try:
os.makedirs(dir2save)
except OSError:
pass
if naive_img is not None:
n = n - 1
if opt.fine_tune:
n = 1
with torch.no_grad():
for idx in range(n):
noise = []
for d in range(depth + 1):
if d == 0:
if opt.fine_tune:
if opt.train_mode == "harmonization":
augmented_image = functions.np2torch(
naive_img, opt)
noise.append(augmented_image)
elif opt.train_mode == "editing":
augmented_image = functions.np2torch(
naive_img, opt)
noise.append(augmented_image + opt.noise_scaling *
functions.generate_noise(
[
opt.nc_im, reals_shapes[d][2],
reals_shapes[d][3]
],
device=opt.device).detach())
else:
if opt.train_mode == "harmonization":
data = {"image": img_to_augment}
augmented = aug.transform(**data)
augmented_image = functions.np2torch(
augmented["image"], opt)
noise.append(augmented_image)
elif opt.train_mode == "editing":
image = functions.shuffle_grid(img_to_augment)
augmented_image = functions.np2torch(image, opt)
noise.append(augmented_image + opt.noise_scaling *
functions.generate_noise(
[
opt.nc_im, reals_shapes[d][2],
reals_shapes[d][3]
],
device=opt.device).detach())
else:
noise.append(
functions.generate_noise(
[opt.nfc, reals_shapes[d][2], reals_shapes[d][3]],
device=opt.device).detach())
sample = netG(noise, reals_shapes, noise_amp)
functions.save_image(
'{}/{}_naive_sample.jpg'.format(dir2save, idx),
augmented_image)
functions.save_image(
'{}/{}_{}_sample.jpg'.format(dir2save, idx, _name),
sample.detach())
augmented_image = imresize_to_shape(augmented_image,
sample.shape[2:], opt)
images.append(augmented_image)
images.append(sample.detach())
if opt.fine_tune:
mask_file_name = '{}_mask{}'.format(opt.naive_img[:-4],
opt.naive_img[-4:])
augmented_image = imresize_to_shape(naive_img_large,
sample.shape[2:], opt)
if os.path.exists(mask_file_name):
mask = get_mask(mask_file_name, augmented_image, opt)
sample_w_mask = (
1 - mask) * augmented_image + mask * sample.detach()
functions.save_image(
'{}/{}_sample_w_mask_{}.jpg'.format(dir2save, _name, iter),
sample_w_mask.detach())
images = torch.cat(
[augmented_image,
sample.detach(), sample_w_mask], 0)
grid = make_grid(images, nrow=3, normalize=True)
writer.add_image('{}_images_{}'.format(_name, depth), grid,
iter)
else:
print(
"Warning: no mask with name {} exists for image {}".format(
mask_file_name, opt.input_name))
print("Only showing results without mask.")
images = torch.cat([augmented_image, sample.detach()], 0)
grid = make_grid(images, nrow=2, normalize=True)
writer.add_image('{}_images_{}'.format(_name, depth), grid,
iter)
functions.save_image(
'{}/{}_sample_{}.jpg'.format(dir2save, _name, iter),
sample.detach())
else:
if naive_img is not None:
noise = []
for d in range(depth + 1):
if d == 0:
if opt.train_mode == "harmonization":
noise.append(functions.np2torch(naive_img, opt))
elif opt.train_mode == "editing":
noise.append(functions.np2torch(naive_img, opt) + opt.noise_scaling * \
functions.generate_noise([opt.nc_im, reals_shapes[d][2],
reals_shapes[d][3]],
device=opt.device).detach())
else:
noise.append(
functions.generate_noise(
[
opt.nfc, reals_shapes[d][2],
reals_shapes[d][3]
],
device=opt.device).detach())
sample = netG(noise, reals_shapes, noise_amp)
_naive_img = imresize_to_shape(naive_img_large,
sample.shape[2:], opt)
images.insert(0, sample.detach())
images.insert(0, _naive_img)
functions.save_image(
'{}/{}_sample_{}.jpg'.format(dir2save, _name, iter),
sample.detach())
mask_file_name = '{}_mask{}'.format(opt.naive_img[:-4],
opt.naive_img[-4:])
if os.path.exists(mask_file_name):
mask = get_mask(mask_file_name, _naive_img, opt)
sample_w_mask = (
1 - mask) * _naive_img + mask * sample.detach()
functions.save_image(
'{}/{}_sample_w_mask_{}.jpg'.format(
dir2save, _name, iter), sample_w_mask)
images = torch.cat(images, 0)
grid = make_grid(images, nrow=4, normalize=True)
writer.add_image('{}_images_{}'.format(_name, depth), grid, iter)
def train(opt):
print("Training model with the following parameters:")
print("\t number of stages: {}".format(opt.train_stages))
print("\t number of concurrently trained stages: {}".format(
opt.train_depth))
print("\t learning rate scaling: {}".format(opt.lr_scale))
print("\t non-linearity: {}".format(opt.activation))
real = functions.read_image(opt)
real = functions.adjust_scales2image(real, opt)
reals = functions.create_reals_pyramid(real, opt)
print("Training on image pyramid: {}".format([r.shape for r in reals]))
print("")
if opt.naive_img != "":
naive_img = functions.read_image_dir(opt.naive_img, opt)
naive_img_large = imresize_to_shape(naive_img, reals[-1].shape[2:],
opt)
naive_img = imresize_to_shape(naive_img, reals[0].shape[2:], opt)
naive_img = functions.convert_image_np(naive_img) * 255.0
else:
naive_img = None
naive_img_large = None
if opt.fine_tune:
img_to_augment = naive_img
else:
img_to_augment = functions.convert_image_np(reals[0]) * 255.0
if opt.train_mode == "editing":
opt.noise_scaling = 0.1
generator = init_G(opt)
if opt.fine_tune:
for _ in range(opt.train_stages - 1):
generator.init_next_stage()
generator.load_state_dict(
torch.load(
'{}/{}/netG.pth'.format(opt.model_dir, opt.train_stages - 1),
map_location="cuda:{}".format(torch.cuda.current_device())))
fixed_noise = []
noise_amp = []
for scale_num in range(opt.start_scale, opt.train_stages):
opt.out_ = functions.generate_dir2save(opt)
opt.outf = '%s/%d' % (opt.out_, scale_num)
try:
os.makedirs(opt.outf)
except OSError:
print(OSError)
pass
functions.save_image('{}/real_scale.jpg'.format(opt.outf),
reals[scale_num])
d_curr = init_D(opt)
if opt.fine_tune:
d_curr.load_state_dict(
torch.load('{}/{}/netD.pth'.format(opt.model_dir,
opt.train_stages - 1),
map_location="cuda:{}".format(
torch.cuda.current_device())))
elif scale_num > 0:
d_curr.load_state_dict(
torch.load('%s/%d/netD.pth' % (opt.out_, scale_num - 1)))
generator.init_next_stage()
writer = SummaryWriter(log_dir=opt.outf)
fixed_noise, noise_amp, generator, d_curr = train_single_scale(
d_curr, generator, reals, img_to_augment, naive_img,
naive_img_large, fixed_noise, noise_amp, opt, scale_num, writer)
torch.save(fixed_noise, '%s/fixed_noise.pth' % (opt.out_))
torch.save(generator, '%s/G.pth' % (opt.out_))
torch.save(reals, '%s/reals.pth' % (opt.out_))
torch.save(noise_amp, '%s/noise_amp.pth' % (opt.out_))
del d_curr
writer.close()
return
