Hands-on Python Guide to Style-based Age Manipulation (SAM) Technique
Table of contentsIntroductionWhat is age transformation?Overview of SAMWhy the name ‘style-based Age Manipulation’?Practical implementation of SAMPrerequisitesPre-trained modelDemo codeOutputGoogle colab notebooks References Introduction Style-based Age Manipulation (SAM) is a method used to perform fine-grained age transformation in digital image processing and computer vision tasks using a single facial image as an input. It was introduced by Yuval Alalu, Or Patashnik and Daneil Cohen-Or of Tel-Aviv University in February 2021 (research paper). This article gives an overview of SAM along with its demonstration using Python code. Before going into the details of SAM, let us first understand what is meant by the ‘age transformation’ task. What is age transformation? Age transformation is a process of representing the change in a person’s appearance across different ages while preserving his identity. In order to model such a process over a single input facial image, the change in head shape and texture must be captured while the identity and other key facial attributes of the input face must be preserved. The complexity of the task increases when modelling lifelong ageing where significant age modification is desired (e.g. from ages 10 to 90 years). To avoid explicit modelling of age transformation, Generative Adversarial Networks (GANs) are largely used for generating images in a data-driven manner, especially on facial images. Overview of SAM SAM is a method for learning a conditional image generation function which can capture the desired change in age but preserve the facial identity. It is an image-to-image translation method i.e. translates a given image of a source domain to a corresponding image of a target domain. It couples the expressiveness of a pre-trained, fixed StyleGAN generator with an encoder architecture. The encoder directly encodes an input facial image into a series of style vectors subject to the desired age shift. These style vectors are then fed into unconditional StyleGAN. The output of StyleGAN represents the desired age transformation. The use of StyleGAN enables leveraging its ability to achieve excellent image quality. A pre-trained, the fixed age regression model is used for generating the latent codes corresponding. The continuous ageing process is formulated as a regression task between the input age and desired target age, providing fine-grained control over image generated by GAN. Why the name ‘style-based Age Manipulation’? Since age transformation in the SAM method is controlled through the intermediate style representations learned by the StyleGAN, it has been named as “Style-based” Age Manipulation. In other words, it does age transformation based on the style of an input facial image. Practical implementation of SAM Prerequisites Python 3Linux or macOSNVIDIA GPU + CUDA CuDNN Pre-trained model The SAM model pre-trained on the FFHQ dataset can be downloaded from here. If you wish to train a SAM model from scratch, the following auxiliary models can be used: pSp Encoder : taken from pixel2style2pixel (an image-to-image translation framework) trained on the FFHQ dataset for StyleGAN inversionFFHQ StyleGAN : StyleGAN model pre-trained on FFHQ taken from rosinality’s repository with 1024×1024 output resolutionIR-SE50 Model : Pretrained IR-SE50 model taken from TreB1eN (used for identity loss)VGG Age Classifier : VGG age classifier from DEX and fine-tuned on the FFHQ-Aging dataset (used for aging loss) Demo code Import the os module to interact with the underlying Operating System import os Define the code directory name CODE_DIR = ‘SAM’ Clone the GitHub repository !git clone https://github.com/yuval-alaluf/SAM.git $CODE_DIR Download Ninja (a small build system focussed on speed) !wget https://github.com/ninja-build/ninja/releases/download/v1.8.2/ ninja-linux.zip Unzip ninja-linux.zip file !sudo unzip ninja-linux.zip -d /usr/local/bin/ Change from default to alternative Python version !sudo update-alternatives --install /usr/bin/ninja ninja /usr/local/bin/ninja 1 --force Change the current working directory to ‘SAM’ os.chdir(f'./{CODE_DIR}') Import the required standard libraries from argparse import Namespace import sys import pprint import numpy as np from PIL import Image import torch import torchvision.transforms as transforms Modify sys.path sys.path.append(".") sys.path.append("..") Import the AgeTransformer class from datasets.augmentations import AgeTransformer Import tensor2im method for tensor-to-image conversion from utils.common import tensor2im Import the pSp class from models.psp import pSp Define experiment type EXPERIMENT_TYPE = 'ffhq_aging' Get wget download command for downloading the desired model and save to directory ../pretrained_models def get_download_model_command(file_id, file_name): current_directory = os.getcwd() save_path = os.path.join(os.path.dirname(current_directory), "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 Define model path MODEL_PATHS = { "ffhq_aging": {"id": "1XyumF6_fdAxFmxpFcmPf-q84LU_22EMC", "name": "sam_ffhq_aging.pt"} } Initialize model path and download command path = MODEL_PATHS[EXPERIMENT_TYPE] download_command = get_download_model_command(file_id=path["id"], file_name=path["name"]) !wget {download_command} Define experiment arguments EXPERIMENT_DATA_ARGS = { "ffhq_aging": { "model_path": "../pretrained_models/sam_ffhq_aging.pt", "image_path": "notebooks/images/1287.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] Initialize the model path model_path = EXPERIMENT_ARGS['model_path'] Load the PyTorch model using torch.load() ckpt = torch.load(model_path, map_location='cpu') Have a look at the model training options opts = ckpt['opts'] pprint.pprint(opts) Update the training options opts['checkpoint_path'] = model_path Load SAM model.Use CUDA tensor types which implement functions like CPU tensors but using GPUs. opts = Namespace(**opts) net = pSp(opts) net.eval() net.cuda() //torch.cuda print('Model successfully loaded') Open the input image and resize it for display image_path = EXPERIMENT_DATA_ARGS[EXPERIMENT_TYPE]["image_path"] original_image = Image.open(image_path).convert("RGB") //Image.open() On executing the above line of code, you will see the input facial image which is as follows: Download Dlib model shape_predictor_68_face_landmarks.dat.bz2 which has been trained on the ibug 300-W dataset. !wget http://dlib.net/files/shape_predictor_68_face_landmarks.dat.bz2 Use bzip2 command for file compression and decompression !bzip2 -dk shape_predictor_68_face_landmarks.dat.bz2 Define function for face alignment def run_alignment(image_path): import dlib from scripts.align_all_parallel import align_face predictor = dlib.shape_predictor("shape_predictor _68_face_landmarks.dat") aligned_image = align_face(filepath=image_path, predictor=predictor) print("Aligned image has shape: {}".format(aligned_image.size)) return aligned_image Align the input face aligned_image = run_alignment(image_path) Resize the aligned face aligned_image.resize((256, 256)) Initialize variables for image transformation img_transforms = EXPERIMENT_ARGS['transform'] input_image = img_transforms(aligned_image) Run the image on multiple target ages target_ages = [0, 10, 20, 30, 40, 50, 60, 70, 80, 90, 100] //0-100 years’ age shift age_transformers = [AgeTransformer(target_age=age) for age in target_ages] Define function to create tensors def run_on_batch(inputs, net): result_batch = net(inputs.to("cuda").float(), randomize_noise=False, resize=False) return result_batch For each age transformed age, concatenate the results to display them side-by-side results = np.array(aligned_image.resize((1024, 1024))) for age_transformer in age_transformers: print(f"Running on target age: {age_transformer.target_age}") with torch.no_grad(): input_image_age = [age_transformer(input_image.cpu()).to('cuda')] input_image_age = torch.stack(input_image_age) //torch.stack result_tensor = run_on_batch(input_image_age, net)[0] result_image = tensor2im(result_tensor) results = np.concatenate([results, result_image], axis=1) Construct image memory from numerical array representation of the output using Image.fromarray() results = Image.fromarray(results) Display the output results Output Source: https://github.com/yuval-alaluf/SAM/blob/master/notebooks/inference_playground.ipynb The above code can also be experimented for different images by changing the image path appropriately. We have tried the code on two more images, the results of which are as follows: Input image 2: Output: Input image 3: Output: Google colab notebooks Demo1Demo2Demo3 References Do you want to have a deep-rooted understanding of the SAM technique? Refer to the following sources: SAM research paperGitHub repositoryStyleGAN research paper
Introduction Style-based Age Manipulation (SAM) is a method used to perform fine-grained age transformation in digital image processing and computer vision tasks using a single facial image as an input. It was introduced by Yuval Alalu, Or Patashnik and Daneil Cohen-Or of Tel-Aviv University in February 2021 (research paper). This article gives an overview of […]
SAM is a technique used to perform fine-grained age transformation using a single facial image as an input.
["face recognition"]
["CUDA", "NumPy", "TPU", "PyTorch", "AI", "Transformers", "computer vision", "RAG", "Colab", "Ray", "face recognition"]
["AI", "computer vision", "Ray", "PyTorch", "Transformers", "Colab", "NumPy", "RAG", "TPU", "CUDA"]
A zealous learner aspiring to advance in the domain of AI/ML. Eager to grasp emerging techniques to get insights from data and hence explore realistic Data Science applications as well.