Spaces:
Runtime error
Runtime error
File size: 4,919 Bytes
c2f174c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 |
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)
# Check if latents exist
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)
|