コード例 #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)
    mask = parsing_mask(args.mask).cuda()
    mask_cpu = parsing_mask(args.mask)
    crop_loss = masked_loss(loss, mask)
    # 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: 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, 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_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, '%s-%s-%d.png' % (image_name_list[img_id], os.path.split(args.mask[:-4])[1], math.floor(time.time()))))
            y_gt_pil.save(os.path.join(args.outputs, '%s-%s.png' % (image_name_list[img_id], os.path.split(args.mask[:-4])[1])))
            # 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()
コード例 #2
0
def saveIMG():
    image_list = image_files(default_load_img_pth)

    for i, images in enumerate(split_to_batches(image_list, 1)):
        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)))

        # 保存结果
        for img_id, image in enumerate(images):
            down = convert_test(image_tensor_list[img_id],
                                factor=16,
                                mode='nearest')
            y_nn_pil = Tensor2PIL(down)
            y_nn_pil.save(
                os.path.join(
                    os.path.join('img_test_output',
                                 '%s-nn.png' % image_name_list[img_id][:-4])))
コード例 #3
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()
コード例 #4
0
ファイル: PmSFC.py プロジェクト: guanyuelee/PmSFC
def main(args):
    os.makedirs(args.outputs, exist_ok=True)
    out_dir, exp_name = create_experiments_directory(args, args.exp_id)
    print(out_dir)
    print(exp_name)
    generator = get_derivable_generator(args.gan_model, args.inversion_type,
                                        args)
    generator.cuda()

    if args.target_images.endswith('.png') or args.target_images.endswith(
            '.jpg'):
        image_list = os.path.abspath(args.target_images)
        image_list = [image_list]
    else:
        image_list = image_files(args.target_images)
    frameSize = MODEL_POOL[args.gan_model]['resolution']
    n_blocks = generator.n_blocks
    print('There are %d blocks in this generator.' % n_blocks)  # 19 for pggan

    latent_space = generator.PGGAN_LATENT[args.layer]
    print('The latent space is ', latent_space)

    with open(args.matrix_dir, 'rb') as file_in:
        matrix = pkl.load(file_in)
        print('Load matrix successfully.')
        print('Matrix shape ', matrix.shape)

    matrix2 = thrC(matrix, args.alpha)
    predict, _ = post_proC(matrix2, args.n_subs, args.d_subs, args.power)
    print(predict)
    p_sum = [sum(predict == k) for k in range(1, args.cluster_numbers, 1)]
    p_sum = np.array(p_sum)
    p_sort = np.argsort(p_sum)[::-1]
    print(p_sum)
    predict_new = predict.copy()
    for i in range(1, args.cluster_numbers, 1):
        predict_new[predict == (p_sort[i - 1] + 1)] = i
    predict = predict_new.copy()
    p_sum = [sum(predict == k) for k in range(1, args.cluster_numbers, 1)]
    print(predict)
    print(p_sum)

    # pre_see the images
    gan_type, image_type = args.gan_model.split("_")
    print('The gan type is %s, and the image type is %s' %
          (gan_type, image_type))
    test_image_dir = os.path.join('./bin/', gan_type, image_type)
    print(test_image_dir)

    files = os.listdir(test_image_dir)
    test_zs = []
    for i in range(len(files)):
        if files[i].endswith('.pkl'):
            with open(os.path.join(test_image_dir, files[i]), 'rb') as file_in:
                test_zs.append(pkl.load(file_in))
    test_zs = torch.from_numpy(
        np.concatenate(test_zs, axis=0).astype(np.float32)).cuda()
    print('Load all testing zs, shape is ', test_zs.size())

    image_number = 3
    sel_idx = np.random.choice(test_zs.shape[0],
                               size=[image_number],
                               replace=False)
    F = generator([test_zs[sel_idx]], which_block=args.layer, pre_model=True)
    features = F.detach().cpu().numpy()

    predict_masks = []
    for i in range(1, args.cluster_numbers + 1, 1):
        mask = torch.from_numpy((predict == i).astype(np.float32)).cuda()
        predict_masks += [mask.reshape((1, -1, 1, 1))]

    for i, images in enumerate(split_to_batches(image_list, 1)):
        # input_size = generator.PGGAN_LATENT[args.layer + 1]
        # print("We are making ", input_size, ". ")
        print('%d: Inverting %d images :' % (i + 1, 1), end='')
        pt_image_str = '%s\n'
        print(pt_image_str % tuple(images))
        image_name_only = images[0].split(".")[0]
        image_name_only = image_name_only.split("/")[-1]
        print(image_name_only)

        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)))
        print("image_name_list", image_name_list)
        print("image_tensor_list, [", image_tensor_list[0].size(), "]")

        y_image = _sigmoid_to_tanh(torch.cat(image_tensor_list, dim=0)).cuda()
        print('image size is ', y_image.size())

        z_estimate = generator.init_value(batch_size=1,
                                          which_layer=0,
                                          init=args.init_type,
                                          z_numbers=args.cluster_numbers *
                                          args.code_per_cluster)
        base_estimate = generator.init_value(batch_size=1,
                                             which_layer=0,
                                             init=args.init_type,
                                             z_numbers=1)

        if args.optimization == 'GD':
            z_optimizer = torch.optim.SGD(z_estimate + base_estimate,
                                          lr=args.lr)
        elif args.optimization == 'Adam':
            z_optimizer = torch.optim.Adam(z_estimate + base_estimate,
                                           lr=args.lr)
        else:
            raise NotImplemented(
                'We don\'t support this type of optimization.')

        for iter in range(args.iterations):
            for estimate in z_estimate:
                estimate.requires_grad = True
            for estimate in base_estimate:
                estimate.requires_grad = True

            features = generator([
                z_estimate[0].reshape(
                    [args.cluster_numbers * args.code_per_cluster, 512, 1, 1])
            ],
                                 which_block=args.layer,
                                 pre_model=True)
            base_feature = generator(
                [base_estimate[0].reshape([1, 512, 1, 1])],
                which_block=args.layer,
                pre_model=True)

            for t in range(args.cluster_numbers * args.code_per_cluster):
                if t == 0:
                    f_mix = features[t].view(*(1, ) +
                                             latent_space) * predict_masks[int(
                                                 t / args.code_per_cluster)]
                else:
                    f_mix = f_mix + features[t].view(
                        *(1, ) + latent_space) * predict_masks[int(
                            t / args.code_per_cluster)]

            f_mix = f_mix + base_feature
            y_estimate = generator([f_mix],
                                   which_block=args.layer,
                                   post_model=True)
            y_raw_estimate = generator([base_feature],
                                       which_block=args.layer,
                                       post_model=True)
            z_optimizer.zero_grad()
            loss = 0.01 * torch.mean(torch.pow(
                y_estimate - y_image, 2.0)) + torch.mean(
                    torch.pow(y_raw_estimate - y_image, 2.0))
            loss.backward()
            z_optimizer.step()

            if iter % args.report_value == 0:
                print('Iter %d, layer %d, loss = %.4f.' %
                      (iter, args.layer, float(loss.item())))

            if iter % args.report_image == 0:
                print('Saving the images.')
                y_estimate_pil = Tensor2PIL(
                    torch.clamp(_tanh_to_sigmoid(y_estimate.detach().cpu()),
                                min=0.0,
                                max=1.0))
                y_estimate_pil.save(
                    os.path.join(
                        out_dir,
                        image_name_only + "_estimate_iter%d.png" % iter))
                y_estimate_pil = Tensor2PIL(
                    torch.clamp(_tanh_to_sigmoid(
                        y_raw_estimate.detach().cpu()),
                                min=0.0,
                                max=1.0))
                y_estimate_pil.save(
                    os.path.join(
                        out_dir,
                        image_name_only + "_raw_estimate_iter%d.png" % iter))
                # add bias added output picture.
                # save all the codes
                codes = []
                for code_idx in range(args.cluster_numbers *
                                      args.code_per_cluster):
                    code_f = generator([z_estimate[0][code_idx]],
                                       which_block=args.layer + 1,
                                       pre_model=True)
                    code_y = generator([code_f],
                                       which_block=args.layer + 1,
                                       post_model=True).detach().cpu()
                    codes.append(
                        torch.clamp(_tanh_to_sigmoid(code_y), min=0, max=1))
                codes = torch.cat(codes, dim=0).detach().cpu()
                torchvision.utils.save_image(
                    codes,
                    os.path.join(
                        out_dir,
                        image_name_only + '_codes_iter%d.png' % (iter)),
                    nrow=(args.cluster_numbers * args.code_per_cluster) // 2)

            if iter % args.report_model == 0:
                print('Save the models')
                save_dict = {
                    "z": z_estimate[0].detach().cpu().numpy(),
                    'matrix': matrix,
                    'layer': args.layer,
                    'predict': predict
                }
                with open(
                        os.path.join(
                            out_dir, 'save_dict_iter_%d_layer_%d.pkl' %
                            (iter, args.layer)), 'wb') as file_out:
                    pkl.dump(save_dict, file_out)
                print('Save the models OK!')
コード例 #5
0
def main(args):
    os.makedirs(args.outputs, exist_ok=True)
    generator = get_derivable_generator(
        args.gan_model, args.inversion_type,
        args)  # pggan_celebahq/stylegan2_ffhq256    PGGAN-Multi-Z/StyleGAN-w+
    loss = get_loss(args.loss_type, args)
    sr_loss = SR_loss(loss, args.down, args.factor)
    # to cuda
    generator.cuda()
    loss.cuda()
    inversion = get_inversion(args.optimization, args)  # SGD
    image_list = image_files(args.target_images)
    frameSize = MODEL_POOL[args.gan_model][
        'resolution']  # gan_model=pggan_celebahq

    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_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,
                                                     sr_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_nn, up_bic = 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)
                    image_cv2 = convert_array_to_images(
                        image.detach().cpu().numpy())[0][:, :, ::-1]
                    video.write(image_cv2)
                video.release()
