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, split=args.split, 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 evaluate(options): config = InferenceConfig(options) config.FITTING_TYPE = options.numAnchorPlanes if options.dataset == '': dataset = PlaneDataset(options, config, split='test', random=False, load_semantics=False) elif options.dataset == 'occlusion': config_dataset = copy.deepcopy(config) config_dataset.OCCLUSION = False dataset = PlaneDataset(options, config_dataset, split='test', random=False, load_semantics=True) elif 'nyu' in options.dataset: dataset = NYUDataset(options, config, split='val', random=False) elif options.dataset == 'synthia': dataset = SynthiaDataset(options, config, split='val', random=False) elif options.dataset == 'kitti': camera = np.zeros(6) camera[0] = 9.842439e+02 camera[1] = 9.808141e+02 camera[2] = 6.900000e+02 camera[3] = 2.331966e+02 camera[4] = 1242 camera[5] = 375 dataset = InferenceDataset( options, config, image_list=glob.glob('../../Data/KITTI/scene_3/*.png'), camera=camera) elif options.dataset == '7scene': camera = np.zeros(6) camera[0] = 519 camera[1] = 519 camera[2] = 320 camera[3] = 240 camera[4] = 640 camera[5] = 480 dataset = InferenceDataset( options, config, image_list=glob.glob('../../Data/SevenScene/scene_3/*.png'), camera=camera) elif options.dataset == 'tanktemple': camera = np.zeros(6) camera[0] = 0.7 camera[1] = 0.7 camera[2] = 0.5 camera[3] = 0.5 camera[4] = 1 camera[5] = 1 dataset = InferenceDataset( options, config, image_list=glob.glob('../../Data/TankAndTemple/scene_4/*.jpg'), camera=camera) elif options.dataset == 'make3d': camera = np.zeros(6) camera[0] = 0.7 camera[1] = 0.7 camera[2] = 0.5 camera[3] = 0.5 camera[4] = 1 camera[5] = 1 dataset = InferenceDataset( options, config, image_list=glob.glob('../../Data/Make3D/*.jpg'), camera=camera) elif options.dataset == 'popup': camera = np.zeros(6) camera[0] = 0.7 camera[1] = 0.7 camera[2] = 0.5 camera[3] = 0.5 camera[4] = 1 camera[5] = 1 dataset = InferenceDataset( options, config, image_list=glob.glob('../../Data/PhotoPopup/*.jpg'), camera=camera) elif options.dataset == 'cross' or options.dataset == 'cross_2': image_list = [ 'test/cross_dataset/' + str(c) + '_image.png' for c in range(12) ] cameras = [] camera = np.zeros(6) camera[0] = 587 camera[1] = 587 camera[2] = 320 camera[3] = 240 camera[4] = 640 camera[5] = 480 for c in range(4): cameras.append(camera) continue camera_kitti = np.zeros(6) camera_kitti[0] = 9.842439e+02 camera_kitti[1] = 9.808141e+02 camera_kitti[2] = 6.900000e+02 camera_kitti[3] = 2.331966e+02 camera_kitti[4] = 1242.0 camera_kitti[5] = 375.0 for c in range(2): cameras.append(camera_kitti) continue camera_synthia = np.zeros(6) camera_synthia[0] = 133.185088 camera_synthia[1] = 134.587036 camera_synthia[2] = 160.000000 camera_synthia[3] = 96.000000 camera_synthia[4] = 320 camera_synthia[5] = 192 for c in range(2): cameras.append(camera_synthia) continue camera_tanktemple = np.zeros(6) camera_tanktemple[0] = 0.7 camera_tanktemple[1] = 0.7 camera_tanktemple[2] = 0.5 camera_tanktemple[3] = 0.5 camera_tanktemple[4] = 1 camera_tanktemple[5] = 1 for c in range(2): cameras.append(camera_tanktemple) continue for c in range(2): cameras.append(camera) continue dataset = InferenceDataset(options, config, image_list=image_list, camera=cameras) elif options.dataset == 'selected': image_list = glob.glob('test/selected_images/*_image_0.png') image_list = [ filename for filename in image_list if '63_image' not in filename and '77_image' not in filename ] + [ filename for filename in image_list if '63_image' in filename or '77_image' in filename ] camera = np.zeros(6) camera[0] = 587 camera[1] = 587 camera[2] = 320 camera[3] = 240 camera[4] = 640 camera[5] = 480 dataset = InferenceDataset(options, config, image_list=image_list, camera=camera) elif options.dataset == 'comparison': image_list = [ 'test/comparison/' + str(index) + '_image_0.png' for index in [65, 11, 24] ] camera = np.zeros(6) camera[0] = 587 camera[1] = 587 camera[2] = 320 camera[3] = 240 camera[4] = 640 camera[5] = 480 dataset = InferenceDataset(options, config, image_list=image_list, camera=camera) elif 'inference' in options.dataset: image_list = glob.glob(options.customDataFolder + '/*.png') + glob.glob(options.customDataFolder + '/*.jpg') if os.path.exists(options.customDataFolder + '/camera.txt'): camera = np.zeros(6) with open(options.customDataFolder + '/camera.txt', 'r') as f: for line in f: values = [ float(token.strip()) for token in line.split(' ') if token.strip() != '' ] for c in range(6): camera[c] = values[c] continue break pass else: camera = [ filename.replace('.png', '.txt').replace('.jpg', '.txt') for filename in image_list ] pass dataset = InferenceDataset(options, config, image_list=image_list, camera=camera) pass print('the number of images', len(dataset)) dataloader = DataLoader(dataset, batch_size=1, shuffle=False, num_workers=1) epoch_losses = [] data_iterator = tqdm(dataloader, total=len(dataset)) specified_suffix = options.suffix with torch.no_grad(): detectors = [] for method in options.methods: if method == 'w': options.suffix = 'pair_' + specified_suffix if specified_suffix != '' else 'pair' detectors.append(('warping', PlaneRCNNDetector(options, config, modelType='pair'))) elif method == 'b': options.suffix = specified_suffix if specified_suffix != '' else '' detectors.append(('basic', PlaneRCNNDetector(options, config, modelType='pair'))) elif method == 'o': options.suffix = 'occlusion_' + specified_suffix if specified_suffix != '' else 'occlusion' detectors.append(('occlusion', PlaneRCNNDetector(options, config, modelType='occlusion'))) elif method == 'p': detectors.append( ('planenet', PlaneNetDetector(options, config))) elif method == 'e': detectors.append( ('planerecover', PlaneRecoverDetector(options, config))) elif method == 't': if 'gt' in options.suffix: detectors.append( ('manhattan_gt', TraditionalDetector(options, config, 'manhattan_gt'))) else: detectors.append( ('manhattan_pred', TraditionalDetector(options, config, 'manhattan_pred'))) pass elif method == 'n': options.suffix = specified_suffix if specified_suffix != '' else '' detectors.append(('non_planar', DepthDetector(options, config, modelType='np'))) elif method == 'r': options.suffix = specified_suffix if specified_suffix != '' else '' detectors.append(('refine', PlaneRCNNDetector(options, config, modelType='refine'))) elif method == 's': options.suffix = specified_suffix if specified_suffix != '' else '' detectors.append( ('refine_single', PlaneRCNNDetector(options, config, modelType='refine_single'))) elif method == 'f': options.suffix = specified_suffix if specified_suffix != '' else '' detectors.append(('final', PlaneRCNNDetector(options, config, modelType='final'))) pass continue pass if not options.debug: for method_name in [detector[0] for detector in detectors]: os.system('rm ' + options.test_dir + '/*_' + method_name + '.png') continue pass all_statistics = [] for name, detector in detectors: statistics = [[], [], [], []] for sampleIndex, sample in enumerate(data_iterator): if options.testingIndex >= 0 and sampleIndex != options.testingIndex: if sampleIndex > options.testingIndex: break continue input_pair = [] camera = sample[30][0].cuda() for indexOffset in [ 0, ]: images, image_metas, rpn_match, rpn_bbox, gt_class_ids, gt_boxes, gt_masks, gt_parameters, gt_depth, extrinsics, planes, gt_segmentation = sample[ indexOffset + 0].cuda(), sample[indexOffset + 1].numpy(), sample[ indexOffset + 2].cuda(), sample[indexOffset + 3].cuda(), sample[ indexOffset + 4].cuda(), sample[indexOffset + 5].cuda(), sample[ indexOffset + 6].cuda(), sample[indexOffset + 7].cuda( ), sample[indexOffset + 8].cuda(), sample[ indexOffset + 9].cuda(), sample[ indexOffset + 10].cuda(), sample[indexOffset + 11].cuda() masks = ( gt_segmentation == torch.arange(gt_segmentation.max() + 1).cuda().view(-1, 1, 1)).float() input_pair.append({ 'image': images, 'depth': gt_depth, 'bbox': gt_boxes, 'extrinsics': extrinsics, 'segmentation': gt_segmentation, 'camera': camera, 'plane': planes[0], 'masks': masks, 'mask': gt_masks }) continue if sampleIndex >= options.numTestingImages: break with torch.no_grad(): detection_pair = detector.detect(sample) pass if options.dataset == 'rob': depth = detection_pair[0]['depth'].squeeze().detach().cpu( ).numpy() os.system('rm ' + image_list[sampleIndex].replace('color', 'depth')) depth_rounded = np.round(depth * 256) depth_rounded[np.logical_or(depth_rounded < 0, depth_rounded >= 256 * 256)] = 0 cv2.imwrite( image_list[sampleIndex].replace('color', 'depth').replace( 'jpg', 'png'), depth_rounded.astype(np.uint16)) continue if 'inference' not in options.dataset: for c in range(len(input_pair)): evaluateBatchDetection( options, config, input_pair[c], detection_pair[c], statistics=statistics, printInfo=options.debug, evaluate_plane=options.dataset == '') continue else: for c in range(len(detection_pair)): np.save( options.test_dir + '/' + str(sampleIndex % 500) + '_plane_parameters_' + str(c) + '.npy', detection_pair[c]['detection'][:, 6:9]) np.save( options.test_dir + '/' + str(sampleIndex % 500) + '_plane_masks_' + str(c) + '.npy', detection_pair[c]['masks'][:, 80:560]) continue pass if sampleIndex < 30 or options.debug or options.dataset != '': visualizeBatchPair(options, config, input_pair, detection_pair, indexOffset=sampleIndex % 500, suffix='_' + name + options.modelType, write_ply=options.testingIndex >= 0, write_new_view=options.testingIndex >= 0 and 'occlusion' in options.suffix) pass if sampleIndex >= options.numTestingImages: break continue if 'inference' not in options.dataset: options.keyname = name printStatisticsDetection(options, statistics) all_statistics.append(statistics) pass continue if 'inference' not in options.dataset: if options.debug and len(detectors) > 1: all_statistics = np.concatenate([ np.arange(len(all_statistics[0][0])).reshape((-1, 1)), ] + [np.array(statistics[3]) for statistics in all_statistics], axis=-1) print(all_statistics.astype(np.int32)) pass if options.testingIndex == -1: np.save('logs/all_statistics.npy', all_statistics) pass pass return
def run(): test_opts = TestOptions().parse() if test_opts.resize_factors is not None: factors = test_opts.resize_factors.split(',') assert len( factors ) == 1, "When running inference, please provide a single downsampling factor!" mixed_path_results = os.path.join( test_opts.exp_dir, 'style_mixing', 'downsampling_{}'.format(test_opts.resize_factors)) else: mixed_path_results = os.path.join(test_opts.exp_dir, 'style_mixing') os.makedirs(mixed_path_results, exist_ok=True) # update test options with options used during training ckpt = torch.load(test_opts.checkpoint_path, map_location='cpu') opts = ckpt['opts'] opts.update(vars(test_opts)) if 'learn_in_w' not in opts: opts['learn_in_w'] = False opts = Namespace(**opts) net = pSp(opts) net.eval() net.cuda() print('Loading dataset for {}'.format(opts.dataset_type)) dataset_args = data_configs.DATASETS[opts.dataset_type] transforms_dict = dataset_args['transforms'](opts).get_transforms() dataset = InferenceDataset( root=opts.data_path, transform=transforms_dict['transform_inference'], opts=opts) dataloader = DataLoader(dataset, batch_size=opts.test_batch_size, shuffle=False, num_workers=int(opts.test_workers), drop_last=True) latent_mask = [int(l) for l in opts.latent_mask.split(",")] if opts.n_images is None: opts.n_images = len(dataset) global_i = 0 for input_batch in tqdm(dataloader): if global_i > opts.n_images: break with torch.no_grad(): input_batch = input_batch.cuda() for image_idx, input_image in enumerate(input_batch): # generate random vectors to inject into input image vecs_to_inject = np.random.randn(opts.n_outputs_to_generate, 512).astype('float32') multi_modal_outputs = [] for vec_to_inject in vecs_to_inject: cur_vec = torch.from_numpy(vec_to_inject).unsqueeze(0).to( "cuda") # get latent vector to inject into our input image _, latent_to_inject = net(cur_vec, input_code=True, return_latents=True) # get output image with injected style vector res = net(input_image.unsqueeze(0).to("cuda").float(), latent_mask=latent_mask, inject_latent=latent_to_inject, alpha=opts.mix_alpha) multi_modal_outputs.append(res[0]) # visualize multi modal outputs input_im_path = dataset.paths[global_i] image = input_batch[image_idx] input_image = log_input_image(image, opts) res = np.array(input_image.resize((256, 256))) for output in multi_modal_outputs: output = tensor2im(output) res = np.concatenate( [res, np.array(output.resize((256, 256)))], axis=1) Image.fromarray(res).save( os.path.join(mixed_path_results, os.path.basename(input_im_path))) global_i += 1
def run(): test_opts = TestOptions().parse() out_path_results = os.path.join(test_opts.exp_dir, 'inference_side_by_side') os.makedirs(out_path_results, exist_ok=True) # update test options with options used during training ckpt = torch.load(test_opts.checkpoint_path, map_location='cpu') opts = ckpt['opts'] opts.update(vars(test_opts)) opts = Namespace(**opts) net = pSp(opts) net.eval() net.cuda() age_transformers = [AgeTransformer(target_age=age) for age in opts.target_age.split(',')] print(f'Loading dataset for {opts.dataset_type}') dataset_args = data_configs.DATASETS[opts.dataset_type] transforms_dict = dataset_args['transforms'](opts).get_transforms() dataset = InferenceDataset(root=opts.data_path, transform=transforms_dict['transform_inference'], opts=opts, return_path=True) dataloader = DataLoader(dataset, batch_size=opts.test_batch_size, shuffle=False, num_workers=int(opts.test_workers), drop_last=False) if opts.n_images is None: opts.n_images = len(dataset) global_time = [] global_i = 0 for input_batch, image_paths in tqdm(dataloader): if global_i >= opts.n_images: break batch_results = {} for idx, age_transformer in enumerate(age_transformers): with torch.no_grad(): input_age_batch = [age_transformer(img.cpu()).to('cuda') for img in input_batch] input_age_batch = torch.stack(input_age_batch) input_cuda = input_age_batch.cuda().float() tic = time.time() result_batch = run_on_batch(input_cuda, net, opts) toc = time.time() global_time.append(toc - tic) resize_amount = (256, 256) if opts.resize_outputs else (1024, 1024) for i in range(len(input_batch)): result = tensor2im(result_batch[i]) im_path = image_paths[i] input_im = log_image(input_batch[i], opts) if im_path not in batch_results.keys(): batch_results[im_path] = np.array(input_im.resize(resize_amount)) batch_results[im_path] = np.concatenate([batch_results[im_path], np.array(result.resize(resize_amount))], axis=1) for im_path, res in batch_results.items(): image_name = os.path.basename(im_path) im_save_path = os.path.join(out_path_results, image_name) Image.fromarray(np.array(res)).save(im_save_path) global_i += 1
def run(): test_opts = TestOptions().parse() if test_opts.resize_factors is not None: assert len( test_opts.resize_factors.split(',') ) == 1, "When running inference, provide a single downsampling factor!" out_path_results = os.path.join( test_opts.exp_dir, 'inference_results', 'downsampling_{}'.format(test_opts.resize_factors)) out_path_coupled = os.path.join( test_opts.exp_dir, 'inference_coupled', 'downsampling_{}'.format(test_opts.resize_factors)) else: out_path_results = os.path.join(test_opts.exp_dir, 'inference_results') out_path_coupled = os.path.join(test_opts.exp_dir, 'inference_coupled') os.makedirs(out_path_results, exist_ok=True) os.makedirs(out_path_coupled, exist_ok=True) # update test options with options used during training ckpt = torch.load(test_opts.checkpoint_path, map_location='cpu') opts = ckpt['opts'] opts.update(vars(test_opts)) if 'learn_in_w' not in opts: opts['learn_in_w'] = False opts = Namespace(**opts) net = pSp(opts) net.eval() net.cuda() print('Loading dataset for {}'.format(opts.dataset_type)) dataset_args = data_configs.DATASETS[opts.dataset_type] transforms_dict = dataset_args['transforms'](opts).get_transforms() dataset = InferenceDataset( root=opts.data_path, transform=transforms_dict['transform_inference'], opts=opts) dataloader = DataLoader(dataset, batch_size=opts.test_batch_size, shuffle=False, num_workers=int(opts.test_workers), drop_last=True) if opts.n_images is None: opts.n_images = len(dataset) global_i = 0 global_time = [] for input_batch in tqdm(dataloader): if global_i >= opts.n_images: break with torch.no_grad(): input_cuda = input_batch.cuda().float() tic = time.time() result_batch = run_on_batch(input_cuda, net, opts) toc = time.time() global_time.append(toc - tic) for i in range(opts.test_batch_size): result = tensor2im(result_batch[i]) im_path = dataset.paths[global_i] if opts.couple_outputs or global_i % 100 == 0: input_im = log_input_image(input_batch[i], opts) resize_amount = (256, 256) if opts.resize_outputs else (1024, 1024) if opts.resize_factors is not None: # for super resolution, save the original, down-sampled, and output source = Image.open(im_path) res = np.concatenate([ np.array(source.resize(resize_amount)), np.array( input_im.resize(resize_amount, resample=Image.NEAREST)), np.array(result.resize(resize_amount)) ], axis=1) else: # otherwise, save the original and output res = np.concatenate([ np.array(input_im.resize(resize_amount)), np.array(result.resize(resize_amount)) ], axis=1) Image.fromarray(res).save( os.path.join(out_path_coupled, os.path.basename(im_path))) im_save_path = os.path.join(out_path_results, os.path.basename(im_path)) Image.fromarray(np.array(result)).save(im_save_path) global_i += 1 stats_path = os.path.join(opts.exp_dir, 'stats.txt') result_str = 'Runtime {:.4f}+-{:.4f}'.format(np.mean(global_time), np.std(global_time)) print(result_str) with open(stats_path, 'w') as f: f.write(result_str)
def run(): test_opts = TestOptions().parse() out_path_coupled = os.path.join(test_opts.exp_dir, 'inference_coupled') os.makedirs(out_path_coupled, exist_ok=True) # update test options with options used during training ckpt = torch.load(test_opts.checkpoint_path, map_location='cpu') opts = ckpt['opts'] opts.update(vars(test_opts)) opts = Namespace(**opts) net = pSp(opts) net.eval() net.cuda() print('Loading dataset for {}'.format(opts.dataset_type)) dataset_args = data_configs.DATASETS[opts.dataset_type] transforms_dict = dataset_args['transforms'](opts).get_transforms() dataset = InferenceDataset(root=opts.data_path, transform=transforms_dict['transform_inference'], opts=opts) dataloader = DataLoader(dataset, batch_size=opts.test_batch_size, shuffle=False, num_workers=int(opts.test_workers), drop_last=False) if opts.n_images is None: opts.n_images = len(dataset) # get the image corresponding to the latent average avg_image = net(net.latent_avg.unsqueeze(0), input_code=True, randomize_noise=False, return_latents=False, average_code=True)[0] avg_image = avg_image.to('cuda').float().detach() if opts.dataset_type == "cars_encode": avg_image = avg_image[:, 32:224, :] tensor2im(avg_image).save(os.path.join(opts.exp_dir, 'avg_image.jpg')) if opts.dataset_type == "cars_encode": resize_amount = (256, 192) if opts.resize_outputs else (512, 384) else: resize_amount = (256, 256) if opts.resize_outputs else (opts.output_size, opts.output_size) global_i = 0 global_time = [] for input_batch in tqdm(dataloader): if global_i >= opts.n_images: break with torch.no_grad(): input_cuda = input_batch.cuda().float() tic = time.time() result_batch, result_latents = run_on_batch(input_cuda, net, opts, avg_image) toc = time.time() global_time.append(toc - tic) for i in range(input_batch.shape[0]): results = [tensor2im(result_batch[i][iter_idx]) for iter_idx in range(opts.n_iters_per_batch)] im_path = dataset.paths[global_i] # save step-by-step results side-by-side input_im = tensor2im(input_batch[i]) res = np.array(results[0].resize(resize_amount)) for idx, result in enumerate(results[1:]): res = np.concatenate([res, np.array(result.resize(resize_amount))], axis=1) res = np.concatenate([res, input_im.resize(resize_amount)], axis=1) Image.fromarray(res).save(os.path.join(out_path_coupled, os.path.basename(im_path))) global_i += 1 stats_path = os.path.join(opts.exp_dir, 'stats.txt') result_str = 'Runtime {:.4f}+-{:.4f}'.format(np.mean(global_time), np.std(global_time)) print(result_str) with open(stats_path, 'w') as f: f.write(result_str)
def run(): test_opts = TestOptions().parse() out_path_results = os.path.join(test_opts.exp_dir, 'inference_results') os.makedirs(out_path_results, exist_ok=True) # load model used for initializing encoder bootstrapping ckpt = torch.load(test_opts.model_1_checkpoint_path, map_location='cpu') opts = ckpt['opts'] opts.update(vars(test_opts)) opts['checkpoint_path'] = test_opts.model_1_checkpoint_path opts = Namespace(**opts) if opts.encoder_type in ENCODER_TYPES['pSp']: net1 = pSp(opts) else: net1 = e4e(opts) net1.eval() net1.cuda() # load model used for translating input image after initialization ckpt = torch.load(test_opts.model_2_checkpoint_path, map_location='cpu') opts = ckpt['opts'] opts.update(vars(test_opts)) opts['checkpoint_path'] = test_opts.model_2_checkpoint_path opts = Namespace(**opts) if opts.encoder_type in ENCODER_TYPES['pSp']: net2 = pSp(opts) else: net2 = e4e(opts) net2.eval() net2.cuda() print('Loading dataset for {}'.format(opts.dataset_type)) dataset_args = data_configs.DATASETS[opts.dataset_type] transforms_dict = dataset_args['transforms'](opts).get_transforms() dataset = InferenceDataset( root=opts.data_path, transform=transforms_dict['transform_inference'], opts=opts) dataloader = DataLoader(dataset, batch_size=opts.test_batch_size, shuffle=False, num_workers=int(opts.test_workers), drop_last=False) if opts.n_images is None: opts.n_images = len(dataset) # get the image corresponding to the latent average avg_image = get_average_image(net1, opts) resize_amount = (256, 256) if opts.resize_outputs else (opts.output_size, opts.output_size) global_i = 0 global_time = [] for input_batch in tqdm(dataloader): if global_i >= opts.n_images: break with torch.no_grad(): input_cuda = input_batch.cuda().float() tic = time.time() result_batch = run_on_batch(input_cuda, net1, net2, opts, avg_image) toc = time.time() global_time.append(toc - tic) for i in range(input_batch.shape[0]): results = [ tensor2im(result_batch[i][iter_idx]) for iter_idx in range(opts.n_iters_per_batch + 1) ] im_path = dataset.paths[global_i] input_im = tensor2im(input_batch[i]) # save step-by-step results side-by-side res = np.array(results[0].resize(resize_amount)) for idx, result in enumerate(results[1:]): res = np.concatenate( [res, np.array(result.resize(resize_amount))], axis=1) res = np.concatenate([res, input_im.resize(resize_amount)], axis=1) Image.fromarray(res).save( os.path.join(out_path_results, os.path.basename(im_path))) global_i += 1 stats_path = os.path.join(opts.exp_dir, 'stats.txt') result_str = 'Runtime {:.4f}+-{:.4f}'.format(np.mean(global_time), np.std(global_time)) print(result_str) with open(stats_path, 'w') as f: f.write(result_str)
def run(): test_opts = TestOptions().parse() out_path_results = os.path.join(test_opts.exp_dir, 'inference_results') out_path_coupled = os.path.join(test_opts.exp_dir, 'inference_coupled') os.makedirs(out_path_results, exist_ok=True) os.makedirs(out_path_coupled, exist_ok=True) # update test options with options used during training ckpt = torch.load(test_opts.checkpoint_path, map_location='cpu') opts = ckpt['opts'] opts.update(vars(test_opts)) opts = Namespace(**opts) net = pSp(opts) net.eval() net.cuda() age_transformers = [ AgeTransformer(target_age=age) for age in opts.target_age.split(',') ] print(f'Loading dataset for {opts.dataset_type}') dataset_args = data_configs.DATASETS[opts.dataset_type] transforms_dict = dataset_args['transforms'](opts).get_transforms() dataset = InferenceDataset( root=opts.data_path, transform=transforms_dict['transform_inference'], opts=opts) dataloader = DataLoader(dataset, batch_size=opts.test_batch_size, shuffle=False, num_workers=int(opts.test_workers), drop_last=False) if opts.n_images is None: opts.n_images = len(dataset) global_time = [] for age_transformer in age_transformers: print(f"Running on target age: {age_transformer.target_age}") global_i = 0 for input_batch in tqdm(dataloader): if global_i >= opts.n_images: break with torch.no_grad(): input_age_batch = [ age_transformer(img.cpu()).to('cuda') for img in input_batch ] input_age_batch = torch.stack(input_age_batch) input_cuda = input_age_batch.cuda().float() tic = time.time() result_batch = run_on_batch(input_cuda, net, opts) toc = time.time() global_time.append(toc - tic) for i in range(len(input_batch)): result = tensor2im(result_batch[i]) im_path = dataset.paths[global_i] if opts.couple_outputs or global_i % 100 == 0: input_im = log_image(input_batch[i], opts) resize_amount = ( 256, 256) if opts.resize_outputs else (1024, 1024) res = np.concatenate([ np.array(input_im.resize(resize_amount)), np.array(result.resize(resize_amount)) ], axis=1) age_out_path_coupled = os.path.join( out_path_coupled, age_transformer.target_age) os.makedirs(age_out_path_coupled, exist_ok=True) Image.fromarray(res).save( os.path.join(age_out_path_coupled, os.path.basename(im_path))) age_out_path_results = os.path.join( out_path_results, age_transformer.target_age) os.makedirs(age_out_path_results, exist_ok=True) image_name = os.path.basename(im_path) im_save_path = os.path.join(age_out_path_results, image_name) Image.fromarray(np.array( result.resize(resize_amount))).save(im_save_path) global_i += 1 stats_path = os.path.join(opts.exp_dir, 'stats.txt') result_str = 'Runtime {:.4f}+-{:.4f}'.format(np.mean(global_time), np.std(global_time)) print(result_str) with open(stats_path, 'w') as f: f.write(result_str)
def run(): test_opts = TestOptions().parse() out_path_results = os.path.join(test_opts.exp_dir, 'reference_guided_inference') os.makedirs(out_path_results, exist_ok=True) # update test options with options used during training ckpt = torch.load(test_opts.checkpoint_path, map_location='cpu') opts = ckpt['opts'] opts.update(vars(test_opts)) opts = Namespace(**opts) net = pSp(opts) net.eval() net.cuda() age_transformers = [ AgeTransformer(target_age=age) for age in opts.target_age.split(',') ] print(f'Loading dataset for {opts.dataset_type}') dataset_args = data_configs.DATASETS[opts.dataset_type] transforms_dict = dataset_args['transforms'](opts).get_transforms() source_dataset = InferenceDataset( root=opts.data_path, transform=transforms_dict['transform_inference'], opts=opts) source_dataloader = DataLoader(source_dataset, batch_size=opts.test_batch_size, shuffle=False, num_workers=int(opts.test_workers), drop_last=False) ref_dataset = InferenceDataset( paths_list=opts.ref_images_paths_file, transform=transforms_dict['transform_inference'], opts=opts) ref_dataloader = DataLoader(ref_dataset, batch_size=1, shuffle=False, num_workers=1, drop_last=False) if opts.n_images is None: opts.n_images = len(source_dataset) for age_transformer in age_transformers: target_age = age_transformer.target_age print(f"Running on target age: {target_age}") age_save_path = os.path.join(out_path_results, str(target_age)) os.makedirs(age_save_path, exist_ok=True) global_i = 0 for i, source_batch in enumerate(tqdm(source_dataloader)): if global_i >= opts.n_images: break results_per_source = {idx: [] for idx in range(len(source_batch))} with torch.no_grad(): for ref_batch in ref_dataloader: source_batch = source_batch.cuda().float() ref_batch = ref_batch.cuda().float() source_input_age_batch = [ age_transformer(img.cpu()).to('cuda') for img in source_batch ] source_input_age_batch = torch.stack( source_input_age_batch) # compute w+ of ref images to be injected for style-mixing ref_latents = net.pretrained_encoder( ref_batch) + net.latent_avg # run age transformation on source images with style-mixing res_batch_mixed = run_on_batch( source_input_age_batch, net, opts, latent_to_inject=ref_latents) # store results for idx in range(len(source_batch)): results_per_source[idx].append( [ref_batch[0], res_batch_mixed[idx]]) # save results resize_amount = (256, 256) if opts.resize_outputs else (1024, 1024) for image_idx, image_results in results_per_source.items(): input_im_path = source_dataset.paths[global_i] image = source_batch[image_idx] input_image = log_image(image, opts) # initialize results image ref_inputs = np.zeros_like( input_image.resize(resize_amount)) mixing_results = np.array( input_image.resize(resize_amount)) for ref_idx in range(len(image_results)): ref_input, mixing_result = image_results[ref_idx] ref_input = log_image(ref_input, opts) mixing_result = log_image(mixing_result, opts) # append current results ref_inputs = np.concatenate([ ref_inputs, np.array(ref_input.resize(resize_amount)) ], axis=1) mixing_results = np.concatenate([ mixing_results, np.array(mixing_result.resize(resize_amount)) ], axis=1) res = np.concatenate([ref_inputs, mixing_results], axis=0) save_path = os.path.join(age_save_path, os.path.basename(input_im_path)) Image.fromarray(res).save(save_path) global_i += 1
def run(): """ This script can be used to perform inversion and editing. Please note that this script supports editing using only the ReStyle-e4e model and currently supports editing using three edit directions found using InterFaceGAN (age, smile, and pose) on the faces domain. For performing the edits please provide the arguments `--edit_directions` and `--factor_ranges`. For example, setting these values to be `--edit_directions=age,smile,pose` and `--factor_ranges=5,5,5` will use a lambda range between -5 and 5 for each of the attributes. These should be comma-separated lists of the same length. You may get better results by playing around with the factor ranges for each edit. """ test_opts = TestOptions().parse() out_path_results = os.path.join(test_opts.exp_dir, 'editing_results') out_path_coupled = os.path.join(test_opts.exp_dir, 'editing_coupled') os.makedirs(out_path_results, exist_ok=True) os.makedirs(out_path_coupled, exist_ok=True) # update test options with options used during training ckpt = torch.load(test_opts.checkpoint_path, map_location='cpu') opts = ckpt['opts'] opts.update(vars(test_opts)) opts = Namespace(**opts) net = e4e(opts) net.eval() net.cuda() print('Loading dataset for {}'.format(opts.dataset_type)) if opts.dataset_type != "ffhq_encode": raise ValueError( "Editing script only supports edits on the faces domain!") dataset_args = data_configs.DATASETS[opts.dataset_type] transforms_dict = dataset_args['transforms'](opts).get_transforms() dataset = InferenceDataset( root=opts.data_path, transform=transforms_dict['transform_inference'], opts=opts) dataloader = DataLoader(dataset, batch_size=opts.test_batch_size, shuffle=False, num_workers=int(opts.test_workers), drop_last=False) if opts.n_images is None: opts.n_images = len(dataset) latent_editor = LatentEditor(net.decoder) opts.edit_directions = opts.edit_directions.split(',') opts.factor_ranges = [ int(factor) for factor in opts.factor_ranges.split(',') ] if len(opts.edit_directions) != len(opts.factor_ranges): raise ValueError( "Invalid edit directions and factor ranges. Please provide a single factor range for each" f"edit direction. Given: {opts.edit_directions} and {opts.factor_ranges}" ) avg_image = get_average_image(net, opts) global_i = 0 global_time = [] for input_batch in tqdm(dataloader): if global_i >= opts.n_images: break with torch.no_grad(): input_cuda = input_batch.cuda().float() tic = time.time() result_batch = edit_batch(input_cuda, net, avg_image, latent_editor, opts) toc = time.time() global_time.append(toc - tic) resize_amount = (256, 256) if opts.resize_outputs else (opts.output_size, opts.output_size) for i in range(input_batch.shape[0]): im_path = dataset.paths[global_i] results = result_batch[i] inversion = results.pop('inversion') input_im = tensor2im(input_batch[i]) all_edit_results = [] for edit_name, edit_res in results.items(): res = np.array( input_im.resize(resize_amount)) # set the input image res = np.concatenate( [res, np.array(inversion.resize(resize_amount))], axis=1) # set the inversion for result in edit_res: res = np.concatenate( [res, np.array(result.resize(resize_amount))], axis=1) res_im = Image.fromarray(res) all_edit_results.append(res_im) edit_save_dir = os.path.join(out_path_results, edit_name) os.makedirs(edit_save_dir, exist_ok=True) res_im.save( os.path.join(edit_save_dir, os.path.basename(im_path))) # save final concatenated result if all factor ranges are equal if opts.factor_ranges.count(opts.factor_ranges[0]) == len( opts.factor_ranges): coupled_res = np.concatenate(all_edit_results, axis=0) im_save_path = os.path.join(out_path_coupled, os.path.basename(im_path)) Image.fromarray(coupled_res).save(im_save_path) global_i += 1 stats_path = os.path.join(opts.exp_dir, 'stats.txt') result_str = 'Runtime {:.4f}+-{:.4f}'.format(np.mean(global_time), np.std(global_time)) print(result_str) with open(stats_path, 'w') as f: f.write(result_str)