lfs

.gitattributes

@@ -0,0 +1 @@
+*.dat filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,2 @@
+docs/
+output.*

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2021 Yuval Alaluf
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -0,0 +1,308 @@
+# Only a Matter of Style: Age Transformation Using a Style-Based Regression Model (SIGGRAPH 2021)
+
+> The  task of age transformation illustrates the change of an individual's appearance over time. Accurately modeling this complex transformation over an input facial image is extremely challenging as it requires making convincing and possibly large changes to facial features and head shape, while still preserving the input identity. In this work, we present an image-to-image translation method that learns to directly encode real facial images into the latent space of a pre-trained unconditional GAN (e.g., StyleGAN) subject to a given aging shift. We employ a pre-trained age regression network used to explicitly guide the encoder to generate the latent codes corresponding to the desired age. In this formulation, our method approaches the continuous aging process as a regression task between the input age and desired target age, providing fine-grained control on the generated image. Moreover, unlike other approaches that operate solely in the latent space using a prior on the path controlling age, our method learns a more disentangled, non-linear path. We demonstrate that the end-to-end nature of our approach, coupled with the rich semantic latent space of StyleGAN, allows for further editing of the generated images. Qualitative and quantitative evaluations show the advantages of our method compared to state-of-the-art approaches.
+
+<a href="https://arxiv.org/abs/2102.02754"><img src="https://img.shields.io/badge/arXiv-2008.00951-b31b1b.svg" height=22.5></a>
+<a href="https://opensource.org/licenses/MIT"><img src="https://img.shields.io/badge/License-MIT-yellow.svg" height=22.5></a>
+
+<a href="https://www.youtube.com/watch?v=zDTUbtmUbG8"><img src="https://img.shields.io/static/v1?label=Two Minute Papers&message=SAM Video&color=red" height=22.5></a>  
+<a href="https://youtu.be/X_pYC_LtBFw"><img src="https://img.shields.io/static/v1?label=SIGGRAPH 2021 &message=5 Minute Video&color=red" height=22.5></a>  
+<a href="https://replicate.ai/yuval-alaluf/sam"><img src="https://img.shields.io/static/v1?label=Replicate&message=Demo and Docker Image&color=darkgreen" height=22.5></a>
+
+
+Inference Notebook: &nbsp;<a href="http://colab.research.google.com/github/yuval-alaluf/SAM/blob/master/notebooks/inference_playground.ipynb"><img src="https://colab.research.google.com/assets/colab-badge.svg" height=22.5></a>  
+Animation Notebook: <a href="http://colab.research.google.com/github/yuval-alaluf/SAM/blob/master/notebooks/animation_inference_playground.ipynb"><img src="https://colab.research.google.com/assets/colab-badge.svg" height=22.5></a>
+
+
+<p align="center">
+<img src="docs/teaser.jpeg" width="800px"/>
+</p>
+
+## Description   
+Official Implementation of our Style-based Age Manipulation (SAM) paper for both training and evaluation. SAM 
+allows modeling fine-grained age transformation using a single input facial image
+
+<p align="center">
+<img src="docs/2195.jpg" width="800px"/>
+<img src="docs/1936.jpg" width="800px"/>
+</p>
+
+
+## Getting Started
+### Prerequisites
+- Linux or macOS
+- NVIDIA GPU + CUDA CuDNN (CPU may be possible with some modifications, but is not inherently supported)
+- Python 3
+
+### Installation
+- Dependencies:  
+We recommend running this repository using [Anaconda](https://docs.anaconda.com/anaconda/install/). 
+All dependencies for defining the environment are provided in `environment/sam_env.yaml`.
+
+## Pretrained Models
+Please download the pretrained aging model from the following links.
+
+| Path | Description
+| :--- | :----------
+|[SAM](https://drive.google.com/file/d/1XyumF6_fdAxFmxpFcmPf-q84LU_22EMC/view?usp=sharing)  | SAM trained on the FFHQ dataset for age transformation.
+
+You can run this to download it to the right place:
+
+```
+mkdir pretrained_models
+pip install gdown
+gdown "https://drive.google.com/u/0/uc?id=1XyumF6_fdAxFmxpFcmPf-q84LU_22EMC&export=download" -O pretrained_models/sam_ffhq_aging.pt
+wget "https://github.com/italojs/facial-landmarks-recognition/raw/master/shape_predictor_68_face_landmarks.dat"
+```
+
+In addition, we provide various auxiliary models needed for training your own SAM model from scratch.  
+This includes the pretrained pSp encoder model for generating the encodings of the input image and the aging classifier 
+used to compute the aging loss during training.
+
+| Path | Description
+| :--- | :----------
+|[pSp Encoder](https://drive.google.com/file/d/1bMTNWkh5LArlaWSc_wa8VKyq2V42T2z0/view?usp=sharing) | pSp taken from [pixel2style2pixel](https://github.com/eladrich/pixel2style2pixel) trained on the FFHQ dataset for StyleGAN inversion.
+|[FFHQ StyleGAN](https://drive.google.com/file/d/1EM87UquaoQmk17Q8d5kYIAHqu0dkYqdT/view?usp=sharing) | StyleGAN model pretrained on FFHQ taken from [rosinality](https://github.com/rosinality/stylegan2-pytorch) with 1024x1024 output resolution.
+|[IR-SE50 Model](https://drive.google.com/file/d/1KW7bjndL3QG3sxBbZxreGHigcCCpsDgn/view?usp=sharing) | Pretrained IR-SE50 model taken from [TreB1eN](https://github.com/TreB1eN/InsightFace_Pytorch) for use in our ID loss during training.
+|[VGG Age Classifier](https://drive.google.com/file/d/1atzjZm_dJrCmFWCqWlyspSpr3nI6Evsh/view?usp=sharing) | VGG age classifier from DEX and fine-tuned on the FFHQ-Aging dataset for use in our aging loss
+
+By default, we assume that all auxiliary models are downloaded and saved to the directory `pretrained_models`. 
+However, you may use your own paths by changing the necessary values in `configs/path_configs.py`. 
+
+## Training
+### Preparing your Data
+Please refer to `configs/paths_config.py` to define the necessary data paths and model paths for training and inference.   
+Then, refer to `configs/data_configs.py` to define the source/target data paths for the train and test sets as well as the 
+transforms to be used for training and inference.
+    
+As an example, we can first go to `configs/paths_config.py` and define:
+``` 
+dataset_paths = {
+    'ffhq': '/path/to/ffhq/images256x256'
+    'celeba_test': '/path/to/CelebAMask-HQ/test_img',
+}
+```
+Then, in `configs/data_configs.py`, we define:
+``` 
+DATASETS = {
+	'ffhq_aging': {
+		'transforms': transforms_config.AgingTransforms,
+		'train_source_root': dataset_paths['ffhq'],
+		'train_target_root': dataset_paths['ffhq'],
+		'test_source_root': dataset_paths['celeba_test'],
+		'test_target_root': dataset_paths['celeba_test'],
+	}
+}
+``` 
+When defining the datasets for training and inference, we will use the values defined in the above dictionary.
+
+
+### Training SAM
+The main training script can be found in `scripts/train.py`.   
+Intermediate training results are saved to `opts.exp_dir`. This includes checkpoints, train outputs, and test outputs.  
+Additionally, if you have tensorboard installed, you can visualize tensorboard logs in `opts.exp_dir/logs`.
+
+Training SAM with the settings used in the paper can be done by running the following command:
+```
+python scripts/train.py \
+--dataset_type=ffhq_aging \
+--exp_dir=/path/to/experiment \
+--workers=6 \
+--batch_size=6 \
+--test_batch_size=6 \
+--test_workers=6 \
+--val_interval=2500 \
+--save_interval=10000 \
+--start_from_encoded_w_plus \
+--id_lambda=0.1 \
+--lpips_lambda=0.1 \
+--lpips_lambda_aging=0.1 \
+--lpips_lambda_crop=0.6 \
+--l2_lambda=0.25 \
+--l2_lambda_aging=0.25 \
+--l2_lambda_crop=1 \
+--w_norm_lambda=0.005 \
+--aging_lambda=5 \
+--cycle_lambda=1 \
+--input_nc=4 \
+--target_age=uniform_random \
+--use_weighted_id_loss
+```
+
+### Additional Notes
+- See `options/train_options.py` for all training-specific flags. 
+- Note that using the flag `--start_from_encoded_w_plus` requires you to specify the path to the pretrained pSp encoder.  
+    By default, this path is taken from `configs.paths_config.model_paths['pretrained_psp']`.
+- If you wish to resume from a specific checkpoint (e.g. a pretrained SAM model), you may do so using `--checkpoint_path`.
+
+
+## Notebooks 
+### Inference Notebook
+To help visualize the results of SAM we provide a Jupyter notebook found in `notebooks/inference_playground.ipynb`.   
+The notebook will download the pretrained aging model and run inference on the images found in `notebooks/images`.  
+
+In addition, [Replicate](https://replicate.ai/) have created a demo for SAM where you can easily upload an image and run SAM on a desired set of ages! Check
+out the demo [here](https://replicate.ai/yuval-alaluf/sam).
+
+### MP4 Notebook
+To show full lifespan results using SAM we provide an additional notebook `notebooks/animation_inference_playground.ipynb` that will 
+run aging on multiple ages between 0 and 100 and interpolate between the results to display full aging. 
+The results will be saved as an MP4 files in `notebooks/animations` showing the aging and de-aging results.
+
+## Testing
+### Inference
+Having trained your model or if you're using a pretrained SAM model, you can use `scripts/inference.py` to run inference
+on a set of images.   
+For example, 
+```
+python scripts/inference.py \
+--exp_dir=/path/to/experiment \
+--checkpoint_path=experiment/checkpoints/best_model.pt \
+--data_path=/path/to/test_data \
+--test_batch_size=4 \
+--test_workers=4 \
+--couple_outputs
+--target_age=0,10,20,30,40,50,60,70,80
+```
+Additional notes to consider: 
+- During inference, the options used during training are loaded from the saved checkpoint and are then updated using the 
+test options passed to the inference script.
+- Adding the flag `--couple_outputs` will save an additional image containing the input and output images side-by-side in the sub-directory
+`inference_coupled`. Otherwise, only the output image is saved to the sub-directory `inference_results`.
+- In the above example, we will run age transformation with target ages 0,10,...,80.
+    - The results of each target age are saved to the sub-directories `inference_results/TARGET_AGE` and `inference_coupled/TARGET_AGE`.
+- By default, the images will be saved at resolution of 1024x1024, the original output size of StyleGAN. 
+    - If you wish to save outputs resized to resolutions of 256x256, you can do so by adding the flag `--resize_outputs`.
+
+### Side-by-Side Inference 
+The above inference script will save each aging result in a different sub-directory for each target age. Sometimes, 
+however, it is more convenient to save all aging results of a given input side-by-side like the following: 
+
+<p align="center">
+<img src="docs/866.jpg" width="800px"/>
+</p>
+
+To do so, we provide a script `inference_side_by_side.py` that works in a similar manner as the regular inference script: 
+```
+python scripts/inference_side_by_side.py \
+--exp_dir=/path/to/experiment \
+--checkpoint_path=experiment/checkpoints/best_model.pt \
+--data_path=/path/to/test_data \
+--test_batch_size=4 \
+--test_workers=4 \
+--target_age=0,10,20,30,40,50,60,70,80
+```
+Here, all aging results 0,10,...,80 will be save side-by-side with the original input image.
+
+### Reference-Guided Inference
+In the paper, we demonstrated how one can perform style-mixing on the fine-level style inputs with a reference image
+to control global features such as hair color. For example, 
+
+<p align="center">
+<img src="docs/1005_style_mixing.jpg" width="800px"/>
+</p>
+
+To perform style mixing using reference images, we provide the script `reference_guided_inference.py`. Here, 
+we first perform aging using the specified target age(s). Then, style mixing is performed using the specified 
+reference images and the specified layers. For example, one can run: 
+```
+python scripts/reference_guided_inference.py \
+--exp_dir=/path/to/experiment \
+--checkpoint_path=experiment/checkpoints/best_model.pt \
+--data_path=/path/to/test_data \
+--test_batch_size=4 \
+--test_workers=4 \
+--ref_images_paths_file=/path/to/ref_list.txt \
+--latent_mask=8,9 \
+--target_age=50,60,70,80
+```
+Here, the reference images should be specified in the file defined by `--ref_images_paths_file` and should have the 
+following format:
+```
+/path/to/reference/1.jpg
+/path/to/reference/2.jpg
+/path/to/reference/3.jpg
+/path/to/reference/4.jpg
+/path/to/reference/5.jpg
+```
+In the above example, we will aging using 4 different target ages. For each target age, we first transform the 
+test samples defined by `--data_path` and then perform style mixing on layers 8,9 defined by `--latent_mask`. 
+The results of each target age are saved in its own sub-directory.
+
+### Style Mixing 
+Instead of performing style mixing using a reference image, you can perform style mixing using randomly generated 
+w latent vectors by running the script `style_mixing.py`. This script works in a similar manner to the reference
+guided inference except you do not need to specify the `--ref_images_paths_file` flag.
+
+## Repository structure
+| Path | Description <img width=200>
+| :--- | :---
+| SAM | Repository root folder
+| &boxvr;&nbsp; configs | Folder containing configs defining model/data paths and data transforms
+| &boxvr;&nbsp; criteria | Folder containing various loss criterias for training
+| &boxvr;&nbsp; datasets | Folder with various dataset objects and augmentations
+| &boxvr;&nbsp; docs | Folder containing images displayed in the README
+| &boxvr;&nbsp; environment | Folder containing Anaconda environment used in our experiments
+| &boxvr; models | Folder containing all the models and training objects
+| &boxv;&nbsp; &boxvr;&nbsp; encoders | Folder containing various architecture implementations
+| &boxv;&nbsp; &boxvr;&nbsp; stylegan2 | StyleGAN2 model from [rosinality](https://github.com/rosinality/stylegan2-pytorch)
+| &boxv;&nbsp; &boxvr;&nbsp; psp.py | Implementation of pSp encoder
+| &boxv;&nbsp; &boxur;&nbsp; dex_vgg.py | Implementation of DEX VGG classifier used in computation of aging loss
+| &boxvr;&nbsp; notebook | Folder with jupyter notebook containing SAM inference playground
+| &boxvr;&nbsp; options | Folder with training and test command-line options
+| &boxvr;&nbsp; scripts | Folder with running scripts for training and inference
+| &boxvr;&nbsp; training | Folder with main training logic and Ranger implementation from [lessw2020](https://github.com/lessw2020/Ranger-Deep-Learning-Optimizer)
+| &boxvr;&nbsp; utils | Folder with various utility functions
+| <img width=300> | <img>
+
+
+## Credits
+**StyleGAN2 model and implementation:**  
+https://github.com/rosinality/stylegan2-pytorch  
+Copyright (c) 2019 Kim Seonghyeon  
+License (MIT) https://github.com/rosinality/stylegan2-pytorch/blob/master/LICENSE  
+
+**IR-SE50 model and implementations:**  
+https://github.com/TreB1eN/InsightFace_Pytorch  
+Copyright (c) 2018 TreB1eN  
+License (MIT) https://github.com/TreB1eN/InsightFace_Pytorch/blob/master/LICENSE  
+
+**Ranger optimizer implementation:**  
+https://github.com/lessw2020/Ranger-Deep-Learning-Optimizer   
+License (Apache License 2.0) https://github.com/lessw2020/Ranger-Deep-Learning-Optimizer/blob/master/LICENSE  
+
+**LPIPS model and implementation:**  
+https://github.com/S-aiueo32/lpips-pytorch  
+Copyright (c) 2020, Sou Uchida  
+License (BSD 2-Clause) https://github.com/S-aiueo32/lpips-pytorch/blob/master/LICENSE  
+
+**DEX VGG model and implementation:**  
+https://github.com/InterDigitalInc/HRFAE  
+Copyright (c) 2020, InterDigital R&D France  
+https://github.com/InterDigitalInc/HRFAE/blob/master/LICENSE.txt
+
+**pSp model and implementation:**   
+https://github.com/eladrich/pixel2style2pixel  
+Copyright (c) 2020 Elad Richardson, Yuval Alaluf  
+https://github.com/eladrich/pixel2style2pixel/blob/master/LICENSE
+
+## Acknowledgments
+This code borrows heavily from [pixel2style2pixel](https://github.com/eladrich/pixel2style2pixel)
+
+## Citation
+If you use this code for your research, please cite our paper <a href="https://arxiv.org/abs/2102.02754">Only a Matter of Style: Age Transformation Using a Style-Based Regression Model</a>:
+
+```
+@article{alaluf2021matter,
+    author = {Alaluf, Yuval and Patashnik, Or and Cohen-Or, Daniel},
+    title = {Only a Matter of Style: Age Transformation Using a Style-Based Regression Model},
+    journal = {ACM Trans. Graph.},
+    issue_date = {August 2021},
+    volume = {40},
+    number = {4},
+    year = {2021},
+    articleno = {45},
+    publisher = {Association for Computing Machinery},
+    url = {https://doi.org/10.1145/3450626.3459805}
+}
+```

cog.yaml

@@ -0,0 +1,24 @@
+image: "r8.im/yuval-alaluf/sam"
+build:
+  gpu: true
+  python_version: "3.8"
+  system_packages:
+    - "cmake"
+    - "libgl1-mesa-glx"
+    - "libglib2.0-0"
+    - "ninja-build"
+  python_packages:
+    - "Pillow==8.3.1"
+    - "cmake==3.21.1"
+    - "dlib==19.22.1"
+    - "imageio==2.9.0"
+    - "ipython==7.21.0"
+    - "matplotlib==3.1.3"
+    - "numpy==1.21.1"
+    - "opencv-python==4.5.3.56"
+    - "scipy==1.4.1"
+    - "tensorboard==2.2.1"
+    - "torch==1.8.0"
+    - "torchvision==0.9.0"
+    - "tqdm==4.42.1"
+predict: "predict.py:Predictor"

configs/data_configs.py

@@ -0,0 +1,13 @@
+from configs import transforms_config
+from configs.paths_config import dataset_paths
+
+
+DATASETS = {
+	'ffhq_aging': {
+		'transforms': transforms_config.AgingTransforms,
+		'train_source_root': dataset_paths['ffhq'],
+		'train_target_root': dataset_paths['ffhq'],
+		'test_source_root': dataset_paths['celeba_test'],
+		'test_target_root': dataset_paths['celeba_test'],
+	}
+}

configs/paths_config.py

@@ -0,0 +1,12 @@
+dataset_paths = {
+    'celeba_test': '',
+    'ffhq': '',
+}
+
+model_paths = {
+    'pretrained_psp_encoder': 'pretrained_models/psp_ffhq_encode.pt',
+    'ir_se50': 'pretrained_models/model_ir_se50.pth',
+    'stylegan_ffhq': 'pretrained_models/stylegan2-ffhq-config-f.pt',
+    'shape_predictor': 'shape_predictor_68_face_landmarks.dat',
+    'age_predictor': 'pretrained_models/dex_age_classifier.pth'
+}

configs/transforms_config.py

@@ -0,0 +1,37 @@
+from abc import abstractmethod
+import torchvision.transforms as transforms
+
+
+class TransformsConfig(object):
+
+	def __init__(self, opts):
+		self.opts = opts
+
+	@abstractmethod
+	def get_transforms(self):
+		pass
+
+
+class AgingTransforms(TransformsConfig):
+
+	def __init__(self, opts):
+		super(AgingTransforms, self).__init__(opts)
+
+	def get_transforms(self):
+		transforms_dict = {
+			'transform_gt_train': transforms.Compose([
+				transforms.Resize((256, 256)),
+				transforms.RandomHorizontalFlip(0.5),
+				transforms.ToTensor(),
+				transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])]),
+			'transform_source': None,
+			'transform_test': transforms.Compose([
+				transforms.Resize((256, 256)),
+				transforms.ToTensor(),
+				transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])]),
+			'transform_inference': transforms.Compose([
+				transforms.Resize((256, 256)),
+				transforms.ToTensor(),
+				transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])
+		}
+		return transforms_dict

criteria/aging_loss.py

@@ -0,0 +1,59 @@
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+from configs.paths_config import model_paths
+from models.dex_vgg import VGG
+
+
+class AgingLoss(nn.Module):
+
+    def __init__(self, opts):
+        super(AgingLoss, self).__init__()
+        self.age_net = VGG()
+        ckpt = torch.load(model_paths['age_predictor'], map_location="cpu")['state_dict']
+        ckpt = {k.replace('-', '_'): v for k, v in ckpt.items()}
+        self.age_net.load_state_dict(ckpt)
+        self.age_net.cuda()
+        self.age_net.eval()
+        self.min_age = 0
+        self.max_age = 100
+        self.opts = opts
+
+    def __get_predicted_age(self, age_pb):
+        predict_age_pb = F.softmax(age_pb)
+        predict_age = torch.zeros(age_pb.size(0)).type_as(predict_age_pb)
+        for i in range(age_pb.size(0)):
+            for j in range(age_pb.size(1)):
+                predict_age[i] += j * predict_age_pb[i][j]
+        return predict_age
+
+    def extract_ages(self, x):
+        x = F.interpolate(x, size=(224, 224), mode='bilinear')
+        predict_age_pb = self.age_net(x)['fc8']
+        predicted_age = self.__get_predicted_age(predict_age_pb)
+        return predicted_age
+
+    def forward(self, y_hat, y, target_ages, id_logs, label=None):
+        n_samples = y.shape[0]
+
+        if id_logs is None:
+            id_logs = []
+
+        input_ages = self.extract_ages(y) / 100.
+        output_ages = self.extract_ages(y_hat) / 100.
+
+        for i in range(n_samples):
+            # if id logs for the same exists, update the dictionary
+            if len(id_logs) > i:
+                id_logs[i].update({f'input_age_{label}': float(input_ages[i]) * 100,
+                                   f'output_age_{label}': float(output_ages[i]) * 100,
+                                   f'target_age_{label}': float(target_ages[i]) * 100})
+            # otherwise, create a new entry for the sample
+            else:
+                id_logs.append({f'input_age_{label}': float(input_ages[i]) * 100,
+                                f'output_age_{label}': float(output_ages[i]) * 100,
+                                f'target_age_{label}': float(target_ages[i]) * 100})
+
+        loss = F.mse_loss(output_ages, target_ages)
+        return loss, id_logs

criteria/id_loss.py

@@ -0,0 +1,55 @@
+import torch
+from torch import nn
+from configs.paths_config import model_paths
+from models.encoders.model_irse import Backbone
+
+
+class IDLoss(nn.Module):
+    def __init__(self):
+        super(IDLoss, self).__init__()
+        print('Loading ResNet ArcFace')
+        self.facenet = Backbone(input_size=112, num_layers=50, drop_ratio=0.6, mode='ir_se')
+        self.facenet.load_state_dict(torch.load(model_paths['ir_se50']))
+        self.face_pool = torch.nn.AdaptiveAvgPool2d((112, 112))
+        self.facenet.eval()
+
+    def extract_feats(self, x):
+        x = x[:, :, 35:223, 32:220]  # Crop interesting region
+        x = self.face_pool(x)
+        x_feats = self.facenet(x)
+        return x_feats
+
+    def forward(self, y_hat, y, x, label=None, weights=None):
+        n_samples = x.shape[0]
+        x_feats = self.extract_feats(x)
+        y_feats = self.extract_feats(y)
+        y_hat_feats = self.extract_feats(y_hat)
+        y_feats = y_feats.detach()
+        total_loss = 0
+        sim_improvement = 0
+        id_logs = []
+        count = 0
+        for i in range(n_samples):
+            diff_target = y_hat_feats[i].dot(y_feats[i])
+            diff_input = y_hat_feats[i].dot(x_feats[i])
+            diff_views = y_feats[i].dot(x_feats[i])
+
+            if label is None:
+                id_logs.append({'diff_target': float(diff_target),
+                                'diff_input': float(diff_input),
+                                'diff_views': float(diff_views)})
+            else:
+                id_logs.append({f'diff_target_{label}': float(diff_target),
+                                f'diff_input_{label}': float(diff_input),
+                                f'diff_views_{label}': float(diff_views)})
+
+            loss = 1 - diff_target
+            if weights is not None:
+                loss = weights[i] * loss
+
+            total_loss += loss
+            id_diff = float(diff_target) - float(diff_views)
+            sim_improvement += id_diff
+            count += 1
+
+        return total_loss / count, sim_improvement / count, id_logs

criteria/lpips/lpips.py

@@ -0,0 +1,35 @@
+import torch
+import torch.nn as nn
+
+from criteria.lpips.networks import get_network, LinLayers
+from criteria.lpips.utils import get_state_dict
+
+
+class LPIPS(nn.Module):
+    r"""Creates a criterion that measures
+    Learned Perceptual Image Patch Similarity (LPIPS).
+    Arguments:
+        net_type (str): the network type to compare the features:
+                        'alex' | 'squeeze' | 'vgg'. Default: 'alex'.
+        version (str): the version of LPIPS. Default: 0.1.
+    """
+    def __init__(self, net_type: str = 'alex', version: str = '0.1'):
+
+        assert version in ['0.1'], 'v0.1 is only supported now'
+
+        super(LPIPS, self).__init__()
+
+        # pretrained network
+        self.net = get_network(net_type).to("cuda")
+
+        # linear layers
+        self.lin = LinLayers(self.net.n_channels_list).to("cuda")
+        self.lin.load_state_dict(get_state_dict(net_type, version))
+
+    def forward(self, x: torch.Tensor, y: torch.Tensor):
+        feat_x, feat_y = self.net(x), self.net(y)
+
+        diff = [(fx - fy) ** 2 for fx, fy in zip(feat_x, feat_y)]
+        res = [l(d).mean((2, 3), True) for d, l in zip(diff, self.lin)]
+
+        return torch.sum(torch.cat(res, 0)) / x.shape[0]

criteria/lpips/networks.py

@@ -0,0 +1,96 @@
+from typing import Sequence
+
+from itertools import chain
+
+import torch
+import torch.nn as nn
+from torchvision import models
+
+from criteria.lpips.utils import normalize_activation
+
+
+def get_network(net_type: str):
+    if net_type == 'alex':
+        return AlexNet()
+    elif net_type == 'squeeze':
+        return SqueezeNet()
+    elif net_type == 'vgg':
+        return VGG16()
+    else:
+        raise NotImplementedError('choose net_type from [alex, squeeze, vgg].')
+
+
+class LinLayers(nn.ModuleList):
+    def __init__(self, n_channels_list: Sequence[int]):
+        super(LinLayers, self).__init__([
+            nn.Sequential(
+                nn.Identity(),
+                nn.Conv2d(nc, 1, 1, 1, 0, bias=False)
+            ) for nc in n_channels_list
+        ])
+
+        for param in self.parameters():
+            param.requires_grad = False
+
+
+class BaseNet(nn.Module):
+    def __init__(self):
+        super(BaseNet, self).__init__()
+
+        # register buffer
+        self.register_buffer(
+            'mean', torch.Tensor([-.030, -.088, -.188])[None, :, None, None])
+        self.register_buffer(
+            'std', torch.Tensor([.458, .448, .450])[None, :, None, None])
+
+    def set_requires_grad(self, state: bool):
+        for param in chain(self.parameters(), self.buffers()):
+            param.requires_grad = state
+
+    def z_score(self, x: torch.Tensor):
+        return (x - self.mean) / self.std
+
+    def forward(self, x: torch.Tensor):
+        x = self.z_score(x)
+
+        output = []
+        for i, (_, layer) in enumerate(self.layers._modules.items(), 1):
+            x = layer(x)
+            if i in self.target_layers:
+                output.append(normalize_activation(x))
+            if len(output) == len(self.target_layers):
+                break
+        return output
+
+
+class SqueezeNet(BaseNet):
+    def __init__(self):
+        super(SqueezeNet, self).__init__()
+
+        self.layers = models.squeezenet1_1(True).features
+        self.target_layers = [2, 5, 8, 10, 11, 12, 13]
+        self.n_channels_list = [64, 128, 256, 384, 384, 512, 512]
+
+        self.set_requires_grad(False)
+
+
+class AlexNet(BaseNet):
+    def __init__(self):
+        super(AlexNet, self).__init__()
+
+        self.layers = models.alexnet(True).features
+        self.target_layers = [2, 5, 8, 10, 12]
+        self.n_channels_list = [64, 192, 384, 256, 256]
+
+        self.set_requires_grad(False)
+
+
+class VGG16(BaseNet):
+    def __init__(self):
+        super(VGG16, self).__init__()
+
+        self.layers = models.vgg16(True).features
+        self.target_layers = [4, 9, 16, 23, 30]
+        self.n_channels_list = [64, 128, 256, 512, 512]
+
+        self.set_requires_grad(False)

criteria/lpips/utils.py

@@ -0,0 +1,30 @@
+from collections import OrderedDict
+
+import torch
+
+
+def normalize_activation(x, eps=1e-10):
+    norm_factor = torch.sqrt(torch.sum(x ** 2, dim=1, keepdim=True))
+    return x / (norm_factor + eps)
+
+
+def get_state_dict(net_type: str = 'alex', version: str = '0.1'):
+    # build url
+    url = 'https://raw.githubusercontent.com/richzhang/PerceptualSimilarity/' \
+        + f'master/lpips/weights/v{version}/{net_type}.pth'
+
+    # download
+    old_state_dict = torch.hub.load_state_dict_from_url(
+        url, progress=True,
+        map_location=None if torch.cuda.is_available() else torch.device('cpu')
+    )
+
+    # rename keys
+    new_state_dict = OrderedDict()
+    for key, val in old_state_dict.items():
+        new_key = key
+        new_key = new_key.replace('lin', '')
+        new_key = new_key.replace('model.', '')
+        new_state_dict[new_key] = val
+
+    return new_state_dict

criteria/w_norm.py

@@ -0,0 +1,14 @@
+import torch
+from torch import nn
+
+
+class WNormLoss(nn.Module):
+
+	def __init__(self, opts):
+		super(WNormLoss, self).__init__()
+		self.opts = opts
+
+	def forward(self, latent, latent_avg=None):
+		if self.opts.start_from_latent_avg or self.opts.start_from_encoded_w_plus:
+			latent = latent - latent_avg
+		return torch.sum(latent.norm(2, dim=(1, 2))) / latent.shape[0]

datasets/augmentations.py

@@ -0,0 +1,24 @@
+import numpy as np
+import torch
+
+
+class AgeTransformer(object):
+
+	def __init__(self, target_age):
+		self.target_age = target_age
+
+	def __call__(self, img):
+		img = self.add_aging_channel(img)
+		return img
+
+	def add_aging_channel(self, img):
+		target_age = self.__get_target_age()
+		target_age = int(target_age) / 100  # normalize aging amount to be in range [-1,1]
+		img = torch.cat((img, target_age * torch.ones((1, img.shape[1], img.shape[2]))))
+		return img
+
+	def __get_target_age(self):
+		if self.target_age == "uniform_random":
+			return np.random.randint(low=0., high=101, size=1)[0]
+		else:
+			return self.target_age

datasets/images_dataset.py

@@ -0,0 +1,33 @@
+from torch.utils.data import Dataset
+from PIL import Image
+from utils import data_utils
+
+
+class ImagesDataset(Dataset):
+
+	def __init__(self, source_root, target_root, opts, target_transform=None, source_transform=None):
+		self.source_paths = sorted(data_utils.make_dataset(source_root))
+		self.target_paths = sorted(data_utils.make_dataset(target_root))
+		self.source_transform = source_transform
+		self.target_transform = target_transform
+		self.opts = opts
+
+	def __len__(self):
+		return len(self.source_paths)
+
+	def __getitem__(self, index):
+		from_path = self.source_paths[index]
+		from_im = Image.open(from_path)
+		from_im = from_im.convert('RGB') if self.opts.label_nc == 0 else from_im.convert('L')
+
+		to_path = self.target_paths[index]
+		to_im = Image.open(to_path).convert('RGB')
+		if self.target_transform:
+			to_im = self.target_transform(to_im)
+
+		if self.source_transform:
+			from_im = self.source_transform(from_im)
+		else:
+			from_im = to_im
+
+		return from_im, to_im

datasets/inference_dataset.py

@@ -0,0 +1,29 @@
+from torch.utils.data import Dataset
+from PIL import Image
+from utils import data_utils
+
+
+class InferenceDataset(Dataset):
+
+	def __init__(self, root=None, paths_list=None, opts=None, transform=None, return_path=False):
+		if paths_list is None:
+			self.paths = sorted(data_utils.make_dataset(root))
+		else:
+			self.paths = data_utils.make_dataset_from_paths_list(paths_list)
+		self.transform = transform
+		self.opts = opts
+		self.return_path = return_path
+
+	def __len__(self):
+		return len(self.paths)
+
+	def __getitem__(self, index):
+		from_path = self.paths[index]
+		from_im = Image.open(from_path)
+		from_im = from_im.convert('RGB') if self.opts.label_nc == 0 else from_im.convert('L')
+		if self.transform:
+			from_im = self.transform(from_im)
+		if self.return_path:
+			return from_im, from_path
+		else:
+			return from_im

environment/sam_env.yaml

@@ -0,0 +1,36 @@
+name: sam_env
+channels:
+  - conda-forge
+  - defaults
+dependencies:
+  - _libgcc_mutex=0.1=main
+  - ca-certificates=2020.4.5.1=hecc5488_0
+  - certifi=2020.4.5.1=py36h9f0ad1d_0
+  - libedit=3.1.20181209=hc058e9b_0
+  - libffi=3.2.1=hd88cf55_4
+  - libgcc-ng=9.1.0=hdf63c60_0
+  - libstdcxx-ng=9.1.0=hdf63c60_0
+  - ncurses=6.2=he6710b0_1
+  - ninja=1.10.0=hc9558a2_0
+  - openssl=1.1.1g=h516909a_0
+  - pip=20.0.2=py36_3
+  - python=3.6.7=h0371630_0
+  - python_abi=3.6=1_cp36m
+  - readline=7.0=h7b6447c_5
+  - setuptools=46.4.0=py36_0
+  - sqlite=3.31.1=h62c20be_1
+  - tk=8.6.8=hbc83047_0
+  - wheel=0.34.2=py36_0
+  - xz=5.2.5=h7b6447c_0
+  - zlib=1.2.11=h7b6447c_3
+  - pip:
+    - scipy==1.4.1
+    - matplotlib==3.2.1
+    - tqdm==4.46.0
+    - numpy==1.18.4
+    - opencv-python==4.2.0.34
+    - pillow==7.1.2
+    - tensorboard==2.2.1
+    - torch==1.6.0
+    - torchvision==0.4.2
+prefix: ~/anaconda3/envs/sam_env

licenses/LICENSE_InterDigitalInc

@@ -0,0 +1,150 @@
+LIMITED SOFTWARE EVALUATION LICENSE AGREEMENT 
+
+
+ 
+This Limited Software Evaluation License Agreement (the “Agreement”) is entered into as of April 9th 2020, (“Effective Date”)  
+
+The following limited software evaluation license agreement (“the Agreement”) constitute an agreement between you (the “licensee”) and InterDigital R&D France, a French company existing and organized under the laws of France with its registered offices located at 975 avenue des champs blancs 35510 Cesson-Sévigné, FRANCE (hereinafter “InterDigital”)
+This Agreement governs the download and use of the Software (as defined below). Your use of the Software is subject to the terms and conditions set forth in this Agreement. By installing, using, accessing or copying the Software, you hereby irrevocably accept the terms and conditions of this Agreement. If you do not accept all or parts of the terms and conditions of this Agreement you cannot install, use, access nor copy the Software
+
+    Article 1.	Definitions
+
+“Affiliate” as used herein shall mean any entity that, directly or indirectly, through one or more intermediates, is controlled by, controls, or is under common control with InterDigital or The Licensee, as the case may be.  For purposes of this definition only, the term “control” means the possession of the power to direct or cause the direction of the management and policies of an entity, whether by ownership of voting stock or partnership interest, by contract, or otherwise, including direct or indirect ownership of more than fifty percent (50%) of the voting interest in the entity in question.
+
+“Authorized Purpose” means any use of the Software for research on the Software and evaluation of the Software exclusively, and academic research using the Software without any commercial use. For the avoidance of doubt, a commercial use includes, but is not limited to:
+-	using the Software in advertisements of any kind,
+-	licensing or selling of the Software,
+-	use the Software to provide any service to any third Party
+-	use the Software to develop a competitive product of the Software
+
+“Documentation” means textual materials delivered by InterDigital to the Licensee pursuant to this Agreement relating to the Software, in written or electronic format, including but not limited to: technical reference manuals, technical notes, user manuals, and application guides.
+
+“Limited Period” means the life of the copyright owned by InterDigital on the Software in each and every country where such copyright would exist.
+
+“Intellectual Property Rights” means all copyrights, trademarks, trade secrets, patents, mask works and other intellectual property rights recognized in any jurisdiction worldwide, including all applications and registrations with respect thereto.
+
+"Open Source software" shall mean any software, including where appropriate, any and all modifications, derivative works, enhancements, upgrades, improvements, fixed bugs, and/or statically linked to the source code of such software, released under a free software license, that requires as a condition of royalty-free usage, copy, modification and/or redistribution of the Open Source Software to: 
+•	Redistribute the Open Source Software royalty-free, and/or;
+•	Redistribute the Open Source Software under the same license/distribution terms as those contained in the open source or free software license under which it has originally been released and/or;
+•	Release to the public, disclose or otherwise make available the source code of the Open Source Software.
+
+For purposes of the Agreement, by means of example and without limitation, any software that is released or distributed under any of the following licenses shall be qualified as Open Source Software: (A) GNU General Public License (GPL), (B) GNU Lesser/Library GPL (LGPL), (C) the Artistic License, (D) the Mozilla Public License, (E) the Common Public License, (F) the Sun Community Source License (SCSL), (G) the Sun Industry Standards Source License (SISSL), (H) BSD License, (I) MIT License, (J) Apache Software License, (K) Open SSL License, (L) IBM Public License, (M) Open Software License.
+
+“Software” means any computer programming code, in object and/or source version, and related Documentation delivered by InterDigital to the Licensee pursuant to this Agreement as described in Exhibit A attached and incorporated herein by reference.
+
+    Article 2.	License
+
+InterDigital grants Licensee a free, worldwide, non-exclusive, license on copyright owned on the Software to download, use, modify and reproduce solely for the Authorized Purpose for the Limited Period. 
+
+The Licensee shall not pay any royalty, license fee or maintenance fee, or other fee of any nature under this Agreement.
+
+The Licensee shall have the right to correct, adapt, modify, reverse engineer, disassemble, decompile and any action leading to the transformation of Software provided that such action is made to accomplish the Authorized Purpose. 
+
+Licensee shall have the right to make a demonstration of the Software, provided that it is in the Purpose and provided that Licensee shall maintain control of the Software at all time. This includes the control of any computer or server on which the Software is installed: no third party shall have access to such computer or server under any circumstances. No computer nor server containing the Software will be left in the possession of any third Party. 
+
+    Article 3.	Restrictions on use of the Software
+  
+Licensee shall not remove, obscure or modify any copyright, trademark or other proprietary rights notices, marks or labels contained on or within the Software, falsify or delete any author attributions, legal notices or other labels of the origin or source of the material.
+
+Licensee shall not have the right to distribute the Software, either modified or not, to any third Party.
+
+The rights granted here above do not include any rights to automatically obtain any upgrade or update of the Software, acquired or otherwise made available by InterDigital. Such deliverance shall be discussed on a case by case basis by the Parties.
+
+    Article 4.	Ownership
+
+Title to and ownership of the Software, the Documentation  and/or any Intellectual Property Right protecting the Software or/and the Documentation shall, at all times, remain with InterDigital. Licensee agrees that except for the rights granted on copyright on the Software set forth in Section 2 above, in no event does anything in this Agreement grant, provide or convey any other rights, immunities or interest in or to any Intellectual Property Rights (including especially patents) of InterDigital or any of its Affiliates whether by implication, estoppel or otherwise.
+
+
+    Article 5.	Publication/Communication
+
+Any publication or oral communication resulting from the use of the Software shall be elaborated in good faith and shall not be driven by a deliberate will to denigrate InterDigital or any of its products. In any publication and on any support joined to an oral communication (for instance a PowerPoint document) resulting from the use of the Software, the following statement shall be inserted:
+
+“HRFAE is an InterDigital product”
+
+And in any publication, the latest publication about the software shall be properly cited. The latest publication currently is:
+"Arxiv preprint (ref to come shortly)”
+
+In any oral communication resulting from the use of the Software, the Licensee shall orally indicate that the Software is InterDigital’s property.
+
+    Article 6.	No Warranty - Disclaimer
+
+THE SOFTWARE AND DOCUMENTATION ARE PROVIDED TO LICENSEE ON AN “AS IS” BASIS. INTERDIGITAL MAKES NO WARRANTY THAT THE LICENSED TECHNOLOGY WILL OPERATE ON ANY PARTICULAR HARDWARE, PLATFORM, OR ENVIRONMENT. THERE IS NO WARRANTY THAT THE OPERATION OF THE LICENSED TECHNOLOGY SHALL BE UNINTERRUPTED, WITHOUT BUGS OR ERROR FREE. THE SOFTWARE AND DOCUMENTATION ARE PROVIDED HEREUNDER WITHOUT WARRANTY OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO ANY IMPLIED LIABILITIES AND WARRANTIES OF NONINFRINGEMENT OF INTELLECTUAL PROPERTY, FREEDOM FROM INHERENT DEFECTS, CONFORMITY TO A SAMPLE OR MODEL, MERCHANTABILITY, FITNESS AND/OR SUITABILITY FOR A SPECIFIC OR GENERAL PURPOSE AND THOSE ARISING BY STATUTE OR BY LAW, OR FROM A CAUSE OF DEALING OR USAGE OF TRADE.
+ 
+InterDigital shall not be obliged to perform any modifications, derivative works, enhancements, upgrades, updates or improvements of the Software or to fix any bug that could arise.
+
+Hence, the Licensee uses the Software at his own cost, risks and responsibility. InterDigital shall not be liable for any damage that could arise to Licensee by using the Software, either in accordance with this Agreement or not.
+
+InterDigital shall not be liable for any consequential or indirect losses, including any indirect loss of profits, revenues, business, and/or anticipated savings, whether or not in the contemplation of the Parties at the time of entering into the Agreement unless expressly set out in the Agreement, or arising from gross negligence, willful misconduct or fraud. 
+
+Licensee agrees that it will defend, indemnify and hold harmless InterDigital and its Affiliates against any and all losses, damages, costs and expenses arising from a breach by the Licensee of any of its obligations or representations hereunder, including, without limitation, any third party, and/or any claims in connection with any such breach and/or any use of the Software, including any claim from third party arising from access, use or any other activity in relation to this Software.
+
+The Licensee shall not make any warranty, representation, or commitment on behalf of InterDigital to any other third party.
+
+    Article 7.	Open Source Software 
+
+InterDigital hereby notifies the Licensee, and the Licensee hereby acknowledges and accepts, that the Software contains Open Source Software. The list of such Open Source Software is enclosed in exhibit B and the relevant license are contained at the root of the Software when downloaded. Hence, the Licensee shall comply with such license and agree on its terms on at its own risks.
+
+The Licensee hereby represents, warrants and covenants to InterDigital that The Licensee’s use of the Software shall not result in the Contamination of all or part of the Software, directly or indirectly, or of any Intellectual Property of InterDigital or its Affiliates.
+
+Contamination effect shall mean that the licensing terms under which one Open Source software, distinct from the Software, is released would also apply, by viral effect, to the software to which such Open Source software is linked to, combined with or otherwise connected to. 
+ 
+    Article 8.	No Future Contract Obligation
+
+Neither this Agreement nor the furnishing of the Software, nor any other Confidential Information shall be construed to obligate either party to: (a) enter into any further agreement or negotiation concerning the deployment of the Software; (b) refrain from entering into any agreement or negotiation with any other third party regarding the same or any other subject matter; or (c) refrain from pursuing its business in whatever manner it elects even if this involves competing with the other party.  
+
+    Article 9.	Term and Termination
+
+This Agreement shall terminate at the end of the Limited Period, unless earlier terminated by either party on the ground of material breach by the other party, which breach is not remedied after thirty (30) days advance written notice, specifying the breach with reasonable particularity and referencing this Agreement.
+
+    Article 10.	 General Provisions
+
+12.1 Severability.  If any provision of this Agreement shall be held to be in contravention of applicable law, this Agreement shall be construed as if such provision were not a part thereof, and in all other respects the terms hereof shall remain in full force and effect. 
+
+12.2 Governing Law.  Regardless of the place of execution, delivery, performance or any other aspect of this Agreement, this Agreement and all of the rights of the parties under this Agreement shall be governed by, construed under and enforced in accordance with the substantive law of the France without regard to conflicts of law principles. In case of a dispute that could not be settled amicably, the courts of Nanterre shall be exclusively competent. 
+
+12.3 Survival.  The provisions of articles 1, 3, 4, 6, 7, 9, 10.2 and 10.6 shall survive termination of this Agreement.
+12.4 Assignment. InterDigital may assign this license to any third Party. Such assignment will be announced on the website as defined in article 5. Licensee may not assign this agreement to any third party without the previous written agreement from InterDigital.
+
+12.5 Entire Agreement.  This Agreement constitutes the entire agreement between the parties hereto with respect to the subject matter hereof and supersedes any prior agreements or understanding. 
+
+12.6 Notices.  To have legal effect, notices must be provided by registered or certified mail, return receipt requested, to the representatives of InterDigital at the following address:
+
+InterDigital
+Legal Dept
+975 avenue des champs blancs
+35510 Cesson-Sévigné
+FRANCE
+
+=======================================================================
+
+Exhibit A
+Software
+
+
+The Software is comprised of the following software and Documentation:
+
+-	README.md file that explains the content of the software and the procedure to use it.
+-	Source python files, as well as pretrained models
+ 
+=======================================================================
+
+Exhibit B
+Open Source licenses
+
+
+PIL                 http://www.pythonware.com/products/pil/license.htm
+
+numpy               https://numpy.org/license.html
+
+tensorboardX        https://github.com/lanpa/tensorboardX/blob/master/LICENSE
+
+pytorch             https://github.com/pytorch/pytorch/blob/master/LICENSE
+
+torchvision         https://github.com/pytorch/vision/blob/master/LICENSE
+
+tensorboard_logger  https://github.com/TeamHG-Memex/tensorboard_logger/blob/master/LICENSE
+
+argparse            https://github.com/ThomasWaldmann/argparse/blob/master/LICENSE.txt
+
+yaml                https://github.com/yaml/pyyaml/blob/master/LICENSE
+

licenses/LICENSE_S-aiueo32

@@ -0,0 +1,25 @@
+BSD 2-Clause License
+
+Copyright (c) 2020, Sou Uchida
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright notice, this
+   list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright notice,
+   this list of conditions and the following disclaimer in the documentation
+   and/or other materials provided with the distribution.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

licenses/LICENSE_TreB1eN

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2018 TreB1eN
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

licenses/LICENSE_eladrich

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2020 Elad Richardson, Yuval Alaluf
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

licenses/LICENSE_lessw2020

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

licenses/LICENSE_rosinality

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2019 Kim Seonghyeon
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

models/dex_vgg.py

@@ -0,0 +1,65 @@
+import torch.nn as nn
+import torch.nn.functional as F
+
+"""
+VGG implementation from [InterDigitalInc](https://github.com/InterDigitalInc/HRFAE/blob/master/nets.py)
+"""
+
+class VGG(nn.Module):
+    def __init__(self, pool='max'):
+        super(VGG, self).__init__()
+        # vgg modules
+        self.conv1_1 = nn.Conv2d(3, 64, kernel_size=3, padding=1)
+        self.conv1_2 = nn.Conv2d(64, 64, kernel_size=3, padding=1)
+        self.conv2_1 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
+        self.conv2_2 = nn.Conv2d(128, 128, kernel_size=3, padding=1)
+        self.conv3_1 = nn.Conv2d(128, 256, kernel_size=3, padding=1)
+        self.conv3_2 = nn.Conv2d(256, 256, kernel_size=3, padding=1)
+        self.conv3_3 = nn.Conv2d(256, 256, kernel_size=3, padding=1)
+        self.conv4_1 = nn.Conv2d(256, 512, kernel_size=3, padding=1)
+        self.conv4_2 = nn.Conv2d(512, 512, kernel_size=3, padding=1)
+        self.conv4_3 = nn.Conv2d(512, 512, kernel_size=3, padding=1)
+        self.conv5_1 = nn.Conv2d(512, 512, kernel_size=3, padding=1)
+        self.conv5_2 = nn.Conv2d(512, 512, kernel_size=3, padding=1)
+        self.conv5_3 = nn.Conv2d(512, 512, kernel_size=3, padding=1)
+        self.fc6 = nn.Linear(25088, 4096, bias=True)
+        self.fc7 = nn.Linear(4096, 4096, bias=True)
+        self.fc8_101 = nn.Linear(4096, 101, bias=True)
+        if pool == 'max':
+            self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
+            self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
+            self.pool3 = nn.MaxPool2d(kernel_size=2, stride=2)
+            self.pool4 = nn.MaxPool2d(kernel_size=2, stride=2)
+            self.pool5 = nn.MaxPool2d(kernel_size=2, stride=2)
+        elif pool == 'avg':
+            self.pool1 = nn.AvgPool2d(kernel_size=2, stride=2)
+            self.pool2 = nn.AvgPool2d(kernel_size=2, stride=2)
+            self.pool3 = nn.AvgPool2d(kernel_size=2, stride=2)
+            self.pool4 = nn.AvgPool2d(kernel_size=2, stride=2)
+            self.pool5 = nn.AvgPool2d(kernel_size=2, stride=2)
+
+    def forward(self, x):
+        out = {}
+        out['r11'] = F.relu(self.conv1_1(x))
+        out['r12'] = F.relu(self.conv1_2(out['r11']))
+        out['p1'] = self.pool1(out['r12'])
+        out['r21'] = F.relu(self.conv2_1(out['p1']))
+        out['r22'] = F.relu(self.conv2_2(out['r21']))
+        out['p2'] = self.pool2(out['r22'])
+        out['r31'] = F.relu(self.conv3_1(out['p2']))
+        out['r32'] = F.relu(self.conv3_2(out['r31']))
+        out['r33'] = F.relu(self.conv3_3(out['r32']))
+        out['p3'] = self.pool3(out['r33'])
+        out['r41'] = F.relu(self.conv4_1(out['p3']))
+        out['r42'] = F.relu(self.conv4_2(out['r41']))
+        out['r43'] = F.relu(self.conv4_3(out['r42']))
+        out['p4'] = self.pool4(out['r43'])
+        out['r51'] = F.relu(self.conv5_1(out['p4']))
+        out['r52'] = F.relu(self.conv5_2(out['r51']))
+        out['r53'] = F.relu(self.conv5_3(out['r52']))
+        out['p5'] = self.pool5(out['r53'])
+        out['p5'] = out['p5'].view(out['p5'].size(0), -1)
+        out['fc6'] = F.relu(self.fc6(out['p5']))
+        out['fc7'] = F.relu(self.fc7(out['fc6']))
+        out['fc8'] = self.fc8_101(out['fc7'])
+        return out

models/encoders/helpers.py

@@ -0,0 +1,119 @@
+from collections import namedtuple
+import torch
+from torch.nn import Conv2d, BatchNorm2d, PReLU, ReLU, Sigmoid, MaxPool2d, AdaptiveAvgPool2d, Sequential, Module
+
+"""
+ArcFace implementation from [TreB1eN](https://github.com/TreB1eN/InsightFace_Pytorch)
+"""
+
+
+class Flatten(Module):
+	def forward(self, input):
+		return input.view(input.size(0), -1)
+
+
+def l2_norm(input, axis=1):
+	norm = torch.norm(input, 2, axis, True)
+	output = torch.div(input, norm)
+	return output
+
+
+class Bottleneck(namedtuple('Block', ['in_channel', 'depth', 'stride'])):
+	""" A named tuple describing a ResNet block. """
+
+
+def get_block(in_channel, depth, num_units, stride=2):
+	return [Bottleneck(in_channel, depth, stride)] + [Bottleneck(depth, depth, 1) for i in range(num_units - 1)]
+
+
+def get_blocks(num_layers):
+	if num_layers == 50:
+		blocks = [
+			get_block(in_channel=64, depth=64, num_units=3),
+			get_block(in_channel=64, depth=128, num_units=4),
+			get_block(in_channel=128, depth=256, num_units=14),
+			get_block(in_channel=256, depth=512, num_units=3)
+		]
+	elif num_layers == 100:
+		blocks = [
+			get_block(in_channel=64, depth=64, num_units=3),
+			get_block(in_channel=64, depth=128, num_units=13),
+			get_block(in_channel=128, depth=256, num_units=30),
+			get_block(in_channel=256, depth=512, num_units=3)
+		]
+	elif num_layers == 152:
+		blocks = [
+			get_block(in_channel=64, depth=64, num_units=3),
+			get_block(in_channel=64, depth=128, num_units=8),
+			get_block(in_channel=128, depth=256, num_units=36),
+			get_block(in_channel=256, depth=512, num_units=3)
+		]
+	else:
+		raise ValueError("Invalid number of layers: {}. Must be one of [50, 100, 152]".format(num_layers))
+	return blocks
+
+
+class SEModule(Module):
+	def __init__(self, channels, reduction):
+		super(SEModule, self).__init__()
+		self.avg_pool = AdaptiveAvgPool2d(1)
+		self.fc1 = Conv2d(channels, channels // reduction, kernel_size=1, padding=0, bias=False)
+		self.relu = ReLU(inplace=True)
+		self.fc2 = Conv2d(channels // reduction, channels, kernel_size=1, padding=0, bias=False)
+		self.sigmoid = Sigmoid()
+
+	def forward(self, x):
+		module_input = x
+		x = self.avg_pool(x)
+		x = self.fc1(x)
+		x = self.relu(x)
+		x = self.fc2(x)
+		x = self.sigmoid(x)
+		return module_input * x
+
+
+class bottleneck_IR(Module):
+	def __init__(self, in_channel, depth, stride):
+		super(bottleneck_IR, self).__init__()
+		if in_channel == depth:
+			self.shortcut_layer = MaxPool2d(1, stride)
+		else:
+			self.shortcut_layer = Sequential(
+				Conv2d(in_channel, depth, (1, 1), stride, bias=False),
+				BatchNorm2d(depth)
+			)
+		self.res_layer = Sequential(
+			BatchNorm2d(in_channel),
+			Conv2d(in_channel, depth, (3, 3), (1, 1), 1, bias=False), PReLU(depth),
+			Conv2d(depth, depth, (3, 3), stride, 1, bias=False), BatchNorm2d(depth)
+		)
+
+	def forward(self, x):
+		shortcut = self.shortcut_layer(x)
+		res = self.res_layer(x)
+		return res + shortcut
+
+
+class bottleneck_IR_SE(Module):
+	def __init__(self, in_channel, depth, stride):
+		super(bottleneck_IR_SE, self).__init__()
+		if in_channel == depth:
+			self.shortcut_layer = MaxPool2d(1, stride)
+		else:
+			self.shortcut_layer = Sequential(
+				Conv2d(in_channel, depth, (1, 1), stride, bias=False),
+				BatchNorm2d(depth)
+			)
+		self.res_layer = Sequential(
+			BatchNorm2d(in_channel),
+			Conv2d(in_channel, depth, (3, 3), (1, 1), 1, bias=False),
+			PReLU(depth),
+			Conv2d(depth, depth, (3, 3), stride, 1, bias=False),
+			BatchNorm2d(depth),
+			SEModule(depth, 16)
+		)
+
+	def forward(self, x):
+		shortcut = self.shortcut_layer(x)
+		res = self.res_layer(x)
+		return res + shortcut

models/encoders/model_irse.py

@@ -0,0 +1,48 @@
+from torch.nn import Linear, Conv2d, BatchNorm1d, BatchNorm2d, PReLU, Dropout, Sequential, Module
+from models.encoders.helpers import get_blocks, Flatten, bottleneck_IR, bottleneck_IR_SE, l2_norm
+
+"""
+Modified Backbone implementation from [TreB1eN](https://github.com/TreB1eN/InsightFace_Pytorch)
+"""
+
+
+class Backbone(Module):
+	def __init__(self, input_size, num_layers, mode='ir', drop_ratio=0.4, affine=True):
+		super(Backbone, self).__init__()
+		assert input_size in [112, 224], "input_size should be 112 or 224"
+		assert num_layers in [50, 100, 152], "num_layers should be 50, 100 or 152"
+		assert mode in ['ir', 'ir_se'], "mode should be ir or ir_se"
+		blocks = get_blocks(num_layers)
+		if mode == 'ir':
+			unit_module = bottleneck_IR
+		elif mode == 'ir_se':
+			unit_module = bottleneck_IR_SE
+		self.input_layer = Sequential(Conv2d(3, 64, (3, 3), 1, 1, bias=False),
+									  BatchNorm2d(64),
+									  PReLU(64))
+		if input_size == 112:
+			self.output_layer = Sequential(BatchNorm2d(512),
+			                               Dropout(drop_ratio),
+			                               Flatten(),
+			                               Linear(512 * 7 * 7, 512),
+			                               BatchNorm1d(512, affine=affine))
+		else:
+			self.output_layer = Sequential(BatchNorm2d(512),
+			                               Dropout(drop_ratio),
+			                               Flatten(),
+			                               Linear(512 * 14 * 14, 512),
+			                               BatchNorm1d(512, affine=affine))
+
+		modules = []
+		for block in blocks:
+			for bottleneck in block:
+				modules.append(unit_module(bottleneck.in_channel,
+										   bottleneck.depth,
+										   bottleneck.stride))
+		self.body = Sequential(*modules)
+
+	def forward(self, x):
+		x = self.input_layer(x)
+		x = self.body(x)
+		x = self.output_layer(x)
+		return l2_norm(x)

models/encoders/psp_encoders.py

@@ -0,0 +1,114 @@
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch import nn
+from torch.nn import Conv2d, BatchNorm2d, PReLU, Sequential, Module
+
+from models.encoders.helpers import get_blocks, bottleneck_IR, bottleneck_IR_SE
+from models.stylegan2.model import EqualLinear
+
+
+class GradualStyleBlock(Module):
+    def __init__(self, in_c, out_c, spatial):
+        super(GradualStyleBlock, self).__init__()
+        self.out_c = out_c
+        self.spatial = spatial
+        num_pools = int(np.log2(spatial))
+        modules = []
+        modules += [Conv2d(in_c, out_c, kernel_size=3, stride=2, padding=1), nn.LeakyReLU()]
+        for i in range(num_pools - 1):
+            modules += [
+                Conv2d(out_c, out_c, kernel_size=3, stride=2, padding=1), nn.LeakyReLU()
+            ]
+        self.convs = nn.Sequential(*modules)
+        self.linear = EqualLinear(out_c, out_c, lr_mul=1)
+
+    def forward(self, x):
+        x = self.convs(x)
+        x = x.view(-1, self.out_c)
+        x = self.linear(x)
+        return x
+
+
+class GradualStyleEncoder(Module):
+    def __init__(self, num_layers, mode='ir', n_styles=18, opts=None):
+        super(GradualStyleEncoder, self).__init__()
+        assert num_layers in [50, 100, 152], 'num_layers should be 50,100, or 152'
+        assert mode in ['ir', 'ir_se'], 'mode should be ir or ir_se'
+        blocks = get_blocks(num_layers)
+        if mode == 'ir':
+            unit_module = bottleneck_IR
+        elif mode == 'ir_se':
+            unit_module = bottleneck_IR_SE
+        self.input_layer = Sequential(Conv2d(opts.input_nc, 64, (3, 3), 1, 1, bias=False),
+                                      BatchNorm2d(64),
+                                      PReLU(64))
+        modules = []
+        for block in blocks:
+            for bottleneck in block:
+                modules.append(unit_module(bottleneck.in_channel,
+                                           bottleneck.depth,
+                                           bottleneck.stride))
+        self.body = Sequential(*modules)
+
+        self.styles = nn.ModuleList()
+        self.style_count = n_styles
+        self.coarse_ind = 3
+        self.middle_ind = 7
+        for i in range(self.style_count):
+            if i < self.coarse_ind:
+                style = GradualStyleBlock(512, 512, 16)
+            elif i < self.middle_ind:
+                style = GradualStyleBlock(512, 512, 32)
+            else:
+                style = GradualStyleBlock(512, 512, 64)
+            self.styles.append(style)
+        self.latlayer1 = nn.Conv2d(256, 512, kernel_size=1, stride=1, padding=0)
+        self.latlayer2 = nn.Conv2d(128, 512, kernel_size=1, stride=1, padding=0)
+
+    def _upsample_add(self, x, y):
+        '''Upsample and add two feature maps.
+		Args:
+		  x: (Variable) top feature map to be upsampled.
+		  y: (Variable) lateral feature map.
+		Returns:
+		  (Variable) added feature map.
+		Note in PyTorch, when input size is odd, the upsampled feature map
+		with `F.upsample(..., scale_factor=2, mode='nearest')`
+		maybe not equal to the lateral feature map size.
+		e.g.
+		original input size: [N,_,15,15] ->
+		conv2d feature map size: [N,_,8,8] ->
+		upsampled feature map size: [N,_,16,16]
+		So we choose bilinear upsample which supports arbitrary output sizes.
+		'''
+        _, _, H, W = y.size()
+        return F.interpolate(x, size=(H, W), mode='bilinear', align_corners=True) + y
+
+    def forward(self, x):
+        x = self.input_layer(x)
+
+        latents = []
+        modulelist = list(self.body._modules.values())
+        for i, l in enumerate(modulelist):
+            x = l(x)
+            if i == 6:
+                c1 = x
+            elif i == 20:
+                c2 = x
+            elif i == 23:
+                c3 = x
+
+        for j in range(self.coarse_ind):
+            latents.append(self.styles[j](c3))
+
+        p2 = self._upsample_add(c3, self.latlayer1(c2))
+        for j in range(self.coarse_ind, self.middle_ind):
+            latents.append(self.styles[j](p2))
+
+        p1 = self._upsample_add(p2, self.latlayer2(c1))
+        for j in range(self.middle_ind, self.style_count):
+            latents.append(self.styles[j](p1))
+
+        out = torch.stack(latents, dim=1)
+        return out

models/psp.py

@@ -0,0 +1,131 @@
+"""
+This file defines the core research contribution
+"""
+import copy
+from argparse import Namespace
+
+import torch
+from torch import nn
+import math
+
+from configs.paths_config import model_paths
+from models.encoders import psp_encoders
+from models.stylegan2.model import Generator
+
+
+class pSp(nn.Module):
+
+	def __init__(self, opts):
+		super(pSp, self).__init__()
+		self.set_opts(opts)
+		self.n_styles = int(math.log(self.opts.output_size, 2)) * 2 - 2
+		# Define architecture
+		self.encoder = self.set_encoder()
+		self.decoder = Generator(self.opts.output_size, 512, 8)
+		self.face_pool = torch.nn.AdaptiveAvgPool2d((256, 256))
+		# Load weights if needed
+		self.load_weights()
+
+	def set_encoder(self):
+		return psp_encoders.GradualStyleEncoder(50, 'ir_se', self.n_styles, self.opts)
+
+	def load_weights(self):
+		if self.opts.checkpoint_path is not None:
+			print(f'Loading SAM from checkpoint: {self.opts.checkpoint_path}')
+			ckpt = torch.load(self.opts.checkpoint_path, map_location='cpu')
+			self.encoder.load_state_dict(self.__get_keys(ckpt, 'encoder'), strict=False)
+			self.decoder.load_state_dict(self.__get_keys(ckpt, 'decoder'), strict=True)
+			if self.opts.start_from_encoded_w_plus:
+				self.pretrained_encoder = self.__get_pretrained_psp_encoder()
+				self.pretrained_encoder.load_state_dict(self.__get_keys(ckpt, 'pretrained_encoder'), strict=True)
+			self.__load_latent_avg(ckpt)
+		else:
+			print('Loading encoders weights from irse50!')
+			encoder_ckpt = torch.load(model_paths['ir_se50'])
+			# Transfer the RGB input of the irse50 network to the first 3 input channels of SAM's encoder
+			if self.opts.input_nc != 3:
+				shape = encoder_ckpt['input_layer.0.weight'].shape
+				altered_input_layer = torch.randn(shape[0], self.opts.input_nc, shape[2], shape[3], dtype=torch.float32)
+				altered_input_layer[:, :3, :, :] = encoder_ckpt['input_layer.0.weight']
+				encoder_ckpt['input_layer.0.weight'] = altered_input_layer
+			self.encoder.load_state_dict(encoder_ckpt, strict=False)
+			print(f'Loading decoder weights from pretrained path: {self.opts.stylegan_weights}')
+			ckpt = torch.load(self.opts.stylegan_weights)
+			self.decoder.load_state_dict(ckpt['g_ema'], strict=True)
+			self.__load_latent_avg(ckpt, repeat=self.n_styles)
+			if self.opts.start_from_encoded_w_plus:
+				self.pretrained_encoder = self.__load_pretrained_psp_encoder()
+				self.pretrained_encoder.eval()
+
+	def forward(self, x, resize=True, latent_mask=None, input_code=False, randomize_noise=True,
+				inject_latent=None, return_latents=False, alpha=None, input_is_full=False):
+		if input_code:
+			codes = x
+		else:
+			codes = self.encoder(x)
+			# normalize with respect to the center of an average face
+			if self.opts.start_from_latent_avg:
+				codes = codes + self.latent_avg
+			# normalize with respect to the latent of the encoded image of pretrained pSp encoder
+			elif self.opts.start_from_encoded_w_plus:
+				with torch.no_grad():
+					encoded_latents = self.pretrained_encoder(x[:, :-1, :, :])
+					encoded_latents = encoded_latents + self.latent_avg
+				codes = codes + encoded_latents
+
+		if latent_mask is not None:
+			for i in latent_mask:
+				if inject_latent is not None:
+					if alpha is not None:
+						codes[:, i] = alpha * inject_latent[:, i] + (1 - alpha) * codes[:, i]
+					else:
+						codes[:, i] = inject_latent[:, i]
+				else:
+					codes[:, i] = 0
+
+		input_is_latent = (not input_code) or (input_is_full)
+		images, result_latent = self.decoder([codes],
+											 input_is_latent=input_is_latent,
+											 randomize_noise=randomize_noise,
+											 return_latents=return_latents)
+
+		if resize:
+			images = self.face_pool(images)
+
+		if return_latents:
+			return images, result_latent
+		else:
+			return images
+
+	def set_opts(self, opts):
+		self.opts = opts
+
+	def __load_latent_avg(self, ckpt, repeat=None):
+		if 'latent_avg' in ckpt:
+			self.latent_avg = ckpt['latent_avg'].to(self.opts.device)
+			if repeat is not None:
+				self.latent_avg = self.latent_avg.repeat(repeat, 1)
+		else:
+			self.latent_avg = None
+
+	def __get_pretrained_psp_encoder(self):
+		opts_encoder = vars(copy.deepcopy(self.opts))
+		opts_encoder['input_nc'] = 3
+		opts_encoder = Namespace(**opts_encoder)
+		encoder = psp_encoders.GradualStyleEncoder(50, 'ir_se', self.n_styles, opts_encoder)
+		return encoder
+
+	def __load_pretrained_psp_encoder(self):
+		print(f'Loading pSp encoder from checkpoint: {self.opts.pretrained_psp_path}')
+		ckpt = torch.load(self.opts.pretrained_psp_path, map_location='cpu')
+		encoder_ckpt = self.__get_keys(ckpt, name='encoder')
+		encoder = self.__get_pretrained_psp_encoder()
+		encoder.load_state_dict(encoder_ckpt, strict=False)
+		return encoder
+
+	@staticmethod
+	def __get_keys(d, name):
+		if 'state_dict' in d:
+			d = d['state_dict']
+		d_filt = {k[len(name) + 1:]: v for k, v in d.items() if k[:len(name)] == name}
+		return d_filt