コード例 #6
0
ファイル: inversion_methods.py プロジェクト: guanyuelee/PmSFC
    def invert(self,
               generator,
               gt_image,
               loss_function,
               batch_size=1,
               video=True,
               init=(),
               z_iterations=3000,
               args=None,
               out_dir=None,
               image_name=None):
        input_size_list = generator.input_size()
        if len(init) == 0:
            if generator.init is False:
                latent_estimate = []
                for input_size in input_size_list:
                    if self.init_type == 'Zero':
                        print("Zero")
                        latent_estimate.append(
                            torch.zeros((batch_size, ) + input_size).cuda())
                    elif self.init_type == 'Normal':
                        latent_estimate.append(
                            torch.randn((batch_size, ) + input_size).cuda())
            else:
                latent_estimate = list(generator.init_value(batch_size))
        else:
            assert len(init) == len(
                input_size_list), 'Please check the number of init value'
            latent_estimate = init

        for latent in latent_estimate:
            latent.requires_grad = True  # 19 PGGAN - 0: 1   4 - 18: 1e-8 x 10 = 1e-6
        print(latent_estimate[0].requires_grad)
        optimizer = self.optimizer(latent_estimate, lr=self.lr)

        history = []
        # Opt
        for i in tqdm(range(z_iterations)):
            y_estimate = generator(latent_estimate)
            if i % args.report_image == 0:
                if out_dir is not None:
                    y_estimate_pil = Tensor2PIL(
                        torch.clamp(_tanh_to_sigmoid(
                            y_estimate.detach().cpu()),
                                    min=0.0,
                                    max=1.0))
                    y_estimate_pil.save(
                        os.path.join(
                            out_dir,
                            image_name.split('.')[0] + "_" + str(i) +
                            "_layer_" + str(args.composing_layer) + ".png"))
                    zs = latent_estimate[0].detach().cpu().numpy()
                    ps = latent_estimate[1].detach().cpu().numpy()
                    save_dict = {'zs': zs, 'ps': ps}
                    with open(
                            os.path.join(
                                out_dir,
                                image_name.split('.')[0] + "_" + str(i) +
                                "_layer_" + str(args.composing_layer) +
                                ".pkl"), 'wb') as file_out:
                        pkl.dump(save_dict, file_out)

            optimizer.zero_grad()
            loss = loss_function(y_estimate, gt_image)
            loss.backward()
            optimizer.step()
            if video:
                history.append(copy.deepcopy(latent_estimate))

        return latent_estimate, history
コード例 #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)
    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.')
コード例 #8
0
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()
コード例 #9
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()
コード例 #10
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.')