A Residual-Based StyleGAN Encoder Via Iterative Refinement
A Generative Adversarial Networks, in short, GAN is an approach to generative modeling using deep neural networks methods such as convolutional neural networks. Those are effective in generating high-quality images. Generative modeling is an unsupervised task of machine learning that involves automatically discovering and learning the regularities and patterns in input data so that the model can be used to output new examples that could be more similar to input data. The Style generative adversarial network, in short StyleGAN, is an extension of GAN architecture that posses large changes to the generator model, including the use of mapping network to map points in latent space(latent space nothing but the representation of compressed data) to an intermediate latent space, the use of intermediate latent space to control the style at each point in the generator model and introduction of noise as a source of variation at each point in the generator model. Thus, the model is capable of generating high-quality photos of faces and offers control over the style of the generated images. In this post, we are going to discuss the Restyle Encoder. Architecture and Method: The author of the Restyle Encoder proposes a novel architecture, a fast iterative method of image inversion into the latent space of pretrained StyleGAN generators that archives state of the art quality at lower inference time. The core idea of this innovation is to start from the average latent vector in W+ and predict an offset that would make the generated image look more like the target image, then repeat this step with a new image and latent vector as a starting point. With this approach, a perfect inversion can be achieved in about just ten steps. The below figure is taken from the official research paper, which clearly explains the architecture with example. For detailed information, you can check the paper here. Source: Architecture of Restyle Encoder Official Code Implementation The Restyle encoder offers various experiments, i.e. it can infer human faces, cars, horses, wild animals, churches and toonify the input image. Just we need to select on which we want to perform inference. Note: The below code is taken from official code resource. Clone the repository to colab: import os os.chdir('/content') CODE_DIR = 'restyle-encoder' !git clone https://github.com/yuval-alaluf/restyle-encoder.git $CODE_DIR !wget https://github.com/ninja-build/ninja/releases/download/v1.8.2/ninja-linux.zip !sudo unzip ninja-linux.zip -d /usr/local/bin/ !sudo update-alternatives --install /usr/bin/ninja ninja /usr/local/bin/ninja 1 --force Import all dependencies: from argparse import Namespace import time import os import sys import pprint import numpy as np from PIL import Image import torch import torchvision.transforms as transforms sys.path.append(".") sys.path.append("..") from utils.common import tensor2im from models.psp import pSp from models.e4e import e4e %load_ext autoreload %autoreload 2 Select the type of experiment you want to perform as listed below #@title Select which experiment you wish to perform inference on: { run: "auto" } experiment_type = 'ffhq_encode' #@param ['ffhq_encode', 'cars_encode', 'church_encode', 'horse_encode', 'afhq_wild_encode', 'toonify'] Download the pretrained models The author has provided a pre-trained model for each example listed above. The below code shows the command to download the pre-trained model for the desired experiment. def get_download_model_command(file_id, file_name): """ Get wget download command for downloading the desired model and save to directory ../pretrained_models. """ current_directory = os.getcwd() save_path = os.path.join(os.path.dirname(current_directory), CODE_DIR, "pretrained_models") if not os.path.exists(save_path): os.makedirs(save_path) url = r"""wget --load-cookies /tmp/cookies.txt "https://docs.google.com/uc?export=download&confirm=$(wget --quiet --save-cookies /tmp/cookies.txt --keep-session-cookies --no-check-certificate 'https://docs.google.com/uc?export=download&id={FILE_ID}' -O- | sed -rn 's/.*confirm=([0-9A-Za-z_]+).*/\1\n/p')&id={FILE_ID}" -O {SAVE_PATH}/{FILE_NAME} && rm -rf /tmp/cookies.txt""".format(FILE_ID=file_id, FILE_NAME=file_name, SAVE_PATH=save_path) return url MODEL_PATHS = { "ffhq_encode": {"id": "1sw6I2lRIB0MpuJkpc8F5BJiSZrc0hjfE", "name": "restyle_psp_ffhq_encode.pt"}, "cars_encode": {"id": "1zJHqHRQ8NOnVohVVCGbeYMMr6PDhRpPR", "name": "restyle_psp_cars_encode.pt"}, "church_encode": {"id": "1bcxx7mw-1z7dzbJI_z7oGpWG1oQAvMaD", "name": "restyle_psp_church_encode.pt"}, "horse_encode": {"id": "19_sUpTYtJmhSAolKLm3VgI-ptYqd-hgY", "name": "restyle_e4e_horse_encode.pt"}, "afhq_wild_encode": {"id": "1GyFXVTNDUw3IIGHmGS71ChhJ1Rmslhk7", "name": "restyle_psp_afhq_wild_encode.pt"}, "toonify": {"id": "1GtudVDig59d4HJ_8bGEniz5huaTSGO_0", "name": "restyle_psp_toonify.pt"} } path = MODEL_PATHS[experiment_type] download_command = get_download_model_command(file_id=path["id"], file_name=path["name"]) Define the inference parameters Inference parameters are nothing but the loading of the particular model and some desired transformation, defining path to pre-trained model and path of images. EXPERIMENT_DATA_ARGS = { "ffhq_encode": { "model_path": "pretrained_models/restyle_psp_ffhq_encode.pt", "image_path": "notebooks/images/face_img.jpg", "transform": transforms.Compose([ transforms.Resize((256, 256)), transforms.ToTensor(), transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])]) }, "cars_encode": { "model_path": "pretrained_models/restyle_psp_cars_encode.pt", "image_path": "notebooks/images/car_img.jpg", "transform": transforms.Compose([ transforms.Resize((192, 256)), transforms.ToTensor(), transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])]) }, "church_encode": { "model_path": "pretrained_models/restyle_psp_church_encode.pt", "image_path": "notebooks/images/church_img.jpg", "transform": transforms.Compose([ transforms.Resize((256, 256)), transforms.ToTensor(), transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])]) }, "horse_encode": { "model_path": "pretrained_models/restyle_e4e_horse_encode.pt", "image_path": "notebooks/images/horse_img.jpg", "transform": transforms.Compose([ transforms.Resize((256, 256)), transforms.ToTensor(), transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])]) }, "afhq_wild_encode": { "model_path": "pretrained_models/restyle_psp_afhq_wild_encode.pt", "image_path": "notebooks/images/afhq_wild_img.jpg", "transform": transforms.Compose([ transforms.Resize((256, 256)), transforms.ToTensor(), transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])]) }, "toonify": { "model_path": "pretrained_models/restyle_psp_toonify.pt", "image_path": "notebooks/images/toonify_img.jpg", "transform": transforms.Compose([ transforms.Resize((256, 256)), transforms.ToTensor(), transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])]) }, } EXPERIMENT_ARGS = EXPERIMENT_DATA_ARGS[experiment_type] Load the pre-trained model model_path = EXPERIMENT_ARGS['model_path'] ckpt = torch.load(model_path, map_location='cpu') opts = ckpt['opts'] pprint.pprint(opts) # update the training options opts['checkpoint_path'] = model_path opts = Namespace(**opts) if experiment_type == 'horse_encode': net = e4e(opts) else: net = pSp(opts) net.eval() net.cuda() print('Model successfully loaded!') Visualise the image for face experiment image_path = EXPERIMENT_DATA_ARGS[experiment_type]["image_path"] original_image = Image.open(image_path).convert("RGB") Perform the Inference Generate the image corresponding to the average latent code because it will be used in an iterative refinement process. img_transforms = EXPERIMENT_ARGS['transform'] transformed_image = img_transforms(input_image) def get_avg_image(net): 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 experiment_type == "cars_encode": avg_image = avg_image[:, 32:224, :] return avg_image Now run the inference. opts.n_iters_per_batch = 5 opts.resize_outputs = False  # generate outputs at full resolution from utils.inference_utils import run_on_batch with torch.no_grad(): avg_image = get_avg_image(net) tic = time.time() result_batch, result_latents = run_on_batch(transformed_image.unsqueeze(0).cuda(), net, opts, avg_image) toc = time.time() print('Inference took {:.4f} seconds.'.format(toc - tic)) Visualise the output result side by side 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) def get_coupled_results(result_batch, transformed_image): """ Visualize output images from left to right (the input image is on the right) """ result_tensors = result_batch[0]  # there's one image in our batch result_images = [tensor2im(result_tensors[iter_idx]) for iter_idx in range(opts.n_iters_per_batch)] input_im = tensor2im(transformed_image) res = np.array(result_images[0].resize(resize_amount)) for idx, result in enumerate(result_images[1:]): res = np.concatenate([res, np.array(result.resize(resize_amount))], axis=1) res = np.concatenate([res, input_im.resize(resize_amount)], axis=1) res = Image.fromarray(res) return res res = get_coupled_results(result_batch, transformed_image) res Note that the original image is on the rightmost side. Just by changing the experiment type earlier in the code, we can perform toonification of the given image shows the result as below; Conclusion In this article, we have seen the basic understanding of GAN and StyleGAN. Later on, the novel architecture called Restyle Encoder, which performs inference on various categories, shows the two inferences: normal facial inference and toonification, which is the coolest thing. The quality of the output images can be improved by changing the no of interactions per batch. References Link for colab implementationOfficial GitHub RepositoryResearch paper
The Style generative adversarial network, in short StyleGAN, is an extension of GAN architecture that posses large…
A Generative Adversarial Networks, in short, GAN is an approach to generative modeling using deep neural networks methods such as...
["AI Tool", "GAN", "image synthesis", "StyleGAN"]
["CUDA", "NumPy", "machine learning", "TPU", "AI", "neural network", "StyleGAN", "Colab", "Ray", "RAG", "image synthesis", "GAN", "AI Tool", "R"]
["AI", "machine learning", "neural network", "Ray", "Colab", "NumPy", "RAG", "TPU", "CUDA", "R"]
Vijaysinh is an enthusiast in machine learning and deep learning. He is skilled in ML algorithms, data manipulation, handling and visualization, model building.