def main(args):
os.makedirs(args.outputs, exist_ok=True)
generator = get_derivable_generator(args.gan_model, args.inversion_type,
args)
loss = get_loss(args.loss_type, args)
generator.cuda()
loss.cuda()
inversion = get_inversion(args.optimization, args)
image_list = image_files(args.target_images)
frameSize = MODEL_POOL[args.gan_model]['resolution']
for i, images in enumerate(split_to_batches(image_list, 1)):
print('%d: Inverting %d images :' % (i + 1, 1), end='')
pt_image_str = '%s\n'
print(pt_image_str % tuple(images))
image_name_list = []
image_tensor_list = []
for image in images:
image_name_list.append(os.path.split(image)[1])
image_tensor_list.append(_add_batch_one(load_as_tensor(image)))
y_gt = _sigmoid_to_tanh(torch.cat(image_tensor_list, dim=0)).cuda()
# Invert
latent_estimates, history = inversion.invert(generator,
y_gt,
loss,
batch_size=1,
video=args.video)
# Get Images
y_estimate_list = torch.split(torch.clamp(_tanh_to_sigmoid(
generator(latent_estimates)),
min=0.,
max=1.).cpu(),
1,
dim=0)
# Save
for img_id, image in enumerate(images):
y_estimate_pil = Tensor2PIL(y_estimate_list[img_id])
y_estimate_pil.save(
os.path.join(args.outputs, image_name_list[img_id]))
# Create video
if args.video:
print('Create GAN-Inversion video.')
video = cv2.VideoWriter(filename=os.path.join(
args.outputs,
'%s_inversion.avi' % image_name_list[img_id]),
fourcc=cv2.VideoWriter_fourcc(*'MJPG'),
fps=args.fps,
frameSize=(frameSize, frameSize))
print('Save frames.')
for i, sample in enumerate(history):
image = generator(sample)
image_cv2 = convert_array_to_images(
image.detach().cpu().numpy())[0][:, :, ::-1]
video.write(image_cv2)
video.release()
def main(args):
os.makedirs(args.outputs, exist_ok=True)
generator = get_derivable_generator(args.gan_model, args.inversion_type,
args) # 生成器
loss = get_loss(args.loss_type, args)
sr_loss = SR_loss(loss, args.down, args.factor) # SR计算loss的方式
# to cuda
generator.cuda()
loss.cuda()
inversion = get_inversion(args.optimization, args)
image_list = image_files(args.target_images)
frameSize = MODEL_POOL[args.gan_model]['resolution']
for i, images in enumerate(split_to_batches(image_list, 1)):
print('%d: Super-resolving %d images ' % (i + 1, 1), end='')
pt_image_str = '%s\n'
print(pt_image_str % tuple(images))
image_name_list = []
image_tensor_list = []
for image in images:
image_name_list.append(os.path.split(image)[1])
# image = _add_batch_one(load_as_tensor(image))
image = convert2target(_add_batch_one(load_as_tensor(image)),
'nearest') # 更改size使得适用于更多类型的分辨率图像
image_tensor_list.append(image)
# print("add..: ", _add_batch_one(load_as_tensor(image)).size()) # torch.Size([1, 3, 64, 64])
y_gt = _sigmoid_to_tanh(torch.cat(image_tensor_list, dim=0)).cuda()
# Invert
latent_estimates, history = inversion.invert(generator,
y_gt,
sr_loss,
batch_size=BATCH_SIZE,
video=args.video)
# Get Images
# 将optimizer优化好的latent_estimates再放入generator中生成图像
# 并且将batch_size那一列的数据去除
y_estimate_list = torch.split(torch.clamp(_tanh_to_sigmoid(
generator(latent_estimates)),
min=0.,
max=1.).cpu(),
1,
dim=0)
# 保存结果
for img_id, image in enumerate(images):
# up_nn, up_bic, down = downsample_images(image_tensor_list[img_id], factor=args.factor, mode=args.down)
# y_nn_pil = Tensor2PIL(up_nn) # 低分辨率化后的图像
y_estimate_pil = Tensor2PIL(y_estimate_list[img_id])
y_estimate_pil.save(
os.path.join(
os.path.join(args.outputs,
'%s.png' % image_name_list[img_id][:-4])))
#y_nn_pil.save(os.path.join(os.path.join(args.outputs, '%s-nn.png' % image_name_list[img_id][:-4])))
# Create video
if args.video:
print('Create GAN-Inversion video.')
video = cv2.VideoWriter(filename=os.path.join(
args.outputs, '%s_sr.avi' % image_name_list[img_id][:-4]),
fourcc=cv2.VideoWriter_fourcc(*'MJPG'),
fps=args.fps,
frameSize=(frameSize, frameSize))
print('Save frames.')
for i, sample in enumerate(history):
image = generator(
sample
) # 用generator从history（保存的训练中的estimate_latent的值）中生成图像
image_cv2 = convert_array_to_images(
image.detach().cpu().numpy())[0][:, :, ::-1]
video.write(image_cv2)
video.release()
def run(args):
os.makedirs(args.outputs, exist_ok=True) # 生成输出路径文件夹，存在则跳过
# 生成器
generator = get_derivable_generator(args.gan_model, args.inversion_type,
args)
loss = get_loss(args.loss_type, args) # 损失函数
generator.cuda() # pytorch需要手动放入GPU进行运算
loss.cuda()
inversion = get_inversion(args.optimization, args)
image_list = image_files(args.target_images) # 获取输入图片路径
frameSize = MODEL_POOL[args.gan_model]['resolution'] # 获取图像分辨率
# 按照batch大小分批处理图像
for i, images in enumerate(split_to_batches(image_list, 1)):
print('%d: Inverting %d images :' % (i + 1, 1), end='')
# pt_image_str = '%s\n'
print('%s\n' % tuple(images))
image_name_list = []
image_tensor_list = []
for image in images:
image_name_list.append(os.path.split(image)[1])
image_tensor_list.append(_add_batch_one(load_as_tensor(image)))
# torch.cat(tensors, dim=0, out=None) → Tensor
# tensors (sequence of Tensors) – any python sequence of tensors of the same type. Non-empty tensors provided must have the same shape, except in the cat dimension.
# dim (int, optional) – the dimension over which the tensors are concatenated
# out (Tensor, optional) – the output tensor.
y_gt = _sigmoid_to_tanh(
torch.cat(image_tensor_list,
dim=0)).cuda() # 在维度0上连接所有的tensor并且将值域映射到[-1, 1]
# 逆映射, 生成图像tensor
latent_estimates, history = inversion.invert(generator,
y_gt,
loss,
batch_size=BATCH_SIZE,
video=args.video)
# 将值域从[-1,1]映射到[0,1], 使用torch.clamp()进一步保证值域在[0,1]
y_estimate_list = torch.split(torch.clamp(_tanh_to_sigmoid(
generator(latent_estimates)),
min=0.,
max=1.).cpu(),
1,
dim=0)
# Save
for img_id, image in enumerate(images):
y_estimate_pil = Tensor2PIL(
y_estimate_list[img_id]) # 从tensor转化为PIL image并保存
y_estimate_pil.save(
os.path.join(args.outputs, image_name_list[img_id]))
# Create video
if args.video:
print('Create GAN-Inversion video.')
video = cv2.VideoWriter(filename=os.path.join(
args.outputs,
'%s_inversion.avi' % image_name_list[img_id]),
fourcc=cv2.VideoWriter_fourcc(*'MJPG'),
fps=args.fps,
frameSize=(frameSize, frameSize))
print('Save frames.')
for i, sample in enumerate(history):
image = generator(sample)
image_cv2 = convert_array_to_images(
image.detach().cpu().numpy())[0][:, :, ::-1]
video.write(image_cv2)
video.release()
def main(args):
os.makedirs(args.outputs + '/input', exist_ok=True)
os.makedirs(args.outputs + '/GT', exist_ok=True)
os.makedirs(args.outputs + '/mGANoutput', exist_ok=True)
with open(args.outputs + '/mGANargs.txt', 'w') as f:
json.dump(args.__dict__, f, indent=2)
generator = get_derivable_generator(args.gan_model, args.inversion_type,
args)
loss = get_loss(args.loss_type, args)
mask = parsing_mask('code/mganprior/masks/' + args.mask).cuda()
mask_cpu = parsing_mask('code/mganprior/masks/' + args.mask)
crop_loss = masked_loss(loss, mask)
generator.cuda()
loss.cuda()
inversion = get_inversion(args.optimization, args)
image_list = image_files(args.target_images)
if len(image_list) > 300:
print('Limiting the image set to 300.')
image_list = image_list[:300]
frameSize = MODEL_POOL[args.gan_model]['resolution']
start_time = time.time()
for i, images in enumerate(split_to_batches(image_list, 1)):
print('%d: Processing %d images :' % (i, 1), end='')
pt_image_str = '%s\n'
print(pt_image_str % tuple(images))
image_name_list = []
image_tensor_list = []
for image in images:
image_name_list.append(os.path.split(image)[1])
image_tensor_list.append(_add_batch_one(load_as_tensor(image)))
y_gt = _sigmoid_to_tanh(torch.cat(image_tensor_list, dim=0)).cuda()
# Invert
if args.varmask:
os.makedirs(args.outputs + '/mask', exist_ok=True)
mask_cpu = get_var_mask(y_gt.shape[-2:], args.min_p, args.max_p,
args.width_mean, args.width_var)
mask = mask_cpu.cuda()
save_image(mask,
os.path.join(args.outputs + '/mask/%d%s' % (i, '.png')))
crop_loss = masked_loss(loss, mask)
latent_estimates, history = inversion.invert(generator,
y_gt,
crop_loss,
batch_size=1,
video=args.video)
# Get Images
y_estimate_list = torch.split(torch.clamp(_tanh_to_sigmoid(
generator(latent_estimates)),
min=0.,
max=1.).cpu(),
1,
dim=0)
# Save
for img_id, image in enumerate(images):
y_RGB = Tensor2PIL(image_tensor_list[img_id])
y_RGB.save(args.outputs + '/GT/%d%s' %
(i, image_name_list[img_id][-4:]))
y_gt_pil = Tensor2PIL(
mask_images(image_tensor_list[img_id], mask_cpu))
y_estimate_pil = Tensor2PIL(y_estimate_list[img_id])
y_estimate_pil.save(
os.path.join(args.outputs + '/mGANoutput/%d%s' %
(i, image_name_list[img_id][-4:])))
y_gt_pil.save(
os.path.join(args.outputs + '/input/%d%s' %
(i, image_name_list[img_id][-4:])))
# Create video
if args.video:
print('Create GAN-Inversion video.')
video = cv2.VideoWriter(filename=os.path.join(
args.outputs, '%s_inpainting_%s.avi' %
(image_name_list[img_id][:-4], os.path.split(
args.mask[:-4])[1])),
fourcc=cv2.VideoWriter_fourcc(*'MJPG'),
fps=args.fps,
frameSize=(frameSize, frameSize))
print('Save frames.')
for i, sample in enumerate(history):
image = generator(sample)
image_cv2 = convert_array_to_images(
image.detach().cpu().numpy())[0][:, :, ::-1]
video.write(image_cv2)
video.release()
print(f'{(time.time()-start_time)/60:.2f}', 'minutes taken in total;',
f'{(time.time()-start_time)/60/len(image_list):.2f}', 'per image.')
def main(args):
generator = get_derivable_generator(args.gan_model, args.inversion_type,
args)
loss = get_loss(args.loss_type, args)
loss.cuda()
generator.cuda()
inversion = get_inversion(args.optimization, args)
os.makedirs(args.outputs, exist_ok=True)
save_manipulate_dir = os.path.join(args.outputs, args.attribute_name)
images_list = image_files(args.target_images)
for i, images in enumerate(images_list):
print('%d: Processing images ' % (i + 1), end='')
image_name = os.path.split(images)[1]
print(image_name)
image_name_list = []
image_tensor_list = []
image_name_list.append(os.path.split(images)[1])
image_tensor_list.append(_add_batch_one(load_as_tensor(images)))
y_gt = _sigmoid_to_tanh(torch.cat(image_tensor_list, dim=0)).cuda()
# Invert
latent_estimates, history = inversion.invert(generator,
y_gt,
loss,
batch_size=1)
image_manipulate_dir = os.path.join(save_manipulate_dir,
image_name[:-4])
os.makedirs(image_manipulate_dir, exist_ok=True)
wp = latent_estimates[0].cpu().detach().numpy()
mask = latent_estimates[1].cpu().detach().numpy()
# Visualize results with given w+ latent vector.
if args.original:
print('Save inversion.')
image = generator(latent_estimates)
image_cv2 = convert_array_to_images(image.detach().cpu().numpy())
cv2.imwrite(
os.path.join(image_manipulate_dir, 'original_inversion.png'),
image_cv2[0][:, :, ::-1])
boundary, bias = get_boundary(
os.path.join(
BOUNDARY_DIR,
'pggan_celebahq_%s_boundary.npy' % args.attribute_name))
wp_list = get_interpolated_wp(wp,
boundary,
max_step=args.max_step,
num_frames=args.fps * args.duration)
# Create video for attribute manipulation with given w+ latent_vector.
if args.video:
print('Create attribute manipulation video.')
video = cv2.VideoWriter(filename=os.path.join(
image_manipulate_dir,
'%s_manipulate.avi' % args.attribute_name),
fourcc=cv2.VideoWriter_fourcc(*'MJPG'),
fps=args.fps,
frameSize=(1024, 1024))
print('Save frames.')
for i, sample in enumerate(wp_list):
image = generator([
torch.from_numpy(sample).view((1, ) + sample.shape).cuda(),
torch.from_numpy(mask).cuda()
])
image_cv2 = convert_array_to_images(
image.detach().cpu().numpy())[0][:, :, ::-1]
cv2.imwrite(os.path.join(image_manipulate_dir, '%d.png' % i),
image_cv2)
video.write(image_cv2)
video.release()