| | import argparse |
| |
|
| | import torch |
| | import numpy as np |
| | import sys |
| | import os |
| | import dlib |
| |
|
| | sys.path.append(".") |
| | sys.path.append("..") |
| |
|
| | from configs import data_configs, paths_config |
| | from datasets.inference_dataset import InferenceDataset |
| | from torch.utils.data import DataLoader |
| | from utils.model_utils import setup_model |
| | from utils.common import tensor2im |
| | from utils.alignment import align_face |
| | from PIL import Image |
| |
|
| |
|
| | def main(args): |
| | net, opts = setup_model(args.ckpt, device) |
| | is_cars = 'cars_' in opts.dataset_type |
| | generator = net.decoder |
| | generator.eval() |
| | args, data_loader = setup_data_loader(args, opts) |
| |
|
| | |
| | latents_file_path = os.path.join(args.save_dir, 'latents.pt') |
| | if os.path.exists(latents_file_path): |
| | latent_codes = torch.load(latents_file_path).to(device) |
| | else: |
| | latent_codes = get_all_latents(net, data_loader, args.n_sample, is_cars=is_cars) |
| | torch.save(latent_codes, latents_file_path) |
| |
|
| | if not args.latents_only: |
| | generate_inversions(args, generator, latent_codes, is_cars=is_cars) |
| |
|
| |
|
| | def setup_data_loader(args, opts): |
| | dataset_args = data_configs.DATASETS[opts.dataset_type] |
| | transforms_dict = dataset_args['transforms'](opts).get_transforms() |
| | images_path = args.images_dir if args.images_dir is not None else dataset_args['test_source_root'] |
| | print(f"images path: {images_path}") |
| | align_function = None |
| | if args.align: |
| | align_function = run_alignment |
| | test_dataset = InferenceDataset(root=images_path, |
| | transform=transforms_dict['transform_test'], |
| | preprocess=align_function, |
| | opts=opts) |
| |
|
| | data_loader = DataLoader(test_dataset, |
| | batch_size=args.batch, |
| | shuffle=False, |
| | num_workers=2, |
| | drop_last=True) |
| |
|
| | print(f'dataset length: {len(test_dataset)}') |
| |
|
| | if args.n_sample is None: |
| | args.n_sample = len(test_dataset) |
| | return args, data_loader |
| |
|
| |
|
| | def get_latents(net, x, is_cars=False): |
| | codes = net.encoder(x) |
| | if net.opts.start_from_latent_avg: |
| | if codes.ndim == 2: |
| | codes = codes + net.latent_avg.repeat(codes.shape[0], 1, 1)[:, 0, :] |
| | else: |
| | codes = codes + net.latent_avg.repeat(codes.shape[0], 1, 1) |
| | if codes.shape[1] == 18 and is_cars: |
| | codes = codes[:, :16, :] |
| | return codes |
| |
|
| |
|
| | def get_all_latents(net, data_loader, n_images=None, is_cars=False): |
| | all_latents = [] |
| | i = 0 |
| | with torch.no_grad(): |
| | for batch in data_loader: |
| | if n_images is not None and i > n_images: |
| | break |
| | x = batch |
| | inputs = x.to(device).float() |
| | latents = get_latents(net, inputs, is_cars) |
| | all_latents.append(latents) |
| | i += len(latents) |
| | return torch.cat(all_latents) |
| |
|
| |
|
| | def save_image(img, save_dir, idx): |
| | result = tensor2im(img) |
| | im_save_path = os.path.join(save_dir, f"{idx:05d}.jpg") |
| | Image.fromarray(np.array(result)).save(im_save_path) |
| |
|
| |
|
| | @torch.no_grad() |
| | def generate_inversions(args, g, latent_codes, is_cars): |
| | print('Saving inversion images') |
| | inversions_directory_path = os.path.join(args.save_dir, 'inversions') |
| | os.makedirs(inversions_directory_path, exist_ok=True) |
| | for i in range(args.n_sample): |
| | imgs, _ = g([latent_codes[i].unsqueeze(0)], input_is_latent=True, randomize_noise=False, return_latents=True) |
| | if is_cars: |
| | imgs = imgs[:, :, 64:448, :] |
| | save_image(imgs[0], inversions_directory_path, i + 1) |
| |
|
| |
|
| | def run_alignment(image_path): |
| | predictor = dlib.shape_predictor(paths_config.model_paths['shape_predictor']) |
| | aligned_image = align_face(filepath=image_path, predictor=predictor) |
| | print("Aligned image has shape: {}".format(aligned_image.size)) |
| | return aligned_image |
| |
|
| |
|
| | if __name__ == "__main__": |
| | device = "cuda" |
| |
|
| | parser = argparse.ArgumentParser(description="Inference") |
| | parser.add_argument("--images_dir", type=str, default=None, |
| | help="The directory of the images to be inverted") |
| | parser.add_argument("--save_dir", type=str, default=None, |
| | help="The directory to save the latent codes and inversion images. (default: images_dir") |
| | parser.add_argument("--batch", type=int, default=1, help="batch size for the generator") |
| | parser.add_argument("--n_sample", type=int, default=None, help="number of the samples to infer.") |
| | parser.add_argument("--latents_only", action="store_true", help="infer only the latent codes of the directory") |
| | parser.add_argument("--align", action="store_true", help="align face images before inference") |
| | parser.add_argument("ckpt", metavar="CHECKPOINT", help="path to generator checkpoint") |
| |
|
| | args = parser.parse_args() |
| | main(args) |
| |
|
