Spaces:
Running
Running
| import argparse | |
| import torch | |
| import os | |
| from Project.configs import data_configs | |
| from Project.datasets.inference_dataset import InferenceDataset | |
| from torch.utils.data import DataLoader | |
| from Project.utils.model_utils import setup_model | |
| def main(args,device): | |
| net, opts ,latent_avg= setup_model(args.ckpt, device) | |
| is_cars = 'cars_' in opts.dataset_type | |
| args, data_loader = setup_data_loader(args, opts) | |
| # Check if latents exist | |
| latents_file_path = os.path.join(args.save_dir, 'latents.pt') | |
| latent_codes = get_all_latents(net, device,data_loader,latent_avg, args.n_sample, is_cars=is_cars) | |
| torch.save(latent_codes, latents_file_path) | |
| 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 | |
| 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=0, | |
| 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,latent_avg, is_cars=False): | |
| input = {net.get_inputs()[0].name: to_numpy(x)} | |
| codes = net.run(None,input) | |
| codes=torch.from_numpy(codes[0]) | |
| if codes.ndim == 2: | |
| codes = codes + latent_avg.repeat(codes.shape[0], 1, 1)[:, 0, :] | |
| else: | |
| codes = codes + latent_avg.repeat(codes.shape[0], 1, 1) | |
| return codes | |
| def get_all_latents(net, device ,data_loader, latent_avg,n_images=None, is_cars=False): | |
| all_latents = [] | |
| with torch.no_grad(): | |
| for batch in data_loader: | |
| x = batch | |
| inputs = x.float() | |
| print(inputs.shape) | |
| latents = get_latents(net, inputs,latent_avg, is_cars) | |
| all_latents.append(latents) | |
| return torch.cat(all_latents) | |
| #@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(min(args.n_sample, len(latent_codes))): | |
| # 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 | |
| def to_numpy(tensor): | |
| return tensor.cpu().numpy() | |
| def inference(): | |
| device = "cpu" | |
| parser = argparse.ArgumentParser(description="Inference") | |
| parser.add_argument("--images_dir", type=str, default='static/img_aligned', | |
| help="The directory of the images to be inverted") | |
| parser.add_argument("--save_dir", type=str, default='static/latents', | |
| 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",default=True, help="infer only the latent codes of the directory") | |
| parser.add_argument("--align", action="store_true",default=False,help="align face images before inference") | |
| parser.add_argument("--ckpt", default='Project/pretrained_models/e4e_ffhq_encode.pt',help="path to generator checkpoint") | |
| args = parser.parse_args() | |
| main(args,device) | |
| if __name__=="__main__": | |
| inference() |