new commit

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 AIRI - Artificial Intelligence Research Institute
+
+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,169 @@
+# HairFastGAN: Realistic and Robust Hair Transfer with a Fast Encoder-Based Approach
+
+<a href="https://arxiv.org/abs/2404.01094"><img src="https://img.shields.io/badge/arXiv-2404.01094-b31b1b.svg" height=22.5></a>
+<a href="https://huggingface.co/spaces/AIRI-Institute/HairFastGAN"><img src="https://huggingface.co/datasets/huggingface/badges/resolve/main/open-in-hf-spaces-md.svg" height=22.5></a>
+<a href="https://colab.research.google.com/#fileId=https://huggingface.co/AIRI-Institute/HairFastGAN/blob/main/notebooks/HairFast_inference.ipynb"><img src="https://colab.research.google.com/assets/colab-badge.svg" height=22.5></a>
+[![License](https://img.shields.io/github/license/AIRI-Institute/al_toolbox)](./LICENSE)
+
+
+> Our paper addresses the complex task of transferring a hairstyle from a reference image to an input photo for virtual hair try-on. This task is challenging due to the need to adapt to various photo poses, the sensitivity of hairstyles, and the lack of objective metrics. The current state of the art hairstyle transfer methods use an optimization process for different parts of the approach, making them inexcusably slow. At the same time, faster encoder-based models are of very low quality because they either operate in StyleGAN's W+ space or use other low-dimensional image generators. Additionally, both approaches have a problem with hairstyle transfer when the source pose is very different from the target pose, because they either don't consider the pose at all or deal with it inefficiently. In our paper, we present the HairFast model, which uniquely solves these problems and achieves high resolution, near real-time performance, and superior reconstruction compared to optimization problem-based methods. Our solution includes a new architecture operating in the FS latent space of StyleGAN, an enhanced inpainting approach, and improved encoders for better alignment, color transfer, and a new encoder for post-processing. The effectiveness of our approach is demonstrated on realism metrics after random hairstyle transfer and reconstruction when the original hairstyle is transferred. In the most difficult scenario of transferring both shape and color of a hairstyle from different images, our method performs in less than a second on the Nvidia V100.
+> 
+
+<p align="center">
+  <img src="docs/assets/logo.webp" alt="Teaser"/>
+  <br>
+The proposed HairFast framework allows to edit a hairstyle on an arbitrary photo based on an example from other photos. Here we have an example of how the method works by transferring a hairstyle from one photo and a hair color from another.
+</p>
+
+## Updates
+
+- [25/09/2024] 🎉🎉🎉 HairFastGAN has been accepted by [NeurIPS 2024](https://nips.cc/virtual/2024/poster/93397).
+- [24/05/2024] 🌟🌟🌟 Release of the [official demo](https://huggingface.co/spaces/AIRI-Institute/HairFastGAN) on Hugging Face 🤗.
+- [01/04/2024] 🔥🔥🔥 HairFastGAN release.
+  
+## Prerequisites
+You need following hardware and python version to run our method.
+- Linux
+- NVIDIA GPU + CUDA CuDNN
+- Python 3.10
+- PyTorch 1.13.1+
+
+## Installation
+
+* Clone this repo:
+```bash
+git clone https://github.com/AIRI-Institute/HairFastGAN
+cd HairFastGAN
+```
+
+* Download all pretrained models:
+```bash
+git clone https://huggingface.co/AIRI-Institute/HairFastGAN
+cd HairFastGAN && git lfs pull && cd ..
+mv HairFastGAN/pretrained_models pretrained_models
+mv HairFastGAN/input input
+rm -rf HairFastGAN
+```
+
+* Setting the environment
+
+**Option 1 [recommended]**, install [Poetry](https://python-poetry.org/docs/) and then:
+```bash
+poetry install
+```
+
+**Option 2**, just install the dependencies in your environment:
+```bash
+pip install -r requirements.txt
+```
+
+## Inference
+You can use `main.py` to run the method, either for a single run or for a batch of experiments.
+
+* An example of running a single experiment:
+
+```
+python main.py --face_path=6.png --shape_path=7.png --color_path=8.png \
+    --input_dir=input --result_path=output/result.png
+```
+
+* To run the batch version, first create an image triples file (face/shape/color):
+```
+cat > example.txt << EOF
+6.png 7.png 8.png
+8.png 4.jpg 5.jpg
+EOF
+```
+
+And now you can run the method:
+```
+python main.py --file_path=example.txt --input_dir=input --output_dir=output
+```
+
+* You can use HairFast in the code directly:
+
+```python
+from hair_swap import HairFast, get_parser
+
+# Init HairFast
+hair_fast = HairFast(get_parser().parse_args([]))
+
+# Inference
+result = hair_fast(face_img, shape_img, color_img)
+```
+
+See the code for input parameters and output formats.
+
+* Alternatively, you can use our [Colab Notebook](https://colab.research.google.com/#fileId=https://huggingface.co/AIRI-Institute/HairFastGAN/blob/main/notebooks/HairFast_inference.ipynb) to prepare the environment, download the code, pretrained weights, and allow you to run experiments with a convenient form.
+
+* You can also try our method on the [Hugging Face demo](https://huggingface.co/spaces/AIRI-Institute/HairFastGAN) 🤗.
+
+## Scripts
+
+There is a list of scripts below, see arguments via --help for details.
+
+| Path                                    | Description <img width=200>
+|:----------------------------------------| :---
+| scripts/align_face.py                   | Processing of raw photos for inference
+| scripts/fid_metric.py                   | Metrics calculation
+| scripts/rotate_gen.py                   | Dataset generation for rotate encoder training
+| scripts/blending_gen.py                 | Dataset generation for color encoder training
+| scripts/pp_gen.py                       | Dataset generation for refinement encoder training
+| scripts/rotate_train.py                 | Rotate encoder training
+| scripts/blending_train.py               | Color encoder training
+| scripts/pp_train.py                     | Refinement encoder training
+
+
+## Training
+For training, you need to generate a dataset and then run the scripts for training. See the scripts section above.
+
+We use [Weights & Biases](https://wandb.ai/home) to track experiments. Before training, you should put your W&B API key into the `WANDB_KEY` environment variable.
+
+## Method diagram
+
+<p align="center">
+  <img src="https://huggingface.co/AIRI-Institute/HairFastGAN/resolve/main/docs/assets/diagram.webp" alt="Diagram"/>
+  <br>
+Overview of HairFast: the images first pass through the Pose alignment module, which generates a pose-aligned face mask with the desired hair shape. Then we transfer the desired hairstyle shape using Shape alignment and the desired hair color using Color alignment. In the last step, Refinement alignment returns the lost details of the original image where they are needed.
+</p>
+
+## Repository structure
+
+    .
+    ├── 📂 datasets                   # Implementation of torch datasets for inference
+    ├── 📂 docs                       # Folder with method diagram and teaser
+    ├── 📂 models                     # Folder containting all the models
+    │   ├── ...
+    │   ├── 📄 Embedding.py               # Implementation of Embedding module
+    │   ├── 📄 Alignment.py               # Implementation of Pose and Shape alignment modules
+    │   ├── 📄 Blending.py                # Implementation of Color and Refinement alignment modules
+    │   ├── 📄 Encoders.py                # Implementation of encoder architectures
+    │   └── 📄 Net.py                     # Implementation of basic models
+    │
+    ├── 📂 losses                     # Folder containing various loss criterias for training
+    ├── 📂 scripts                    # Folder with various scripts
+    ├── 📂 utils                      # Folder with utility functions
+    │
+    ├── 📜 poetry.lock                # Records exact dependency versions.
+    ├── 📜 pyproject.toml             # Poetry configuration for dependencies.
+    ├── 📜 requirements.txt           # Lists required Python packages.
+    ├── 📄 hair_swap.py               # Implementation of the HairFast main class
+    └── 📄 main.py                    # Script for inference
+
+## References & Acknowledgments
+
+The repository was started from [Barbershop](https://github.com/ZPdesu/Barbershop).
+
+The code [CtrlHair](https://github.com/XuyangGuo/CtrlHair), [SEAN](https://github.com/ZPdesu/SEAN), [HairCLIP](https://github.com/wty-ustc/HairCLIP), [FSE](https://github.com/InterDigitalInc/FeatureStyleEncoder), [E4E](https://github.com/omertov/encoder4editing) and [STAR](https://github.com/ZhenglinZhou/STAR) was also used.
+
+## Citation
+
+If you use this code for your research, please cite our paper:
+```
+@article{nikolaev2024hairfastgan,
+  title={HairFastGAN: Realistic and Robust Hair Transfer with a Fast Encoder-Based Approach},
+  author={Nikolaev, Maxim and Kuznetsov, Mikhail and Vetrov, Dmitry and Alanov, Aibek},
+  journal={arXiv preprint arXiv:2404.01094},
+  year={2024}
+}
+```

datasets/image_dataset.py

@@ -0,0 +1,29 @@
+import torch
+from torch.utils.data import Dataset
+
+
+class ImagesDataset(Dataset):
+    def __init__(self, images: dict[torch.Tensor, list[str]] | list[torch.Tensor]):
+        if isinstance(images, list):
+            images = dict.fromkeys(images)
+
+        self.images = list(images)
+        self.names = list(images.values())
+
+    def __len__(self):
+        return len(self.images)
+
+    def __getitem__(self, index):
+        image = self.images[index]
+
+        if image.dtype is torch.uint8:
+            image = image / 255
+
+        names = self.names[index]
+        return image, names
+
+
+def image_collate(batch):
+    images = torch.stack([item[0] for item in batch])
+    names = [item[1] for item in batch]
+    return images, names

hair_swap.py

@@ -0,0 +1,139 @@
+import argparse
+import typing as tp
+from collections import defaultdict
+from functools import wraps
+from pathlib import Path
+
+import numpy as np
+import torch
+import torchvision.transforms.functional as F
+from PIL import Image
+from torchvision.io import read_image, ImageReadMode
+
+from models.Alignment import Alignment
+from models.Blending import Blending
+from models.Embedding import Embedding
+from models.Net import Net
+from utils.image_utils import equal_replacer
+from utils.seed import seed_setter
+from utils.shape_predictor import align_face
+from utils.time import bench_session
+
+TImage = tp.TypeVar('TImage', torch.Tensor, Image.Image, np.ndarray)
+TPath = tp.TypeVar('TPath', Path, str)
+TReturn = tp.TypeVar('TReturn', torch.Tensor, tuple[torch.Tensor, ...])
+
+
+class HairFast:
+    """
+    HairFast implementation with hairstyle transfer interface
+    """
+
+    def __init__(self, args):
+        self.args = args
+        self.net = Net(self.args)
+        self.embed = Embedding(args, net=self.net)
+        self.align = Alignment(args, self.embed.get_e4e_embed, net=self.net)
+        self.blend = Blending(args, net=self.net)
+
+    @seed_setter
+    @bench_session
+    def __swap_from_tensors(self, face: torch.Tensor, shape: torch.Tensor, color: torch.Tensor,
+                            **kwargs) -> torch.Tensor:
+        images_to_name = defaultdict(list)
+        for image, name in zip((face, shape, color), ('face', 'shape', 'color')):
+            images_to_name[image].append(name)
+
+        # Embedding stage
+        name_to_embed = self.embed.embedding_images(images_to_name, **kwargs)
+
+        # Alignment stage
+        align_shape = self.align.align_images('face', 'shape', name_to_embed, **kwargs)
+
+        # Shape Module stage for blending
+        if shape is not color:
+            align_color = self.align.shape_module('face', 'color', name_to_embed, **kwargs)
+        else:
+            align_color = align_shape
+
+        # Blending and Post Process stage
+        final_image = self.blend.blend_images(align_shape, align_color, name_to_embed, **kwargs)
+        return final_image
+
+    def swap(self, face_img: TImage | TPath, shape_img: TImage | TPath, color_img: TImage | TPath,
+             benchmark=False, align=False, seed=None, exp_name=None, **kwargs) -> TReturn:
+        """
+        Run HairFast on the input images to transfer hair shape and color to the desired images.
+        :param face_img:  face image in Tensor, PIL Image, array or file path format
+        :param shape_img: shape image in Tensor, PIL Image, array or file path format
+        :param color_img: color image in Tensor, PIL Image, array or file path format
+        :param benchmark: starts counting the speed of the session
+        :param align:     for arbitrary photos crops images to faces
+        :param seed:      fixes seed for reproducibility, default 3407
+        :param exp_name:  used as a folder name when 'save_all' model is enabled
+        :return:          returns the final image as a Tensor
+        """
+        images: list[torch.Tensor] = []
+        path_to_images: dict[TPath, torch.Tensor] = {}
+
+        for img in (face_img, shape_img, color_img):
+            if isinstance(img, (torch.Tensor, Image.Image, np.ndarray)):
+                if not isinstance(img, torch.Tensor):
+                    img = F.to_tensor(img)
+            elif isinstance(img, (Path, str)):
+                path_img = img
+                if path_img not in path_to_images:
+                    path_to_images[path_img] = read_image(str(path_img), mode=ImageReadMode.RGB)
+                img = path_to_images[path_img]
+            else:
+                raise TypeError(f'Unsupported image format {type(img)}')
+
+            images.append(img)
+
+        if align:
+            images = align_face(images)
+        images = equal_replacer(images)
+
+        final_image = self.__swap_from_tensors(*images, seed=seed, benchmark=benchmark, exp_name=exp_name, **kwargs)
+
+        if align:
+            return final_image, *images
+        return final_image
+
+    @wraps(swap)
+    def __call__(self, *args, **kwargs):
+        return self.swap(*args, **kwargs)
+
+
+def get_parser():
+    parser = argparse.ArgumentParser(description='HairFast')
+
+    # I/O arguments
+    parser.add_argument('--save_all_dir', type=Path, default=Path('output'),
+                        help='the directory to save the latent codes and inversion images')
+
+    # StyleGAN2 setting
+    parser.add_argument('--size', type=int, default=1024)
+    parser.add_argument('--ckpt', type=str, default="pretrained_models/StyleGAN/ffhq.pt")
+    parser.add_argument('--channel_multiplier', type=int, default=2)
+    parser.add_argument('--latent', type=int, default=512)
+    parser.add_argument('--n_mlp', type=int, default=8)
+
+    # Arguments
+    parser.add_argument('--device', type=str, default='cuda')
+    parser.add_argument('--batch_size', type=int, default=3, help='batch size for encoding images')
+    parser.add_argument('--save_all', action='store_true', help='save and print mode information')
+
+    # HairFast setting
+    parser.add_argument('--mixing', type=float, default=0.95, help='hair blending in alignment')
+    parser.add_argument('--smooth', type=int, default=5, help='dilation and erosion parameter')
+    parser.add_argument('--rotate_checkpoint', type=str, default='pretrained_models/Rotate/rotate_best.pth')
+    parser.add_argument('--blending_checkpoint', type=str, default='pretrained_models/Blending/checkpoint.pth')
+    parser.add_argument('--pp_checkpoint', type=str, default='pretrained_models/PostProcess/pp_model.pth')
+    return parser
+
+
+if __name__ == '__main__':
+    model_args = get_parser()
+    args = model_args.parse_args()
+    hair_fast = HairFast(args)

inference_server.py

@@ -0,0 +1,5 @@
+if __name__ == "__main__":
+    server = grpc.server(...)
+    ...
+    server.start()
+    server.wait_for_termination()

losses/lpips/__init__.py

@@ -0,0 +1,160 @@
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+from skimage.metrics import structural_similarity
+import torch
+from torch.autograd import Variable
+
+from ..lpips import dist_model
+
+class PerceptualLoss(torch.nn.Module):
+    def __init__(self, model='net-lin', net='alex', colorspace='rgb', spatial=False, use_gpu=True, gpu_ids=[0]): # VGG using our perceptually-learned weights (LPIPS metric)
+    # def __init__(self, model='net', net='vgg', use_gpu=True): # "default" way of using VGG as a perceptual loss
+        super(PerceptualLoss, self).__init__()
+        print('Setting up Perceptual loss...')
+        self.use_gpu = use_gpu
+        self.spatial = spatial
+        self.gpu_ids = gpu_ids
+        self.model = dist_model.DistModel()
+        self.model.initialize(model=model, net=net, use_gpu=use_gpu, colorspace=colorspace, spatial=self.spatial, gpu_ids=gpu_ids)
+        print('...[%s] initialized'%self.model.name())
+        print('...Done')
+
+    def forward(self, pred, target, normalize=False):
+        """
+        Pred and target are Variables.
+        If normalize is True, assumes the images are between [0,1] and then scales them between [-1,+1]
+        If normalize is False, assumes the images are already between [-1,+1]
+
+        Inputs pred and target are Nx3xHxW
+        Output pytorch Variable N long
+        """
+
+        if normalize:
+            target = 2 * target  - 1
+            pred = 2 * pred  - 1
+
+        return self.model.forward(target, pred)
+
+def normalize_tensor(in_feat,eps=1e-10):
+    norm_factor = torch.sqrt(torch.sum(in_feat**2,dim=1,keepdim=True))
+    return in_feat/(norm_factor+eps)
+
+def l2(p0, p1, range=255.):
+    return .5*np.mean((p0 / range - p1 / range)**2)
+
+def psnr(p0, p1, peak=255.):
+    return 10*np.log10(peak**2/np.mean((1.*p0-1.*p1)**2))
+
+def dssim(p0, p1, range=255.):
+    return (1 - structural_similarity(p0, p1, data_range=range, multichannel=True)) / 2.
+
+def rgb2lab(in_img,mean_cent=False):
+    from skimage import color
+    img_lab = color.rgb2lab(in_img)
+    if(mean_cent):
+        img_lab[:,:,0] = img_lab[:,:,0]-50
+    return img_lab
+
+def tensor2np(tensor_obj):
+    # change dimension of a tensor object into a numpy array
+    return tensor_obj[0].cpu().float().numpy().transpose((1,2,0))
+
+def np2tensor(np_obj):
+     # change dimenion of np array into tensor array
+    return torch.Tensor(np_obj[:, :, :, np.newaxis].transpose((3, 2, 0, 1)))
+
+def tensor2tensorlab(image_tensor,to_norm=True,mc_only=False):
+    # image tensor to lab tensor
+    from skimage import color
+
+    img = tensor2im(image_tensor)
+    img_lab = color.rgb2lab(img)
+    if(mc_only):
+        img_lab[:,:,0] = img_lab[:,:,0]-50
+    if(to_norm and not mc_only):
+        img_lab[:,:,0] = img_lab[:,:,0]-50
+        img_lab = img_lab/100.
+
+    return np2tensor(img_lab)
+
+def tensorlab2tensor(lab_tensor,return_inbnd=False):
+    from skimage import color
+    import warnings
+    warnings.filterwarnings("ignore")
+
+    lab = tensor2np(lab_tensor)*100.
+    lab[:,:,0] = lab[:,:,0]+50
+
+    rgb_back = 255.*np.clip(color.lab2rgb(lab.astype('float')),0,1)
+    if(return_inbnd):
+        # convert back to lab, see if we match
+        lab_back = color.rgb2lab(rgb_back.astype('uint8'))
+        mask = 1.*np.isclose(lab_back,lab,atol=2.)
+        mask = np2tensor(np.prod(mask,axis=2)[:,:,np.newaxis])
+        return (im2tensor(rgb_back),mask)
+    else:
+        return im2tensor(rgb_back)
+
+def rgb2lab(input):
+    from skimage import color
+    return color.rgb2lab(input / 255.)
+
+def tensor2im(image_tensor, imtype=np.uint8, cent=1., factor=255./2.):
+    image_numpy = image_tensor[0].cpu().float().numpy()
+    image_numpy = (np.transpose(image_numpy, (1, 2, 0)) + cent) * factor
+    return image_numpy.astype(imtype)
+
+def im2tensor(image, imtype=np.uint8, cent=1., factor=255./2.):
+    return torch.Tensor((image / factor - cent)
+                        [:, :, :, np.newaxis].transpose((3, 2, 0, 1)))
+
+def tensor2vec(vector_tensor):
+    return vector_tensor.data.cpu().numpy()[:, :, 0, 0]
+
+def voc_ap(rec, prec, use_07_metric=False):
+    """ ap = voc_ap(rec, prec, [use_07_metric])
+    Compute VOC AP given precision and recall.
+    If use_07_metric is true, uses the
+    VOC 07 11 point method (default:False).
+    """
+    if use_07_metric:
+        # 11 point metric
+        ap = 0.
+        for t in np.arange(0., 1.1, 0.1):
+            if np.sum(rec >= t) == 0:
+                p = 0
+            else:
+                p = np.max(prec[rec >= t])
+            ap = ap + p / 11.
+    else:
+        # correct AP calculation
+        # first append sentinel values at the end
+        mrec = np.concatenate(([0.], rec, [1.]))
+        mpre = np.concatenate(([0.], prec, [0.]))
+
+        # compute the precision envelope
+        for i in range(mpre.size - 1, 0, -1):
+            mpre[i - 1] = np.maximum(mpre[i - 1], mpre[i])
+
+        # to calculate area under PR curve, look for points
+        # where X axis (recall) changes value
+        i = np.where(mrec[1:] != mrec[:-1])[0]
+
+        # and sum (\Delta recall) * prec
+        ap = np.sum((mrec[i + 1] - mrec[i]) * mpre[i + 1])
+    return ap
+
+def tensor2im(image_tensor, imtype=np.uint8, cent=1., factor=255./2.):
+# def tensor2im(image_tensor, imtype=np.uint8, cent=1., factor=1.):
+    image_numpy = image_tensor[0].cpu().float().numpy()
+    image_numpy = (np.transpose(image_numpy, (1, 2, 0)) + cent) * factor
+    return image_numpy.astype(imtype)
+
+def im2tensor(image, imtype=np.uint8, cent=1., factor=255./2.):
+# def im2tensor(image, imtype=np.uint8, cent=1., factor=1.):
+    return torch.Tensor((image / factor - cent)
+                        [:, :, :, np.newaxis].transpose((3, 2, 0, 1)))

losses/lpips/base_model.py

@@ -0,0 +1,58 @@
+import os
+import numpy as np
+import torch
+from torch.autograd import Variable
+from pdb import set_trace as st
+from IPython import embed
+
+class BaseModel():
+    def __init__(self):
+        pass;
+        
+    def name(self):
+        return 'BaseModel'
+
+    def initialize(self, use_gpu=True, gpu_ids=[0]):
+        self.use_gpu = use_gpu
+        self.gpu_ids = gpu_ids
+
+    def forward(self):
+        pass
+
+    def get_image_paths(self):
+        pass
+
+    def optimize_parameters(self):
+        pass
+
+    def get_current_visuals(self):
+        return self.input
+
+    def get_current_errors(self):
+        return {}
+
+    def save(self, label):
+        pass
+
+    # helper saving function that can be used by subclasses
+    def save_network(self, network, path, network_label, epoch_label):
+        save_filename = '%s_net_%s.pth' % (epoch_label, network_label)
+        save_path = os.path.join(path, save_filename)
+        torch.save(network.state_dict(), save_path)
+
+    # helper loading function that can be used by subclasses
+    def load_network(self, network, network_label, epoch_label):
+        save_filename = '%s_net_%s.pth' % (epoch_label, network_label)
+        save_path = os.path.join(self.save_dir, save_filename)
+        print('Loading network from %s'%save_path)
+        network.load_state_dict(torch.load(save_path))
+
+    def update_learning_rate():
+        pass
+
+    def get_image_paths(self):
+        return self.image_paths
+
+    def save_done(self, flag=False):
+        np.save(os.path.join(self.save_dir, 'done_flag'),flag)
+        np.savetxt(os.path.join(self.save_dir, 'done_flag'),[flag,],fmt='%i')

losses/lpips/dist_model.py

@@ -0,0 +1,284 @@
+
+from __future__ import absolute_import
+
+import sys
+import numpy as np
+import torch
+from torch import nn
+import os
+from collections import OrderedDict
+from torch.autograd import Variable
+import itertools
+from .base_model import BaseModel
+from scipy.ndimage import zoom
+import fractions
+import functools
+import skimage.transform
+from tqdm import tqdm
+
+from IPython import embed
+
+from . import networks_basic as networks
+from losses import lpips as util
+
+class DistModel(BaseModel):
+    def name(self):
+        return self.model_name
+
+    def initialize(self, model='net-lin', net='alex', colorspace='Lab', pnet_rand=False, pnet_tune=False, model_path=None,
+            use_gpu=True, printNet=False, spatial=False, 
+            is_train=False, lr=.0001, beta1=0.5, version='0.1', gpu_ids=[0]):
+        '''
+        INPUTS
+            model - ['net-lin'] for linearly calibrated network
+                    ['net'] for off-the-shelf network
+                    ['L2'] for L2 distance in Lab colorspace
+                    ['SSIM'] for ssim in RGB colorspace
+            net - ['squeeze','alex','vgg']
+            model_path - if None, will look in weights/[NET_NAME].pth
+            colorspace - ['Lab','RGB'] colorspace to use for L2 and SSIM
+            use_gpu - bool - whether or not to use a GPU
+            printNet - bool - whether or not to print network architecture out
+            spatial - bool - whether to output an array containing varying distances across spatial dimensions
+            spatial_shape - if given, output spatial shape. if None then spatial shape is determined automatically via spatial_factor (see below).
+            spatial_factor - if given, specifies upsampling factor relative to the largest spatial extent of a convolutional layer. if None then resized to size of input images.
+            spatial_order - spline order of filter for upsampling in spatial mode, by default 1 (bilinear).
+            is_train - bool - [True] for training mode
+            lr - float - initial learning rate
+            beta1 - float - initial momentum term for adam
+            version - 0.1 for latest, 0.0 was original (with a bug)
+            gpu_ids - int array - [0] by default, gpus to use
+        '''
+        BaseModel.initialize(self, use_gpu=use_gpu, gpu_ids=gpu_ids)
+
+        self.model = model
+        self.net = net
+        self.is_train = is_train
+        self.spatial = spatial
+        self.gpu_ids = gpu_ids
+        self.model_name = '%s [%s]'%(model,net)
+
+        if(self.model == 'net-lin'): # pretrained net + linear layer
+            self.net = networks.PNetLin(pnet_rand=pnet_rand, pnet_tune=pnet_tune, pnet_type=net,
+                use_dropout=True, spatial=spatial, version=version, lpips=True)
+            kw = {}
+            if not use_gpu:
+                kw['map_location'] = 'cpu'
+            if(model_path is None):
+                import inspect
+                model_path = os.path.abspath(os.path.join(inspect.getfile(self.initialize), '..', 'weights/v%s/%s.pth'%(version,net)))
+
+            if(not is_train):
+                print('Loading model from: %s'%model_path)
+                self.net.load_state_dict(torch.load(model_path, **kw), strict=False)
+
+        elif(self.model=='net'): # pretrained network
+            self.net = networks.PNetLin(pnet_rand=pnet_rand, pnet_type=net, lpips=False)
+        elif(self.model in ['L2','l2']):
+            self.net = networks.L2(use_gpu=use_gpu,colorspace=colorspace) # not really a network, only for testing
+            self.model_name = 'L2'
+        elif(self.model in ['DSSIM','dssim','SSIM','ssim']):
+            self.net = networks.DSSIM(use_gpu=use_gpu,colorspace=colorspace)
+            self.model_name = 'SSIM'
+        else:
+            raise ValueError("Model [%s] not recognized." % self.model)
+
+        self.parameters = list(self.net.parameters())
+
+        if self.is_train: # training mode
+            # extra network on top to go from distances (d0,d1) => predicted human judgment (h*)
+            self.rankLoss = networks.BCERankingLoss()
+            self.parameters += list(self.rankLoss.net.parameters())
+            self.lr = lr
+            self.old_lr = lr
+            self.optimizer_net = torch.optim.Adam(self.parameters, lr=lr, betas=(beta1, 0.999))
+        else: # test mode
+            self.net.eval()
+
+        if(use_gpu):
+            self.net.to(gpu_ids[0])
+            self.net = torch.nn.DataParallel(self.net, device_ids=gpu_ids)
+            if(self.is_train):
+                self.rankLoss = self.rankLoss.to(device=gpu_ids[0]) # just put this on GPU0
+
+        if(printNet):
+            print('---------- Networks initialized -------------')
+            networks.print_network(self.net)
+            print('-----------------------------------------------')
+
+    def forward(self, in0, in1, retPerLayer=False):
+        ''' Function computes the distance between image patches in0 and in1
+        INPUTS
+            in0, in1 - torch.Tensor object of shape Nx3xXxY - image patch scaled to [-1,1]
+        OUTPUT
+            computed distances between in0 and in1
+        '''
+
+        return self.net.forward(in0, in1, retPerLayer=retPerLayer)
+
+    # ***** TRAINING FUNCTIONS *****
+    def optimize_parameters(self):
+        self.forward_train()
+        self.optimizer_net.zero_grad()
+        self.backward_train()
+        self.optimizer_net.step()
+        self.clamp_weights()
+
+    def clamp_weights(self):
+        for module in self.net.modules():
+            if(hasattr(module, 'weight') and module.kernel_size==(1,1)):
+                module.weight.data = torch.clamp(module.weight.data,min=0)
+
+    def set_input(self, data):
+        self.input_ref = data['ref']
+        self.input_p0 = data['p0']
+        self.input_p1 = data['p1']
+        self.input_judge = data['judge']
+
+        if(self.use_gpu):
+            self.input_ref = self.input_ref.to(device=self.gpu_ids[0])
+            self.input_p0 = self.input_p0.to(device=self.gpu_ids[0])
+            self.input_p1 = self.input_p1.to(device=self.gpu_ids[0])
+            self.input_judge = self.input_judge.to(device=self.gpu_ids[0])
+
+        self.var_ref = Variable(self.input_ref,requires_grad=True)
+        self.var_p0 = Variable(self.input_p0,requires_grad=True)
+        self.var_p1 = Variable(self.input_p1,requires_grad=True)
+
+    def forward_train(self): # run forward pass
+        # print(self.net.module.scaling_layer.shift)
+        # print(torch.norm(self.net.module.net.slice1[0].weight).item(), torch.norm(self.net.module.lin0.model[1].weight).item())
+
+        self.d0 = self.forward(self.var_ref, self.var_p0)
+        self.d1 = self.forward(self.var_ref, self.var_p1)
+        self.acc_r = self.compute_accuracy(self.d0,self.d1,self.input_judge)
+
+        self.var_judge = Variable(1.*self.input_judge).view(self.d0.size())
+
+        self.loss_total = self.rankLoss.forward(self.d0, self.d1, self.var_judge*2.-1.)
+
+        return self.loss_total
+
+    def backward_train(self):
+        torch.mean(self.loss_total).backward()
+
+    def compute_accuracy(self,d0,d1,judge):
+        ''' d0, d1 are Variables, judge is a Tensor '''
+        d1_lt_d0 = (d1<d0).cpu().data.numpy().flatten()
+        judge_per = judge.cpu().numpy().flatten()
+        return d1_lt_d0*judge_per + (1-d1_lt_d0)*(1-judge_per)
+
+    def get_current_errors(self):
+        retDict = OrderedDict([('loss_total', self.loss_total.data.cpu().numpy()),
+                            ('acc_r', self.acc_r)])
+
+        for key in retDict.keys():
+            retDict[key] = np.mean(retDict[key])
+
+        return retDict
+
+    def get_current_visuals(self):
+        zoom_factor = 256/self.var_ref.data.size()[2]
+
+        ref_img = util.tensor2im(self.var_ref.data)
+        p0_img = util.tensor2im(self.var_p0.data)
+        p1_img = util.tensor2im(self.var_p1.data)
+
+        ref_img_vis = zoom(ref_img,[zoom_factor, zoom_factor, 1],order=0)
+        p0_img_vis = zoom(p0_img,[zoom_factor, zoom_factor, 1],order=0)
+        p1_img_vis = zoom(p1_img,[zoom_factor, zoom_factor, 1],order=0)
+
+        return OrderedDict([('ref', ref_img_vis),
+                            ('p0', p0_img_vis),
+                            ('p1', p1_img_vis)])
+
+    def save(self, path, label):
+        if(self.use_gpu):
+            self.save_network(self.net.module, path, '', label)
+        else:
+            self.save_network(self.net, path, '', label)
+        self.save_network(self.rankLoss.net, path, 'rank', label)
+
+    def update_learning_rate(self,nepoch_decay):
+        lrd = self.lr / nepoch_decay
+        lr = self.old_lr - lrd
+
+        for param_group in self.optimizer_net.param_groups:
+            param_group['lr'] = lr
+
+        print('update lr [%s] decay: %f -> %f' % (type,self.old_lr, lr))
+        self.old_lr = lr
+
+def score_2afc_dataset(data_loader, func, name=''):
+    ''' Function computes Two Alternative Forced Choice (2AFC) score using
+        distance function 'func' in dataset 'data_loader'
+    INPUTS
+        data_loader - CustomDatasetDataLoader object - contains a TwoAFCDataset inside
+        func - callable distance function - calling d=func(in0,in1) should take 2
+            pytorch tensors with shape Nx3xXxY, and return numpy array of length N
+    OUTPUTS
+        [0] - 2AFC score in [0,1], fraction of time func agrees with human evaluators
+        [1] - dictionary with following elements
+            d0s,d1s - N arrays containing distances between reference patch to perturbed patches 
+            gts - N array in [0,1], preferred patch selected by human evaluators
+                (closer to "0" for left patch p0, "1" for right patch p1,
+                "0.6" means 60pct people preferred right patch, 40pct preferred left)
+            scores - N array in [0,1], corresponding to what percentage function agreed with humans
+    CONSTS
+        N - number of test triplets in data_loader
+    '''
+
+    d0s = []
+    d1s = []
+    gts = []
+
+    for data in tqdm(data_loader.load_data(), desc=name):
+        d0s+=func(data['ref'],data['p0']).data.cpu().numpy().flatten().tolist()
+        d1s+=func(data['ref'],data['p1']).data.cpu().numpy().flatten().tolist()
+        gts+=data['judge'].cpu().numpy().flatten().tolist()
+
+    d0s = np.array(d0s)
+    d1s = np.array(d1s)
+    gts = np.array(gts)
+    scores = (d0s<d1s)*(1.-gts) + (d1s<d0s)*gts + (d1s==d0s)*.5
+
+    return(np.mean(scores), dict(d0s=d0s,d1s=d1s,gts=gts,scores=scores))
+
+def score_jnd_dataset(data_loader, func, name=''):
+    ''' Function computes JND score using distance function 'func' in dataset 'data_loader'
+    INPUTS
+        data_loader - CustomDatasetDataLoader object - contains a JNDDataset inside
+        func - callable distance function - calling d=func(in0,in1) should take 2
+            pytorch tensors with shape Nx3xXxY, and return pytorch array of length N
+    OUTPUTS
+        [0] - JND score in [0,1], mAP score (area under precision-recall curve)
+        [1] - dictionary with following elements
+            ds - N array containing distances between two patches shown to human evaluator
+            sames - N array containing fraction of people who thought the two patches were identical
+    CONSTS
+        N - number of test triplets in data_loader
+    '''
+
+    ds = []
+    gts = []
+
+    for data in tqdm(data_loader.load_data(), desc=name):
+        ds+=func(data['p0'],data['p1']).data.cpu().numpy().tolist()
+        gts+=data['same'].cpu().numpy().flatten().tolist()
+
+    sames = np.array(gts)
+    ds = np.array(ds)
+
+    sorted_inds = np.argsort(ds)
+    ds_sorted = ds[sorted_inds]
+    sames_sorted = sames[sorted_inds]
+
+    TPs = np.cumsum(sames_sorted)
+    FPs = np.cumsum(1-sames_sorted)
+    FNs = np.sum(sames_sorted)-TPs
+
+    precs = TPs/(TPs+FPs)
+    recs = TPs/(TPs+FNs)
+    score = util.voc_ap(recs,precs)
+
+    return(score, dict(ds=ds,sames=sames))

losses/lpips/networks_basic.py

@@ -0,0 +1,187 @@
+
+from __future__ import absolute_import
+
+import sys
+import torch
+import torch.nn as nn
+import torch.nn.init as init
+from torch.autograd import Variable
+import numpy as np
+from pdb import set_trace as st
+from skimage import color
+from IPython import embed
+from . import pretrained_networks as pn
+
+from losses import lpips as util
+
+def spatial_average(in_tens, keepdim=True):
+    return in_tens.mean([2,3],keepdim=keepdim)
+
+def upsample(in_tens, out_H=64): # assumes scale factor is same for H and W
+    in_H = in_tens.shape[2]
+    scale_factor = 1.*out_H/in_H
+
+    return nn.Upsample(scale_factor=scale_factor, mode='bilinear', align_corners=False)(in_tens)
+
+# Learned perceptual metric
+class PNetLin(nn.Module):
+    def __init__(self, pnet_type='vgg', pnet_rand=False, pnet_tune=False, use_dropout=True, spatial=False, version='0.1', lpips=True):
+        super(PNetLin, self).__init__()
+
+        self.pnet_type = pnet_type
+        self.pnet_tune = pnet_tune
+        self.pnet_rand = pnet_rand
+        self.spatial = spatial
+        self.lpips = lpips
+        self.version = version
+        self.scaling_layer = ScalingLayer()
+
+        if(self.pnet_type in ['vgg','vgg16']):
+            net_type = pn.vgg16
+            self.chns = [64,128,256,512,512]
+        elif(self.pnet_type=='alex'):
+            net_type = pn.alexnet
+            self.chns = [64,192,384,256,256]
+        elif(self.pnet_type=='squeeze'):
+            net_type = pn.squeezenet
+            self.chns = [64,128,256,384,384,512,512]
+        self.L = len(self.chns)
+
+        self.net = net_type(pretrained=not self.pnet_rand, requires_grad=self.pnet_tune)
+
+        if(lpips):
+            self.lin0 = NetLinLayer(self.chns[0], use_dropout=use_dropout)
+            self.lin1 = NetLinLayer(self.chns[1], use_dropout=use_dropout)
+            self.lin2 = NetLinLayer(self.chns[2], use_dropout=use_dropout)
+            self.lin3 = NetLinLayer(self.chns[3], use_dropout=use_dropout)
+            self.lin4 = NetLinLayer(self.chns[4], use_dropout=use_dropout)
+            self.lins = [self.lin0,self.lin1,self.lin2,self.lin3,self.lin4]
+            if(self.pnet_type=='squeeze'): # 7 layers for squeezenet
+                self.lin5 = NetLinLayer(self.chns[5], use_dropout=use_dropout)
+                self.lin6 = NetLinLayer(self.chns[6], use_dropout=use_dropout)
+                self.lins+=[self.lin5,self.lin6]
+
+    def forward(self, in0, in1, retPerLayer=False):
+        # v0.0 - original release had a bug, where input was not scaled
+        in0_input, in1_input = (self.scaling_layer(in0), self.scaling_layer(in1)) if self.version=='0.1' else (in0, in1)
+        outs0, outs1 = self.net.forward(in0_input), self.net.forward(in1_input)
+        feats0, feats1, diffs = {}, {}, {}
+
+        for kk in range(self.L):
+            feats0[kk], feats1[kk] = util.normalize_tensor(outs0[kk]), util.normalize_tensor(outs1[kk])
+            diffs[kk] = (feats0[kk]-feats1[kk])**2
+
+        if(self.lpips):
+            if(self.spatial):
+                res = [upsample(self.lins[kk].model(diffs[kk]), out_H=in0.shape[2]) for kk in range(self.L)]
+            else:
+                res = [spatial_average(self.lins[kk].model(diffs[kk]), keepdim=True) for kk in range(self.L)]
+        else:
+            if(self.spatial):
+                res = [upsample(diffs[kk].sum(dim=1,keepdim=True), out_H=in0.shape[2]) for kk in range(self.L)]
+            else:
+                res = [spatial_average(diffs[kk].sum(dim=1,keepdim=True), keepdim=True) for kk in range(self.L)]
+
+        val = res[0]
+        for l in range(1,self.L):
+            val += res[l]
+        
+        if(retPerLayer):
+            return (val, res)
+        else:
+            return val
+
+class ScalingLayer(nn.Module):
+    def __init__(self):
+        super(ScalingLayer, self).__init__()
+        self.register_buffer('shift', torch.Tensor([-.030,-.088,-.188])[None,:,None,None])
+        self.register_buffer('scale', torch.Tensor([.458,.448,.450])[None,:,None,None])
+
+    def forward(self, inp):
+        return (inp - self.shift) / self.scale
+
+
+class NetLinLayer(nn.Module):
+    ''' A single linear layer which does a 1x1 conv '''
+    def __init__(self, chn_in, chn_out=1, use_dropout=False):
+        super(NetLinLayer, self).__init__()
+
+        layers = [nn.Dropout(),] if(use_dropout) else []
+        layers += [nn.Conv2d(chn_in, chn_out, 1, stride=1, padding=0, bias=False),]
+        self.model = nn.Sequential(*layers)
+
+
+class Dist2LogitLayer(nn.Module):
+    ''' takes 2 distances, puts through fc layers, spits out value between [0,1] (if use_sigmoid is True) '''
+    def __init__(self, chn_mid=32, use_sigmoid=True):
+        super(Dist2LogitLayer, self).__init__()
+
+        layers = [nn.Conv2d(5, chn_mid, 1, stride=1, padding=0, bias=True),]
+        layers += [nn.LeakyReLU(0.2,True),]
+        layers += [nn.Conv2d(chn_mid, chn_mid, 1, stride=1, padding=0, bias=True),]
+        layers += [nn.LeakyReLU(0.2,True),]
+        layers += [nn.Conv2d(chn_mid, 1, 1, stride=1, padding=0, bias=True),]
+        if(use_sigmoid):
+            layers += [nn.Sigmoid(),]
+        self.model = nn.Sequential(*layers)
+
+    def forward(self,d0,d1,eps=0.1):
+        return self.model.forward(torch.cat((d0,d1,d0-d1,d0/(d1+eps),d1/(d0+eps)),dim=1))
+
+class BCERankingLoss(nn.Module):
+    def __init__(self, chn_mid=32):
+        super(BCERankingLoss, self).__init__()
+        self.net = Dist2LogitLayer(chn_mid=chn_mid)
+        # self.parameters = list(self.net.parameters())
+        self.loss = torch.nn.BCELoss()
+
+    def forward(self, d0, d1, judge):
+        per = (judge+1.)/2.
+        self.logit = self.net.forward(d0,d1)
+        return self.loss(self.logit, per)
+
+# L2, DSSIM metrics
+class FakeNet(nn.Module):
+    def __init__(self, use_gpu=True, colorspace='Lab'):
+        super(FakeNet, self).__init__()
+        self.use_gpu = use_gpu
+        self.colorspace=colorspace
+
+class L2(FakeNet):
+
+    def forward(self, in0, in1, retPerLayer=None):
+        assert(in0.size()[0]==1) # currently only supports batchSize 1
+
+        if(self.colorspace=='RGB'):
+            (N,C,X,Y) = in0.size()
+            value = torch.mean(torch.mean(torch.mean((in0-in1)**2,dim=1).view(N,1,X,Y),dim=2).view(N,1,1,Y),dim=3).view(N)
+            return value
+        elif(self.colorspace=='Lab'):
+            value = util.l2(util.tensor2np(util.tensor2tensorlab(in0.data,to_norm=False)), 
+                util.tensor2np(util.tensor2tensorlab(in1.data,to_norm=False)), range=100.).astype('float')
+            ret_var = Variable( torch.Tensor((value,) ) )
+            if(self.use_gpu):
+                ret_var = ret_var.cuda()
+            return ret_var
+
+class DSSIM(FakeNet):
+
+    def forward(self, in0, in1, retPerLayer=None):
+        assert(in0.size()[0]==1) # currently only supports batchSize 1
+
+        if(self.colorspace=='RGB'):
+            value = util.dssim(1.*util.tensor2im(in0.data), 1.*util.tensor2im(in1.data), range=255.).astype('float')
+        elif(self.colorspace=='Lab'):
+            value = util.dssim(util.tensor2np(util.tensor2tensorlab(in0.data,to_norm=False)), 
+                util.tensor2np(util.tensor2tensorlab(in1.data,to_norm=False)), range=100.).astype('float')
+        ret_var = Variable( torch.Tensor((value,) ) )
+        if(self.use_gpu):
+            ret_var = ret_var.cuda()
+        return ret_var
+
+def print_network(net):
+    num_params = 0
+    for param in net.parameters():
+        num_params += param.numel()
+    print('Network',net)
+    print('Total number of parameters: %d' % num_params)

losses/lpips/pretrained_networks.py

@@ -0,0 +1,181 @@
+from collections import namedtuple
+import torch
+from torchvision import models as tv
+from IPython import embed
+
+class squeezenet(torch.nn.Module):
+    def __init__(self, requires_grad=False, pretrained=True):
+        super(squeezenet, self).__init__()
+        pretrained_features = tv.squeezenet1_1(pretrained=pretrained).features
+        self.slice1 = torch.nn.Sequential()
+        self.slice2 = torch.nn.Sequential()
+        self.slice3 = torch.nn.Sequential()
+        self.slice4 = torch.nn.Sequential()
+        self.slice5 = torch.nn.Sequential()
+        self.slice6 = torch.nn.Sequential()
+        self.slice7 = torch.nn.Sequential()
+        self.N_slices = 7
+        for x in range(2):
+            self.slice1.add_module(str(x), pretrained_features[x])
+        for x in range(2,5):
+            self.slice2.add_module(str(x), pretrained_features[x])
+        for x in range(5, 8):
+            self.slice3.add_module(str(x), pretrained_features[x])
+        for x in range(8, 10):
+            self.slice4.add_module(str(x), pretrained_features[x])
+        for x in range(10, 11):
+            self.slice5.add_module(str(x), pretrained_features[x])
+        for x in range(11, 12):
+            self.slice6.add_module(str(x), pretrained_features[x])
+        for x in range(12, 13):
+            self.slice7.add_module(str(x), pretrained_features[x])
+        if not requires_grad:
+            for param in self.parameters():
+                param.requires_grad = False
+
+    def forward(self, X):
+        h = self.slice1(X)
+        h_relu1 = h
+        h = self.slice2(h)
+        h_relu2 = h
+        h = self.slice3(h)
+        h_relu3 = h
+        h = self.slice4(h)
+        h_relu4 = h
+        h = self.slice5(h)
+        h_relu5 = h
+        h = self.slice6(h)
+        h_relu6 = h
+        h = self.slice7(h)
+        h_relu7 = h
+        vgg_outputs = namedtuple("SqueezeOutputs", ['relu1','relu2','relu3','relu4','relu5','relu6','relu7'])
+        out = vgg_outputs(h_relu1,h_relu2,h_relu3,h_relu4,h_relu5,h_relu6,h_relu7)
+
+        return out
+
+
+class alexnet(torch.nn.Module):
+    def __init__(self, requires_grad=False, pretrained=True):
+        super(alexnet, self).__init__()
+        alexnet_pretrained_features = tv.alexnet(pretrained=pretrained).features
+        self.slice1 = torch.nn.Sequential()
+        self.slice2 = torch.nn.Sequential()
+        self.slice3 = torch.nn.Sequential()
+        self.slice4 = torch.nn.Sequential()
+        self.slice5 = torch.nn.Sequential()
+        self.N_slices = 5
+        for x in range(2):
+            self.slice1.add_module(str(x), alexnet_pretrained_features[x])
+        for x in range(2, 5):
+            self.slice2.add_module(str(x), alexnet_pretrained_features[x])
+        for x in range(5, 8):
+            self.slice3.add_module(str(x), alexnet_pretrained_features[x])
+        for x in range(8, 10):
+            self.slice4.add_module(str(x), alexnet_pretrained_features[x])
+        for x in range(10, 12):
+            self.slice5.add_module(str(x), alexnet_pretrained_features[x])
+        if not requires_grad:
+            for param in self.parameters():
+                param.requires_grad = False
+
+    def forward(self, X):
+        h = self.slice1(X)
+        h_relu1 = h
+        h = self.slice2(h)
+        h_relu2 = h
+        h = self.slice3(h)
+        h_relu3 = h
+        h = self.slice4(h)
+        h_relu4 = h
+        h = self.slice5(h)
+        h_relu5 = h
+        alexnet_outputs = namedtuple("AlexnetOutputs", ['relu1', 'relu2', 'relu3', 'relu4', 'relu5'])
+        out = alexnet_outputs(h_relu1, h_relu2, h_relu3, h_relu4, h_relu5)
+
+        return out
+
+class vgg16(torch.nn.Module):
+    def __init__(self, requires_grad=False, pretrained=True):
+        super(vgg16, self).__init__()
+        vgg_pretrained_features = tv.vgg16(pretrained=pretrained).features
+        self.slice1 = torch.nn.Sequential()
+        self.slice2 = torch.nn.Sequential()
+        self.slice3 = torch.nn.Sequential()
+        self.slice4 = torch.nn.Sequential()
+        self.slice5 = torch.nn.Sequential()
+        self.N_slices = 5
+        for x in range(4):
+            self.slice1.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(4, 9):
+            self.slice2.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(9, 16):
+            self.slice3.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(16, 23):
+            self.slice4.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(23, 30):
+            self.slice5.add_module(str(x), vgg_pretrained_features[x])
+        if not requires_grad:
+            for param in self.parameters():
+                param.requires_grad = False
+
+    def forward(self, X):
+        h = self.slice1(X)
+        h_relu1_2 = h
+        h = self.slice2(h)
+        h_relu2_2 = h
+        h = self.slice3(h)
+        h_relu3_3 = h
+        h = self.slice4(h)
+        h_relu4_3 = h
+        h = self.slice5(h)
+        h_relu5_3 = h
+        vgg_outputs = namedtuple("VggOutputs", ['relu1_2', 'relu2_2', 'relu3_3', 'relu4_3', 'relu5_3'])
+        out = vgg_outputs(h_relu1_2, h_relu2_2, h_relu3_3, h_relu4_3, h_relu5_3)
+
+        return out
+
+
+
+class resnet(torch.nn.Module):
+    def __init__(self, requires_grad=False, pretrained=True, num=18):
+        super(resnet, self).__init__()
+        if(num==18):
+            self.net = tv.resnet18(pretrained=pretrained)
+        elif(num==34):
+            self.net = tv.resnet34(pretrained=pretrained)
+        elif(num==50):
+            self.net = tv.resnet50(pretrained=pretrained)
+        elif(num==101):
+            self.net = tv.resnet101(pretrained=pretrained)
+        elif(num==152):
+            self.net = tv.resnet152(pretrained=pretrained)
+        self.N_slices = 5
+
+        self.conv1 = self.net.conv1
+        self.bn1 = self.net.bn1
+        self.relu = self.net.relu
+        self.maxpool = self.net.maxpool
+        self.layer1 = self.net.layer1
+        self.layer2 = self.net.layer2
+        self.layer3 = self.net.layer3
+        self.layer4 = self.net.layer4
+
+    def forward(self, X):
+        h = self.conv1(X)
+        h = self.bn1(h)
+        h = self.relu(h)
+        h_relu1 = h
+        h = self.maxpool(h)
+        h = self.layer1(h)
+        h_conv2 = h
+        h = self.layer2(h)
+        h_conv3 = h
+        h = self.layer3(h)
+        h_conv4 = h
+        h = self.layer4(h)
+        h_conv5 = h
+
+        outputs = namedtuple("Outputs", ['relu1','conv2','conv3','conv4','conv5'])
+        out = outputs(h_relu1, h_conv2, h_conv3, h_conv4, h_conv5)
+
+        return out

losses/masked_lpips/__init__.py

@@ -0,0 +1,199 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+from skimage.metrics import structural_similarity
+import torch
+from torch.autograd import Variable
+
+from ..masked_lpips import dist_model
+
+
+class PerceptualLoss(torch.nn.Module):
+    def __init__(
+        self,
+        model="net-lin",
+        net="alex",
+        vgg_blocks=[1, 2, 3, 4, 5],
+        colorspace="rgb",
+        spatial=False,
+        use_gpu=True,
+        gpu_ids=[0],
+    ):  # VGG using our perceptually-learned weights (LPIPS metric)
+        # def __init__(self, model='net', net='vgg', use_gpu=True): # "default" way of using VGG as a perceptual loss
+        super(PerceptualLoss, self).__init__()
+        print("Setting up Perceptual loss...")
+        self.use_gpu = use_gpu
+        self.spatial = spatial
+        self.gpu_ids = gpu_ids
+        self.model = dist_model.DistModel()
+        self.model.initialize(
+            model=model,
+            net=net,
+            vgg_blocks=vgg_blocks,
+            use_gpu=use_gpu,
+            colorspace=colorspace,
+            spatial=self.spatial,
+            gpu_ids=gpu_ids,
+        )
+        print("...[%s] initialized" % self.model.name())
+        print("...Done")
+
+    def forward(self, pred, target, mask=None, normalize=False):
+        """
+        Pred and target are Variables.
+        If normalize is True, assumes the images are between [0,1] and then scales them between [-1,+1]
+        If normalize is False, assumes the images are already between [-1,+1]
+
+        Inputs pred and target are Nx3xHxW
+        Output pytorch Variable N long
+        """
+
+        if normalize:
+            target = 2 * target - 1
+            pred = 2 * pred - 1
+
+        return self.model.forward(target, pred, mask=mask)
+
+
+def normalize_tensor(in_feat, eps=1e-10):
+    # takes care of masked tensors implicitly.
+    norm_factor = torch.sqrt(torch.sum(in_feat ** 2, dim=1, keepdim=True))
+    return in_feat / (norm_factor + eps)
+
+
+def l2(p0, p1, range=255.0):
+    return 0.5 * np.mean((p0 / range - p1 / range) ** 2)
+
+
+def psnr(p0, p1, peak=255.0):
+    return 10 * np.log10(peak ** 2 / np.mean((1.0 * p0 - 1.0 * p1) ** 2))
+
+
+def dssim(p0, p1, range=255.0):
+    return (1 - compare_ssim(p0, p1, data_range=range, multichannel=True)) / 2.0
+
+
+def rgb2lab(in_img, mean_cent=False):
+    from skimage import color
+
+    img_lab = color.rgb2lab(in_img)
+    if mean_cent:
+        img_lab[:, :, 0] = img_lab[:, :, 0] - 50
+    return img_lab
+
+
+def tensor2np(tensor_obj):
+    # change dimension of a tensor object into a numpy array
+    return tensor_obj[0].cpu().float().numpy().transpose((1, 2, 0))
+
+
+def np2tensor(np_obj):
+    # change dimenion of np array into tensor array
+    return torch.Tensor(np_obj[:, :, :, np.newaxis].transpose((3, 2, 0, 1)))
+
+
+def tensor2tensorlab(image_tensor, to_norm=True, mc_only=False):
+    # image tensor to lab tensor
+    from skimage import color
+
+    img = tensor2im(image_tensor)
+    img_lab = color.rgb2lab(img)
+    if mc_only:
+        img_lab[:, :, 0] = img_lab[:, :, 0] - 50
+    if to_norm and not mc_only:
+        img_lab[:, :, 0] = img_lab[:, :, 0] - 50
+        img_lab = img_lab / 100.0
+
+    return np2tensor(img_lab)
+
+
+def tensorlab2tensor(lab_tensor, return_inbnd=False):
+    from skimage import color
+    import warnings
+
+    warnings.filterwarnings("ignore")
+
+    lab = tensor2np(lab_tensor) * 100.0
+    lab[:, :, 0] = lab[:, :, 0] + 50
+
+    rgb_back = 255.0 * np.clip(color.lab2rgb(lab.astype("float")), 0, 1)
+    if return_inbnd:
+        # convert back to lab, see if we match
+        lab_back = color.rgb2lab(rgb_back.astype("uint8"))
+        mask = 1.0 * np.isclose(lab_back, lab, atol=2.0)
+        mask = np2tensor(np.prod(mask, axis=2)[:, :, np.newaxis])
+        return (im2tensor(rgb_back), mask)
+    else:
+        return im2tensor(rgb_back)
+
+
+def rgb2lab(input):
+    from skimage import color
+
+    return color.rgb2lab(input / 255.0)
+
+
+def tensor2im(image_tensor, imtype=np.uint8, cent=1.0, factor=255.0 / 2.0):
+    image_numpy = image_tensor[0].cpu().float().numpy()
+    image_numpy = (np.transpose(image_numpy, (1, 2, 0)) + cent) * factor
+    return image_numpy.astype(imtype)
+
+
+def im2tensor(image, imtype=np.uint8, cent=1.0, factor=255.0 / 2.0):
+    return torch.Tensor(
+        (image / factor - cent)[:, :, :, np.newaxis].transpose((3, 2, 0, 1))
+    )
+
+
+def tensor2vec(vector_tensor):
+    return vector_tensor.data.cpu().numpy()[:, :, 0, 0]
+
+
+def voc_ap(rec, prec, use_07_metric=False):
+    """ap = voc_ap(rec, prec, [use_07_metric])
+    Compute VOC AP given precision and recall.
+    If use_07_metric is true, uses the
+    VOC 07 11 point method (default:False).
+    """
+    if use_07_metric:
+        # 11 point metric
+        ap = 0.0
+        for t in np.arange(0.0, 1.1, 0.1):
+            if np.sum(rec >= t) == 0:
+                p = 0
+            else:
+                p = np.max(prec[rec >= t])
+            ap = ap + p / 11.0
+    else:
+        # correct AP calculation
+        # first append sentinel values at the end
+        mrec = np.concatenate(([0.0], rec, [1.0]))
+        mpre = np.concatenate(([0.0], prec, [0.0]))
+
+        # compute the precision envelope
+        for i in range(mpre.size - 1, 0, -1):
+            mpre[i - 1] = np.maximum(mpre[i - 1], mpre[i])
+
+        # to calculate area under PR curve, look for points
+        # where X axis (recall) changes value
+        i = np.where(mrec[1:] != mrec[:-1])[0]
+
+        # and sum (\Delta recall) * prec
+        ap = np.sum((mrec[i + 1] - mrec[i]) * mpre[i + 1])
+    return ap
+
+
+def tensor2im(image_tensor, imtype=np.uint8, cent=1.0, factor=255.0 / 2.0):
+    # def tensor2im(image_tensor, imtype=np.uint8, cent=1., factor=1.):
+    image_numpy = image_tensor[0].cpu().float().numpy()
+    image_numpy = (np.transpose(image_numpy, (1, 2, 0)) + cent) * factor
+    return image_numpy.astype(imtype)
+
+
+def im2tensor(image, imtype=np.uint8, cent=1.0, factor=255.0 / 2.0):
+    # def im2tensor(image, imtype=np.uint8, cent=1., factor=1.):
+    return torch.Tensor(
+        (image / factor - cent)[:, :, :, np.newaxis].transpose((3, 2, 0, 1))
+    )

losses/masked_lpips/base_model.py

@@ -0,0 +1,65 @@
+import os
+import numpy as np
+import torch
+from torch.autograd import Variable
+from pdb import set_trace as st
+from IPython import embed
+
+
+class BaseModel:
+    def __init__(self):
+        pass
+
+    def name(self):
+        return "BaseModel"
+
+    def initialize(self, use_gpu=True, gpu_ids=[0]):
+        self.use_gpu = use_gpu
+        self.gpu_ids = gpu_ids
+
+    def forward(self):
+        pass
+
+    def get_image_paths(self):
+        pass
+
+    def optimize_parameters(self):
+        pass
+
+    def get_current_visuals(self):
+        return self.input
+
+    def get_current_errors(self):
+        return {}
+
+    def save(self, label):
+        pass
+
+    # helper saving function that can be used by subclasses
+    def save_network(self, network, path, network_label, epoch_label):
+        save_filename = "%s_net_%s.pth" % (epoch_label, network_label)
+        save_path = os.path.join(path, save_filename)
+        torch.save(network.state_dict(), save_path)
+
+    # helper loading function that can be used by subclasses
+    def load_network(self, network, network_label, epoch_label):
+        save_filename = "%s_net_%s.pth" % (epoch_label, network_label)
+        save_path = os.path.join(self.save_dir, save_filename)
+        print("Loading network from %s" % save_path)
+        network.load_state_dict(torch.load(save_path))
+
+    def update_learning_rate():
+        pass
+
+    def get_image_paths(self):
+        return self.image_paths
+
+    def save_done(self, flag=False):
+        np.save(os.path.join(self.save_dir, "done_flag"), flag)
+        np.savetxt(
+            os.path.join(self.save_dir, "done_flag"),
+            [
+                flag,
+            ],
+            fmt="%i",
+        )

losses/masked_lpips/dist_model.py

@@ -0,0 +1,325 @@
+from __future__ import absolute_import
+
+import sys
+import numpy as np
+import torch
+from torch import nn
+import os
+from collections import OrderedDict
+from torch.autograd import Variable
+import itertools
+from .base_model import BaseModel
+from scipy.ndimage import zoom
+import fractions
+import functools
+import skimage.transform
+from tqdm import tqdm
+
+from IPython import embed
+
+from . import networks_basic as netw
+from losses import masked_lpips as util
+
+
+class DistModel(BaseModel):
+    def name(self):
+        return self.model_name
+
+    def initialize(
+        self,
+        model="net-lin",
+        net="alex",
+        vgg_blocks=[1, 2, 3, 4, 5],
+        colorspace="Lab",
+        pnet_rand=False,
+        pnet_tune=False,
+        model_path=None,
+        use_gpu=True,
+        printNet=False,
+        spatial=False,
+        is_train=False,
+        lr=0.0001,
+        beta1=0.5,
+        version="0.1",
+        gpu_ids=[0],
+    ):
+        """
+        INPUTS
+            model - ['net-lin'] for linearly calibrated network
+                    ['net'] for off-the-shelf network
+                    ['L2'] for L2 distance in Lab colorspace
+                    ['SSIM'] for ssim in RGB colorspace
+            net - ['squeeze','alex','vgg']
+            model_path - if None, will look in weights/[NET_NAME].pth
+            colorspace - ['Lab','RGB'] colorspace to use for L2 and SSIM
+            use_gpu - bool - whether or not to use a GPU
+            printNet - bool - whether or not to print network architecture out
+            spatial - bool - whether to output an array containing varying distances across spatial dimensions
+            spatial_shape - if given, output spatial shape. if None then spatial shape is determined automatically via spatial_factor (see below).
+            spatial_factor - if given, specifies upsampling factor relative to the largest spatial extent of a convolutional layer. if None then resized to size of input images.
+            spatial_order - spline order of filter for upsampling in spatial mode, by default 1 (bilinear).
+            is_train - bool - [True] for training mode
+            lr - float - initial learning rate
+            beta1 - float - initial momentum term for adam
+            version - 0.1 for latest, 0.0 was original (with a bug)
+            gpu_ids - int array - [0] by default, gpus to use
+        """
+        BaseModel.initialize(self, use_gpu=use_gpu, gpu_ids=gpu_ids)
+
+        self.model = model
+        self.net = net
+        self.is_train = is_train
+        self.spatial = spatial
+        self.gpu_ids = gpu_ids
+        self.model_name = "%s [%s]" % (model, net)
+
+        if self.model == "net-lin":  # pretrained net + linear layer
+            self.net = netw.PNetLin(
+                pnet_rand=pnet_rand,
+                pnet_tune=pnet_tune,
+                pnet_type=net,
+                use_dropout=True,
+                spatial=spatial,
+                version=version,
+                lpips=True,
+                vgg_blocks=vgg_blocks,
+            )
+            kw = {}
+            if not use_gpu:
+                kw["map_location"] = "cpu"
+            if model_path is None:
+                import inspect
+
+                model_path = os.path.abspath(
+                    os.path.join(
+                        inspect.getfile(self.initialize),
+                        "..",
+                        "weights/v%s/%s.pth" % (version, net),
+                    )
+                )
+
+            if not is_train:
+                print("Loading model from: %s" % model_path)
+                self.net.load_state_dict(torch.load(model_path, **kw), strict=False)
+
+        elif self.model == "net":  # pretrained network
+            self.net = netw.PNetLin(pnet_rand=pnet_rand, pnet_type=net, lpips=False)
+        elif self.model in ["L2", "l2"]:
+            self.net = netw.L2(
+                use_gpu=use_gpu, colorspace=colorspace
+            )  # not really a network, only for testing
+            self.model_name = "L2"
+        elif self.model in ["DSSIM", "dssim", "SSIM", "ssim"]:
+            self.net = netw.DSSIM(use_gpu=use_gpu, colorspace=colorspace)
+            self.model_name = "SSIM"
+        else:
+            raise ValueError("Model [%s] not recognized." % self.model)
+
+        self.parameters = list(self.net.parameters())
+
+        if self.is_train:  # training mode
+            # extra network on top to go from distances (d0,d1) => predicted human judgment (h*)
+            self.rankLoss = netw.BCERankingLoss()
+            self.parameters += list(self.rankLoss.net.parameters())
+            self.lr = lr
+            self.old_lr = lr
+            self.optimizer_net = torch.optim.Adam(
+                self.parameters, lr=lr, betas=(beta1, 0.999)
+            )
+        else:  # test mode
+            self.net.eval()
+
+        if use_gpu:
+            self.net.to(gpu_ids[0])
+            self.net = torch.nn.DataParallel(self.net, device_ids=gpu_ids)
+            if self.is_train:
+                self.rankLoss = self.rankLoss.to(
+                    device=gpu_ids[0]
+                )  # just put this on GPU0
+
+        if printNet:
+            print("---------- Networks initialized -------------")
+            netw.print_network(self.net)
+            print("-----------------------------------------------")
+
+    def forward(self, in0, in1, mask=None, retPerLayer=False):
+        """Function computes the distance between image patches in0 and in1
+        INPUTS
+            in0, in1 - torch.Tensor object of shape Nx3xXxY - image patch scaled to [-1,1]
+        OUTPUT
+            computed distances between in0 and in1
+        """
+
+        return self.net.forward(in0, in1, mask=mask, retPerLayer=retPerLayer)
+
+    # ***** TRAINING FUNCTIONS *****
+    def optimize_parameters(self):
+        self.forward_train()
+        self.optimizer_net.zero_grad()
+        self.backward_train()
+        self.optimizer_net.step()
+        self.clamp_weights()
+
+    def clamp_weights(self):
+        for module in self.net.modules():
+            if hasattr(module, "weight") and module.kernel_size == (1, 1):
+                module.weight.data = torch.clamp(module.weight.data, min=0)
+
+    def set_input(self, data):
+        self.input_ref = data["ref"]
+        self.input_p0 = data["p0"]
+        self.input_p1 = data["p1"]
+        self.input_judge = data["judge"]
+
+        if self.use_gpu:
+            self.input_ref = self.input_ref.to(device=self.gpu_ids[0])
+            self.input_p0 = self.input_p0.to(device=self.gpu_ids[0])
+            self.input_p1 = self.input_p1.to(device=self.gpu_ids[0])
+            self.input_judge = self.input_judge.to(device=self.gpu_ids[0])
+
+        self.var_ref = Variable(self.input_ref, requires_grad=True)
+        self.var_p0 = Variable(self.input_p0, requires_grad=True)
+        self.var_p1 = Variable(self.input_p1, requires_grad=True)
+
+    def forward_train(self):  # run forward pass
+        # print(self.net.module.scaling_layer.shift)
+        # print(torch.norm(self.net.module.net.slice1[0].weight).item(), torch.norm(self.net.module.lin0.model[1].weight).item())
+
+        self.d0 = self.forward(self.var_ref, self.var_p0)
+        self.d1 = self.forward(self.var_ref, self.var_p1)
+        self.acc_r = self.compute_accuracy(self.d0, self.d1, self.input_judge)
+
+        self.var_judge = Variable(1.0 * self.input_judge).view(self.d0.size())
+
+        self.loss_total = self.rankLoss.forward(
+            self.d0, self.d1, self.var_judge * 2.0 - 1.0
+        )
+
+        return self.loss_total
+
+    def backward_train(self):
+        torch.mean(self.loss_total).backward()
+
+    def compute_accuracy(self, d0, d1, judge):
+        """ d0, d1 are Variables, judge is a Tensor """
+        d1_lt_d0 = (d1 < d0).cpu().data.numpy().flatten()
+        judge_per = judge.cpu().numpy().flatten()
+        return d1_lt_d0 * judge_per + (1 - d1_lt_d0) * (1 - judge_per)
+
+    def get_current_errors(self):
+        retDict = OrderedDict(
+            [("loss_total", self.loss_total.data.cpu().numpy()), ("acc_r", self.acc_r)]
+        )
+
+        for key in retDict.keys():
+            retDict[key] = np.mean(retDict[key])
+
+        return retDict
+
+    def get_current_visuals(self):
+        zoom_factor = 256 / self.var_ref.data.size()[2]
+
+        ref_img = util.tensor2im(self.var_ref.data)
+        p0_img = util.tensor2im(self.var_p0.data)
+        p1_img = util.tensor2im(self.var_p1.data)
+
+        ref_img_vis = zoom(ref_img, [zoom_factor, zoom_factor, 1], order=0)
+        p0_img_vis = zoom(p0_img, [zoom_factor, zoom_factor, 1], order=0)
+        p1_img_vis = zoom(p1_img, [zoom_factor, zoom_factor, 1], order=0)
+
+        return OrderedDict(
+            [("ref", ref_img_vis), ("p0", p0_img_vis), ("p1", p1_img_vis)]
+        )
+
+    def save(self, path, label):
+        if self.use_gpu:
+            self.save_network(self.net.module, path, "", label)
+        else:
+            self.save_network(self.net, path, "", label)
+        self.save_network(self.rankLoss.net, path, "rank", label)
+
+    def update_learning_rate(self, nepoch_decay):
+        lrd = self.lr / nepoch_decay
+        lr = self.old_lr - lrd
+
+        for param_group in self.optimizer_net.param_groups:
+            param_group["lr"] = lr
+
+        print("update lr [%s] decay: %f -> %f" % (type, self.old_lr, lr))
+        self.old_lr = lr
+
+
+def score_2afc_dataset(data_loader, func, name=""):
+    """Function computes Two Alternative Forced Choice (2AFC) score using
+        distance function 'func' in dataset 'data_loader'
+    INPUTS
+        data_loader - CustomDatasetDataLoader object - contains a TwoAFCDataset inside
+        func - callable distance function - calling d=func(in0,in1) should take 2
+            pytorch tensors with shape Nx3xXxY, and return numpy array of length N
+    OUTPUTS
+        [0] - 2AFC score in [0,1], fraction of time func agrees with human evaluators
+        [1] - dictionary with following elements
+            d0s,d1s - N arrays containing distances between reference patch to perturbed patches
+            gts - N array in [0,1], preferred patch selected by human evaluators
+                (closer to "0" for left patch p0, "1" for right patch p1,
+                "0.6" means 60pct people preferred right patch, 40pct preferred left)
+            scores - N array in [0,1], corresponding to what percentage function agreed with humans
+    CONSTS
+        N - number of test triplets in data_loader
+    """
+
+    d0s = []
+    d1s = []
+    gts = []
+
+    for data in tqdm(data_loader.load_data(), desc=name):
+        d0s += func(data["ref"], data["p0"]).data.cpu().numpy().flatten().tolist()
+        d1s += func(data["ref"], data["p1"]).data.cpu().numpy().flatten().tolist()
+        gts += data["judge"].cpu().numpy().flatten().tolist()
+
+    d0s = np.array(d0s)
+    d1s = np.array(d1s)
+    gts = np.array(gts)
+    scores = (d0s < d1s) * (1.0 - gts) + (d1s < d0s) * gts + (d1s == d0s) * 0.5
+
+    return (np.mean(scores), dict(d0s=d0s, d1s=d1s, gts=gts, scores=scores))
+
+
+def score_jnd_dataset(data_loader, func, name=""):
+    """Function computes JND score using distance function 'func' in dataset 'data_loader'
+    INPUTS
+        data_loader - CustomDatasetDataLoader object - contains a JNDDataset inside
+        func - callable distance function - calling d=func(in0,in1) should take 2
+            pytorch tensors with shape Nx3xXxY, and return pytorch array of length N
+    OUTPUTS
+        [0] - JND score in [0,1], mAP score (area under precision-recall curve)
+        [1] - dictionary with following elements
+            ds - N array containing distances between two patches shown to human evaluator
+            sames - N array containing fraction of people who thought the two patches were identical
+    CONSTS
+        N - number of test triplets in data_loader
+    """
+
+    ds = []
+    gts = []
+
+    for data in tqdm(data_loader.load_data(), desc=name):
+        ds += func(data["p0"], data["p1"]).data.cpu().numpy().tolist()
+        gts += data["same"].cpu().numpy().flatten().tolist()
+
+    sames = np.array(gts)
+    ds = np.array(ds)
+
+    sorted_inds = np.argsort(ds)
+    ds_sorted = ds[sorted_inds]
+    sames_sorted = sames[sorted_inds]
+
+    TPs = np.cumsum(sames_sorted)
+    FPs = np.cumsum(1 - sames_sorted)
+    FNs = np.sum(sames_sorted) - TPs
+
+    precs = TPs / (TPs + FPs)
+    recs = TPs / (TPs + FNs)
+    score = util.voc_ap(recs, precs)
+
+    return (score, dict(ds=ds, sames=sames))

losses/masked_lpips/networks_basic.py

@@ -0,0 +1,331 @@
+from __future__ import absolute_import
+
+import sys
+import torch
+import torch.nn as nn
+import torch.nn.init as init
+from torch.autograd import Variable
+from torch.nn import functional as F
+import numpy as np
+from pdb import set_trace as st
+from skimage import color
+from IPython import embed
+from . import pretrained_networks as pn
+
+from losses import masked_lpips as util
+
+
+def spatial_average(in_tens, mask=None, keepdim=True):
+    if mask is None:
+        return in_tens.mean([2, 3], keepdim=keepdim)
+    else:
+        in_tens = in_tens * mask
+
+        # sum masked_in_tens across spatial dims
+        in_tens = in_tens.sum([2, 3], keepdim=keepdim)
+        in_tens = in_tens / torch.sum(mask)
+
+        return in_tens
+
+
+def upsample(in_tens, out_H=64):  # assumes scale factor is same for H and W
+    in_H = in_tens.shape[2]
+    scale_factor = 1.0 * out_H / in_H
+
+    return nn.Upsample(scale_factor=scale_factor, mode="bilinear", align_corners=False)(
+        in_tens
+    )
+
+
+# Learned perceptual metric
+class PNetLin(nn.Module):
+    def __init__(
+        self,
+        pnet_type="vgg",
+        pnet_rand=False,
+        pnet_tune=False,
+        use_dropout=True,
+        spatial=False,
+        version="0.1",
+        lpips=True,
+        vgg_blocks=[1, 2, 3, 4, 5]
+    ):
+        super(PNetLin, self).__init__()
+
+        self.pnet_type = pnet_type
+        self.pnet_tune = pnet_tune
+        self.pnet_rand = pnet_rand
+        self.spatial = spatial
+        self.lpips = lpips
+        self.version = version
+        self.scaling_layer = ScalingLayer()
+
+        if self.pnet_type in ["vgg", "vgg16"]:
+            net_type = pn.vgg16
+            self.blocks = vgg_blocks
+            self.chns = []
+            self.chns = [64, 128, 256, 512, 512]
+
+        elif self.pnet_type == "alex":
+            net_type = pn.alexnet
+            self.chns = [64, 192, 384, 256, 256]
+        elif self.pnet_type == "squeeze":
+            net_type = pn.squeezenet
+            self.chns = [64, 128, 256, 384, 384, 512, 512]
+        self.L = len(self.chns)
+
+        self.net = net_type(pretrained=not self.pnet_rand, requires_grad=self.pnet_tune)
+
+        if lpips:
+            self.lin0 = NetLinLayer(self.chns[0], use_dropout=use_dropout)
+            self.lin1 = NetLinLayer(self.chns[1], use_dropout=use_dropout)
+            self.lin2 = NetLinLayer(self.chns[2], use_dropout=use_dropout)
+            self.lin3 = NetLinLayer(self.chns[3], use_dropout=use_dropout)
+            self.lin4 = NetLinLayer(self.chns[4], use_dropout=use_dropout)
+            self.lins = [self.lin0, self.lin1, self.lin2, self.lin3, self.lin4]
+            #self.lins = [self.lin0, self.lin1, self.lin2, self.lin3, self.lin4]
+            if self.pnet_type == "squeeze":  # 7 layers for squeezenet
+                self.lin5 = NetLinLayer(self.chns[5], use_dropout=use_dropout)
+                self.lin6 = NetLinLayer(self.chns[6], use_dropout=use_dropout)
+                self.lins += [self.lin5, self.lin6]
+
+    def forward(self, in0, in1, mask=None, retPerLayer=False):
+        # blocks: list of layer names
+
+        # v0.0 - original release had a bug, where input was not scaled
+        in0_input, in1_input = (
+            (self.scaling_layer(in0), self.scaling_layer(in1))
+            if self.version == "0.1"
+            else (in0, in1)
+        )
+        outs0, outs1 = self.net.forward(in0_input), self.net.forward(in1_input)
+        feats0, feats1, diffs = {}, {}, {}
+        
+        # prepare list of masks at different resolutions
+        if mask is not None:
+            masks = []
+            if len(mask.shape) == 3:
+                mask = torch.unsqueeze(mask, axis=0)  # 4D
+
+            for kk in range(self.L):
+                N, C, H, W = outs0[kk].shape
+                mask = F.interpolate(mask, size=(H, W), mode="nearest")
+                masks.append(mask)
+
+        """
+        outs0 has 5 feature maps 
+        1. [1, 64, 256, 256]
+        2. [1, 128, 128, 128]
+        3. [1, 256, 64, 64]
+        4. [1, 512, 32, 32]
+        5. [1, 512, 16, 16]
+        """
+        for kk in range(self.L):
+            feats0[kk], feats1[kk] = (
+                util.normalize_tensor(outs0[kk]),
+                util.normalize_tensor(outs1[kk]),
+            )
+            diffs[kk] = (feats0[kk] - feats1[kk]) ** 2
+        
+        if self.lpips:
+            if self.spatial:
+                res = [
+                    upsample(self.lins[kk].model(diffs[kk]), out_H=in0.shape[2])
+                    for kk in range(self.L)
+                ]
+            else:
+                # NOTE: this block is used
+                # self.lins has 5 elements, where each element is a layer of LIN
+                """
+                self.lin0 = NetLinLayer(self.chns[0], use_dropout=use_dropout)
+                self.lin1 = NetLinLayer(self.chns[1], use_dropout=use_dropout)
+                self.lin2 = NetLinLayer(self.chns[2], use_dropout=use_dropout)
+                self.lin3 = NetLinLayer(self.chns[3], use_dropout=use_dropout)
+                self.lin4 = NetLinLayer(self.chns[4], use_dropout=use_dropout)
+                self.lins = [self.lin0,self.lin1,self.lin2,self.lin3,self.lin4]
+                """
+
+                # NOTE:
+                # Each lins is applying a 1x1 conv on the spatial tensor to output 1 channel
+                # Therefore, to prevent this problem, we can simply mask out the activations
+                # in the spatial_average block. Right now, spatial_average does a spatial mean.
+                # We can mask out the tensor and then consider only on pixels for the mean op.
+                res = [
+                    spatial_average(
+                        self.lins[kk].model(diffs[kk]),
+                        mask=masks[kk] if mask is not None else None,
+                        keepdim=True,
+                    )
+                    for kk in range(self.L)
+                ]
+        else:
+            if self.spatial:
+                res = [
+                    upsample(diffs[kk].sum(dim=1, keepdim=True), out_H=in0.shape[2])
+                    for kk in range(self.L)
+                ]
+            else:
+                res = [
+                    spatial_average(diffs[kk].sum(dim=1, keepdim=True), keepdim=True)
+                    for kk in range(self.L)
+                ]
+        
+        '''
+        val = res[0]
+        for l in range(1, self.L):
+            val += res[l]
+        '''
+        
+        val = 0.0
+        for l in range(self.L):
+            # l is going to run from 0 to 4
+            # check if (l + 1), i.e., [1 -> 5] in self.blocks, then count the loss
+            if str(l + 1) in self.blocks:
+                val += res[l]
+
+        if retPerLayer:
+            return (val, res)
+        else:
+            return val
+
+
+class ScalingLayer(nn.Module):
+    def __init__(self):
+        super(ScalingLayer, self).__init__()
+        self.register_buffer(
+            "shift", torch.Tensor([-0.030, -0.088, -0.188])[None, :, None, None]
+        )
+        self.register_buffer(
+            "scale", torch.Tensor([0.458, 0.448, 0.450])[None, :, None, None]
+        )
+
+    def forward(self, inp):
+        return (inp - self.shift) / self.scale
+
+
+class NetLinLayer(nn.Module):
+    """ A single linear layer which does a 1x1 conv """
+
+    def __init__(self, chn_in, chn_out=1, use_dropout=False):
+        super(NetLinLayer, self).__init__()
+
+        layers = (
+            [
+                nn.Dropout(),
+            ]
+            if (use_dropout)
+            else []
+        )
+        layers += [
+            nn.Conv2d(chn_in, chn_out, 1, stride=1, padding=0, bias=False),
+        ]
+        self.model = nn.Sequential(*layers)
+
+
+class Dist2LogitLayer(nn.Module):
+    """ takes 2 distances, puts through fc layers, spits out value between [0,1] (if use_sigmoid is True) """
+
+    def __init__(self, chn_mid=32, use_sigmoid=True):
+        super(Dist2LogitLayer, self).__init__()
+
+        layers = [
+            nn.Conv2d(5, chn_mid, 1, stride=1, padding=0, bias=True),
+        ]
+        layers += [
+            nn.LeakyReLU(0.2, True),
+        ]
+        layers += [
+            nn.Conv2d(chn_mid, chn_mid, 1, stride=1, padding=0, bias=True),
+        ]
+        layers += [
+            nn.LeakyReLU(0.2, True),
+        ]
+        layers += [
+            nn.Conv2d(chn_mid, 1, 1, stride=1, padding=0, bias=True),
+        ]
+        if use_sigmoid:
+            layers += [
+                nn.Sigmoid(),
+            ]
+        self.model = nn.Sequential(*layers)
+
+    def forward(self, d0, d1, eps=0.1):
+        return self.model.forward(
+            torch.cat((d0, d1, d0 - d1, d0 / (d1 + eps), d1 / (d0 + eps)), dim=1)
+        )
+
+
+class BCERankingLoss(nn.Module):
+    def __init__(self, chn_mid=32):
+        super(BCERankingLoss, self).__init__()
+        self.net = Dist2LogitLayer(chn_mid=chn_mid)
+        # self.parameters = list(self.net.parameters())
+        self.loss = torch.nn.BCELoss()
+
+    def forward(self, d0, d1, judge):
+        per = (judge + 1.0) / 2.0
+        self.logit = self.net.forward(d0, d1)
+        return self.loss(self.logit, per)
+
+
+# L2, DSSIM metrics
+class FakeNet(nn.Module):
+    def __init__(self, use_gpu=True, colorspace="Lab"):
+        super(FakeNet, self).__init__()
+        self.use_gpu = use_gpu
+        self.colorspace = colorspace
+
+
+class L2(FakeNet):
+    def forward(self, in0, in1, retPerLayer=None):
+        assert in0.size()[0] == 1  # currently only supports batchSize 1
+
+        if self.colorspace == "RGB":
+            (N, C, X, Y) = in0.size()
+            value = torch.mean(
+                torch.mean(
+                    torch.mean((in0 - in1) ** 2, dim=1).view(N, 1, X, Y), dim=2
+                ).view(N, 1, 1, Y),
+                dim=3,
+            ).view(N)
+            return value
+        elif self.colorspace == "Lab":
+            value = util.l2(
+                util.tensor2np(util.tensor2tensorlab(in0.data, to_norm=False)),
+                util.tensor2np(util.tensor2tensorlab(in1.data, to_norm=False)),
+                range=100.0,
+            ).astype("float")
+            ret_var = Variable(torch.Tensor((value,)))
+            if self.use_gpu:
+                ret_var = ret_var.cuda()
+            return ret_var
+
+
+class DSSIM(FakeNet):
+    def forward(self, in0, in1, retPerLayer=None):
+        assert in0.size()[0] == 1  # currently only supports batchSize 1
+
+        if self.colorspace == "RGB":
+            value = util.dssim(
+                1.0 * util.tensor2im(in0.data),
+                1.0 * util.tensor2im(in1.data),
+                range=255.0,
+            ).astype("float")
+        elif self.colorspace == "Lab":
+            value = util.dssim(
+                util.tensor2np(util.tensor2tensorlab(in0.data, to_norm=False)),
+                util.tensor2np(util.tensor2tensorlab(in1.data, to_norm=False)),
+                range=100.0,
+            ).astype("float")
+        ret_var = Variable(torch.Tensor((value,)))
+        if self.use_gpu:
+            ret_var = ret_var.cuda()
+        return ret_var
+
+
+def print_network(net):
+    num_params = 0
+    for param in net.parameters():
+        num_params += param.numel()
+    print("Network", net)
+    print("Total number of parameters: %d" % num_params)

losses/masked_lpips/pretrained_networks.py

@@ -0,0 +1,190 @@
+from collections import namedtuple
+import torch
+from torchvision import models as tv
+from IPython import embed
+
+
+class squeezenet(torch.nn.Module):
+    def __init__(self, requires_grad=False, pretrained=True):
+        super(squeezenet, self).__init__()
+        pretrained_features = tv.squeezenet1_1(pretrained=pretrained).features
+        self.slice1 = torch.nn.Sequential()
+        self.slice2 = torch.nn.Sequential()
+        self.slice3 = torch.nn.Sequential()
+        self.slice4 = torch.nn.Sequential()
+        self.slice5 = torch.nn.Sequential()
+        self.slice6 = torch.nn.Sequential()
+        self.slice7 = torch.nn.Sequential()
+        self.N_slices = 7
+        for x in range(2):
+            self.slice1.add_module(str(x), pretrained_features[x])
+        for x in range(2, 5):
+            self.slice2.add_module(str(x), pretrained_features[x])
+        for x in range(5, 8):
+            self.slice3.add_module(str(x), pretrained_features[x])
+        for x in range(8, 10):
+            self.slice4.add_module(str(x), pretrained_features[x])
+        for x in range(10, 11):
+            self.slice5.add_module(str(x), pretrained_features[x])
+        for x in range(11, 12):
+            self.slice6.add_module(str(x), pretrained_features[x])
+        for x in range(12, 13):
+            self.slice7.add_module(str(x), pretrained_features[x])
+        if not requires_grad:
+            for param in self.parameters():
+                param.requires_grad = False
+
+    def forward(self, X):
+        h = self.slice1(X)
+        h_relu1 = h
+        h = self.slice2(h)
+        h_relu2 = h
+        h = self.slice3(h)
+        h_relu3 = h
+        h = self.slice4(h)
+        h_relu4 = h
+        h = self.slice5(h)
+        h_relu5 = h
+        h = self.slice6(h)
+        h_relu6 = h
+        h = self.slice7(h)
+        h_relu7 = h
+        vgg_outputs = namedtuple(
+            "SqueezeOutputs",
+            ["relu1", "relu2", "relu3", "relu4", "relu5", "relu6", "relu7"],
+        )
+        out = vgg_outputs(h_relu1, h_relu2, h_relu3, h_relu4, h_relu5, h_relu6, h_relu7)
+
+        return out
+
+
+class alexnet(torch.nn.Module):
+    def __init__(self, requires_grad=False, pretrained=True):
+        super(alexnet, self).__init__()
+        alexnet_pretrained_features = tv.alexnet(pretrained=pretrained).features
+        self.slice1 = torch.nn.Sequential()
+        self.slice2 = torch.nn.Sequential()
+        self.slice3 = torch.nn.Sequential()
+        self.slice4 = torch.nn.Sequential()
+        self.slice5 = torch.nn.Sequential()
+        self.N_slices = 5
+        for x in range(2):
+            self.slice1.add_module(str(x), alexnet_pretrained_features[x])
+        for x in range(2, 5):
+            self.slice2.add_module(str(x), alexnet_pretrained_features[x])
+        for x in range(5, 8):
+            self.slice3.add_module(str(x), alexnet_pretrained_features[x])
+        for x in range(8, 10):
+            self.slice4.add_module(str(x), alexnet_pretrained_features[x])
+        for x in range(10, 12):
+            self.slice5.add_module(str(x), alexnet_pretrained_features[x])
+        if not requires_grad:
+            for param in self.parameters():
+                param.requires_grad = False
+
+    def forward(self, X):
+        h = self.slice1(X)
+        h_relu1 = h
+        h = self.slice2(h)
+        h_relu2 = h
+        h = self.slice3(h)
+        h_relu3 = h
+        h = self.slice4(h)
+        h_relu4 = h
+        h = self.slice5(h)
+        h_relu5 = h
+        alexnet_outputs = namedtuple(
+            "AlexnetOutputs", ["relu1", "relu2", "relu3", "relu4", "relu5"]
+        )
+        out = alexnet_outputs(h_relu1, h_relu2, h_relu3, h_relu4, h_relu5)
+
+        return out
+
+
+class vgg16(torch.nn.Module):
+    def __init__(self, requires_grad=False, pretrained=True):
+        super(vgg16, self).__init__()
+        vgg_pretrained_features = tv.vgg16(pretrained=pretrained).features
+        self.slice1 = torch.nn.Sequential()
+        self.slice2 = torch.nn.Sequential()
+        self.slice3 = torch.nn.Sequential()
+        self.slice4 = torch.nn.Sequential()
+        self.slice5 = torch.nn.Sequential()
+        self.N_slices = 5
+        for x in range(4):
+            self.slice1.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(4, 9):
+            self.slice2.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(9, 16):
+            self.slice3.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(16, 23):
+            self.slice4.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(23, 30):
+            self.slice5.add_module(str(x), vgg_pretrained_features[x])
+        if not requires_grad:
+            for param in self.parameters():
+                param.requires_grad = False
+
+    def forward(self, X):
+        h = self.slice1(X)
+        h_relu1_2 = h
+        h = self.slice2(h)
+        h_relu2_2 = h
+        h = self.slice3(h)
+        h_relu3_3 = h
+        h = self.slice4(h)
+        h_relu4_3 = h
+        h = self.slice5(h)
+        h_relu5_3 = h
+
+        vgg_outputs = namedtuple(
+            "VggOutputs", ["relu1_2", "relu2_2", "relu3_3", "relu4_3", "relu5_3"]
+        )
+        out = vgg_outputs(h_relu1_2, h_relu2_2, h_relu3_3, h_relu4_3, h_relu5_3)
+
+        return out
+
+
+class resnet(torch.nn.Module):
+    def __init__(self, requires_grad=False, pretrained=True, num=18):
+        super(resnet, self).__init__()
+        if num == 18:
+            self.net = tv.resnet18(pretrained=pretrained)
+        elif num == 34:
+            self.net = tv.resnet34(pretrained=pretrained)
+        elif num == 50:
+            self.net = tv.resnet50(pretrained=pretrained)
+        elif num == 101:
+            self.net = tv.resnet101(pretrained=pretrained)
+        elif num == 152:
+            self.net = tv.resnet152(pretrained=pretrained)
+        self.N_slices = 5
+
+        self.conv1 = self.net.conv1
+        self.bn1 = self.net.bn1
+        self.relu = self.net.relu
+        self.maxpool = self.net.maxpool
+        self.layer1 = self.net.layer1
+        self.layer2 = self.net.layer2
+        self.layer3 = self.net.layer3
+        self.layer4 = self.net.layer4
+
+    def forward(self, X):
+        h = self.conv1(X)
+        h = self.bn1(h)
+        h = self.relu(h)
+        h_relu1 = h
+        h = self.maxpool(h)
+        h = self.layer1(h)
+        h_conv2 = h
+        h = self.layer2(h)
+        h_conv3 = h
+        h = self.layer3(h)
+        h_conv4 = h
+        h = self.layer4(h)
+        h_conv5 = h
+
+        outputs = namedtuple("Outputs", ["relu1", "conv2", "conv3", "conv4", "conv5"])
+        out = outputs(h_relu1, h_conv2, h_conv3, h_conv4, h_conv5)
+
+        return out

losses/pp_losses.py

@@ -0,0 +1,677 @@
+from dataclasses import dataclass
+
+import torch.nn as nn
+import torch.nn.functional as F
+from torchvision import transforms as T
+
+from utils.bicubic import BicubicDownSample
+
+normalize = T.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
+
+@dataclass
+class DefaultPaths:
+    psp_path: str = "pretrained_models/psp_ffhq_encode.pt"
+    ir_se50_path: str = "pretrained_models/ArcFace/ir_se50.pth"
+    stylegan_weights: str = "pretrained_models/stylegan2-ffhq-config-f.pt"
+    stylegan_car_weights: str = "pretrained_models/stylegan2-car-config-f-new.pkl"
+    stylegan_weights_pkl: str = (
+        "pretrained_models/stylegan2-ffhq-config-f.pkl"
+    )
+    arcface_model_path: str = "pretrained_models/ArcFace/backbone_ir50.pth"
+    moco: str = "pretrained_models/moco_v2_800ep_pretrain.pt"
+    
+
+from collections import namedtuple
+from torch.nn import (
+    Conv2d,
+    BatchNorm2d,
+    PReLU,
+    ReLU,
+    Sigmoid,
+    MaxPool2d,
+    AdaptiveAvgPool2d,
+    Sequential,
+    Module,
+    Dropout,
+    Linear,
+    BatchNorm1d,
+)
+
+"""
+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
+
+
+"""
+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)
+
+
+def IR_50(input_size):
+    """Constructs a ir-50 model."""
+    model = Backbone(input_size, num_layers=50, mode="ir", drop_ratio=0.4, affine=False)
+    return model
+
+
+def IR_101(input_size):
+    """Constructs a ir-101 model."""
+    model = Backbone(
+        input_size, num_layers=100, mode="ir", drop_ratio=0.4, affine=False
+    )
+    return model
+
+
+def IR_152(input_size):
+    """Constructs a ir-152 model."""
+    model = Backbone(
+        input_size, num_layers=152, mode="ir", drop_ratio=0.4, affine=False
+    )
+    return model
+
+
+def IR_SE_50(input_size):
+    """Constructs a ir_se-50 model."""
+    model = Backbone(
+        input_size, num_layers=50, mode="ir_se", drop_ratio=0.4, affine=False
+    )
+    return model
+
+
+def IR_SE_101(input_size):
+    """Constructs a ir_se-101 model."""
+    model = Backbone(
+        input_size, num_layers=100, mode="ir_se", drop_ratio=0.4, affine=False
+    )
+    return model
+
+
+def IR_SE_152(input_size):
+    """Constructs a ir_se-152 model."""
+    model = Backbone(
+        input_size, num_layers=152, mode="ir_se", drop_ratio=0.4, affine=False
+    )
+    return model
+
+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(DefaultPaths.ir_se50_path))
+        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):
+        n_samples = y.shape[0]
+        y_feats = self.extract_feats(y)
+        y_hat_feats = self.extract_feats(y_hat)
+        y_feats = y_feats.detach()
+        loss = 0
+        count = 0
+        for i in range(n_samples):
+            diff_target = y_hat_feats[i].dot(y_feats[i])
+            loss += 1 - diff_target
+            count += 1
+
+        return loss / count
+    
+class FeatReconLoss(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.loss_fn = nn.MSELoss()
+
+    def forward(self, recon_1, recon_2):
+        return self.loss_fn(recon_1, recon_2).mean()
+    
+class EncoderAdvLoss:
+    def __call__(self, fake_preds):
+        loss_G_adv = F.softplus(-fake_preds).mean()
+        return loss_G_adv
+
+class AdvLoss:
+    def __init__(self, coef=0.0):
+        self.coef = coef
+
+    def __call__(self, disc, real_images, generated_images):
+        fake_preds = disc(generated_images, None)
+        real_preds = disc(real_images, None)
+        loss = self.d_logistic_loss(real_preds, fake_preds)
+
+        return {'disc adv': loss}
+
+    def d_logistic_loss(self, real_preds, fake_preds):
+        real_loss = F.softplus(-real_preds)
+        fake_loss = F.softplus(fake_preds)
+
+        return (real_loss.mean() + fake_loss.mean()) / 2
+
+from models.face_parsing.model import BiSeNet, seg_mean, seg_std
+    
+class DiceLoss(nn.Module):
+    def __init__(self, gamma=2):
+        super().__init__()
+        self.gamma = gamma
+        self.seg = BiSeNet(n_classes=16)
+        self.seg.to('cuda')
+        self.seg.load_state_dict(torch.load('pretrained_models/BiSeNet/seg.pth'))
+        for param in self.seg.parameters():
+            param.requires_grad = False
+        self.seg.eval()
+        self.downsample_512 = BicubicDownSample(factor=2)
+    
+    def calc_landmark(self, x):
+        IM = (self.downsample_512(x) - seg_mean) / seg_std
+        out, _, _ = self.seg(IM)
+        return out
+
+    def dice_loss(self, input, target):
+        smooth = 1.
+
+        iflat = input.view(input.size(0), -1)
+        tflat = target.view(target.size(0), -1)
+        intersection = (iflat * tflat).sum(dim=1)
+        
+        fn = torch.sum((tflat * (1-iflat))**self.gamma, dim=1)
+        fp = torch.sum(((1-tflat) * iflat)**self.gamma, dim=1)
+
+        return 1 - ((2. * intersection + smooth) /
+                  (iflat.sum(dim=1) + tflat.sum(dim=1) + fn + fp + smooth))
+    
+    def __call__(self, in_logit, tg_logit):
+        probs1 = F.softmax(in_logit, dim=1)
+        probs2 = F.softmax(tg_logit, dim=1)
+        return self.dice_loss(probs1, probs2).mean()
+        
+
+from typing import Sequence
+
+from itertools import chain
+
+import torch
+import torch.nn as nn
+from torchvision import models
+
+
+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([-0.030, -0.088, -0.188])[None, :, None, None]
+        )
+        self.register_buffer(
+            "std", torch.Tensor([0.458, 0.448, 0.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)
+
+    
+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
+    
+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]
+
+class LPIPSLoss(LPIPS):
+    pass
+    
+class LPIPSScaleLoss(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.loss_fn = LPIPSLoss()
+
+    def forward(self, x, y):
+        out = 0
+        for res in [256, 128, 64]:
+            x_scale = F.interpolate(x, size=(res, res), mode="bilinear", align_corners=False)
+            y_scale = F.interpolate(y, size=(res, res), mode="bilinear", align_corners=False)
+            out += self.loss_fn.forward(x_scale, y_scale).mean()
+        return out
+    
+class SyntMSELoss(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.loss_fn = nn.MSELoss()
+
+    def forward(self, im1, im2):
+        return self.loss_fn(im1, im2).mean()
+    
+class R1Loss:
+    def __init__(self, coef=10.0):
+        self.coef = coef
+
+    def __call__(self, disc, real_images):
+        real_images.requires_grad = True
+
+        real_preds = disc(real_images, None)
+        real_preds = real_preds.view(real_images.size(0), -1)
+        real_preds = real_preds.mean(dim=1).unsqueeze(1)
+        r1_loss = self.d_r1_loss(real_preds, real_images)
+
+        loss_D_R1 = self.coef / 2 * r1_loss * 16 + 0 * real_preds[0]
+        return {'disc r1 loss': loss_D_R1}
+
+    def d_r1_loss(self, real_pred, real_img):
+        (grad_real,) = torch.autograd.grad(
+            outputs=real_pred.sum(), inputs=real_img, create_graph=True
+        )
+        grad_penalty = grad_real.pow(2).reshape(grad_real.shape[0], -1).sum(1).mean()
+
+        return grad_penalty
+
+
+class DilatedMask:
+    def __init__(self, kernel_size=5):
+        self.kernel_size = kernel_size
+        
+        cords_x = torch.arange(0, kernel_size).view(1, -1).expand(kernel_size, -1) - kernel_size // 2
+        cords_y = cords_x.clone().permute(1, 0)
+        self.kernel = torch.as_tensor((cords_x ** 2 + cords_y ** 2) <= (kernel_size // 2) ** 2, dtype=torch.float).view(1, 1, kernel_size, kernel_size).cuda()
+        self.kernel /= self.kernel.sum()
+    
+    def __call__(self, mask):
+        smooth_mask = F.conv2d(mask, self.kernel, padding=self.kernel_size // 2)
+        return smooth_mask ** 0.25
+
+
+class LossBuilder:
+    def __init__(self, losses_dict, device='cuda'):
+        self.losses_dict = losses_dict
+        self.device = device
+        
+        self.EncoderAdvLoss = EncoderAdvLoss()
+        self.AdvLoss = AdvLoss()
+        self.R1Loss = R1Loss()
+        self.FeatReconLoss = FeatReconLoss().to(device).eval()
+        self.IDLoss = IDLoss().to(device).eval()
+        self.LPIPS = LPIPSScaleLoss().to(device).eval()
+        self.SyntMSELoss = SyntMSELoss().to(device).eval()
+        self.downsample_256 = BicubicDownSample(factor=4)
+        
+    def CalcAdvLoss(self, disc, gen_F):
+        fake_preds_F = disc(gen_F, None)
+        
+        return {'adv': self.losses_dict['adv'] * self.EncoderAdvLoss(fake_preds_F)}
+    
+    def CalcDisLoss(self, disc, real_images, generated_images):
+        return self.AdvLoss(disc, real_images, generated_images)
+    
+    def CalcR1Loss(self, disc, real_images):
+        return self.R1Loss(disc, real_images)
+        
+    def __call__(self, source, target, target_mask, HT_E, gen_w, F_w, gen_F, F_gen, **kwargs):
+        losses = {}
+        
+        gen_w_256 = self.downsample_256(gen_w)
+        gen_F_256 = self.downsample_256(gen_F)
+        
+        # ID loss
+        losses['rec id'] = self.losses_dict['id'] * (self.IDLoss(normalize(source), gen_w_256) + self.IDLoss(normalize(source), gen_F_256))
+
+        # Feat Recons Loss
+        losses['rec feat_rec'] = self.losses_dict['feat_rec'] * self.FeatReconLoss(F_w.detach(), F_gen)
+        
+        # LPIPS loss
+        losses['rec lpips_scale'] = self.losses_dict['lpips_scale'] * (self.LPIPS(normalize(source), gen_w_256) + self.LPIPS(normalize(source), gen_F_256))
+        
+        # Synt loss
+        # losses['l2_synt'] = self.losses_dict['l2_synt'] * self.SyntMSELoss(target * HT_E, (gen_F_256 + 1) / 2 * HT_E)
+        
+        return losses
+
+    
+class LossBuilderMulti(LossBuilder):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.DiceLoss = DiceLoss().to(kwargs.get('device', 'cuda')).eval()
+        self.dilated = DilatedMask(25)
+        
+    def __call__(self, source, target, target_mask, HT_E, gen_w, F_w, gen_F, F_gen, **kwargs):
+        losses = {}
+        
+        gen_w_256 = self.downsample_256(gen_w)
+        gen_F_256 = self.downsample_256(gen_F)
+        
+        # Dice loss
+        with torch.no_grad():
+            target_512 = F.interpolate(target, size=(512, 512), mode='bilinear').clip(0, 1)
+            seg_target = self.DiceLoss.calc_landmark(target_512)
+            seg_target = F.interpolate(seg_target, size=(256, 256), mode='nearest')
+        seg_gen = F.interpolate(self.DiceLoss.calc_landmark((gen_F + 1) / 2), size=(256, 256), mode='nearest')
+        
+        losses['DiceLoss'] = self.losses_dict['landmark'] * self.DiceLoss(seg_gen, seg_target)
+        
+        # ID loss
+        losses['id'] = self.losses_dict['id'] * (self.IDLoss(normalize(source) * target_mask, gen_w_256 * target_mask) +
+                                                 self.IDLoss(normalize(source) * target_mask, gen_F_256 * target_mask))
+
+        # Feat Recons loss
+        losses['feat_rec'] = self.losses_dict['feat_rec'] * self.FeatReconLoss(F_w.detach(), F_gen)
+        
+        # LPIPS loss
+        losses['lpips_face'] = 0.5 * self.losses_dict['lpips_scale'] * (self.LPIPS(normalize(source) * target_mask, gen_w_256 * target_mask) +
+                                                                         self.LPIPS(normalize(source) * target_mask, gen_F_256 * target_mask))
+        losses['lpips_hair'] = 0.5 * self.losses_dict['lpips_scale'] * (self.LPIPS(normalize(target) * HT_E, gen_w_256 * HT_E) +
+                                                                          self.LPIPS(normalize(target) * HT_E, gen_F_256 * HT_E))
+                                                                          
+        # Inpaint loss
+        if self.losses_dict['inpaint'] != 0.:
+            M_Inp = (1 - target_mask) * (1 - HT_E)
+            Smooth_M = self.dilated(M_Inp)
+            losses['inpaint'] = 0.5 * self.losses_dict['inpaint'] * self.LPIPS(normalize(target) * Smooth_M, gen_F_256 * Smooth_M)
+            losses['inpaint'] += 0.5 * self.losses_dict['inpaint'] * self.LPIPS(gen_w_256.detach() * Smooth_M * (1 - HT_E), gen_F_256 * Smooth_M * (1 - HT_E))
+        
+        return losses

losses/style/custom_loss.py

@@ -0,0 +1,67 @@
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+
+mse_loss = nn.MSELoss(reduction="mean")
+
+
+def custom_loss(x, y, mask=None, loss_type="l2", include_bkgd=True):
+    """
+    x, y: [N, C, H, W]
+    Computes L1/L2 loss
+
+    if include_bkgd is True:
+        use traditional MSE and L1 loss
+    else:
+        mask out background info using :mask
+        normalize loss with #1's in mask
+    """
+    if include_bkgd:
+        # perform simple mse or l1 loss
+        if loss_type == "l2":
+            loss_rec = mse_loss(x, y)
+        elif loss_type == "l1":
+            loss_rec = F.l1_loss(x, y)
+
+        return loss_rec
+
+    Nx, Cx, Hx, Wx = x.shape
+    Nm, Cm, Hm, Wm = mask.shape
+    mask = prepare_mask(x, mask)
+
+    x_reshape = torch.reshape(x, [Nx, -1])
+    y_reshape = torch.reshape(y, [Nx, -1])
+    mask_reshape = torch.reshape(mask, [Nx, -1])
+
+    if loss_type == "l2":
+        diff = (x_reshape - y_reshape) ** 2
+    elif loss_type == "l1":
+        diff = torch.abs(x_reshape - y_reshape)
+
+    # diff: [N, Cx * Hx * Wx]
+    # set elements in diff to 0 using mask
+    masked_diff = diff * mask_reshape
+    sum_diff = torch.sum(masked_diff, axis=-1)
+    # count non-zero elements; add :mask_reshape elements
+    norm_count = torch.sum(mask_reshape, axis=-1)
+    diff_norm = sum_diff / (norm_count + 1.0)
+
+    loss_rec = torch.mean(diff_norm)
+
+    return loss_rec
+
+
+def prepare_mask(x, mask):
+    """
+    Make mask similar to x.
+    Mask contains values in [0, 1].
+    Adjust channels and spatial dimensions.
+    """
+    Nx, Cx, Hx, Wx = x.shape
+    Nm, Cm, Hm, Wm = mask.shape
+    if Cm == 1:
+        mask = mask.repeat(1, Cx, 1, 1)
+
+    mask = F.interpolate(mask, scale_factor=Hx / Hm, mode="nearest")
+
+    return mask

losses/style/style_loss.py

@@ -0,0 +1,113 @@
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+
+import os
+
+from losses.style.custom_loss import custom_loss, prepare_mask
+from losses.style.vgg_activations import VGG16_Activations, VGG19_Activations, Vgg_face_dag
+
+
+class StyleLoss(nn.Module):
+    def __init__(self, VGG16_ACTIVATIONS_LIST=[21], normalize=False, distance="l2"):
+
+        super(StyleLoss, self).__init__()
+
+        self.vgg16_act = VGG16_Activations(VGG16_ACTIVATIONS_LIST)
+        self.vgg16_act.eval()
+
+        self.normalize = normalize
+        self.distance = distance
+
+    def get_features(self, model, x):
+
+        return model(x)
+
+    def mask_features(self, x, mask):
+
+        mask = prepare_mask(x, mask)
+        return x * mask
+
+    def gram_matrix(self, x):
+        """
+        :x is an activation tensor
+        """
+        N, C, H, W = x.shape
+        x = x.view(N * C, H * W)
+        G = torch.mm(x, x.t())
+
+        return G.div(N * H * W * C)
+
+    def cal_style(self, model, x, x_hat, mask1=None, mask2=None):
+        # Get features from the model for x and x_hat
+        with torch.no_grad():
+            act_x = self.get_features(model, x)
+        for layer in range(0, len(act_x)):
+            act_x[layer].detach_()
+
+        act_x_hat = self.get_features(model, x_hat)
+
+        loss = 0.0
+        for layer in range(0, len(act_x)):
+
+            # mask features if present
+            if mask1 is not None:
+                feat_x = self.mask_features(act_x[layer], mask1)
+            else:
+                feat_x = act_x[layer]
+            if mask2 is not None:
+                feat_x_hat = self.mask_features(act_x_hat[layer], mask2)
+            else:
+                feat_x_hat = act_x_hat[layer]
+
+            """
+            import ipdb; ipdb.set_trace()
+            fx = feat_x[0, ...].detach().cpu().numpy()
+            fx = (fx - fx.min()) / (fx.max() - fx.min())
+            fx = fx * 255.
+            fxhat = feat_x_hat[0, ...].detach().cpu().numpy()
+            fxhat = (fxhat - fxhat.min()) / (fxhat.max() - fxhat.min())
+            fxhat = fxhat * 255
+            from PIL import Image
+            import numpy as np
+            for idx, img in enumerate(fx):
+                img = fx[idx, ...]
+                img = img.astype(np.uint8)
+                img = Image.fromarray(img)
+                img.save('plot/feat_x/{}.png'.format(str(idx)))
+                img = fxhat[idx, ...]
+                img = img.astype(np.uint8)
+                img = Image.fromarray(img)
+                img.save('plot/feat_x_hat/{}.png'.format(str(idx)))
+            import ipdb; ipdb.set_trace()
+            """
+
+            # compute Gram matrix for x and x_hat
+            G_x = self.gram_matrix(feat_x)
+            G_x_hat = self.gram_matrix(feat_x_hat)
+
+            # compute layer wise loss and aggregate
+            loss += custom_loss(
+                G_x, G_x_hat, mask=None, loss_type=self.distance, include_bkgd=True
+            )
+
+        loss = loss / len(act_x)
+
+        return loss
+
+    def forward(self, x, x_hat, mask1=None, mask2=None):
+        x = x.cuda()
+        x_hat = x_hat.cuda()
+
+        # resize images to 256px resolution
+        N, C, H, W = x.shape
+        upsample2d = nn.Upsample(
+            scale_factor=256 / H, mode="bilinear", align_corners=True
+        )
+
+        x = upsample2d(x)
+        x_hat = upsample2d(x_hat)
+
+        loss = self.cal_style(self.vgg16_act, x, x_hat, mask1=mask1, mask2=mask2)
+
+        return loss

losses/style/vgg_activations.py

@@ -0,0 +1,187 @@
+import torch
+import torch.nn as nn
+
+import torchvision.models as models
+
+
+def requires_grad(model, flag=True):
+    for p in model.parameters():
+        p.requires_grad = flag
+
+
+class VGG16_Activations(nn.Module):
+    def __init__(self, feature_idx):
+        super(VGG16_Activations, self).__init__()
+        vgg16 = models.vgg16(pretrained=True)
+        features = list(vgg16.features)
+        self.features = nn.ModuleList(features).eval()
+        self.layer_id_list = feature_idx
+
+    def forward(self, x):
+        activations = []
+        for i, model in enumerate(self.features):
+            x = model(x)
+            if i in self.layer_id_list:
+                activations.append(x)
+
+        return activations
+
+
+class VGG19_Activations(nn.Module):
+    def __init__(self, feature_idx, requires_grad=False):
+        super(VGG19_Activations, self).__init__()
+        vgg19 = models.vgg19(pretrained=True)
+        requires_grad(vgg19, flag=False)
+        features = list(vgg19.features)
+        self.features = nn.ModuleList(features).eval()
+        self.layer_id_list = feature_idx
+
+    def forward(self, x):
+        activations = []
+        for i, model in enumerate(self.features):
+            x = model(x)
+            if i in self.layer_id_list:
+                activations.append(x)
+
+        return activations
+
+
+# http://www.robots.ox.ac.uk/~albanie/models/pytorch-mcn/vgg_face_dag.py
+class Vgg_face_dag(nn.Module):
+    def __init__(self):
+        super(Vgg_face_dag, self).__init__()
+        self.meta = {
+            "mean": [129.186279296875, 104.76238250732422, 93.59396362304688],
+            "std": [1, 1, 1],
+            "imageSize": [224, 224, 3],
+        }
+        self.conv1_1 = nn.Conv2d(
+            3, 64, kernel_size=[3, 3], stride=(1, 1), padding=(1, 1)
+        )
+        self.relu1_1 = nn.ReLU(inplace=True)
+        self.conv1_2 = nn.Conv2d(
+            64, 64, kernel_size=[3, 3], stride=(1, 1), padding=(1, 1)
+        )
+        self.relu1_2 = nn.ReLU(inplace=True)
+        self.pool1 = nn.MaxPool2d(
+            kernel_size=[2, 2], stride=[2, 2], padding=0, dilation=1, ceil_mode=False
+        )
+        self.conv2_1 = nn.Conv2d(
+            64, 128, kernel_size=[3, 3], stride=(1, 1), padding=(1, 1)
+        )
+        self.relu2_1 = nn.ReLU(inplace=True)
+        self.conv2_2 = nn.Conv2d(
+            128, 128, kernel_size=[3, 3], stride=(1, 1), padding=(1, 1)
+        )
+        self.relu2_2 = nn.ReLU(inplace=True)
+        self.pool2 = nn.MaxPool2d(
+            kernel_size=[2, 2], stride=[2, 2], padding=0, dilation=1, ceil_mode=False
+        )
+        self.conv3_1 = nn.Conv2d(
+            128, 256, kernel_size=[3, 3], stride=(1, 1), padding=(1, 1)
+        )
+        self.relu3_1 = nn.ReLU(inplace=True)
+        self.conv3_2 = nn.Conv2d(
+            256, 256, kernel_size=[3, 3], stride=(1, 1), padding=(1, 1)
+        )
+        self.relu3_2 = nn.ReLU(inplace=True)
+        self.conv3_3 = nn.Conv2d(
+            256, 256, kernel_size=[3, 3], stride=(1, 1), padding=(1, 1)
+        )
+        self.relu3_3 = nn.ReLU(inplace=True)
+        self.pool3 = nn.MaxPool2d(
+            kernel_size=[2, 2], stride=[2, 2], padding=0, dilation=1, ceil_mode=False
+        )
+        self.conv4_1 = nn.Conv2d(
+            256, 512, kernel_size=[3, 3], stride=(1, 1), padding=(1, 1)
+        )
+        self.relu4_1 = nn.ReLU(inplace=True)
+        self.conv4_2 = nn.Conv2d(
+            512, 512, kernel_size=[3, 3], stride=(1, 1), padding=(1, 1)
+        )
+        self.relu4_2 = nn.ReLU(inplace=True)
+        self.conv4_3 = nn.Conv2d(
+            512, 512, kernel_size=[3, 3], stride=(1, 1), padding=(1, 1)
+        )
+        self.relu4_3 = nn.ReLU(inplace=True)
+        self.pool4 = nn.MaxPool2d(
+            kernel_size=[2, 2], stride=[2, 2], padding=0, dilation=1, ceil_mode=False
+        )
+        self.conv5_1 = nn.Conv2d(
+            512, 512, kernel_size=[3, 3], stride=(1, 1), padding=(1, 1)
+        )
+        self.relu5_1 = nn.ReLU(inplace=True)
+        self.conv5_2 = nn.Conv2d(
+            512, 512, kernel_size=[3, 3], stride=(1, 1), padding=(1, 1)
+        )
+        self.relu5_2 = nn.ReLU(inplace=True)
+        self.conv5_3 = nn.Conv2d(
+            512, 512, kernel_size=[3, 3], stride=(1, 1), padding=(1, 1)
+        )
+        self.relu5_3 = nn.ReLU(inplace=True)
+        self.pool5 = nn.MaxPool2d(
+            kernel_size=[2, 2], stride=[2, 2], padding=0, dilation=1, ceil_mode=False
+        )
+        self.fc6 = nn.Linear(in_features=25088, out_features=4096, bias=True)
+        self.relu6 = nn.ReLU(inplace=True)
+        self.dropout6 = nn.Dropout(p=0.5)
+        self.fc7 = nn.Linear(in_features=4096, out_features=4096, bias=True)
+        self.relu7 = nn.ReLU(inplace=True)
+        self.dropout7 = nn.Dropout(p=0.5)
+        self.fc8 = nn.Linear(in_features=4096, out_features=2622, bias=True)
+
+    def forward(self, x):
+        activations = []
+        x1 = self.conv1_1(x)
+        activations.append(x1)
+
+        x2 = self.relu1_1(x1)
+        x3 = self.conv1_2(x2)
+        x4 = self.relu1_2(x3)
+        x5 = self.pool1(x4)
+        x6 = self.conv2_1(x5)
+        activations.append(x6)
+
+        x7 = self.relu2_1(x6)
+        x8 = self.conv2_2(x7)
+        x9 = self.relu2_2(x8)
+        x10 = self.pool2(x9)
+        x11 = self.conv3_1(x10)
+        activations.append(x11)
+
+        x12 = self.relu3_1(x11)
+        x13 = self.conv3_2(x12)
+        x14 = self.relu3_2(x13)
+        x15 = self.conv3_3(x14)
+        x16 = self.relu3_3(x15)
+        x17 = self.pool3(x16)
+        x18 = self.conv4_1(x17)
+        activations.append(x18)
+
+        x19 = self.relu4_1(x18)
+        x20 = self.conv4_2(x19)
+        x21 = self.relu4_2(x20)
+        x22 = self.conv4_3(x21)
+        x23 = self.relu4_3(x22)
+        x24 = self.pool4(x23)
+        x25 = self.conv5_1(x24)
+        activations.append(x25)
+
+        """
+        x26 = self.relu5_1(x25)
+        x27 = self.conv5_2(x26)
+        x28 = self.relu5_2(x27)
+        x29 = self.conv5_3(x28)
+        x30 = self.relu5_3(x29)
+        x31_preflatten = self.pool5(x30)
+        x31 = x31_preflatten.view(x31_preflatten.size(0), -1)
+        x32 = self.fc6(x31)
+        x33 = self.relu6(x32)
+        x34 = self.dropout6(x33)
+        x35 = self.fc7(x34)
+        x36 = self.relu7(x35)
+        x37 = self.dropout7(x36)
+        x38 = self.fc8(x37)
+        """
+
+        return activations

losses/vgg_loss.py

@@ -0,0 +1,51 @@
+import torch
+import torchvision
+import torch.nn as nn
+
+class VGG19(torch.nn.Module):
+    def __init__(self, requires_grad=False):
+        super().__init__()
+        vgg_pretrained_features = torchvision.models.vgg19(pretrained=True).features
+        self.slice1 = torch.nn.Sequential()
+        self.slice2 = torch.nn.Sequential()
+        self.slice3 = torch.nn.Sequential()
+        self.slice4 = torch.nn.Sequential()
+        self.slice5 = torch.nn.Sequential()
+        for x in range(2):
+            self.slice1.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(2, 7):
+            self.slice2.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(7, 12):
+            self.slice3.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(12, 21):
+            self.slice4.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(21, 30):
+            self.slice5.add_module(str(x), vgg_pretrained_features[x])
+        if not requires_grad:
+            for param in self.parameters():
+                param.requires_grad = False
+
+    def forward(self, X):
+        h_relu1 = self.slice1(X)
+        h_relu2 = self.slice2(h_relu1)
+        h_relu3 = self.slice3(h_relu2)
+        h_relu4 = self.slice4(h_relu3)
+        h_relu5 = self.slice5(h_relu4)
+        out = [h_relu1, h_relu2, h_relu3, h_relu4, h_relu5]
+        return out
+
+
+# Perceptual loss that uses a pretrained VGG network
+class VGGLoss(nn.Module):
+    def __init__(self):
+        super(VGGLoss, self).__init__()
+        self.vgg = VGG19().cuda()
+        self.criterion = nn.L1Loss()
+        self.weights = [1.0 / 32, 1.0 / 16, 1.0 / 8, 1.0 / 4, 1.0]
+
+    def forward(self, x, y):
+        x_vgg, y_vgg = self.vgg(x), self.vgg(y)
+        loss = 0
+        for i in range(len(x_vgg)):
+            loss += self.weights[i] * self.criterion(x_vgg[i], y_vgg[i].detach())
+        return loss

main.py

@@ -0,0 +1,80 @@
+import argparse
+import os
+import sys
+from pathlib import Path
+
+from torchvision.utils import save_image
+from tqdm.auto import tqdm
+
+from hair_swap import HairFast, get_parser
+
+
+def main(model_args, args):
+    hair_fast = HairFast(model_args)
+
+    experiments: list[str | tuple[str, str, str]] = []
+    if args.file_path is not None:
+        with open(args.file_path, 'r') as file:
+            experiments.extend(file.readlines())
+
+    if all(path is not None for path in (args.face_path, args.shape_path, args.color_path)):
+        experiments.append((args.face_path, args.shape_path, args.color_path))
+
+    for exp in tqdm(experiments):
+        if isinstance(exp, str):
+            file_1, file_2, file_3 = exp.split()
+        else:
+            file_1, file_2, file_3 = exp
+
+        face_path = args.input_dir / file_1
+        shape_path = args.input_dir / file_2
+        color_path = args.input_dir / file_3
+
+        base_name = '_'.join([path.stem for path in (face_path, shape_path, color_path)])
+        exp_name = base_name if model_args.save_all else None
+
+        if isinstance(exp, str) or args.result_path is None:
+            os.makedirs(args.output_dir, exist_ok=True)
+            output_image_path = args.output_dir / f'{base_name}.png'
+        else:
+            os.makedirs(args.result_path.parent, exist_ok=True)
+            output_image_path = args.result_path
+
+        final_image = hair_fast.swap(face_path, shape_path, color_path, benchmark=args.benchmark, exp_name=exp_name)
+        save_image(final_image, output_image_path)
+
+
+if __name__ == "__main__":
+    model_parser = get_parser()
+    parser = argparse.ArgumentParser(description='HairFast evaluate')
+    parser.add_argument('--input_dir', type=Path, default='', help='The directory of the images to be inverted')
+    parser.add_argument('--benchmark', action='store_true', help='Calculates the speed of the method during the session')
+
+    # Arguments for a set of experiments
+    parser.add_argument('--file_path', type=Path, default=None,
+                        help='File with experiments with the format "face_path.png shape_path.png color_path.png"')
+    parser.add_argument('--output_dir', type=Path, default=Path('output'), help='The directory for final results')
+
+    # Arguments for single experiment
+    parser.add_argument('--face_path', type=Path, default=None, help='Path to the face image')
+    parser.add_argument('--shape_path', type=Path, default=None, help='Path to the shape image')
+    parser.add_argument('--color_path', type=Path, default=None, help='Path to the color image')
+    parser.add_argument('--result_path', type=Path, default=None, help='Path to save the result')
+
+    args, unknown1 = parser.parse_known_args()
+    model_args, unknown2 = model_parser.parse_known_args()
+
+    unknown_args = set(unknown1) & set(unknown2)
+    if unknown_args:
+        file_ = sys.stderr
+        print(f"Unknown arguments: {unknown_args}", file=file_)
+
+        print("\nExpected arguments for the model:", file=file_)
+        model_parser.print_help(file=file_)
+
+        print("\nExpected arguments for evaluate:", file=file_)
+        parser.print_help(file=file_)
+
+        sys.exit(1)
+
+    main(model_args, args)

models/Alignment.py

@@ -0,0 +1,181 @@
+import torch
+import torch.nn.functional as F
+import torchvision.transforms as T
+from torch import nn
+
+from models.CtrlHair.shape_branch.config import cfg as cfg_mask
+from models.CtrlHair.shape_branch.solver import get_hair_face_code, get_new_shape, Solver as SolverMask
+from models.Encoders import RotateModel
+from models.Net import Net, get_segmentation
+from models.sean_codes.models.pix2pix_model import Pix2PixModel, SEAN_OPT, encode_sean, decode_sean
+from utils.image_utils import DilateErosion
+from utils.save_utils import save_vis_mask, save_gen_image, save_latents
+
+
+class Alignment(nn.Module):
+    """
+    Module for transferring the desired hair shape
+    """
+
+    def __init__(self, opts, latent_encoder=None, net=None):
+        super().__init__()
+        self.opts = opts
+        self.latent_encoder = latent_encoder
+        if not net:
+            self.net = Net(self.opts)
+        else:
+            self.net = net
+
+        self.sean_model = Pix2PixModel(SEAN_OPT)
+        self.sean_model.eval()
+
+        solver_mask = SolverMask(cfg_mask, device=self.opts.device, local_rank=-1, training=False)
+        self.mask_generator = solver_mask.gen
+        self.mask_generator.load_state_dict(torch.load('pretrained_models/ShapeAdaptor/mask_generator.pth'))
+
+        self.rotate_model = RotateModel()
+        self.rotate_model.load_state_dict(torch.load(self.opts.rotate_checkpoint)['model_state_dict'])
+        self.rotate_model.to(self.opts.device).eval()
+
+        self.dilate_erosion = DilateErosion(dilate_erosion=self.opts.smooth, device=self.opts.device)
+        self.to_bisenet = T.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
+
+    @torch.inference_mode()
+    def shape_module(self, im_name1: str, im_name2: str, name_to_embed, only_target=True, **kwargs):
+        device = self.opts.device
+
+        # load images
+        img1_in = name_to_embed[im_name1]['image_256']
+        img2_in = name_to_embed[im_name2]['image_256']
+
+        # load latents
+        latent_W_1 = name_to_embed[im_name1]["W"]
+        latent_W_2 = name_to_embed[im_name2]["W"]
+
+        # load masks
+        inp_mask1 = name_to_embed[im_name1]['mask']
+        inp_mask2 = name_to_embed[im_name2]['mask']
+
+        # Rotate stage
+        if img1_in is not img2_in:
+            rotate_to = self.rotate_model(latent_W_2[:, :6], latent_W_1[:, :6])
+            rotate_to = torch.cat((rotate_to, latent_W_2[:, 6:]), dim=1)
+            I_rot, _ = self.net.generator([rotate_to], input_is_latent=True, return_latents=False)
+
+            I_rot_to_seg = ((I_rot + 1) / 2).clip(0, 1)
+            I_rot_to_seg = self.to_bisenet(I_rot_to_seg)
+            rot_mask = get_segmentation(I_rot_to_seg)
+        else:
+            I_rot = None
+            rot_mask = inp_mask2
+
+        # Shape Adaptor
+        if img1_in is not img2_in:
+            face_1, hair_1 = get_hair_face_code(self.mask_generator, inp_mask1[0, 0, ...])
+            face_2, hair_2 = get_hair_face_code(self.mask_generator, rot_mask[0, 0, ...])
+
+            target_mask = get_new_shape(self.mask_generator, face_1, hair_2)[None, None]
+        else:
+            target_mask = inp_mask1
+
+        # Hair mask
+        hair_mask_target = torch.where(target_mask == 13, torch.ones_like(target_mask, device=device),
+                                       torch.zeros_like(target_mask, device=device))
+
+        if self.opts.save_all:
+            exp_name = exp_name if (exp_name := kwargs.get('exp_name')) is not None else ""
+            output_dir = self.opts.save_all_dir / exp_name
+            if I_rot is not None:
+                save_gen_image(output_dir, 'Shape', f'{im_name2}_rotate_to_{im_name1}.png', I_rot)
+            save_vis_mask(output_dir, 'Shape', f'mask_{im_name1}.png', inp_mask1)
+            save_vis_mask(output_dir, 'Shape', f'mask_{im_name2}.png', inp_mask2)
+            save_vis_mask(output_dir, 'Shape', f'mask_{im_name2}_rotate_to_{im_name1}.png', rot_mask)
+            save_vis_mask(output_dir, 'Shape', f'mask_{im_name1}_{im_name2}_target.png', target_mask)
+
+        if only_target:
+            return {'HM_X': hair_mask_target}
+        else:
+            hair_mask1 = torch.where(inp_mask1 == 13, torch.ones_like(inp_mask1, device=device),
+                                     torch.zeros_like(inp_mask1, device=device))
+            hair_mask2 = torch.where(inp_mask2 == 13, torch.ones_like(inp_mask2, device=device),
+                                     torch.zeros_like(inp_mask2, device=device))
+
+            return inp_mask1, hair_mask1, inp_mask2, hair_mask2, target_mask, hair_mask_target
+
+    @torch.inference_mode()
+    def align_images(self, im_name1, im_name2, name_to_embed, **kwargs):
+        # load images
+        img1_in = name_to_embed[im_name1]['image_256']
+        img2_in = name_to_embed[im_name2]['image_256']
+
+        # load latents
+        latent_S_1, latent_F_1 = name_to_embed[im_name1]["S"], name_to_embed[im_name1]["F"]
+        latent_S_2, latent_F_2 = name_to_embed[im_name2]["S"], name_to_embed[im_name2]["F"]
+
+        # Shape Module
+        if img1_in is img2_in:
+            hair_mask_target = self.shape_module(im_name1, im_name2, name_to_embed, only_target=True, **kwargs)['HM_X']
+            return {'latent_F_align': latent_F_1, 'HM_X': hair_mask_target}
+
+        inp_mask1, hair_mask1, inp_mask2, hair_mask2, target_mask, hair_mask_target = (
+            self.shape_module(im_name1, im_name2, name_to_embed, only_target=False, **kwargs)
+        )
+
+        images = torch.cat([img1_in, img2_in], dim=0)
+        labels = torch.cat([inp_mask1, inp_mask2], dim=0)
+
+        # SEAN for inpaint
+        img1_code, img2_code = encode_sean(self.sean_model, images, labels)
+
+        gen1_sean = decode_sean(self.sean_model, img1_code.unsqueeze(0), target_mask)
+        gen2_sean = decode_sean(self.sean_model, img2_code.unsqueeze(0), target_mask)
+
+        # Encoding result in F from E4E
+        enc_imgs = self.latent_encoder([gen1_sean, gen2_sean])
+        intermediate_align, latent_inter = enc_imgs["F"][0].unsqueeze(0), enc_imgs["W"][0].unsqueeze(0)
+        latent_F_out_new, latent_out = enc_imgs["F"][1].unsqueeze(0), enc_imgs["W"][1].unsqueeze(0)
+
+        # Alignment of F space
+        masks = [
+            1 - (1 - hair_mask1) * (1 - hair_mask_target),
+            hair_mask_target,
+            hair_mask2 * hair_mask_target
+        ]
+        masks = torch.cat(masks, dim=0)
+        # masks = T.functional.resize(masks, (1024, 1024), interpolation=T.InterpolationMode.NEAREST)
+
+        dilate, erosion = self.dilate_erosion.mask(masks)
+        free_mask = [
+            dilate[0],
+            erosion[1],
+            erosion[2]
+        ]
+        free_mask = torch.stack(free_mask, dim=0)
+        free_mask_down_32 = F.interpolate(free_mask.float(), size=(32, 32), mode='bicubic')
+        interpolation_low = 1 - free_mask_down_32
+
+        latent_F_align = intermediate_align + interpolation_low[0] * (latent_F_1 - intermediate_align)
+        latent_F_align = latent_F_out_new + interpolation_low[1] * (latent_F_align - latent_F_out_new)
+        latent_F_align = latent_F_2 + interpolation_low[2] * (latent_F_align - latent_F_2)
+
+        if self.opts.save_all:
+            exp_name = exp_name if (exp_name := kwargs.get('exp_name')) is not None else ""
+            output_dir = self.opts.save_all_dir / exp_name
+            save_gen_image(output_dir, 'Align', f'{im_name1}_{im_name2}_SEAN.png', gen1_sean)
+            save_gen_image(output_dir, 'Align', f'{im_name2}_{im_name1}_SEAN.png', gen2_sean)
+
+            img1_e4e = self.net.generator([latent_inter], input_is_latent=True, return_latents=False, start_layer=4,
+                                          end_layer=8, layer_in=intermediate_align)[0]
+            img2_e4e = self.net.generator([latent_out], input_is_latent=True, return_latents=False, start_layer=4,
+                                          end_layer=8, layer_in=latent_F_out_new)[0]
+
+            save_gen_image(output_dir, 'Align', f'{im_name1}_{im_name2}_e4e.png', img1_e4e)
+            save_gen_image(output_dir, 'Align', f'{im_name2}_{im_name1}_e4e.png', img2_e4e)
+
+            gen_im, _ = self.net.generator([latent_S_1], input_is_latent=True, return_latents=False, start_layer=4,
+                                           end_layer=8, layer_in=latent_F_align)
+
+            save_gen_image(output_dir, 'Align', f'{im_name1}_{im_name2}_output.png', gen_im)
+            save_latents(output_dir, 'Align', f'{im_name1}_{im_name2}_F.npz', latent_F_align=latent_F_align)
+
+        return {'latent_F_align': latent_F_align, 'HM_X': hair_mask_target}

models/Blending.py

@@ -0,0 +1,81 @@
+import torch
+from torch import nn
+
+from models.Encoders import ClipBlendingModel, PostProcessModel
+from models.Net import Net
+from utils.bicubic import BicubicDownSample
+from utils.image_utils import DilateErosion
+from utils.save_utils import save_gen_image, save_latents
+
+
+class Blending(nn.Module):
+    """
+    Module for transferring the desired hair color and post processing
+    """
+
+    def __init__(self, opts, net=None):
+        super().__init__()
+        self.opts = opts
+        if net is None:
+            self.net = Net(self.opts)
+        else:
+            self.net = net
+
+        blending_checkpoint = torch.load(self.opts.blending_checkpoint)
+        self.blending_encoder = ClipBlendingModel(blending_checkpoint.get('clip', "ViT-B/32"))
+        self.blending_encoder.load_state_dict(blending_checkpoint['model_state_dict'], strict=False)
+        self.blending_encoder.to(self.opts.device).eval()
+
+        self.post_process = PostProcessModel().to(self.opts.device).eval()
+        self.post_process.load_state_dict(torch.load(self.opts.pp_checkpoint)['model_state_dict'])
+
+        self.dilate_erosion = DilateErosion(dilate_erosion=self.opts.smooth, device=self.opts.device)
+        self.downsample_256 = BicubicDownSample(factor=4)
+
+    @torch.inference_mode()
+    def blend_images(self, align_shape, align_color, name_to_embed, **kwargs):
+        I_1 = name_to_embed['face']['image_norm_256']
+        I_2 = name_to_embed['shape']['image_norm_256']
+        I_3 = name_to_embed['color']['image_norm_256']
+
+        mask_de = self.dilate_erosion.hair_from_mask(
+            torch.cat([name_to_embed[x]['mask'] for x in ['face', 'color']], dim=0)
+        )
+        HM_1D, _ = mask_de[0][0].unsqueeze(0), mask_de[1][0].unsqueeze(0)
+        HM_3D, HM_3E = mask_de[0][1].unsqueeze(0), mask_de[1][1].unsqueeze(0)
+
+        latent_S_1, latent_F_align = name_to_embed['face']['S'], align_shape['latent_F_align']
+        HM_X = align_color['HM_X']
+
+        latent_S_3 = name_to_embed['color']["S"]
+
+        HM_XD, _ = self.dilate_erosion.mask(HM_X)
+        target_mask = (1 - HM_1D) * (1 - HM_3D) * (1 - HM_XD)
+
+        # Blending
+        if I_1 is not I_3 or I_1 is not I_2:
+            S_blend_6_18 = self.blending_encoder(latent_S_1[:, 6:], latent_S_3[:, 6:], I_1 * target_mask, I_3 * HM_3E)
+            S_blend = torch.cat((latent_S_1[:, :6], S_blend_6_18), dim=1)
+        else:
+            S_blend = latent_S_1
+
+        I_blend, _ = self.net.generator([S_blend], input_is_latent=True, return_latents=False, start_layer=4,
+                                        end_layer=8, layer_in=latent_F_align)
+        I_blend_256 = self.downsample_256(I_blend)
+
+        # Post Process
+        S_final, F_final = self.post_process(I_1, I_blend_256)
+        I_final, _ = self.net.generator([S_final], input_is_latent=True, return_latents=False,
+                                         start_layer=5, end_layer=8, layer_in=F_final)
+
+        if self.opts.save_all:
+            exp_name = exp_name if (exp_name := kwargs.get('exp_name')) is not None else ""
+            output_dir = self.opts.save_all_dir / exp_name
+            save_gen_image(output_dir, 'Blending', 'blending.png', I_blend)
+            save_latents(output_dir, 'Blending', 'blending.npz', S_blend=S_blend)
+
+            save_gen_image(output_dir, 'Final', 'final.png', I_final)
+            save_latents(output_dir, 'Final', 'final.npz', S_final=S_final, F_final=F_final)
+
+        final_image = ((I_final[0] + 1) / 2).clip(0, 1)
+        return final_image

models/CtrlHair/.gitignore

@@ -0,0 +1,8 @@
+.idea/
+.DS_Store
+/dataset_info_ctrlhair/
+/external_model_params/
+/model_trained/
+**/__pycache__/**
+
+

models/CtrlHair/README.md

@@ -0,0 +1,270 @@
+# [GAN with Multivariate Disentangling for Controllable Hair Editing (ECCV 2022)](https://github.com/XuyangGuo/xuyangguo.github.io/raw/master/database/CtrlHair/CtrlHair.pdf)
+
+[Xuyang Guo](https://xuyangguo.github.io/), [Meina Kan](http://vipl.ict.ac.cn/homepage/mnkan/Publication/), Tianle Chen, [Shiguang Shan](https://scholar.google.com/citations?user=Vkzd7MIAAAAJ)
+
+![demo1](https://github.com/XuyangGuo/xuyangguo.github.io/blob/master/database/CtrlHair/resources/demo1.gif?raw=true)
+
+![demo2](https://github.com/XuyangGuo/xuyangguo.github.io/blob/master/database/CtrlHair/resources/demo2.gif?raw=true)
+
+![demo3](https://github.com/XuyangGuo/xuyangguo.github.io/blob/master/database/CtrlHair/resources/demo3.gif?raw=true)
+
+## Abstract
+
+> Hair editing is an essential but challenging task in portrait editing considering the complex geometry and material of hair. Existing methods have achieved promising results by editing through a reference photo, user-painted mask, or guiding strokes. However, when a user provides no reference photo or hardly paints a desirable mask, these works fail to edit. Going a further step, we propose an efficiently controllable method that can provide a set of sliding bars to do continuous and fine hair editing. Meanwhile, it also naturally supports discrete editing through a reference photo and user-painted mask. Specifically, we propose a generative adversarial network with a multivariate Gaussian disentangling module. Firstly, an encoder disentangles the hair's major attributes, including color, texture, and shape, to separate latent representations. The latent representation of each attribute is modeled as a standard multivariate Gaussian distribution, to make each dimension of an attribute be changed continuously and finely. Benefiting from the Gaussian distribution, any manual editing including sliding a bar, providing a reference photo, and painting a mask can be easily made, which is flexible and friendly for users to interact with. Finally, with changed latent representations, the decoder outputs a portrait with the edited hair. Experiments show that our method can edit each attribute's dimension continuously and separately. Besides, when editing through reference images and painted masks like existing methods, our method achieves comparable results in terms of FID and visualization. Codes can be found at [https://github.com/XuyangGuo/CtrlHair](https://github.com/XuyangGuo/CtrlHair).
+
+![archi](https://raw.githubusercontent.com/XuyangGuo/xuyangguo.github.io/master/database/CtrlHair/resources/architecture.png)
+
+
+## Installation
+
+Clone this repo.
+
+```bash
+git clone https://github.com/XuyangGuo/CtrlHair.git
+cd CtrlHair
+```
+
+The code requires python 3.6.
+
+We recommend using Anaconda to regulate packages.
+
+Dependencies:
+- PyTorch 1.8.2
+- torchvision, tensorboardX, dlib
+- pillow, pandas, scikit-learn, opencv-python
+- PyQt5, tqdm, addict, dill
+
+Please download all [external trained models](https://drive.google.com/drive/folders/1X0y82o7-JB6nGYIWdbbJa4MTFbtV9bei?usp=sharing) ([using Baidu Netdisk](https://pan.baidu.com/s/1vnldtoT9G-5gfvUjjUSVOA?pwd=1234) with password `1234` alternatively), move it to `./` with the correct path `CtrlHair/external_model_params`. (The directory contains the model parameters of face analysis tools required by our project, including SEAN for face encoding and generator with masks, BiSeNet for face parsing, and 68/81 facial landmark detector.)
+
+## Editing by Pretrained Model
+
+Firstly, refer to the "Installation" section above.
+
+Please download [the pre-trained model of CtrlHair](https://drive.google.com/drive/folders/1opQhmc7ckS3J8qdLii_EMqmxYCcBLznO?usp=sharing) ([using Baidu Netdisk](https://pan.baidu.com/s/1O_Hu5dnk4GwmUrIBDWpwFA?pwd=1234) with password `1234` alternatively), move it to `./` with the correct path `CtrlHair/model_trained`.
+
+#### Editing with UI directly (recommended)
+
+```bash
+python ui/frontend_demo.py
+```
+
+Here are some parameters of it:
+
+- `-g D` Use the gpu `D`. (default is `0`)
+- `-n True|False` Whether the input image need crop. (default is `True`)
+- `--no_blending` Do not use poisson blending as post processing. If not blend, the result image will look slightly different from the input image in some details in non-hair region, but the image quality will be better.
+
+The edited results can be found in `temp_folder/demo_output/out_img.png`, and the edited shape is in `temp_folder/demo_output/input_parsing.png`. The `temp_folder` is created automaticly during running, which could be removed after closing the procedure.
+
+#### Editing with Batch
+
+If you want to edit with a mass batch, or want to achieve editing functions such as interpolation, multi style sampling, and continuous gradient, etc, please use the interfaces of `ui/backend.py/Backend` and code your own python script.`Backend` class is the convenient encapsulation of basic functions of CtrlHair, and there are detailed comments for each function. The `main` scetion in the final of `backend.py` shows an simle example of usage of `Backend`.
+
+## Training New Models
+
+#### Data Preparation
+In addition to the images in the dataset, training also involves many image annotations, including face parsing, facial landmarks, sean codes, color annotations, face rotation angle, a small amount of annotations (Curliness of Hair), etc.
+
+Please download [the dataset information](https://drive.google.com/drive/folders/10p87Mobgueg9rdHyLPkX8xqEIvga6p2C?usp=sharing)  ([using Baidu Netdisk](https://pan.baidu.com/s/1D3D2JqxIR6miCeMNssucKw?pwd=1234) with password `1234` alternatively) that we have partially processed, move it to `./` with the correct path `CtrlHair/dataset_info_ctrlhair`.
+
+Then execute the following scripts sequentially for preprocessing.
+
+Get facial segmentation mask
+```bash
+python dataset_scripts/script_get_mask.py
+```
+
+Get 68 facial landmarks and 81 facial landmarks
+```bash
+python dataset_scripts/script_landmark_detection.py
+```
+
+Get SEAN feature codes of the dataset
+```bash
+python dataset_scripts/script_get_sean_code.py
+```
+
+Get color label of the dataset
+```bash
+python dataset_scripts/script_get_rgb_hsv_label.py
+python dataset_scripts/script_color_var_label.py
+```
+
+After complete processing, for training, the correct directory structure in `CtrlHair/dataset_info_ctrlhair` is as follows:
+
+- `CelebaMask_HQ` / `ffhq`  (if you want to add your own dataset, please regulate them as these two cases)
+  - `images_256`   -> cropped images with the resolution 256.
+  - `label`       -> mask label (0, 1, 2, ..., 20) for each pixel
+  - `angle.csv`    restore face rotation angle of each image
+  - `attr_gender.csv`  restore gender of each image
+- `color_var_stat_dict.pkl`, `rgb_stat_dict.pkl`, `hsv_stat_dict_ordered.pkl` store the label of variance, rgb of hair color, and the hsv distribution
+- `sean_code_dict.pkl`  store the sean feature code of images in dataset
+- `landmark68.pkl`, `landmark81.pkl`  store the facial landmarks of the dataset
+- `manual_label`
+  - `curliness`
+    - `-1.txt`, `1.txt`, `test_1.txt`, `test_-1.txt` labeled data list
+
+#### Training Networks
+In order to better control the parameters in the model, we train the entire model separately and divide it into four parts, including curliness classifier, color encoder, color & texture branch, shape branch.
+
+**1. Train the curliness classifier**
+```bash
+python color_texture_branch/predictor/predictor_train.py -c p002 -g 0
+```
+
+Here are some parameters of it:
+
+- `-g D` Use the gpu `D`. (default is `0`)
+- `-c pxxx` Using the model hyper-parameters config named `pxxx`. Please see the config detail in `color_texture_branch/predictor/predictor_config.py`
+
+The trained model and its tensorboard summary are saved in `model_trained/curliness_classifier`.
+
+**2. Train the color encoder**
+```bash
+python color_texture_branch/predictor/predictor_train.py -c p004 -g 0
+```
+
+The parameters are similar like the curliness classifier.
+
+The trained model and its tensorboard summary are saved in `model_trained/color_encoder`.
+
+**3. Train the color & texture branch**
+```bash
+python color_texture_branch/scripts.py -c 045 -g 0
+```
+
+This part depends on the curliness classifier and the color encoder as seen in `color_texture_branch/config.py`:
+
+```python
+...
+'predictor': {'curliness': 'p002', 'rgb': 'p004'},
+...
+```
+
+Here are some parameters of it:
+
+- `-g D` Use the gpu `D`. (default is `0`)
+- `-c xxx` Using the model hyper-parameters config named `xxx`. Please see the config detail in `color_texture_branch/config.py`
+
+The trained model, its tensorboard, editing results in training are saved in `model_trained/color_texture`.
+
+
+Since the training of texture is unsupervised, we need to find some semantic orthogonal directions after training for sliding bars. Please run:
+```bash
+python color_texture_branch/script_find_direction.py -c xxx
+```
+The parameter `-c xxx` is same as above shape config.
+This process will generate a folder named `direction_find` in the directory `model_trained/color_texture/yourConfigName`,
+where `direction_find/texture_dir_n` stores many random `n`-th directions to be selected, 
+and the corresponding visual changes can be seen in `direction_find/texture_n`. 
+When the choice is decided, move the corresponding `texture_dir_n/xxx.pkl` file to `../texture_dir_used` and rename it as you wish (the pretrained No.045 texture model shows an example).
+Afterthat, run `python color_texture_branch/script_find_direction.py -c xxx` again, and repeat the process until the amount of semantic directions is enough.
+
+
+**4. Train the shape branch**
+
+Shape editing employs a transfer-like training process. Before transferring, we use 68 feature points to pre-align the face of the target hairstyle, so as to achieve a certain degree of face adaptation. In order to speed up this process during training, it is necessary to generate some buffer pools to store the pre-aligned masks of the training set and test set respectively.
+
+Generate buffer pool for testing
+```bash
+python shape_branch/script_adaptor_test_pool.py
+```
+The testing buffer pool will be saved in `dataset_info_ctrlhair/shape_testing_wrap_pool`.
+
+Generate buffer pool for training
+```bash
+python shape_branch/script_adaptor_train_pool.py
+```
+The training buffer pool will be saved in `dataset_info_ctrlhair/shape_training_wrap_pool`.
+
+Note that the `script_adaptor_train_pool.py` process will execute for a very very long time until the setting of maximum number of files for buffering is reached.  
+This process can be performed concurrently with the subsequent shape training process. 
+The training data for shape training are all dynamically picked from this buffer pool.
+
+Training the shape branch model
+```bash
+python shape_branch/scripts.py -c 054 -g 0
+```
+
+Here are some parameters of it:
+
+- `-g D` Use the gpu `D`. (default is `0`)
+- `-c xxx` Using the model hyper-parameters config named `xxx`. Please see the config detail in `shape_branch/config.py`
+
+The trained model, its tensorboard, editing results in training are saved in `model_trained/shape`.
+
+Since the training is unsupervised, we need to find some semantic orthogonal directions after training for sliding bars. Please run:
+```bash
+python shape_branch/script_find_direction.py -c xxx
+```
+The parameter `-c xxx` is same as above shape config.
+The entire usage method is similar to texture, but the folder is changed to the `model_trained/shape` directory (the pretrained No.054 shape model shows an example).
+
+After all the above training, use `python ui/frontend_demo.py` to edit. 
+You can also use interfaces in `ui/backend.py/Backend` to program your editing scripts.
+
+## Training New Models with Your Own Images Dataset
+
+Our method only needs unlabeled face images to augment the dataset, which is convenient and is a strength of CtrlHair.
+
+#### Data Preparation
+
+For your own images dataset, firstly, crop and resize them. Please collect them into a single directory, and modify `root_dir` and `dataset_name` for your dataset in `dataset_scripts/script_crop.py`. Then execute
+```bash
+python dataset_scripts/script_crop.py
+```
+After cropping, the dataset should be cropped at `dataset_info_ctrlhair/your_dataset_name/images_256`. Your dataset should have similar structure like `dataset_info_ctrlhair/ffhq`.
+
+Modify `DATASET_NAME` in `global_value_utils.py ` for your dataset.
+
+Do the same steps as the section "Data Preparation" of "Training New Models" in this README.
+
+Predict face rotation angle and gender for your dataset. 
+This will be used to filter the dataset. 
+You can use tools like [3DDFA](https://github.com/cleardusk/3DDFA) and [deepface](https://github.com/serengil/deepface), then output them to `angle.csv` and `attr_gender.csv` in `dataset_info_ctrlhair/yourdataset` (pandas is recommended for generating csv). `dataset_info_ctrlhair/ffhq` shows a preprocessing example. 
+Sorry for that we don't provide these code. Alternatively, if you don't want to depend and use these filter, please modify `angle_filter` and `gender_filter` to `False` in `common_dataset.py`. 
+
+#### Training Networks
+
+Add and adjust your config in `color_texture_branch/predictor/predictor_config.py`, `color_texture_branch/config.py`, `shape/config.py`.
+
+Do the same steps as the section "Training Networks" of "Training New Models" in this README, but with your config.
+
+Finally, change the `DEFAULT_CONFIG_COLOR_TEXTURE_BRANCH` and `DEFAULT_CONFIG_SHAPE_BRANCH` as yours in `global_value_utils.py`.
+Use `python ui/frontend_demo.py` to edit. Or you can also use interfaces in `ui/backend.py/Backend` to program your editing scripts.
+
+
+## Code Structure
+- `color_texture_branch`: color and texture editing branch
+  - `predictor`: color encoder and curliness classifier
+- `shape_branch`: shape editing branch
+- `ui`: encapsulated backend interfaces and frontend UI
+- `dataset_scripts`: scripts for preprocessing dataset
+- `external_code`: codes of external tools
+- `sean_codes`: modified from SEAN project, which is used for image feature extraction and generation
+- `my_pylib`, `my_torchlib`, `utils`: auxiliary code library
+- `wrap_codes`: used for shape align before shape transfer
+- `dataset_info_ctrlhair`: the root directory of dataset
+- `model_trained`: trained model parameters, tensorboard and visual results during training
+- `external_model_params`: pretrained model parameters used for external codes
+- `imgs`: some example images are provided for testing
+
+## Citation
+If you use this code for your research, please cite our papers.
+```
+@inproceedings{guo2022gan,
+  title={GAN with Multivariate Disentangling for Controllable Hair Editing},
+  author={Guo, Xuyang and Kan, Meina and Chen, Tianle and Shan, Shiguang},
+  booktitle={European Conference on Computer Vision},
+  year={2022},
+  pages={655--670},
+  organization={Springer}
+}
+```
+
+This work is also inspired by our previous work [X. Guo, et al., STD-GAN (CVIU2021)](https://github.com/XuyangGuo/STD-GAN) for instance-level facial attributes editing.
+
+## References & Acknowledgments
+- [ZPdesu / SEAN](https://github.com/ZPdesu/SEAN)
+- [zhhoper / RI_render_DPR](https://github.com/zhhoper/RI_render_DPR)
+- [zllrunning / face-parsing.PyTorch](https://github.com/zllrunning/face-parsing.PyTorch)

models/CtrlHair/color_texture_branch/__init__.py

@@ -0,0 +1,8 @@
+# -*- coding: utf-8 -*-
+
+"""
+# File name:    __init__.py.py
+# Time :        2021/11/15 11:18
+# Author:       xyguoo@163.com
+# Description:  
+"""

models/CtrlHair/color_texture_branch/config.py

@@ -0,0 +1,141 @@
+# -*- coding: utf-8 -*-
+
+"""
+# File name:    config.py
+# Time :        2021/11/17 13:10
+# Author:       xyguoo@163.com
+# Description:  
+"""
+
+import addict  # nesting dict
+import os
+import argparse
+
+from global_value_utils import GLOBAL_DATA_ROOT, DEFAULT_CONFIG_COLOR_TEXTURE_BRANCH
+
+configs = [
+    addict.Dict({
+        "experiment_name": "045__color_texture_final",
+        'lambda_rgb': 0.01,
+        'lambda_pca_std': 0.01,
+        'noise_dim': 8,
+        'filter_female_and_frontal': True,
+        'g_hidden_layer_num': 4,
+        'd_hidden_layer_num': 4,
+        'lambda_moment_1': 0.01,
+        'lambda_moment_2': 0.01,
+        'lambda_cls_curliness': {0: 0.1},
+        'lambda_info_curliness': 1.0,
+        'lambda_info': 1.0,
+        'curliness_dim': 1,
+        'predictor': {'curliness': 'p002', 'rgb': 'p004'},
+        'gan_input_from_encoder_prob': 0.3,
+        'curliness_with_weight': True,
+        'lambda_rec': 1000.0,
+        'lambda_rec_img': {0: 0, 600000: 1000},
+        'gen_mode': 'eigengan',
+        'lambda_orthogonal': 0.1,
+    }),
+]
+
+
+def get_config(configs, config_id):
+    for c in configs:
+        if c.experiment_name.startswith(config_id):
+            check_add_default_value_to_base_cfg(c)
+            return c
+    cfg = addict.Dict({})
+    check_add_default_value_to_base_cfg(cfg)
+    return cfg
+
+
+def check_add_default_value_to_base_cfg(cfg):
+    add_default_value_to_cfg(cfg, 'lr_d', 0.0002)
+    add_default_value_to_cfg(cfg, 'lr_g', 0.0002)
+    add_default_value_to_cfg(cfg, 'beta1', 0.5)
+    add_default_value_to_cfg(cfg, 'beta2', 0.999)
+
+    add_default_value_to_cfg(cfg, 'total_step', 650100)
+    add_default_value_to_cfg(cfg, 'log_step', 10)
+    add_default_value_to_cfg(cfg, 'sample_step', 25000)
+    add_default_value_to_cfg(cfg, 'model_save_step', 20000)
+    add_default_value_to_cfg(cfg, 'sample_batch_size', 32)
+    add_default_value_to_cfg(cfg, 'max_save', 2)
+    add_default_value_to_cfg(cfg, 'vae_var_output', 'var')
+    add_default_value_to_cfg(cfg, 'SEAN_code', 512)
+
+    # Model configuration
+    add_default_value_to_cfg(cfg, 'total_batch_size', 128)
+    add_default_value_to_cfg(cfg, 'g_hidden_layer_num', 4)
+    add_default_value_to_cfg(cfg, 'd_hidden_layer_num', 4)
+    add_default_value_to_cfg(cfg, 'd_noise_hidden_layer_num', 3)
+    add_default_value_to_cfg(cfg, 'g_hidden_dim', 256)
+    add_default_value_to_cfg(cfg, 'd_hidden_dim', 256)
+    add_default_value_to_cfg(cfg, 'gan_type', 'wgan_gp')
+    add_default_value_to_cfg(cfg, 'lambda_gp', 10.0)
+    add_default_value_to_cfg(cfg, 'lambda_adv', 1.0)
+
+    add_default_value_to_cfg(cfg, 'noise_dim', 8)
+    add_default_value_to_cfg(cfg, 'g_norm', 'none')
+    add_default_value_to_cfg(cfg, 'd_norm', 'none')
+    add_default_value_to_cfg(cfg, 'g_activ', 'relu')
+    add_default_value_to_cfg(cfg, 'd_activ', 'lrelu')
+    add_default_value_to_cfg(cfg, 'init_type', 'normal')
+    add_default_value_to_cfg(cfg, 'G_D_train_num', {'G': 1, 'D': 1}, )
+
+    output_root_dir = 'model_trained/color_texture/%s' % cfg['experiment_name']
+    add_default_value_to_cfg(cfg, 'root_dir', output_root_dir)
+    add_default_value_to_cfg(cfg, 'log_dir', output_root_dir + '/logs')
+    add_default_value_to_cfg(cfg, 'model_save_dir', output_root_dir + '/models')
+    add_default_value_to_cfg(cfg, 'sample_dir', output_root_dir + '/samples')
+    try:
+        add_default_value_to_cfg(cfg, 'gpu_num', len(args.gpu.split(',')))
+    except:
+        add_default_value_to_cfg(cfg, 'gpu_num', 1)
+
+    add_default_value_to_cfg(cfg, 'data_root', GLOBAL_DATA_ROOT)
+
+
+def add_default_value_to_cfg(cfg, key, value):
+    if key not in cfg:
+        cfg[key] = value
+
+
+def merge_config_in_place(ori_cfg, new_cfg):
+    for k in new_cfg:
+        ori_cfg[k] = new_cfg[k]
+
+
+def back_process(cfg):
+    cfg.batch_size = cfg.total_batch_size // cfg.gpu_num
+    if cfg.predictor:
+        from color_texture_branch.predictor import predictor_config
+        for ke in cfg.predictor:
+            pred_cfg = predictor_config.get_config(predictor_config.configs, cfg.predictor[ke])
+            predictor_config.back_process(pred_cfg)
+            cfg.predictor[ke] = pred_cfg
+
+    if 'gen_mode' in cfg and cfg.gen_mode is 'eigengan':
+        cfg.subspace_dim = cfg.noise_dim // cfg.g_hidden_layer_num
+
+
+def get_basic_arg_parser():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('-c', '--config', type=str, help='Specify config number', default=DEFAULT_CONFIG_COLOR_TEXTURE_BRANCH)
+    parser.add_argument('-g', '--gpu', type=str, help='Specify GPU number', default='0')
+    parser.add_argument('--local_rank', type=int, default=-1)
+    return parser
+
+
+import sys
+
+if sys.argv[0].endswith('color_texture_branch/scripts.py'):
+    parser = get_basic_arg_parser()
+    args = parser.parse_args()
+    cfg = get_config(configs, args.config)
+    os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
+    back_process(cfg)
+else:
+    cfg = get_config(configs, DEFAULT_CONFIG_COLOR_TEXTURE_BRANCH)
+    # cfg = get_config(configs, '046')
+    back_process(cfg)

models/CtrlHair/color_texture_branch/dataset.py

@@ -0,0 +1,158 @@
+# -*- coding: utf-8 -*-
+
+"""
+# File name:    dataset.py
+# Time :        2021/11/16 21:24
+# Author:       xyguoo@163.com
+# Description:  
+"""
+import os
+
+from common_dataset import DataFilter
+
+import random
+import numpy as np
+from global_value_utils import HAIR_IDX, GLOBAL_DATA_ROOT
+import pickle
+import torch
+
+
+class Dataset:
+    """A general image-attributes dataset class."""
+
+    def valid_hair(self, item):
+        if np.isnan(self.rgb_stat_dict[item][0]).any():  # not have hair
+            return False
+        if (self.sean_code_dict[item][HAIR_IDX] == 0).all():
+            return False
+        if item not in self.color_var_stat_dict:
+            return False
+        return True
+
+    def __init__(self, cfg, rank=0, test_part=0.096):
+        # super().__init__()
+        self.cfg = cfg
+
+        self.random_seed = 7  # Do not change the random seed, which determines the split of scripts and test set
+        self.hair_root = GLOBAL_DATA_ROOT
+        with open(os.path.join(self.hair_root, 'sean_code_dict.pkl'), 'rb') as f:
+            self.sean_code_dict = pickle.load(f)
+        with open(os.path.join(self.hair_root, 'rgb_stat_dict.pkl'), 'rb') as f:
+            self.rgb_stat_dict = pickle.load(f)
+        with open(os.path.join(self.hair_root, 'color_var_stat_dict.pkl'), 'rb') as f:
+            self.color_var_stat_dict = pickle.load(f)
+
+        self.local_rank = rank
+        random.seed(self.random_seed + self.local_rank + 1)
+
+        self.data_list = [dd for dd in list(self.sean_code_dict) if self.valid_hair(dd)]
+        random.shuffle(self.data_list)
+
+        self.data_filter = DataFilter(cfg, test_part)
+
+        self.test_list = []
+        for ll in self.data_filter.test_list:
+            path_part = ll.split('/')
+            self.test_list.append('%s___%s' % (path_part[-3], path_part[-1][:-4]))
+
+        self.train_filter = []
+        for ll in self.data_filter.train_list:
+            path_part = ll.split('/')
+            self.train_filter.append('%s___%s' % (path_part[-3], path_part[-1][:-4]))
+        self.train_filter = set(self.train_filter)
+        self.train_list = [ll for ll in self.data_list if ll not in self.test_list]
+        if cfg.filter_female_and_frontal:
+            self.train_list = [ll for ll in self.train_list if ll in self.train_filter]
+        self.train_set = set(self.train_list)
+
+        ### curliness code
+        self.curliness_hair_list = {}
+        self.curliness_hair_list_test = {}
+        self.curliness_hair_dict = {ke: 0 for ke in self.color_var_stat_dict}
+
+        for label in [-1, 1]:
+            img_file = os.path.join(cfg.data_root, 'manual_label', 'curliness', '%d.txt' % label)
+            with open(img_file, 'r') as f:
+                imgs = [l.strip() for l in f.readlines()]
+            imgs = [ii for ii in imgs if ii in self.train_set]
+            self.curliness_hair_list[label] = imgs
+            for ii in imgs:
+                self.curliness_hair_dict[ii] = label
+
+            img_file = os.path.join(cfg.data_root, 'manual_label', 'curliness', 'test_%d.txt' % label)
+            with open(img_file, 'r') as f:
+                imgs = [l.strip() for l in f.readlines()]
+            self.curliness_hair_list_test[label] = imgs
+            for ii in imgs:
+                self.curliness_hair_dict[ii] = label
+
+    def get_sean_code(self, ke):
+        return self.sean_code_dict[ke]
+
+    def get_list_by_items(self, items):
+        res_code, res_rgb_mean, res_pca_std, res_sean_code, res_curliness = [], [], [], [], []
+        for item in items:
+            code = self.sean_code_dict[item][HAIR_IDX]
+            res_code.append(code)
+            rgb_mean = self.rgb_stat_dict[item][0]
+            res_rgb_mean.append(rgb_mean)
+            pca_std = self.color_var_stat_dict[item]['var_pca']
+            # here the 'var' is 'std'
+            res_pca_std.append(pca_std[..., None])
+            res_sean_code.append(self.get_sean_code(ke=item))
+            res_curliness.append(self.curliness_hair_dict[item])
+        res_code = torch.tensor(np.stack(res_code), dtype=torch.float32)
+        res_rgb_mean = torch.tensor(np.stack(res_rgb_mean), dtype=torch.float32)
+        res_pca_std = torch.tensor(np.stack(res_pca_std), dtype=torch.float32)
+        res_curliness = torch.tensor(np.stack(res_curliness), dtype=torch.int)[..., None]
+        data = {'code': res_code, 'rgb_mean': res_rgb_mean, 'pca_std': res_pca_std, 'items': items,
+                'sean_code': res_sean_code, 'curliness_label': res_curliness}
+        return data
+
+    def get_training_batch(self, batch_size):
+        items = []
+        while len(items) < batch_size:
+            item = random.choice(self.train_list)
+            # if not self.valid_hair(item):
+            #     continue
+            items.append(item)
+        data = self.get_list_by_items(items)
+        return data
+
+    def get_testing_batch(self, batch_size):
+        ptr = 0
+        items = []
+        while len(items) < batch_size:
+            item = self.test_list[ptr]
+            ptr += 1
+            if not self.valid_hair(item):
+                continue
+            items.append(item)
+        data = self.get_list_by_items(items)
+        return data
+
+    def get_curliness_hair(self, labels):
+        labels = labels.cpu().numpy()
+        items = []
+        for label in labels:
+            item_list = self.curliness_hair_list[label[0]]
+            items.append(np.random.choice(item_list))
+        data = self.get_list_by_items(items)
+        return data
+
+    def get_curliness_hair_test(self):
+        return self.get_list_by_items(self.curliness_hair_list_test[-1] + self.curliness_hair_list_test[1])
+
+
+# if __name__ == '__main__':
+#     ds = Dataset(cfg)
+#     resources = ds.get_training_batch(8)
+#     pass
+
+
+# for label in [-1, 1]:
+#     img_dir = os.path.join(cfg.data_root, 'manual_label', 'curliness', 'test_%d' % label)
+#     imgs = [pat[:-4] + '\n' for pat in os.listdir(img_dir)]
+#     imgs.sort()
+#     with open(img_dir + '.txt', 'w') as f:
+#         f.writelines(imgs)

models/CtrlHair/color_texture_branch/model.py

@@ -0,0 +1,159 @@
+# -*- coding: utf-8 -*-
+
+"""
+# File name:    model.py
+# Time :        2021/11/17 15:37
+# Author:       xyguoo@163.com
+# Description:  
+"""
+
+import torch.nn as nn
+from my_torchlib.module import LinearBlock
+import torch
+from torch.nn import init
+
+
+def init_weights(net, init_type='normal', init_gain=0.02):
+    """Initialize network weights.
+
+    Parameters:
+        net (network)   -- network to be initialized
+        init_type (str) -- the name of an initialization method: normal | xavier | kaiming | orthogonal
+        init_gain (float)    -- scaling factor for normal, xavier and orthogonal.
+
+    We use 'normal' in the original pix2pix and CycleGAN paper. But xavier and kaiming might
+    work better for some applications. Feel free to try yourself.
+    """
+
+    def init_func(m):  # define the initialization function
+        classname = m.__class__.__name__
+        if hasattr(m, 'weight') and (classname.find('Conv') != -1 or classname.find('Linear') != -1):
+            if init_type == 'normal':
+                init.normal_(m.weight.data, 0.0, init_gain)
+            elif init_type == 'xavier':
+                init.xavier_normal_(m.weight.data, gain=init_gain)
+            elif init_type == 'kaiming':
+                init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')
+            elif init_type == 'orthogonal':
+                init.orthogonal_(m.weight.data, gain=init_gain)
+            else:
+                raise NotImplementedError('initialization method [%s] is not implemented' % init_type)
+            if hasattr(m, 'bias') and m.bias is not None:
+                init.constant_(m.bias.data, 0.0)
+        elif classname.find(
+                'BatchNorm2d') != -1:  # BatchNorm Layer's weight is not a matrix; only normal distribution applies.
+            init.normal_(m.weight.data, 1.0, init_gain)
+            init.constant_(m.bias.data, 0.0)
+
+    print('initialize network with %s' % init_type)
+    net.apply(init_func)  # apply the initialization function <init_func>
+
+
+class Generator(nn.Module):
+    """Generator network."""
+
+    def __init__(self, cfg):
+        super(Generator, self).__init__()
+        self.cfg = cfg
+
+        input_dim = cfg.noise_dim
+        if cfg.lambda_rgb:
+            input_dim += 3
+        if cfg.lambda_pca_std:
+            input_dim += 1
+        if cfg.lambda_cls_curliness:
+            input_dim += cfg.curliness_dim
+
+        layers = [LinearBlock(input_dim, cfg.g_hidden_dim, cfg.g_norm, activation=cfg.g_activ)]
+        for _ in range(cfg.g_hidden_layer_num - 1):
+            layers.append(LinearBlock(cfg.g_hidden_dim, cfg.g_hidden_dim, cfg.g_norm, activation=cfg.g_activ))
+        layers.append(LinearBlock(cfg.g_hidden_dim, cfg.SEAN_code, 'none', 'none'))
+        self.net = nn.Sequential(*layers)
+
+    def forward(self, data):
+        x = data['noise']
+        if self.cfg.lambda_cls_curliness:
+            x = torch.cat([x, data['noise_curliness']], dim=1)
+        if self.cfg.lambda_rgb:
+            x = torch.cat([x, data['rgb_mean']], dim=1)
+        if self.cfg.lambda_pca_std:
+            x = torch.cat([x, data['pca_std']], dim=1)
+        output = self.net(x)
+        res = {'code': output}
+        return res
+
+
+class Discriminator(nn.Module):
+    """Discriminator network."""
+
+    def __init__(self, cfg):
+        super(Discriminator, self).__init__()
+        self.cfg = cfg
+        layers = [LinearBlock(cfg.SEAN_code, cfg.d_hidden_dim, cfg.d_norm, activation=cfg.d_activ)]
+        for _ in range(cfg.d_hidden_layer_num - 1):
+            layers.append(LinearBlock(cfg.d_hidden_dim, cfg.d_hidden_dim, cfg.d_norm, activation=cfg.d_activ))
+        output_dim = 1 + cfg.noise_dim
+        if cfg.lambda_rgb and 'curliness' not in cfg.predictor:
+            output_dim += 3
+        if cfg.lambda_pca_std:
+            output_dim += 1
+        if cfg.lambda_cls_curliness:
+            output_dim += cfg.curliness_dim
+            if 'curliness' not in cfg.predictor:
+                output_dim += 1
+        layers.append(LinearBlock(cfg.d_hidden_dim, output_dim, 'none', 'none'))
+        self.net = nn.Sequential(*layers)
+        self.input_name = 'code'
+
+    def forward(self, data_in):
+        x = data_in[self.input_name]
+        out = self.net(x)
+        data = {'adv': out[:, [0]]}
+        ptr = 1
+        data['noise'] = out[:, ptr:(ptr + self.cfg.noise_dim)]
+        ptr += self.cfg.noise_dim
+        if self.cfg.lambda_cls_curliness:
+            data['noise_curliness'] = out[:, ptr:(ptr + self.cfg.curliness_dim)]
+            ptr += self.cfg.curliness_dim
+            if not 'curliness' in self.cfg.predictor:
+                data['cls_curliness'] = out[:, ptr: ptr + 1]
+                ptr += 1
+        if self.cfg.lambda_rgb and 'rgb' not in self.cfg.predictor:
+            data['rgb_mean'] = out[:, ptr:ptr + 3]
+            ptr += 3
+        if self.cfg.lambda_pca_std and 'rgb' not in self.cfg.predictor:
+            data['pca_std'] = out[:, ptr:]
+            ptr += 1
+        return data
+
+    def forward_adv_direct(self, x):
+        return self.net(x)[:, [0]]
+
+
+class DiscriminatorNoise(nn.Module):
+    """Discriminator network."""
+
+    def __init__(self, cfg):
+        super(DiscriminatorNoise, self).__init__()
+        self.cfg = cfg
+        input_dim = cfg.noise_dim
+        if cfg.lambda_cls_curliness:
+            input_dim += cfg.curliness_dim
+        layers = [LinearBlock(input_dim, cfg.d_hidden_dim, cfg.d_norm, activation=cfg.d_activ)]
+        for _ in range(cfg.d_noise_hidden_layer_num - 1):
+            layers.append(LinearBlock(cfg.d_hidden_dim, cfg.d_hidden_dim, cfg.d_norm, activation=cfg.d_activ))
+        output_dim = 1
+        layers.append(LinearBlock(cfg.d_hidden_dim, output_dim, 'none', 'none'))
+        self.net = nn.Sequential(*layers)
+        self.input_name = 'noise'
+
+    def forward(self, data_in):
+        x = data_in[self.input_name]
+        if self.cfg.lambda_cls_curliness:
+            x = torch.cat([x, data_in['noise_curliness']], dim=1)
+        out = self.net(x)
+        data = {'adv': out[:, [0]]}
+        return data
+
+    def forward_adv_direct(self, x):
+        return self.net(x)[:, [0]]

models/CtrlHair/color_texture_branch/model_eigengan.py

@@ -0,0 +1,89 @@
+# -*- coding: utf-8 -*-
+
+"""
+# File name:    model_eigengan.py
+# Time :        2021/12/28 21:57
+# Author:       xyguoo@163.com
+# Description:  
+"""
+
+import torch.nn as nn
+import torch
+
+
+class SubspaceLayer(nn.Module):
+    def __init__(self, dim: int, n_basis: int):
+        super().__init__()
+        self.U = nn.Parameter(torch.empty(n_basis, dim), requires_grad=True)
+        nn.init.orthogonal_(self.U)
+        self.L = nn.Parameter(torch.FloatTensor([3 * i for i in range(n_basis, 0, -1)]), requires_grad=True)
+        self.mu = nn.Parameter(torch.zeros(dim), requires_grad=True)
+
+        self.unit_matrix = torch.eye(n_basis, requires_grad=False)
+
+    def forward(self, z):
+        return (self.L * z) @ self.U + self.mu
+
+    def orthogonal_regularizer(self):
+        UUT = self.U @ self.U.t()
+        self.unit_matrix = self.unit_matrix.to(self.U.device)
+        reg = ((UUT - self.unit_matrix) ** 2).mean()
+        return reg
+
+
+class EigenGenerator(nn.Module):
+    def __init__(self, cfg):
+        super(EigenGenerator, self).__init__()
+        self.cfg = cfg
+
+        input_dim = 0
+        if cfg.lambda_rgb:
+            input_dim += 3
+        if cfg.lambda_pca_std:
+            input_dim += 1
+        if cfg.lambda_cls_curliness:
+            input_dim += cfg.curliness_dim
+
+        self.main_layer_in = nn.Linear(input_dim, cfg.g_hidden_dim, bias=True)
+        main_layers_mid = []
+        for _ in range(cfg.g_hidden_layer_num - 1):
+            main_layers_mid.append(nn.Sequential(nn.LeakyReLU(0.2), nn.Linear(cfg.g_hidden_dim,
+                                                                              cfg.g_hidden_dim, bias=True)))
+        main_layers_mid.append(nn.Sequential(nn.LeakyReLU(0.2), nn.Linear(cfg.g_hidden_dim, cfg.SEAN_code)))
+
+        subspaces = []
+        for _ in range(cfg.g_hidden_layer_num):
+            sub = SubspaceLayer(cfg.g_hidden_dim, cfg.subspace_dim)
+            subspaces.append(sub)
+
+        self.main_layer_mid = nn.ModuleList(main_layers_mid)
+        self.subspaces = nn.ModuleList(subspaces)
+
+    def forward(self, data):
+        noise = data['noise']
+        noise = noise.reshape(len(data['noise']), self.cfg.g_hidden_layer_num, self.cfg.subspace_dim)
+
+        input_data = []
+        if self.cfg.lambda_cls_curliness:
+            input_data.append(data['noise_curliness'])
+        if self.cfg.lambda_rgb:
+            input_data.append(data['rgb_mean'])
+        if self.cfg.lambda_pca_std:
+            input_data.append(data['pca_std'])
+
+        x = torch.cat(input_data, dim=1)
+        x_mid = self.main_layer_in(x)
+
+        for layer_idx in range(self.cfg.g_hidden_layer_num):
+            subspace_data = self.subspaces[layer_idx](noise[:, layer_idx, :])
+            x_mid = x_mid + subspace_data
+            x_mid = self.main_layer_mid[layer_idx](x_mid)
+
+        res = {'code': x_mid}
+        return res
+
+    def orthogonal_regularizer_loss(self):
+        loss = 0
+        for s in self.subspaces:
+            loss = loss + s.orthogonal_regularizer()
+        return loss

models/CtrlHair/color_texture_branch/module.py

@@ -0,0 +1,61 @@
+# -*- coding: utf-8 -*-
+
+"""
+# File name:    module.py
+# Time :        2021/11/17 15:38
+# Author:       xyguoo@163.com
+# Description:  
+"""
+
+import torch.nn as nn
+
+
+class LinearBlock(nn.Module):
+
+    def __init__(self, input_dim, output_dim, norm, activation='relu', use_bias=True, leaky_slope=0.2, dropout=0):
+        super(LinearBlock, self).__init__()
+        # initialize fully connected layer
+        self.fc = nn.Linear(input_dim, output_dim, bias=use_bias)
+
+        # initialize normalization
+        norm_dim = output_dim
+        if norm == 'bn':
+            self.norm = nn.BatchNorm1d(norm_dim)
+        elif norm == 'in':
+            self.norm = nn.InstanceNorm1d(norm_dim)
+        elif norm == 'ln':
+            self.norm = nn.LayerNorm(norm_dim)
+        elif norm == 'none':
+            self.norm = None
+        else:
+            assert 0, "Unsupported normalization: {}".format(norm)
+
+        # initialize activation
+        if activation == 'relu':
+            self.activation = nn.ReLU(inplace=True)
+        elif activation == 'lrelu':
+            self.activation = nn.LeakyReLU(leaky_slope, inplace=True)
+        elif activation == 'prelu':
+            self.activation = nn.PReLU()
+        elif activation == 'selu':
+            self.activation = nn.SELU(inplace=True)
+        elif activation == 'tanh':
+            self.activation = nn.Tanh()
+        elif activation == 'none':
+            self.activation = None
+        else:
+            assert 0, "Unsupported activation: {}".format(activation)
+
+        self.dropout = dropout
+        if bool(self.dropout) and self.dropout > 0:
+            self.dropout_layer = nn.Dropout(p=self.dropout)
+
+    def forward(self, x):
+        out = self.fc(x)
+        if self.norm:
+            out = self.norm(out)
+        if self.activation:
+            out = self.activation(out)
+        if bool(self.dropout) and self.dropout > 0:
+            out = self.dropout_layer(out)
+        return out

models/CtrlHair/color_texture_branch/predictor/__init__.py

@@ -0,0 +1,8 @@
+# -*- coding: utf-8 -*-
+
+"""
+# File name:    __init__.py.py
+# Time :        2021/12/14 22:43
+# Author:       xyguoo@163.com
+# Description:  
+"""

models/CtrlHair/color_texture_branch/predictor/predictor_config.py

@@ -0,0 +1,119 @@
+# -*- coding: utf-8 -*-
+
+"""
+# File name:    predictor_config.py
+# Time :        2021/12/14 21:03
+# Author:       xyguoo@163.com
+"""
+
+import addict  # nesting dict
+import os
+import argparse
+import sys
+
+from global_value_utils import GLOBAL_DATA_ROOT
+
+configs = [
+    addict.Dict({
+        "experiment_name": "p002___curliness",
+        'basic_dir': 'model_trained/curliness_classifier',
+        'filter_female_and_frontal': True,
+        'hidden_layer_num': 3,
+        'hidden_dim': 32,
+        'lambda_cls_curliness': {0: 1, 200: 0.1, 400: 0.01, 2500: 0.001},
+        'init_type': 'xavier',
+        'norm': 'bn',
+        'dropout': 0.5,
+        'total_batch_size': 256,
+        'total_step': 7000,
+    }),
+    addict.Dict({
+        "experiment_name": "p004___pca_std",
+        'basic_dir': 'model_trained/color_encoder',
+        'filter_female_and_frontal': True,
+        'hidden_layer_num': 3,
+        'hidden_dim': 256,
+        'lambda_rgb': {0: 1, 7000: 1},
+        'lambda_pca_std': {0: 1, 7000: 1},
+        'init_type': 'xavier',
+        'norm': 'bn',
+        'dropout': 0.2,
+        'total_batch_size': 256,
+        'total_step': 10000,
+    }),
+]
+
+
+def get_config(configs, config_id):
+    for c in configs:
+        if c.experiment_name.startswith(config_id):
+            check_add_default_value_to_base_cfg(c)
+            return c
+
+
+def check_add_default_value_to_base_cfg(cfg):
+    add_default_value_to_cfg(cfg, 'lr', 0.002)
+    add_default_value_to_cfg(cfg, 'beta1', 0.5)
+    add_default_value_to_cfg(cfg, 'beta2', 0.999)
+
+    add_default_value_to_cfg(cfg, 'log_step', 10)
+    add_default_value_to_cfg(cfg, 'model_save_step', 1000)
+    add_default_value_to_cfg(cfg, 'sample_batch_size', 100)
+    add_default_value_to_cfg(cfg, 'max_save', 2)
+    add_default_value_to_cfg(cfg, 'SEAN_code', 512)
+
+    # Model configuration
+    add_default_value_to_cfg(cfg, 'total_batch_size', 64)
+    add_default_value_to_cfg(cfg, 'gan_type', 'wgan_gp')
+
+    add_default_value_to_cfg(cfg, 'norm', 'none')
+    add_default_value_to_cfg(cfg, 'activ', 'lrelu')
+    add_default_value_to_cfg(cfg, 'init_type', 'normal')
+
+    add_default_value_to_cfg(cfg, 'root_dir', '%s/%s' % (cfg.basic_dir, cfg['experiment_name']))
+    add_default_value_to_cfg(cfg, 'log_dir', cfg.root_dir + '/logs')
+    add_default_value_to_cfg(cfg, 'model_save_dir', cfg.root_dir + '/models')
+    add_default_value_to_cfg(cfg, 'sample_dir', cfg.root_dir + '/samples')
+    try:
+        add_default_value_to_cfg(cfg, 'gpu_num', len(args.gpu.split(',')))
+    except:
+        add_default_value_to_cfg(cfg, 'gpu_num', 1)
+
+    add_default_value_to_cfg(cfg, 'data_root', GLOBAL_DATA_ROOT)
+
+
+def add_default_value_to_cfg(cfg, key, value):
+    if key not in cfg:
+        cfg[key] = value
+
+
+def merge_config_in_place(ori_cfg, new_cfg):
+    for k in new_cfg:
+        ori_cfg[k] = new_cfg[k]
+
+
+def back_process(cfg):
+    cfg.batch_size = cfg.total_batch_size // cfg.gpu_num
+    cfg.predict_dict = {}
+    if 'lambda_cls_curliness' in cfg:
+        cfg.predict_dict['cls_curliness'] = 1
+    if 'lambda_rgb' in cfg:
+        cfg.predict_dict['rgb_mean'] = 3
+    if 'lambda_pca_std' in cfg:
+        cfg.predict_dict['pca_std'] = 1
+
+
+def get_basic_arg_parser():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('-c', '--config', type=str, help='Specify config number', default='000')
+    parser.add_argument('-g', '--gpu', type=str, help='Specify GPU number', default='0')
+    parser.add_argument('--local_rank', type=int, default=-1)
+    return parser
+
+
+if sys.argv[0].endswith('predictor_train.py'):
+    parser = get_basic_arg_parser()
+    args = parser.parse_args()
+    cfg = get_config(configs, args.config)
+    os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
+    back_process(cfg)

models/CtrlHair/color_texture_branch/predictor/predictor_model.py

@@ -0,0 +1,41 @@
+# -*- coding: utf-8 -*-
+
+"""
+# File name:    predictor_model.py
+# Time :        2021/12/14 20:47
+# Author:       xyguoo@163.com
+# Description:  
+"""
+
+import torch.nn as nn
+from my_torchlib.module import LinearBlock
+
+
+class Predictor(nn.Module):
+    """Discriminator network."""
+
+    def __init__(self, cfg):
+        super(Predictor, self).__init__()
+        self.cfg = cfg
+        layers = [LinearBlock(cfg.SEAN_code, cfg.hidden_dim, cfg.norm, activation=cfg.activ, dropout=cfg.dropout)]
+        for _ in range(cfg.hidden_layer_num - 1):
+            layers.append(LinearBlock(cfg.hidden_dim, cfg.hidden_dim, cfg.norm, activation=cfg.activ,
+                                      dropout=cfg.dropout))
+
+        output_dim = 0
+        for ke in cfg.predict_dict:
+            output_dim += cfg.predict_dict[ke]
+        layers.append(LinearBlock(cfg.hidden_dim, output_dim, 'none', 'none'))
+        self.net = nn.Sequential(*layers)
+        self.input_name = 'code'
+
+    def forward(self, data_in):
+        x = data_in[self.input_name]
+        out = self.net(x)
+        ptr = 0
+        data = {}
+        for ke in self.cfg.predict_dict:
+            cur_dim = self.cfg.predict_dict[ke]
+            data[ke] = out[:, ptr:(ptr + cur_dim)]
+            ptr += cur_dim
+        return data

models/CtrlHair/color_texture_branch/predictor/predictor_solver.py

@@ -0,0 +1,51 @@
+# -*- coding: utf-8 -*-
+
+"""
+# File name:    predictor_solver.py
+# Time :        2021/12/14 21:57
+# Author:       xyguoo@163.com
+# Description:  
+"""
+
+import torch
+import torch.nn.functional as F
+from color_texture_branch.predictor.predictor_config import cfg
+from color_texture_branch.predictor.predictor_model import Predictor
+from torch.nn.parallel import DistributedDataParallel as DDP
+
+
+class PredictorSolver:
+
+    def __init__(self, cfg, device, local_rank, training=True):
+        self.mse_loss = torch.nn.MSELoss()
+        self.cfg = cfg
+
+        # model
+        self.pred = Predictor(cfg)
+        self.pred.to(device)
+        self.optimizer = torch.optim.Adam(self.pred.parameters(), lr=cfg.lr, betas=(cfg.beta1, cfg.beta2),
+                                          weight_decay=0.00)
+
+        if local_rank >= 0:
+            pDDP = lambda m, find_unused: DDP(m, device_ids=[local_rank], output_device=local_rank,
+                                              find_unused_parameters=False)
+            self.pred = pDDP(self.pred, find_unused=True)
+        self.local_rank = local_rank
+        self.device = device
+
+    def forward(self, data):
+        self.data = data
+        self.pred_res = self.pred(data)
+
+    def forward_d(self, loss_dict):
+        if 'lambda_rgb' in cfg:
+            loss_dict['lambda_rgb'] = self.mse_loss(self.pred_res['rgb_mean'], self.data['rgb_mean'])
+        if 'lambda_pca_std' in cfg:
+            d_pca_std = self.pred_res['pca_std']
+            loss_dict['lambda_pca_std'] = self.mse_loss(d_pca_std, self.data['pca_std'])
+
+    def forward_d_curliness(self, data_curliness, loss_dict):
+        if cfg.lambda_cls_curliness:
+            cls_curliness = self.pred_res['cls_curliness']
+            loss_dict['lambda_cls_curliness'] = F.binary_cross_entropy(
+                torch.sigmoid(cls_curliness), data_curliness['curliness_label'].float() / 2 + 0.5)

models/CtrlHair/color_texture_branch/predictor/predictor_train.py

@@ -0,0 +1,159 @@
+# -*- coding: utf-8 -*-
+
+"""
+# File name:    predictor_train.py
+# Time :        2021/12/14 20:58
+# Author:       xyguoo@163.com
+# Description:  
+"""
+
+import sys
+
+sys.path.append('.')
+
+import tensorboardX
+import torch
+import tqdm
+import numpy as np
+from color_texture_branch.predictor.predictor_config import cfg, args
+from color_texture_branch.dataset import Dataset
+import my_pylib
+# distributed training
+import torch.distributed as dist
+from color_texture_branch.predictor.predictor_solver import PredictorSolver
+from my_torchlib.train_utils import LossUpdater, to_device, train
+import my_torchlib
+from color_texture_branch.model import init_weights
+
+
+def get_total_step():
+    total = 0
+    for key in cfg.iter:
+        total += cfg.iter[key]
+    return total
+
+
+def worker(proc, nprocs, args):
+    local_rank = args.local_rank
+    if local_rank >= 0:
+        torch.cuda.set_device(local_rank)
+        dist.init_process_group(backend='nccl',
+                                init_method='tcp://localhost:%d' % (6030 + int(cfg.experiment_name[:3])),
+                                rank=args.local_rank,
+                                world_size=cfg.gpu_num)
+        print('setup rank %d' % local_rank)
+    device = torch.device('cuda', max(0, local_rank))
+
+    # config
+    out_dir = cfg.root_dir
+
+    # data
+    ds = Dataset(cfg)
+
+    # Loss class
+    solver = PredictorSolver(cfg, device, local_rank, training=True)
+
+    loss_updater = LossUpdater(cfg)
+
+    # load checkpoint
+    ckpt_dir = out_dir + '/checkpoints'
+    if local_rank <= 0:
+        my_pylib.mkdir(out_dir)
+        my_pylib.save_json(out_dir + '/setting_hair.json', cfg, indent=4, separators=(',', ': '))
+        my_pylib.mkdir(ckpt_dir)
+
+    try:
+        ckpt = my_torchlib.load_checkpoint(ckpt_dir)
+        start_step = ckpt['step'] + 1
+        for model_name in ['Predictor']:
+            cur_model = ckpt[model_name]
+            if list(cur_model)[0].startswith('module'):
+                ckpt[model_name] = {kk[7:]: cur_model[kk] for kk in cur_model}
+        solver.pred.load_state_dict(ckpt['Predictor'], strict=True)
+        solver.optimizer.load_state_dict(ckpt['optimizer'])
+        print('Load succeed!')
+    except:
+        print(' [*] No checkpoint!')
+        init_weights(solver.pred, init_type=cfg.init_type)
+        start_step = 1
+
+    # writer
+    if local_rank <= 0:
+        writer = tensorboardX.SummaryWriter(out_dir + '/summaries')
+    else:
+        writer = None
+
+    # start training
+    test_batch = ds.get_testing_batch(100)
+    test_batch_curliness = ds.get_curliness_hair_test()
+
+    to_device(test_batch_curliness, device)
+    to_device(test_batch, device)
+
+    if local_rank >= 0:
+        dist.barrier()
+
+    total_step = cfg.total_step + 2
+    for step in tqdm.tqdm(range(start_step, total_step), total=total_step, initial=start_step, desc='step'):
+        loss_updater.update(step)
+        write_log = (writer and step % 11 == 0)
+
+        loss_dict = {}
+        if 'rgb_mean' in cfg.predict_dict or 'pca_std' in cfg.predict_dict:
+            data = ds.get_training_batch(cfg.batch_size)
+            to_device(data, device)
+            solver.forward(data)
+            solver.forward_d(loss_dict)
+            if write_log:
+                solver.pred.eval()
+                if 'rgb_mean' in cfg.predict_dict:
+                    loss_dict['test_lambda_rgb'] = solver.mse_loss(
+                        solver.pred(test_batch)['rgb_mean'], test_batch['rgb_mean'])
+                    print('rgb loss: %f' % loss_dict['test_lambda_rgb'])
+                if 'pca_std' in cfg.predict_dict:
+                    loss_dict['test_lambda_pca_std'] = solver.mse_loss(
+                        solver.pred(test_batch)['pca_std'], test_batch['pca_std'])
+                    print('pca_std loss: %f' % loss_dict['test_lambda_pca_std'])
+                solver.pred.train()
+
+        if cfg.lambda_cls_curliness:
+            data = {}
+            curliness_label = torch.tensor(np.random.choice([-1, 1], (cfg.batch_size, 1)))
+            data['curliness_label'] = curliness_label
+            data_curliness = ds.get_curliness_hair(curliness_label)
+            to_device(data_curliness, device)
+            solver.forward(data_curliness)
+            solver.forward_d_curliness(data_curliness, loss_dict)
+
+            # validation to show whether over-fit
+            if write_log:
+                solver.pred.eval()
+                logit = solver.pred(test_batch_curliness)['cls_curliness']
+                loss_dict['test_cls_curliness'] = torch.nn.functional.binary_cross_entropy_with_logits(
+                    logit, test_batch_curliness['curliness_label'] / 2 + 0.5)
+
+                print('cls_curliness: %f' % loss_dict['test_cls_curliness'])
+                print('acc %f' % ((logit * test_batch_curliness['curliness_label'] > 0).sum() / logit.shape[0]))
+                solver.pred.train()
+
+        train(cfg, loss_dict, optimizers=[solver.optimizer],
+              step=step, writer=writer, flag='Pred', write_log=write_log)
+
+        if step > 0 and step % cfg.model_save_step == 0:
+            if local_rank <= 0:
+                save_model(step, solver, ckpt_dir)
+            if local_rank >= 0:
+                dist.barrier()
+
+
+def save_model(step, solver, ckpt_dir):
+    save_dic = {'step': step,
+                'Predictor': solver.pred.state_dict(),
+                'optimizer': solver.optimizer.state_dict()}
+    my_torchlib.save_checkpoint(save_dic, '%s/%07d.ckpt' % (ckpt_dir, step), max_keep=cfg.max_save)
+
+
+if __name__ == '__main__':
+    # with torch.autograd.set_detect_anomaly(True):
+    # mp.spawn(worker, nprocs=cfg.gpu_num, args=(cfg.gpu_num, args))
+    worker(proc=None, nprocs=None, args=args)

models/CtrlHair/color_texture_branch/script_find_direction.py

@@ -0,0 +1,74 @@
+# -*- coding: utf-8 -*-
+
+"""
+# File name:    script_find_direction.py
+# Time :        2022/02/28
+# Author:       xyguoo@163.com
+# Description:  
+"""
+
+import sys
+
+sys.path.append('.')
+
+import os
+import tqdm
+
+from ui.backend import Backend
+from util.canvas_grid import Canvas
+import numpy as np
+
+import pickle
+from common_dataset import DataFilter
+from util.imutil import read_rgb, write_rgb
+from color_texture_branch.config import cfg
+from util.find_semantic_direction import get_random_direction
+
+df = DataFilter(cfg)
+be = Backend(2.5, blending=False)
+
+exist_direction = 'model_trained/color_texture/%s' % cfg.experiment_name
+code_dim = cfg.noise_dim
+att_name = 'texture'
+interpolate_num = 6
+max_val = 2.5
+batch = 10
+
+interpolate_values = np.linspace(-max_val, max_val, interpolate_num)
+
+existing_dirs_dir = os.path.join(exist_direction, '%s_dir_used' % att_name)
+
+existing_dirs_list = os.listdir(existing_dirs_dir)
+existing_dirs = []
+for dd in existing_dirs_list:
+    with open(os.path.join(existing_dirs_dir, dd), 'rb') as f:
+        existing_dirs.append(pickle.load(f))
+
+direction_dir = '%s/direction_find/%s_dir_%d' % (exist_direction, att_name, len(existing_dirs) + 1)
+img_gen_dir = '%s/direction_find/%s_%d' % (exist_direction, att_name, len(existing_dirs) + 1)
+for dd in [direction_dir, img_gen_dir]:
+    if not os.path.exists(dd):
+        os.makedirs(dd)
+
+img_list = df.train_list
+
+for dir_idx in tqdm.tqdm(range(0, 300)):
+    rand_dir = get_random_direction(code_dim, existing_dirs)
+    with open('%s/%d.pkl' % (direction_dir, dir_idx,), 'wb') as f:
+        pickle.dump(rand_dir, f)
+    rand_dir = rand_dir.to(be.device)
+
+    canvas = Canvas(batch, interpolate_num + 1)
+    for img_idx, img_file in tqdm.tqdm(enumerate(img_list[:batch])):
+        img = read_rgb(img_file)
+        _, img_parsing = be.set_input_img(img)
+
+        canvas.process_draw_image(img, img_idx, 0)
+
+        for inter_idx in range(interpolate_num):
+            inter_val = interpolate_values[inter_idx]
+            be.continue_change_with_direction(att_name, rand_dir, inter_val)
+
+            out_img = be.output()
+            canvas.process_draw_image(out_img, img_idx, inter_idx + 1)
+    write_rgb('%s/%d.png' % (img_gen_dir, dir_idx), canvas.canvas)

models/CtrlHair/color_texture_branch/solver.py

@@ -0,0 +1,299 @@
+# -*- coding: utf-8 -*-
+
+"""
+# File name:    solver.py
+# Time :        2021/11/17 16:24
+# Author:       xyguoo@163.com
+# Description:  
+"""
+import torch
+import torch.nn.functional as F
+import numpy as np
+from .config import cfg
+from .model import Discriminator, DiscriminatorNoise
+from torch.nn.parallel import DistributedDataParallel as DDP
+import random
+import os
+import cv2
+
+
+# solver
+class Solver:
+
+    def __init__(self, cfg, device, local_rank, training=True):
+        self.mse_loss = torch.nn.MSELoss()
+        self.cfg = cfg
+
+        # model
+        if 'gen_mode' in cfg and cfg.gen_mode is 'eigengan':
+            from color_texture_branch.model_eigengan import EigenGenerator
+            self.gen = EigenGenerator(cfg)
+        else:
+            from color_texture_branch.model import Generator
+            self.gen = Generator(cfg)
+        self.gen.to(device)
+
+        self.dis = Discriminator(cfg)
+        self.dis.to(device)
+
+        if 'curliness' in cfg.predictor:
+            from color_texture_branch.predictor.predictor_model import Predictor
+            self.curliness_model = Predictor(cfg.predictor['curliness'])
+            self.curliness_model.to(device)
+            self.curliness_model.eval()
+
+        if 'rgb' in cfg.predictor:
+            from color_texture_branch.predictor.predictor_model import Predictor
+            self.rgb_model = Predictor(cfg.predictor['rgb'])
+            self.rgb_model.to(device)
+            self.rgb_model.eval()
+
+        if training:
+            self.G_optimizer = torch.optim.Adam(self.gen.parameters(), lr=cfg.lr_d, betas=(cfg.beta1, cfg.beta2),
+                                                weight_decay=0.00)
+            self.D_optimizer = torch.optim.Adam(self.dis.parameters(), lr=cfg.lr_g, betas=(cfg.beta1, cfg.beta2),
+                                                weight_decay=0.00)
+
+            if cfg.lambda_adv_noise:
+                self.dis_noise = DiscriminatorNoise(cfg)
+                self.dis_noise.to(device)
+                self.D_noise_optimizer = torch.optim.Adam(self.dis_noise.parameters(), lr=cfg.lr_g,
+                                                          betas=(cfg.beta1, cfg.beta2), weight_decay=0.00)
+            else:
+                self.dis_noise = None
+        else:
+            self.gen.eval()
+            self.dis.eval()
+
+        if local_rank >= 0:
+            pDDP = lambda m, find_unused: DDP(m, device_ids=[local_rank], output_device=local_rank,
+                                              find_unused_parameters=False)
+            self.gen = pDDP(self.gen, find_unused=True)
+            self.dis = pDDP(self.dis, find_unused=True)
+            if cfg.lambda_adv_noise:
+                self.dis_noise = pDDP(self.dis_noise, find_unused=True)
+        self.local_rank = local_rank
+        self.device = device
+
+    def edit_infer(self, hair_code, data):
+        self.inner_code = self.dis({'code': hair_code})
+        for ke in data:
+            self.inner_code[ke] = data[ke]
+        self.res = self.gen(self.inner_code)
+        return self.res['code']
+
+    def forward(self, data):
+        self.ae_in = {'code': data['code']}
+
+        # rec
+        d_res_real = self.dis(self.ae_in)
+        self.ae_mid = {'noise': d_res_real['noise'], 'rgb_mean': data['rgb_mean'], 'pca_std': data['pca_std']}
+        if cfg.lambda_cls_curliness:
+            self.ae_mid['noise_curliness'] = d_res_real['noise_curliness']
+        self.ae_out = self.gen(self.ae_mid)
+
+        # gan
+        random_list = list(range(data['rgb_mean'].shape[0]))
+        random.shuffle(random_list)
+        self.gan_in = {'rgb_mean': data['rgb_mean'][random_list]}
+        self.gan_in['pca_std'] = data['pca_std'][random_list]
+        random.shuffle(random_list)
+        if cfg.lambda_cls_curliness:
+            self.gan_in['noise_curliness'] = data['noise_curliness'][random_list]
+        self.gan_in['curliness_label'] = data['curliness_label'][random_list]
+        random.shuffle(random_list)
+
+        if self.cfg.gan_input_from_encoder_prob and \
+                random.random() < self.cfg.gan_input_from_encoder_prob:
+            self.gan_in['noise'] = d_res_real['noise'][random_list].detach()
+        else:
+            self.gan_in['noise'] = data['noise'][random_list]
+        self.gan_mid = self.gen(self.gan_in)
+        self.gan_out_fake = self.dis(self.gan_mid)
+        self.gan_out_real = d_res_real
+        self.gan_label = {'rgb_mean': data['rgb_mean'], 'pca_std': data['pca_std'], 'curliness_label': data['curliness_label']}
+        self.real_noise = {'noise': data['noise']}
+        if cfg.lambda_cls_curliness:
+            self.real_noise['noise_curliness'] = data['noise_curliness']
+
+    def forward_g(self, loss_dict):
+        self.forward_general_gen(self.gan_out_fake['adv'], loss_dict)
+
+        loss_dict['lambda_info'] = self.mse_loss(self.gan_out_fake['noise'], self.gan_in['noise'])
+        loss_dict['lambda_rec'] = self.mse_loss(self.ae_out['code'], self.ae_in['code'])
+
+        if 'rgb' in cfg.predictor:
+            p_rgb = self.rgb_model(self.gan_mid)
+        if cfg.lambda_rgb:
+            if 'rgb' in cfg.predictor:
+                d_rgb_mean = p_rgb['rgb_mean']
+            else:
+                d_rgb_mean = self.gan_out_fake['rgb_mean']
+            loss_dict['lambda_rgb'] = self.mse_loss(d_rgb_mean, self.gan_in['rgb_mean'])
+
+        if cfg.lambda_pca_std:
+            if 'rgb' in cfg.predictor:
+                d_pca_std = p_rgb['pca_std']
+            else:
+                d_pca_std = self.gan_out_fake['pca_std']
+            loss_dict['lambda_pca_std'] = self.mse_loss(d_pca_std, self.gan_in['pca_std'])
+
+        if cfg.lambda_cls_curliness:
+            d_noise_curliness = self.gan_out_fake['noise_curliness']
+            loss_dict['lambda_info_curliness'] = self.mse_loss(d_noise_curliness, self.gan_in['noise_curliness'])
+            if 'curliness' in cfg.predictor:
+                cls_curliness = self.curliness_model(self.gan_mid)['cls_curliness']
+            else:
+                cls_curliness = self.gan_out_fake['cls_curliness']
+            if cfg.curliness_with_weight:
+                weights = self.gan_in['noise_curliness'].abs()
+                weights = weights / weights.sum() * weights.shape[0]
+                loss_dict['lambda_cls_curliness'] = F.binary_cross_entropy(torch.sigmoid(cls_curliness),
+                                                                      self.gan_in['curliness_label'].float() / 2 + 0.5,
+                                                                      weight=weights)
+            else:
+                loss_dict['lambda_cls_curliness'] = F.binary_cross_entropy(torch.sigmoid(cls_curliness),
+                                                                      self.gan_in['curliness_label'].float() / 2 + 0.5)
+
+        if 'gen_mode' in cfg and cfg.gen_mode is 'eigengan':
+            loss_dict['lambda_orthogonal'] = self.gen.orthogonal_regularizer_loss()
+
+        for loss_d in [loss_dict]:
+            for ke in loss_d:
+                if np.isnan(np.array(loss_d[ke].detach().cpu())):
+                    print('!!!!!!!!!  %s is nan' % ke)
+                    print(loss_d)
+                    raise Exception()
+
+    @staticmethod
+    def forward_general_gen(dis_res, loss_dict, loss_name_suffix=''):
+        if cfg.gan_type == 'lsgan':
+            loss_dis = torch.mean((dis_res - 1) ** 2)
+        elif cfg.gan_type == 'nsgan':
+            all1 = torch.ones_like(dis_res.data).cuda()
+            loss_dis = torch.mean(F.binary_cross_entropy(torch.sigmoid(dis_res), all1))
+        elif cfg.gan_type == 'wgan_gp':
+            loss_dis = - torch.mean(dis_res)
+        elif cfg.gan_type == 'hinge':
+            loss_dis = -torch.mean(dis_res)
+        elif cfg.gan_type == 'hinge2':
+            loss_dis = torch.mean(torch.max(1 - dis_res, torch.zeros_like(dis_res)))
+        else:
+            raise NotImplementedError()
+        loss_dict['lambda_adv' + loss_name_suffix] = loss_dis
+
+    @staticmethod
+    def forward_general_dis(dis1, dis0, dis_model, loss_dict, input_real, input_fake, loss_name_suffix=''):
+
+        if cfg.gan_type == 'lsgan':
+            loss_dis = torch.mean((dis0 - 0) ** 2) + torch.mean((dis1 - 1) ** 2)
+        elif cfg.gan_type == 'nsgan':
+            all0 = torch.zeros_like(dis0.data).cuda()
+            all1 = torch.ones_like(dis1.data).cuda()
+            loss_dis = torch.mean(F.binary_cross_entropy(torch.sigmoid(dis0), all0) +
+                                  F.binary_cross_entropy(torch.sigmoid(dis1), all1))
+        elif cfg.gan_type == 'wgan_gp':
+            loss_dis = torch.mean(dis0) - torch.mean(dis1)
+        elif cfg.gan_type == 'hinge' or cfg.gan_type == 'hinge2':
+            loss_dis = torch.mean(torch.max(1 - dis1, torch.zeros_like(dis1)))
+            loss_dis += torch.mean(torch.max(1 + dis0, torch.zeros_like(dis1)))
+        else:
+            assert 0, "Unsupported GAN type: {}".format(dis_model.gan_type)
+        loss_dict['lambda_adv' + loss_name_suffix] = loss_dis
+
+        if cfg.gan_type == 'wgan_gp':
+            loss_gp = 0
+            alpha_gp = torch.rand(input_real.size(0), 1, ).type_as(input_real)
+            x_hat = (alpha_gp * input_real + (1 - alpha_gp) * input_fake).requires_grad_(True)
+            out_hat = dis_model.forward_adv_direct(x_hat)
+            # gradient penalty
+            weight = torch.ones(out_hat.size()).type_as(out_hat)
+            dydx = torch.autograd.grad(outputs=out_hat, inputs=x_hat, grad_outputs=weight, retain_graph=True,
+                                       create_graph=True, only_inputs=True)[0]
+            dydx = dydx.contiguous().view(dydx.size(0), -1)
+            dydx_l2norm = torch.sqrt(torch.sum(dydx ** 2, dim=1))
+            loss_gp += torch.mean((dydx_l2norm - 1) ** 2)
+            loss_dict['lambda_gp' + loss_name_suffix] = loss_gp
+
+    def forward_d(self, loss_dict):
+        self.forward_general_dis(self.gan_out_real['adv'], self.gan_out_fake['adv'],
+                                 self.dis, loss_dict, input_real=self.ae_in['code'],
+                                 input_fake=self.gan_mid['code'])
+
+        loss_dict['lambda_info'] = self.mse_loss(self.gan_out_fake['noise'], self.gan_in['noise'])
+        if cfg.lambda_rgb and 'rgb' not in cfg.predictor:
+            loss_dict['lambda_rgb'] = self.mse_loss(self.gan_in['rgb_mean'], self.gan_out_fake['rgb_mean'])
+        loss_dict['lambda_rec'] = self.mse_loss(self.ae_out['code'], self.ae_in['code'])
+        if cfg.lambda_pca_std and 'rgb' not in cfg.predictor:
+            loss_dict['lambda_pca_std'] = self.mse_loss(self.gan_out_real['pca_std'], self.gan_label['pca_std'])
+
+        if cfg.lambda_adv_noise:
+            self.d_noise_res = self.dis_noise(self.ae_mid)
+            self.forward_general_gen(self.d_noise_res['adv'], loss_dict, loss_name_suffix='_noise')
+
+        if cfg.lambda_moment_1 or cfg.lambda_moment_2:
+            if cfg.lambda_cls_curliness:
+                noise_mid = torch.cat([self.ae_mid['noise_curliness'], self.ae_mid['noise']], dim=1)
+            else:
+                noise_mid = self.ae_mid['noise']
+            if cfg.lambda_moment_1:
+                loss_dict['lambda_moment_1'] = (noise_mid.mean(dim=0) ** 2).mean()
+            if cfg.lambda_moment_2:
+                loss_dict['lambda_moment_2'] = (((noise_mid ** 2).mean(dim=0) - 1) ** 2).mean()
+
+        if cfg.lambda_cls_curliness:
+            loss_dict['lambda_info_curliness'] = self.mse_loss(self.gan_out_fake['noise_curliness'], self.gan_in['noise_curliness'])
+
+    def forward_d_curliness(self, data_curliness, loss_dict):
+        d_res = self.dis(data_curliness)
+        cls_curliness = d_res['cls_curliness']
+        loss_dict['lambda_cls_curliness'] = F.binary_cross_entropy(torch.sigmoid(cls_curliness),
+                                                              data_curliness['curliness_label'].float() / 2 + 0.5)
+
+    def forward_adv_noise(self, loss_dict):
+        dis1 = self.dis_noise(self.real_noise)['adv']
+        self.ae_mid['noise'] = self.ae_mid['noise'].detach()
+        if self.cfg.lambda_cls_curliness:
+            self.ae_mid['noise_curliness'] = self.ae_mid['noise_curliness'].detach()
+        self.d_noise_res = self.dis_noise(self.ae_mid)
+        dis0 = self.d_noise_res['adv']
+
+        input_real = self.real_noise['noise']
+        input_fake = self.ae_mid['noise']
+        if self.cfg.lambda_cls_curliness:
+            input_real = torch.cat([input_real, self.real_noise['noise_curliness']], dim=1)
+            input_fake = torch.cat([input_fake, self.ae_mid['noise_curliness']], dim=1)
+
+        self.forward_general_dis(dis1, dis0, self.dis_noise, loss_dict, input_real=input_real,
+                                 input_fake=input_fake, loss_name_suffix='_noise')
+
+    def forward_rec_img(self, data, loss_dict, batch_size=4):
+        from .validation_in_train import he
+        from global_value_utils import HAIR_IDX
+
+        items = data['items']
+
+        rec_loss = 0
+        for idx in range(batch_size):
+            item = items[idx]
+            dataset_name, img_name = item.split('___')
+            parsing_img = cv2.imread(os.path.join(self.cfg.data_root, '%s/label/%s.png') %
+                                     (dataset_name, img_name), cv2.IMREAD_GRAYSCALE)
+            parsing_img = cv2.resize(parsing_img, (256, 256), cv2.INTER_NEAREST)
+
+            sean_code = torch.tensor(data['sean_code'][idx].copy(), dtype=torch.float32).to(self.ae_out['code'].device)
+            sean_code[HAIR_IDX] = self.ae_out['code'][idx, ...]
+
+            render_img = he.gen_img(sean_code[None, ...], parsing_img[None, None, ...])
+
+            input_img = cv2.cvtColor(cv2.imread(os.path.join(self.cfg.data_root, '%s/images_256/%s.png') %
+                                                (dataset_name, img_name)), cv2.COLOR_BGR2RGB)
+            input_img = input_img / 127.5 - 1.0
+            input_img = input_img.transpose(2, 0, 1)
+
+            input_img = torch.tensor(input_img).to(render_img.device)
+            parsing_img = torch.tensor(parsing_img).to(render_img.device)
+
+            rec_loss = rec_loss + ((input_img - render_img)[:, (parsing_img == HAIR_IDX)] ** 2).mean()
+        rec_loss = rec_loss / batch_size
+        loss_dict['lambda_rec_img'] = rec_loss

models/CtrlHair/color_texture_branch/train.py

@@ -0,0 +1,166 @@
+# -*- coding: utf-8 -*-
+
+"""
+# File name:    scripts.py.py
+# Time :        2021/11/17 15:24
+# Author:       xyguoo@163.com
+# Description:  
+"""
+
+import sys
+
+sys.path.append('.')
+
+import tensorboardX
+import torch
+import tqdm
+import numpy as np
+from color_texture_branch.config import cfg, args
+from color_texture_branch.dataset import Dataset
+import my_pylib
+from color_texture_branch.validation_in_train import print_val_save_model
+# distributed training
+import torch.distributed as dist
+from color_texture_branch.solver import Solver
+from my_torchlib.train_utils import LossUpdater, to_device, generate_noise, train
+import my_torchlib
+from color_texture_branch.model import init_weights
+
+
+def get_total_step():
+    total = 0
+    for key in cfg.iter:
+        total += cfg.iter[key]
+    return total
+
+
+def worker(proc, nprocs, args):
+    local_rank = args.local_rank
+    if local_rank >= 0:
+        torch.cuda.set_device(local_rank)
+        dist.init_process_group(backend='nccl',
+                                init_method='tcp://localhost:%d' % (6030 + int(cfg.experiment_name[:3])),
+                                rank=args.local_rank,
+                                world_size=cfg.gpu_num)
+        print('setup rank %d' % local_rank)
+    device = torch.device('cuda', max(0, local_rank))
+
+    # config
+    out_dir = cfg.root_dir
+
+    # data
+    ds = Dataset(cfg)
+
+    loss_updater = LossUpdater(cfg)
+    loss_updater.update(0)
+
+    # Loss class
+    solver = Solver(cfg, device, local_rank=local_rank)
+
+    # load checkpoint
+    ckpt_dir = out_dir + '/checkpoints'
+    if local_rank <= 0:
+        my_pylib.mkdir(out_dir)
+        my_pylib.save_json(out_dir + '/setting_hair.json', cfg, indent=4, separators=(',', ': '))
+        my_pylib.mkdir(ckpt_dir)
+
+    try:
+        ckpt = my_torchlib.load_checkpoint(ckpt_dir)
+        start_step = ckpt['step'] + 1
+        for model_name in ['Model_G', 'Model_D']:
+            cur_model = ckpt[model_name]
+            if list(cur_model)[0].startswith('module'):
+                ckpt[model_name] = {kk[7:]: cur_model[kk] for kk in cur_model}
+        solver.gen.load_state_dict(ckpt['Model_G'], strict=True)
+        solver.dis.load_state_dict(ckpt['Model_D'], strict=True)
+        solver.D_optimizer.load_state_dict(ckpt['D_optimizer'])
+        solver.G_optimizer.load_state_dict(ckpt['G_optimizer'])
+        if cfg.lambda_adv_noise:
+            solver.dis_noise.load_state_dict(ckpt['Model_D_noise'], strict=True)
+            solver.D_noise_optimizer.load_state_dict(ckpt['D_noise_optimizer'])
+        print('Load succeed!')
+    except:
+        print(' [*] No checkpoint!')
+        init_weights(solver.gen, init_type=cfg.init_type)
+        init_weights(solver.dis, init_type=cfg.init_type)
+        if cfg.lambda_adv_noise:
+            init_weights(solver.dis_noise, init_type=cfg.init_type)
+        start_step = 1
+
+    if 'curliness' in cfg.predictor:
+        ckpt = my_torchlib.load_checkpoint(cfg.predictor.curliness.root_dir + '/checkpoints')
+        solver.curliness_model.load_state_dict(ckpt['Predictor'], strict=True)
+
+    if 'rgb' in cfg.predictor:
+        ckpt = my_torchlib.load_checkpoint(cfg.predictor.rgb.root_dir + '/checkpoints')
+        solver.rgb_model.load_state_dict(ckpt['Predictor'], strict=True)
+
+    # writer
+    if local_rank <= 0:
+        writer = tensorboardX.SummaryWriter(out_dir + '/summaries')
+    else:
+        writer = None
+
+    # start training
+    test_batch = ds.get_testing_batch(cfg.sample_batch_size)
+    test_batch_curliness = ds.get_curliness_hair_test()
+
+    to_device(test_batch_curliness, device)
+    to_device(test_batch, device)
+
+    if local_rank >= 0:
+        dist.barrier()
+
+    total_step = cfg.total_step + 2
+    for step in tqdm.tqdm(range(start_step, total_step), total=total_step, initial=start_step, desc='step'):
+        loss_updater.update(step)
+        write_log = (writer and step % 23 == 0)
+
+        for i in range(sum(cfg.G_D_train_num.values())):
+            data = ds.get_training_batch(cfg.batch_size)
+            data['noise'] = generate_noise(cfg.batch_size, cfg.noise_dim)
+            if cfg.lambda_cls_curliness:
+                curliness_label = torch.tensor(np.random.choice([-1, 1], (cfg.batch_size, 1)))
+                data['curliness_label'] = curliness_label
+                data['noise_curliness'] = generate_noise(cfg.batch_size, cfg.curliness_dim, curliness_label)
+            to_device(data, device)
+            loss_dict = {}
+            solver.forward(data)
+
+            if 'lambda_rec_img' in cfg and cfg.lambda_rec_img > 0:
+                solver.forward_rec_img(data, loss_dict)
+
+            if i < cfg.G_D_train_num['D']:
+                solver.forward_d(loss_dict)
+                if cfg.lambda_cls_curliness and not 'curliness' in cfg.predictor:
+                    data_curliness = ds.get_curliness_hair(curliness_label)
+                    to_device(data_curliness, device)
+                    solver.forward_d_curliness(data_curliness, loss_dict)
+
+                    # validation to show whether over-fit
+                    if write_log:
+                        loss_dict['test_cls_curliness'] = torch.nn.functional.binary_cross_entropy_with_logits(
+                            solver.dis(test_batch_curliness)['cls_curliness'], test_batch_curliness['curliness_label'] / 2 + 0.5)
+                if cfg.lambda_rgb and 'rgb' not in cfg.predictor and write_log:
+                    loss_dict['test_lambda_rgb'] = solver.mse_loss(solver.dis(test_batch)['rgb_mean'],
+                                                                   test_batch['rgb_mean'])
+                train(cfg, loss_dict, optimizers=[solver.D_optimizer],
+                      step=step, writer=writer, flag='D', write_log=write_log)
+            else:
+                solver.forward_g(loss_dict)
+                train(cfg, loss_dict, optimizers=[solver.G_optimizer],
+                      step=step, writer=writer, flag='G', write_log=write_log)
+
+                if cfg.lambda_adv_noise:
+                    loss_dict = {}
+                    solver.forward_adv_noise(loss_dict)
+                    train(cfg, loss_dict, optimizers=[solver.D_noise_optimizer], step=step, writer=writer, flag='D_noise',
+                          write_log=write_log)
+
+        print_val_save_model(step, out_dir, solver, test_batch, ckpt_dir, local_rank)
+
+
+if __name__ == '__main__':
+    # with torch.autograd.set_detect_anomaly(True):
+    # mp.spawn(worker, nprocs=cfg.gpu_num, args=(cfg.gpu_num, args))
+    worker(proc=None, nprocs=None, args=args)

models/CtrlHair/color_texture_branch/validation_in_train.py

@@ -0,0 +1,299 @@
+# -*- coding: utf-8 -*-
+
+"""
+# File name:    validation_in_train.py
+# Time :        2021/12/10 12:55
+# Author:       xyguoo@163.com
+# Description:  
+"""
+
+import cv2
+import os
+
+from .config import cfg
+import my_pylib
+import torch
+from my_torchlib.train_utils import tensor2numpy, to_device, generate_noise
+from util.canvas_grid import Canvas
+from global_value_utils import HAIR_IDX
+import copy
+import numpy as np
+import torch.distributed as dist
+import my_torchlib
+from hair_editor import HairEditor
+
+he = HairEditor(load_feature_model=False, load_mask_model=False)
+
+
+def gen_by_sean(sean_code, item):
+    dataset_name, img_name = item.split('___')
+    parsing_img = cv2.imread(os.path.join(cfg.data_root, '%s/label/%s.png') %
+                             (dataset_name, img_name), cv2.IMREAD_GRAYSCALE)
+    parsing_img = cv2.resize(parsing_img, (256, 256), cv2.INTER_NEAREST)
+    return he.gen_img(sean_code[None, ...], parsing_img[None, None, ...])
+
+
+def save_model(step, solver, ckpt_dir):
+    save_dic = {'step': step,
+                'Model_G': solver.gen.state_dict(), 'Model_D': solver.dis.state_dict(),
+                'D_optimizer': solver.D_optimizer.state_dict(), 'G_optimizer': solver.G_optimizer.state_dict()}
+    if cfg.lambda_adv_noise:
+        save_dic['Model_D_noise'] = solver.dis_noise.state_dict()
+        save_dic['D_noise_optimizer'] = solver.D_noise_optimizer.state_dict()
+    my_torchlib.save_checkpoint(save_dic, '%s/%07d.ckpt' % (ckpt_dir, step), max_keep=cfg.max_save)
+
+
+def print_val_save_model(step, out_dir, solver, test_batch, ckpt_dir, local_rank):
+    """
+    :param step:
+    :param validation_data:
+    :param img_size:
+    :param alpha:
+    :return:
+    """
+    if step > 0 and step % cfg.sample_step == 0:
+        gen = solver.gen
+        dis = solver.dis
+        local_rank = solver.local_rank
+        device = solver.device
+
+        save_dir = out_dir + '/sample_training'
+        my_pylib.mkdir(save_dir)
+        gen.eval()
+        dis.eval()
+        # gen.cpu()
+        show_batch_size = 10
+        row_idxs = list(range(show_batch_size))
+        instance_idxs = list(range(show_batch_size))
+
+        with torch.no_grad():
+            items = test_batch['items']
+            decoder_res = dis({'code': test_batch['code'].cuda()})
+            hair_noise = decoder_res['noise'].cpu()
+            test_data = {'noise': hair_noise, 'rgb_mean': test_batch['rgb_mean'],
+                         'pca_std': test_batch['pca_std']}
+            if cfg.lambda_cls_curliness:
+                test_data['noise_curliness'] = decoder_res['noise_curliness'].cpu()
+            to_device(test_data, device)
+            rec_code = gen(test_data)['code']
+
+            # ----------------
+            # generate each noise dim
+            # ----------------
+            # grid_count = 10
+            # lin_space = np.linspace(-3, 3, grid_count)
+            grid_count = 6
+            lin_space = np.linspace(-2.5, 2.5, grid_count)
+            for dim_idx in range(cfg.noise_dim):
+                canvas = Canvas(len(row_idxs), grid_count + 1)
+                for draw_idx, idx in enumerate(row_idxs):
+                    item = items[idx]
+                    dataset_name, img_name = item.split('___')
+                    # generate origin and reconstruction
+                    ori_img = cv2.cvtColor(cv2.imread(os.path.join(cfg.data_root, '%s/images_256/%s.png') %
+                                                      (dataset_name, img_name)), cv2.COLOR_BGR2RGB)
+
+                    canvas.process_draw_image(ori_img, draw_idx, 0)
+
+                for grid_idx in range(grid_count):
+                    temp_noise = hair_noise.clone()
+                    temp_noise[:, dim_idx] = lin_space[grid_idx]
+                    data = copy.deepcopy(test_data)
+                    data['noise'] = temp_noise
+                    to_device(data, device)
+                    code = gen(data)['code']
+
+                    for draw_idx, idx in enumerate(row_idxs):
+                        cur_code = test_batch['sean_code'][idx].copy()
+                        cur_code[HAIR_IDX] = code[idx].cpu().numpy()
+                        item = items[idx]
+                        out_img = gen_by_sean(cur_code, item)
+                        out_img = tensor2numpy(out_img)
+                        canvas.process_draw_image(out_img, draw_idx, grid_idx + 1)
+                if local_rank <= 0:
+                    canvas.write_(os.path.join(save_dir, '%06d_noise_%02d.png' % (step, dim_idx)))
+
+            # -------------
+            # direct transfer all content
+            # -------------
+            canvas = Canvas(len(row_idxs) + 1, len(instance_idxs) + 2)
+            for draw_idx, instance_idx in enumerate(instance_idxs):
+                item = items[instance_idx]
+                dataset_name, img_name = item.split('___')
+                # generate origin and reconstruction
+                ori_img = cv2.cvtColor(cv2.imread(os.path.join(cfg.data_root, '%s/images_256/%s.png') %
+                                                  (dataset_name, img_name)), cv2.COLOR_BGR2RGB)
+                canvas.process_draw_image(ori_img, 0, draw_idx + 2)
+
+            for draw_idx, idx in enumerate(row_idxs):
+                item_row = items[idx]
+                dataset_name, img_name = item_row.split('___')
+                img_row = cv2.cvtColor(cv2.imread(os.path.join(cfg.data_root, '%s/images_256/%s.png') %
+                                                  (dataset_name, img_name)), cv2.COLOR_BGR2RGB)
+                canvas.process_draw_image(img_row, draw_idx + 1, 0)
+                sean_code = test_batch['sean_code'][idx]
+                rec_img = tensor2numpy(gen_by_sean(sean_code, item_row))
+                canvas.process_draw_image(rec_img, draw_idx + 1, 1)
+                for draw_idx2, instance_idx in enumerate(instance_idxs):
+                    cur_code = test_batch['sean_code'][idx].copy()
+                    cur_code[HAIR_IDX] = test_batch['sean_code'][instance_idx][HAIR_IDX]
+                    res_img = gen_by_sean(cur_code, item_row)
+                    res_img = tensor2numpy(res_img)
+                    canvas.process_draw_image(res_img, draw_idx + 1, draw_idx2 + 2)
+            if local_rank <= 0:
+                canvas.write_(os.path.join(save_dir, 'rgb_direct.png'))
+
+            # -----------
+            # random choice
+            # -----------
+            grid_count = 10
+            # generate each noise dim
+            canvas = Canvas(len(row_idxs), grid_count + 2)
+            for draw_idx, idx in enumerate(row_idxs):
+                item = items[idx]
+                dataset_name, img_name = item.split('___')
+                # generate origin and reconstruction
+                ori_img = cv2.cvtColor(cv2.imread(os.path.join(cfg.data_root, '%s/images_256/%s.png') %
+                                                  (dataset_name, img_name)), cv2.COLOR_BGR2RGB)
+
+                canvas.process_draw_image(ori_img, draw_idx, 0)
+                cur_code = test_batch['sean_code'][idx].copy()
+                cur_code[HAIR_IDX] = rec_code[idx].cpu().numpy()
+                rec_img = tensor2numpy(gen_by_sean(cur_code, item))
+                canvas.process_draw_image(rec_img, draw_idx, 1)
+
+            temp_noise = generate_noise(grid_count, cfg.noise_dim)
+            for grid_idx in range(grid_count):
+                data = copy.deepcopy(test_data)
+                data['noise'] = torch.tile(temp_noise[[grid_idx]], [test_batch['rgb_mean'].shape[0], 1])
+                to_device(data, device)
+                code = gen(data)['code']
+
+                for draw_idx, idx in enumerate(row_idxs):
+                    cur_code = test_batch['sean_code'][idx].copy()
+                    cur_code[HAIR_IDX] = code[idx].cpu().numpy()
+                    item = items[idx]
+                    out_img = gen_by_sean(cur_code, item)
+                    out_img = tensor2numpy(out_img)
+                    canvas.process_draw_image(out_img, draw_idx, grid_idx + 2)
+
+            if local_rank <= 0:
+                canvas.write_(os.path.join(save_dir, '%06d_random.png' % step))
+
+            # ------------
+            # generate curliness
+            # ------------
+            if cfg.lambda_cls_curliness:
+                grid_count = 10
+                lin_space = np.linspace(-3, 3, grid_count)
+                canvas = Canvas(len(row_idxs), grid_count + 2)
+                for draw_idx, idx in enumerate(row_idxs):
+                    item = items[idx]
+                    dataset_name, img_name = item.split('___')
+                    # generate origin and reconstruction
+                    ori_img = cv2.cvtColor(cv2.imread(os.path.join(cfg.data_root, '%s/images_256/%s.png') %
+                                                      (dataset_name, img_name)), cv2.COLOR_BGR2RGB)
+                    canvas.process_draw_image(ori_img, draw_idx, 0)
+                    cur_code = test_batch['sean_code'][idx].copy()
+                    cur_code[HAIR_IDX] = rec_code[idx].cpu().numpy()
+                    rec_img = tensor2numpy(gen_by_sean(cur_code, item))
+                    canvas.process_draw_image(rec_img, draw_idx, 1)
+
+                for grid_idx in range(grid_count):
+                    cur_noise_curliness = torch.tensor(lin_space[grid_idx]).reshape([1, 1]).tile([cfg.sample_batch_size, 1]).float()
+                    data = copy.deepcopy(test_data)
+                    data['noise_curliness'] = cur_noise_curliness
+                    to_device(data, device)
+                    code = gen(data)['code']
+                    for draw_idx, idx in enumerate(row_idxs):
+                        cur_code = test_batch['sean_code'][idx].copy()
+                        cur_code[HAIR_IDX] = code[idx].cpu().numpy()
+                        item = items[idx]
+                        out_img = gen_by_sean(cur_code, item)
+                        out_img = tensor2numpy(out_img)
+                        canvas.process_draw_image(out_img, draw_idx, grid_idx + 2)
+                if local_rank <= 0:
+                    canvas.write_(os.path.join(save_dir, '%06d_curliness.png' % step))
+
+            # ------------
+            # generate variance
+            # ------------
+            if cfg.lambda_pca_std:
+                grid_count = 10
+                lin_space = np.linspace(10, 150, grid_count)
+                canvas = Canvas(len(row_idxs), grid_count + 2)
+                for draw_idx, idx in enumerate(row_idxs):
+                    item = items[idx]
+                    dataset_name, img_name = item.split('___')
+                    # generate origin and reconstruction
+                    ori_img = cv2.cvtColor(cv2.imread(os.path.join(cfg.data_root, '%s/images_256/%s.png') %
+                                                      (dataset_name, img_name)), cv2.COLOR_BGR2RGB)
+                    canvas.process_draw_image(ori_img, draw_idx, 0)
+
+                    cur_code = test_batch['sean_code'][idx].copy()
+                    cur_code[HAIR_IDX] = rec_code[idx].cpu().numpy()
+                    rec_img = tensor2numpy(gen_by_sean(cur_code, item))
+                    canvas.process_draw_image(rec_img, draw_idx, 1)
+
+                for grid_idx in range(grid_count):
+                    cur_pca_std = torch.tensor(lin_space[grid_idx]).reshape([1, 1]).tile([cfg.sample_batch_size, 1]).float()
+                    data = copy.deepcopy(test_data)
+                    data['pca_std'] = cur_pca_std
+                    to_device(data, device)
+                    code = gen(data)['code']
+
+                    for draw_idx, idx in enumerate(row_idxs):
+                        cur_code = test_batch['sean_code'][idx].copy()
+                        cur_code[HAIR_IDX] = code[idx].cpu().numpy()
+                        item = items[idx]
+                        out_img = gen_by_sean(cur_code, item)
+                        out_img = tensor2numpy(out_img)
+                        canvas.process_draw_image(out_img, draw_idx, grid_idx + 2)
+                if local_rank <= 0:
+                    canvas.write_(os.path.join(save_dir, '%06d_variance.png' % step))
+
+            # -------------
+            # generate each rgb
+            # -------------
+            canvas = Canvas(len(row_idxs) + 1, len(instance_idxs) + 1)
+            for draw_idx, instance_idx in enumerate(instance_idxs):
+                item = items[instance_idx]
+                dataset_name, img_name = item.split('___')
+                # generate origin and reconstruction
+                ori_img = cv2.cvtColor(cv2.imread(os.path.join(cfg.data_root, '%s/images_256/%s.png') %
+                                                  (dataset_name, img_name)), cv2.COLOR_BGR2RGB)
+                ori_img[5:45, 5:45, :] = test_batch['rgb_mean'][instance_idx].cpu().numpy()
+                canvas.process_draw_image(ori_img, 0, draw_idx + 1)
+            for draw_idx, idx in enumerate(row_idxs):
+                item_row = items[idx]
+                dataset_name, img_name = item_row.split('___')
+                img_row = cv2.cvtColor(cv2.imread(os.path.join(cfg.data_root, '%s/images_256/%s.png') %
+                                                  (dataset_name, img_name)), cv2.COLOR_BGR2RGB)
+                canvas.process_draw_image(img_row, draw_idx + 1, 0)
+            for draw_idx2, instance_idx in enumerate(instance_idxs):
+                color = test_batch['rgb_mean'][[instance_idx]]
+                data = copy.deepcopy(test_data)
+                data['rgb_mean'] = torch.tile(color, [cfg.sample_batch_size, 1])
+                data['pca_std'] = torch.tile(test_batch['pca_std'][[instance_idx]], [cfg.sample_batch_size, 1])
+                to_device(data, device)
+                hair_code = gen(data)['code']
+                for draw_idx, idx in enumerate(row_idxs):
+                    item_row = items[idx]
+                    cur_code = test_batch['sean_code'][idx].copy()
+                    cur_code[HAIR_IDX] = hair_code[idx].cpu().numpy()
+                    res_img = gen_by_sean(cur_code, item_row)
+                    res_img = tensor2numpy(res_img)
+                    canvas.process_draw_image(res_img, draw_idx + 1, draw_idx2 + 1)
+            if local_rank <= 0:
+                canvas.write_(os.path.join(save_dir, '%06d_rgb.png' % step))
+
+        gen.train()
+        dis.train()
+        if local_rank >= 0:
+            dist.barrier()
+
+    if step > 0 and step % cfg.model_save_step == 0:
+        if local_rank <= 0:
+            save_model(step, solver, ckpt_dir)
+        if local_rank >= 0:
+            dist.barrier()

models/CtrlHair/common_dataset.py

@@ -0,0 +1,104 @@
+# -*- coding: utf-8 -*-
+
+"""
+# File name:    common_dataset.py
+# Time :        2021/12/31 19:19
+# Author:       xyguoo@163.com
+# Description:  
+"""
+import os
+import random
+
+import numpy as np
+import pandas as pd
+
+from global_value_utils import GLOBAL_DATA_ROOT, HAIR_IDX, HAT_IDX, DATASET_NAME
+from util.util import path_join_abs
+import cv2
+
+
+class DataFilter:
+    def __init__(self, cfg=None, test_part=0.096):
+        # cfg
+        if cfg is None:
+            from color_texture_branch.config import cfg
+        self.cfg = cfg
+
+        base_dataset_dir = GLOBAL_DATA_ROOT
+
+        dataset_dir = [os.path.join(base_dataset_dir, dn) for dn in DATASET_NAME]
+
+        self.data_dirs = dataset_dir
+
+        self.random_seed = 7
+        random.seed(self.random_seed)
+
+        ####################################################################
+        # Please modify these if you don't want to use these filters
+        angle_filter = True
+        gender_filter = True
+        gender = ['female']
+        # gender = ['male', 'female']
+        gender = set([{'male': 1, 'female': -1}[g] for g in gender])
+        ####################################################################
+
+        self.total_list = []
+        for data_dir in self.data_dirs:
+            img_dir = os.path.join(data_dir, 'images_256')
+
+            if angle_filter:
+                angle_csv = pd.read_csv(os.path.join(data_dir, 'angle.csv'), index_col=0)
+                angle_filter_imgs = list(angle_csv.index[angle_csv['angle'] < 5])
+                cur_list = ['%05d.png' % dd for dd in angle_filter_imgs]
+            else:
+                cur_list = os.listdir(img_dir)
+
+            if gender_filter:
+                attr_filter = pd.read_csv(os.path.join(data_dir, 'attr_gender.csv'))
+                cur_list = [p for p in cur_list if attr_filter.Male[int(p[:-4])] in gender]
+
+            self.total_list += [os.path.join(img_dir, p) for p in cur_list]
+
+        random.shuffle(self.total_list)
+        self.test_start = int(len(self.total_list) * (1 - test_part))
+        self.test_list = self.total_list[self.test_start:]
+        self.train_list = [st for st in self.total_list if st not in self.test_list]
+
+        idx = 0
+        # make sure the area of hair is big enough
+        self.hair_region_threshold = 0.07
+        self.test_face_list = []
+        while len(self.test_face_list) < cfg.sample_batch_size:
+            test_file = self.test_list[idx]
+            if self.valid_face(path_join_abs(test_file, '../..'), test_file[-9:-4]):
+                self.test_face_list.append(test_file)
+            idx += 1
+
+        self.test_hair_list = []
+        while len(self.test_hair_list) < cfg.sample_batch_size:
+            test_file = self.test_list[idx]
+            if self.valid_hair(path_join_abs(test_file, '../..'), test_file[-9:-4]):
+                self.test_hair_list.append(test_file)
+            idx += 1
+
+    def valid_face(self, data_dir, img_idx_str):
+        label_path = os.path.join(data_dir, 'label', img_idx_str + '.png')
+        label_img = cv2.imread(label_path, cv2.IMREAD_GRAYSCALE)
+        hat_region = label_img == HAT_IDX
+
+        if hat_region.mean() > 0.03:
+            return False
+        return True
+
+    def valid_hair(self, data_dir, img_idx_str):
+        label_path = os.path.join(data_dir, 'label', img_idx_str + '.png')
+        label_img = cv2.imread(label_path, cv2.IMREAD_GRAYSCALE)
+        hair_region = label_img == HAIR_IDX
+        hat_region = label_img == HAT_IDX
+
+        if hat_region.mean() > 0.03:
+            return False
+        if hair_region.mean() < self.hair_region_threshold:
+            return False
+        return True
+