| import os | |
| from random import choice | |
| from string import ascii_uppercase | |
| from PIL import Image | |
| from tqdm import tqdm | |
| from scripts.latent_editor_wrapper import LatentEditorWrapper | |
| from evaluation.experiment_setting_creator import ExperimentRunner | |
| import torch | |
| from configs import paths_config, hyperparameters, evaluation_config | |
| from utils.log_utils import save_concat_image, save_single_image | |
| from utils.models_utils import load_tuned_G | |
| class EditComparison: | |
| def __init__(self, save_single_images, save_concatenated_images, run_id): | |
| self.run_id = run_id | |
| self.experiment_creator = ExperimentRunner(run_id) | |
| self.save_single_images = save_single_images | |
| self.save_concatenated_images = save_concatenated_images | |
| self.latent_editor = LatentEditorWrapper() | |
| def save_reconstruction_images(self, image_latents, new_inv_image_latent, new_G, target_image): | |
| if self.save_concatenated_images: | |
| save_concat_image(self.concat_base_dir, image_latents, new_inv_image_latent, new_G, | |
| self.experiment_creator.old_G, | |
| 'rec', | |
| target_image) | |
| if self.save_single_images: | |
| save_single_image(self.single_base_dir, new_inv_image_latent, new_G, 'rec') | |
| target_image.save(f'{self.single_base_dir}/Original.jpg') | |
| def create_output_dirs(self, full_image_name): | |
| output_base_dir_path = f'{paths_config.experiments_output_dir}/{paths_config.input_data_id}/{self.run_id}/{full_image_name}' | |
| os.makedirs(output_base_dir_path, exist_ok=True) | |
| self.concat_base_dir = f'{output_base_dir_path}/concat_images' | |
| self.single_base_dir = f'{output_base_dir_path}/single_images' | |
| os.makedirs(self.concat_base_dir, exist_ok=True) | |
| os.makedirs(self.single_base_dir, exist_ok=True) | |
| def get_image_latent_codes(self, image_name): | |
| image_latents = [] | |
| for method in evaluation_config.evaluated_methods: | |
| if method == 'SG2': | |
| image_latents.append(torch.load( | |
| f'{paths_config.embedding_base_dir}/{paths_config.input_data_id}/' | |
| f'{paths_config.pti_results_keyword}/{image_name}/0.pt')) | |
| else: | |
| image_latents.append(torch.load( | |
| f'{paths_config.embedding_base_dir}/{paths_config.input_data_id}/{method}/{image_name}/0.pt')) | |
| new_inv_image_latent = torch.load( | |
| f'{paths_config.embedding_base_dir}/{paths_config.input_data_id}/{paths_config.pti_results_keyword}/{image_name}/0.pt') | |
| return image_latents, new_inv_image_latent | |
| def save_interfacegan_edits(self, image_latents, new_inv_image_latent, interfacegan_factors, new_G, target_image): | |
| new_w_inv_edits = self.latent_editor.get_single_interface_gan_edits(new_inv_image_latent, | |
| interfacegan_factors) | |
| inv_edits = [] | |
| for latent in image_latents: | |
| inv_edits.append(self.latent_editor.get_single_interface_gan_edits(latent, interfacegan_factors)) | |
| for direction, edits in new_w_inv_edits.items(): | |
| for factor, edit_tensor in edits.items(): | |
| if self.save_concatenated_images: | |
| save_concat_image(self.concat_base_dir, [edits[direction][factor] for edits in inv_edits], | |
| new_w_inv_edits[direction][factor], | |
| new_G, | |
| self.experiment_creator.old_G, | |
| f'{direction}_{factor}', target_image) | |
| if self.save_single_images: | |
| save_single_image(self.single_base_dir, new_w_inv_edits[direction][factor], new_G, | |
| f'{direction}_{factor}') | |
| def save_ganspace_edits(self, image_latents, new_inv_image_latent, factors, new_G, target_image): | |
| new_w_inv_edits = self.latent_editor.get_single_ganspace_edits(new_inv_image_latent, factors) | |
| inv_edits = [] | |
| for latent in image_latents: | |
| inv_edits.append(self.latent_editor.get_single_ganspace_edits(latent, factors)) | |
| for idx in range(len(new_w_inv_edits)): | |
| if self.save_concatenated_images: | |
| save_concat_image(self.concat_base_dir, [edit[idx] for edit in inv_edits], new_w_inv_edits[idx], | |
| new_G, | |
| self.experiment_creator.old_G, | |
| f'ganspace_{idx}', target_image) | |
| if self.save_single_images: | |
| save_single_image(self.single_base_dir, new_w_inv_edits[idx], new_G, | |
| f'ganspace_{idx}') | |
| def run_experiment(self, run_pt, create_other_latents, use_multi_id_training, use_wandb=False): | |
| images_counter = 0 | |
| new_G = None | |
| interfacegan_factors = [val / 2 for val in range(-6, 7) if val != 0] | |
| ganspace_factors = range(-20, 25, 5) | |
| self.experiment_creator.run_experiment(run_pt, create_other_latents, use_multi_id_training, use_wandb) | |
| if use_multi_id_training: | |
| new_G = load_tuned_G(self.run_id, paths_config.multi_id_model_type) | |
| for idx, image_path in tqdm(enumerate(self.experiment_creator.images_paths), | |
| total=len(self.experiment_creator.images_paths)): | |
| if images_counter >= hyperparameters.max_images_to_invert: | |
| break | |
| image_name = image_path.split('.')[0].split('/')[-1] | |
| target_image = Image.open(self.experiment_creator.target_paths[idx]) | |
| if not use_multi_id_training: | |
| new_G = load_tuned_G(self.run_id, image_name) | |
| image_latents, new_inv_image_latent = self.get_image_latent_codes(image_name) | |
| self.create_output_dirs(image_name) | |
| self.save_reconstruction_images(image_latents, new_inv_image_latent, new_G, target_image) | |
| self.save_interfacegan_edits(image_latents, new_inv_image_latent, interfacegan_factors, new_G, target_image) | |
| self.save_ganspace_edits(image_latents, new_inv_image_latent, ganspace_factors, new_G, target_image) | |
| target_image.close() | |
| torch.cuda.empty_cache() | |
| images_counter += 1 | |
| def run_pti_and_full_edit(iid): | |
| evaluation_config.evaluated_methods = ['SG2Plus', 'e4e', 'SG2'] | |
| edit_figure_creator = EditComparison(save_single_images=True, save_concatenated_images=True, | |
| run_id=f'{paths_config.input_data_id}_pti_full_edit_{iid}') | |
| edit_figure_creator.run_experiment(True, True, use_multi_id_training=False, use_wandb=False) | |
| def pti_no_comparison(iid): | |
| evaluation_config.evaluated_methods = [] | |
| edit_figure_creator = EditComparison(save_single_images=True, save_concatenated_images=True, | |
| run_id=f'{paths_config.input_data_id}_pti_no_comparison_{iid}') | |
| edit_figure_creator.run_experiment(True, False, use_multi_id_training=False, use_wandb=False) | |
| def edits_for_existed_experiment(run_id): | |
| evaluation_config.evaluated_methods = ['SG2Plus', 'e4e', 'SG2'] | |
| edit_figure_creator = EditComparison(save_single_images=True, save_concatenated_images=True, | |
| run_id=run_id) | |
| edit_figure_creator.run_experiment(False, True, use_multi_id_training=False, use_wandb=False) | |
| if __name__ == '__main__': | |
| iid = ''.join(choice(ascii_uppercase) for i in range(7)) | |
| pti_no_comparison(iid) | |