Exemplo n.º 1
0
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()
Exemplo n.º 2
0
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()
Exemplo n.º 3
0
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 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()
Exemplo n.º 5
0
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)
    cor_loss = Color_loss(loss)
    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: Processing %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,
                                                     cor_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):
            up_gray = colorization_images(image_tensor_list[img_id])
            y_gray_pil = Tensor2PIL(up_gray, mode='L')
            y_gray_pil.save(
                os.path.join(args.outputs,
                             '%s-%s.png' % (image_name_list[img_id], 'gray')))

            Y_gt = Tensor2PIL(image_tensor_list[img_id],
                              mode='RGB').convert('YCbCr')
            y_estimate_pil = Tensor2PIL(y_estimate_list[img_id],
                                        mode='RGB').convert('YCbCr')

            _, Cb, Cr = y_estimate_pil.split()
            Y, _, _ = Y_gt.split()
            y_colorization = Image.merge('YCbCr', (Y, Cb, Cr))
            y_colorization.convert('RGB').save(
                os.path.join(
                    args.outputs, '%s-%d.png' %
                    (image_name_list[img_id], math.floor(time.time()))))
            # 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 = torch.clamp(_tanh_to_sigmoid(generator(sample)),
                                        min=0.,
                                        max=1.).cpu()
                    image_pil = Tensor2PIL(image, mode='RGB').convert('YCbCr')
                    _, Cb, Cr = image_pil.split()
                    y_colorization = Image.merge('YCbCr',
                                                 (Y, Cb, Cr)).convert('RGB')
                    image_cv2 = cv2.cvtColor(np.asarray(y_colorization),
                                             cv2.COLOR_RGB2BGR)
                    # image_cv2 = cv2.cvtColor(np.asarray(image_pil.convert('RGB')), cv2.COLOR_RGB2BGR)
                    video.write(image_cv2)
                video.release()
Exemplo n.º 6
0
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()
Exemplo n.º 7
0
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)
    cor_loss = Color_loss(loss)
    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
        latent_estimates, history = inversion.invert(generator, y_gt, cor_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):
            up_gray = colorization_images(image_tensor_list[img_id])
            y_gray_pil = Tensor2PIL(up_gray, mode='L')
            y_gray_pil.save(args.outputs + '/input/%d%s' % (i, image_name_list[img_id][-4:]))
            
            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 = Tensor2PIL(image_tensor_list[img_id], mode='RGB').convert('YCbCr')
            y_estimate_pil = Tensor2PIL(y_estimate_list[img_id], mode='RGB').convert('YCbCr')

            _, Cb, Cr = y_estimate_pil.split()
            Y, _, _ = Y_gt.split()
            y_colorization = Image.merge('YCbCr', (Y, Cb, Cr))
            y_colorization.convert('RGB').save(args.outputs + '/mGANoutput/%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_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 = torch.clamp(_tanh_to_sigmoid(generator(sample)), min=0., max=1.).cpu()
                    image_pil = Tensor2PIL(image, mode='RGB').convert('YCbCr')
                    _, Cb, Cr = image_pil.split()
                    y_colorization = Image.merge('YCbCr', (Y, Cb, Cr)).convert('RGB')
                    image_cv2 = cv2.cvtColor(np.asarray(y_colorization), cv2.COLOR_RGB2BGR)
                    # image_cv2 = cv2.cvtColor(np.asarray(image_pil.convert('RGB')), cv2.COLOR_RGB2BGR)
                    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.')