diff --git a/.DS_Store b/.DS_Store new file mode 100644 index 0000000000000000000000000000000000000000..f7535c13a2ef8123341f1f697fc50dbca8cafbe4 Binary files /dev/null and b/.DS_Store differ diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..6b1e3dba671318e3576aea1ef913a5ca3180ed96 --- /dev/null +++ b/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. diff --git a/README.md b/README.md new file mode 100644 index 0000000000000000000000000000000000000000..b92342ea91e82d3d812b5dba0d47eb1da42b70a0 --- /dev/null +++ b/README.md @@ -0,0 +1,169 @@ +# HairFastGAN: Realistic and Robust Hair Transfer with a Fast Encoder-Based Approach + + + + +[![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. +> + +

+ Teaser +
+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. +

+ +## 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 +|:----------------------------------------| :--- +| 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 + +

+ Diagram +
+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. +

+ +## 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} +} +``` diff --git a/app.py b/app.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/datasets/__init__.py b/datasets/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/datasets/image_dataset.py b/datasets/image_dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..4e56368ac37ba10cb1a29b5172699589035fd22d --- /dev/null +++ b/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 diff --git a/docs/assets/diagram.webp b/docs/assets/diagram.webp new file mode 100644 index 0000000000000000000000000000000000000000..84275e3a2b7ab57c8b83139d9245194aab07219b --- /dev/null +++ b/docs/assets/diagram.webp @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:82e7fce0312bcc7931243fb1d834539bc9d24bbd07a87135d89974987b217882 +size 742106 diff --git a/docs/assets/logo.webp b/docs/assets/logo.webp new file mode 100644 index 0000000000000000000000000000000000000000..e0c09a628ed9e6d650287c68468d7a044d85e046 --- /dev/null +++ b/docs/assets/logo.webp @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:3262e841eadcdd85ccdb9519e6ac7cf0f486857fc32a5bd59eb363594ad2b66d +size 426760 diff --git a/hair_swap.py b/hair_swap.py new file mode 100644 index 0000000000000000000000000000000000000000..18e044919d4ade45f9b6b621756ab0da3b2e57b9 --- /dev/null +++ b/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) diff --git a/inference_server.py b/inference_server.py new file mode 100644 index 0000000000000000000000000000000000000000..0f35d3c0d64915d9cd8f8a3db3322610192d4d5b --- /dev/null +++ b/inference_server.py @@ -0,0 +1,5 @@ +if __name__ == "__main__": + server = grpc.server(...) + ... + server.start() + server.wait_for_termination() diff --git a/losses/.DS_Store b/losses/.DS_Store new file mode 100644 index 0000000000000000000000000000000000000000..924c6ff93402ee0e429fcc6d9e7f887b7cfee310 Binary files /dev/null and b/losses/.DS_Store differ diff --git a/losses/__init__.py b/losses/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/losses/lpips/__init__.py b/losses/lpips/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..2dd73e7b248bd9bbc814df6210f14f95fc5045ae --- /dev/null +++ b/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))) diff --git a/losses/lpips/base_model.py b/losses/lpips/base_model.py new file mode 100644 index 0000000000000000000000000000000000000000..8de1d16f0c7fa52d8067139abc6e769e96d0a6a1 --- /dev/null +++ b/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') diff --git a/losses/lpips/dist_model.py b/losses/lpips/dist_model.py new file mode 100644 index 0000000000000000000000000000000000000000..6c69380084f412a7f8d3d4e64466d251c4e6f19e --- /dev/null +++ b/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 %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= 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)) + ) diff --git a/losses/masked_lpips/base_model.py b/losses/masked_lpips/base_model.py new file mode 100644 index 0000000000000000000000000000000000000000..20b3b3458bd4a6751708e3ec501f1a59c619a1c4 --- /dev/null +++ b/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", + ) diff --git a/losses/masked_lpips/dist_model.py b/losses/masked_lpips/dist_model.py new file mode 100644 index 0000000000000000000000000000000000000000..ddee6c38c14a43e4171ac485b86483918492e84c --- /dev/null +++ b/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)) diff --git a/losses/masked_lpips/networks_basic.py b/losses/masked_lpips/networks_basic.py new file mode 100644 index 0000000000000000000000000000000000000000..ea81e39c338bc13c7597f57e260e35612b8d2aab --- /dev/null +++ b/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) diff --git a/losses/masked_lpips/pretrained_networks.py b/losses/masked_lpips/pretrained_networks.py new file mode 100644 index 0000000000000000000000000000000000000000..c251390679737ebf7ae279cd93571d369b41796e --- /dev/null +++ b/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 diff --git a/losses/pp_losses.py b/losses/pp_losses.py new file mode 100644 index 0000000000000000000000000000000000000000..29d2d5694a748f9a77d74774fbcc10dd8a7efee8 --- /dev/null +++ b/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 diff --git a/losses/style/__init__.py b/losses/style/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/losses/style/custom_loss.py b/losses/style/custom_loss.py new file mode 100644 index 0000000000000000000000000000000000000000..cceb6775b1df8bf5768a1f6614461054b86287f8 --- /dev/null +++ b/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 diff --git a/losses/style/style_loss.py b/losses/style/style_loss.py new file mode 100644 index 0000000000000000000000000000000000000000..750aa6523151a62b4cf9a6f42dbc240239f10093 --- /dev/null +++ b/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 diff --git a/losses/style/vgg_activations.py b/losses/style/vgg_activations.py new file mode 100644 index 0000000000000000000000000000000000000000..42ef45a5c585cfa4ddf2a10186046887a7edb255 --- /dev/null +++ b/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 diff --git a/losses/vgg_loss.py b/losses/vgg_loss.py new file mode 100644 index 0000000000000000000000000000000000000000..4194c6b660008868d852e290faf63f34b3894656 --- /dev/null +++ b/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 diff --git a/main.py b/main.py new file mode 100644 index 0000000000000000000000000000000000000000..9ec1b92ec705000941d68e7cfb30fcd691834efc --- /dev/null +++ b/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) diff --git a/models/.DS_Store b/models/.DS_Store new file mode 100644 index 0000000000000000000000000000000000000000..a9e2fa569183c06b0259e640bec0059bf7055292 Binary files /dev/null and b/models/.DS_Store differ diff --git a/models/Alignment.py b/models/Alignment.py new file mode 100644 index 0000000000000000000000000000000000000000..062fafc43f924a556112b5bc3c937b38b8337a8c --- /dev/null +++ b/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} diff --git a/models/Blending.py b/models/Blending.py new file mode 100644 index 0000000000000000000000000000000000000000..e7dcc87a7a1bdfc39a1b7f03be945fb06f791801 --- /dev/null +++ b/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 diff --git a/models/CtrlHair/.gitignore b/models/CtrlHair/.gitignore new file mode 100644 index 0000000000000000000000000000000000000000..ed7a7d64d933d04668e59c7ea3b88c1460e6a766 --- /dev/null +++ b/models/CtrlHair/.gitignore @@ -0,0 +1,8 @@ +.idea/ +.DS_Store +/dataset_info_ctrlhair/ +/external_model_params/ +/model_trained/ +**/__pycache__/** + + diff --git a/models/CtrlHair/README.md b/models/CtrlHair/README.md new file mode 100644 index 0000000000000000000000000000000000000000..6111f3531fe612e0bab5efb53e7ad388cbb198a4 --- /dev/null +++ b/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) diff --git a/models/CtrlHair/color_texture_branch/__init__.py b/models/CtrlHair/color_texture_branch/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..01eede24e15d0c9bb2f34b35ae02d518c810373f --- /dev/null +++ b/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: +""" diff --git a/models/CtrlHair/color_texture_branch/config.py b/models/CtrlHair/color_texture_branch/config.py new file mode 100644 index 0000000000000000000000000000000000000000..b161a1d470aa3b5ed9613a2976f6bd3538251c81 --- /dev/null +++ b/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) diff --git a/models/CtrlHair/color_texture_branch/dataset.py b/models/CtrlHair/color_texture_branch/dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..7af6f199871c8f550644ff907508f463ba3ad38d --- /dev/null +++ b/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) diff --git a/models/CtrlHair/color_texture_branch/model.py b/models/CtrlHair/color_texture_branch/model.py new file mode 100644 index 0000000000000000000000000000000000000000..4ced35665953cb023054dd21642518dff73c956a --- /dev/null +++ b/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 + + +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]] diff --git a/models/CtrlHair/color_texture_branch/model_eigengan.py b/models/CtrlHair/color_texture_branch/model_eigengan.py new file mode 100644 index 0000000000000000000000000000000000000000..58b90369f67c51f049293332aee7051e452e0fc4 --- /dev/null +++ b/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 diff --git a/models/CtrlHair/color_texture_branch/module.py b/models/CtrlHair/color_texture_branch/module.py new file mode 100644 index 0000000000000000000000000000000000000000..ada8738fb429c453bb936cbef88f1b972ee51abc --- /dev/null +++ b/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 diff --git a/models/CtrlHair/color_texture_branch/predictor/__init__.py b/models/CtrlHair/color_texture_branch/predictor/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d04f75c257e1c83b1e1920c2a9fc22a464669da4 --- /dev/null +++ b/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: +""" diff --git a/models/CtrlHair/color_texture_branch/predictor/predictor_config.py b/models/CtrlHair/color_texture_branch/predictor/predictor_config.py new file mode 100644 index 0000000000000000000000000000000000000000..ec1980d5ec9adf2c32107bb6627761a1758c16e8 --- /dev/null +++ b/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) diff --git a/models/CtrlHair/color_texture_branch/predictor/predictor_model.py b/models/CtrlHair/color_texture_branch/predictor/predictor_model.py new file mode 100644 index 0000000000000000000000000000000000000000..c1b948b62daf474d95a175217f4af333e3516e42 --- /dev/null +++ b/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 diff --git a/models/CtrlHair/color_texture_branch/predictor/predictor_solver.py b/models/CtrlHair/color_texture_branch/predictor/predictor_solver.py new file mode 100644 index 0000000000000000000000000000000000000000..812f1469641eae4aa8dcd83adea96c5dd4eea76c --- /dev/null +++ b/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) diff --git a/models/CtrlHair/color_texture_branch/predictor/predictor_train.py b/models/CtrlHair/color_texture_branch/predictor/predictor_train.py new file mode 100644 index 0000000000000000000000000000000000000000..dab04b3012f5aaea1979fd852f3fbef4de13cceb --- /dev/null +++ b/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) diff --git a/models/CtrlHair/color_texture_branch/script_find_direction.py b/models/CtrlHair/color_texture_branch/script_find_direction.py new file mode 100644 index 0000000000000000000000000000000000000000..f8b99fc16656545ef9938a1870135a81fcdb9e67 --- /dev/null +++ b/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) diff --git a/models/CtrlHair/color_texture_branch/solver.py b/models/CtrlHair/color_texture_branch/solver.py new file mode 100644 index 0000000000000000000000000000000000000000..58b3ee17579729437c249cd171649be4bbd47f3c --- /dev/null +++ b/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 diff --git a/models/CtrlHair/color_texture_branch/train.py b/models/CtrlHair/color_texture_branch/train.py new file mode 100644 index 0000000000000000000000000000000000000000..9797b9063f94b848a5a748fdc1e05dda2e00ccc7 --- /dev/null +++ b/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) diff --git a/models/CtrlHair/color_texture_branch/validation_in_train.py b/models/CtrlHair/color_texture_branch/validation_in_train.py new file mode 100644 index 0000000000000000000000000000000000000000..8c2c9b460c3b091849d8b148e6f40caf6c0d5a52 --- /dev/null +++ b/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() diff --git a/models/CtrlHair/common_dataset.py b/models/CtrlHair/common_dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..c825ab92540c6c16a1fc63797c397293aecab210 --- /dev/null +++ b/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 + diff --git a/models/CtrlHair/dataset_scripts/__init__.py b/models/CtrlHair/dataset_scripts/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..6f506a5ea11d8716ff556d6868d67965bcb34bdf --- /dev/null +++ b/models/CtrlHair/dataset_scripts/__init__.py @@ -0,0 +1,8 @@ +# -*- coding: utf-8 -*- + +""" +# File name: __init__.py +# Time : 2022/07/15 +# Author: xyguoo@163.com +# Description: +""" diff --git a/models/CtrlHair/dataset_scripts/script_crop.py b/models/CtrlHair/dataset_scripts/script_crop.py new file mode 100644 index 0000000000000000000000000000000000000000..71fcb12d1642ebe5365d934f2406e4ed1f236afd --- /dev/null +++ b/models/CtrlHair/dataset_scripts/script_crop.py @@ -0,0 +1,46 @@ +# -*- coding: utf-8 -*- + +""" +# File name: script_crop.py +# Time : 2022/07/17 +# Author: xyguoo@163.com +# Description: +""" + +import sys +sys.path.append('.') + +from external_code.crop import recreate_aligned_images +from global_value_utils import GLOBAL_DATA_ROOT + +import os +from external_code.landmarks_util import predictor_dict, detector +import numpy as np +import cv2 + +predictor_68 = predictor_dict[68] + +############################################## +# Please input your dataset dir +root_dir = 'your/dataset/original/images' +dataset_name = 'your_dataset_name' +############################################## + +dataset_dir = os.path.join(GLOBAL_DATA_ROOT, dataset_name) +out_dir = os.path.join(dataset_dir, 'images_256') +if not os.path.exists(out_dir): + os.makedirs(out_dir) + +files = os.listdir(root_dir) +files.sort() +for face_path in files: + face_img_bgr = cv2.imread(os.path.join(root_dir, face_path)) + face_img_rgb = cv2.cvtColor(face_img_bgr, cv2.COLOR_BGR2RGB) + face_img_rgb = cv2.resize(face_img_rgb, dsize=(face_img_rgb.shape[1], face_img_rgb.shape[0])) + face_bbox = detector(face_img_rgb, 0) + face_lm_68 = np.array([[p.x, p.y] for p in predictor_68(face_img_bgr, face_bbox[0]).parts()]) + + face_img_pil, _ = recreate_aligned_images(face_img_rgb, face_lm_68, output_size=256) + + img_np = np.array(face_img_pil) + cv2.imwrite(os.path.join(out_dir, face_path), cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)) \ No newline at end of file diff --git a/models/CtrlHair/dataset_scripts/script_get_color_var_label.py b/models/CtrlHair/dataset_scripts/script_get_color_var_label.py new file mode 100644 index 0000000000000000000000000000000000000000..2c7fd6c6c5b57a1024cb9ee002418e4788c22ee4 --- /dev/null +++ b/models/CtrlHair/dataset_scripts/script_get_color_var_label.py @@ -0,0 +1,94 @@ +# -*- coding: utf-8 -*- + +""" +# File name: script_get_color_var_label.py +# Time : 2021/11/29 22:10 +# Author: xyguoo@163.com +# Description: +""" +import os +import sys +sys.path.append('.') + +from dataset_scripts.utils import merge_pickle_dir_to_dict +import cv2 +import tqdm + +import pickle +from common_dataset import DataFilter +from global_value_utils import HAIR_IDX, GLOBAL_DATA_ROOT, DATASET_NAME +import numpy as np + + +data_name = DATASET_NAME + +root_dir = GLOBAL_DATA_ROOT +imgs_sub_dir = 'images_256' +target_dir = os.path.join(root_dir, 'hair_info_all_dataset/color_var_stat') + +ds = DataFilter() +ds.total_list.sort() + +if not os.path.exists(target_dir): + os.makedirs(target_dir) + +path_list = [] +for d in data_name: + data_dir = os.path.join(root_dir, d, imgs_sub_dir) + path_list += [os.path.join(data_dir, pp) for pp in os.listdir(data_dir)] + +res_dict = {} +for img_path in tqdm.tqdm(path_list[:], total=len(path_list)): + for dd in data_name: + if img_path.find(dd) != -1: + dataset_name = dd + break + else: + raise NotImplementedError + base_name = os.path.basename(img_path) + hair_path = os.path.join(root_dir, dataset_name, imgs_sub_dir, base_name) + hair_img = cv2.cvtColor(cv2.imread(hair_path), cv2.COLOR_BGR2RGB) + hair_parsing = cv2.imread(os.path.join(root_dir, dataset_name, 'label', base_name), cv2.IMREAD_GRAYSCALE) + mask_img = cv2.resize(hair_parsing.astype('uint8'), hair_img.shape[:2], interpolation=cv2.INTER_NEAREST) + + hair_mask = (mask_img == HAIR_IDX).astype('uint8') + kernel13 = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, ksize=(19, 19)) + hair_mask = cv2.erode(hair_mask, kernel13, iterations=1) + + points = hair_img[hair_mask.astype('bool')] + + # points color + color_var_stat = [] + var_rgb = (points / 255).var(axis=0) + if len(points) > 5: + points_hsv = cv2.cvtColor(points[None, ...], cv2.COLOR_RGB2HSV) + points_hsv = points_hsv / np.array([180, 255, 255]) + var_hsv = points_hsv.var(axis=(0, 1)) + + points_hls = cv2.cvtColor(points[None, ...], cv2.COLOR_RGB2HLS) + points_hls = points_hls / np.array([180, 255, 255]) + var_hls = points_hls.var(axis=(0, 1)) + + points_lab = cv2.cvtColor(points[None, ...], cv2.COLOR_RGB2LAB) + points_lab = points_lab / np.array([255, 100, 100]) + var_lab = points_lab.var(axis=(0, 1)) + + points_yuv = cv2.cvtColor(points[None, ...], cv2.COLOR_RGB2YUV) + points_yuv = points_yuv / np.array([255, 200, 200]) + var_yuv = points_yuv.var(axis=(0, 1)) + color_var_stat = {'var_rgb': var_rgb, 'var_hsv': var_hsv, 'var_hls': var_hls, + 'var_yuv': var_yuv} + + from sklearn.decomposition import PCA + pca = PCA(n_components=2) + pca.fit(points) + pca_dim = pca.transform(points)[:, 0] + color_var_stat['var_pca'] = pca_dim.std() + color_var_stat['var_pca_mean'] = pca.mean_ + color_var_stat['var_pca_comp'] = pca.components_ + + target_file = os.path.join(target_dir, '%s___%s.pkl' % (dataset_name, base_name[:-4])) + with open(target_file, 'wb') as f: + pickle.dump(color_var_stat, f) + +merge_pickle_dir_to_dict(target_dir, os.path.join(root_dir, 'color_var_stat_dict.pkl')) \ No newline at end of file diff --git a/models/CtrlHair/dataset_scripts/script_get_mask.py b/models/CtrlHair/dataset_scripts/script_get_mask.py new file mode 100644 index 0000000000000000000000000000000000000000..114cea8adf40068e0efafaaddc0141bc3b37ff88 --- /dev/null +++ b/models/CtrlHair/dataset_scripts/script_get_mask.py @@ -0,0 +1,71 @@ +#!/usr/bin/python +# -*- encoding: utf-8 -*- + +import sys +sys.path.append('.') + +import tqdm + +import os +import os.path as osp +import numpy as np +import cv2 +from global_value_utils import GLOBAL_DATA_ROOT, PARSING_COLOR_LIST, DATASET_NAME +from util.imutil import read_rgb, write_rgb +from external_code.face_parsing.my_parsing_util import FaceParsing + + +data_name = [d for d in DATASET_NAME if d != 'CelebaMask_HQ'] + +def makedir(pat): + if not os.path.exists(pat): + os.makedirs(pat) + + +def vis_parsing_maps(im, parsing_anno, stride, save_im, save_path, img_path): + # Colors for all 20 parts + + label_path = os.path.join(save_path, 'label') + vis_path = os.path.join(save_path, 'vis') + makedir(pat=label_path) + makedir(pat=vis_path) + + im = np.array(im) + vis_im = im.copy().astype(np.uint8) + vis_parsing_anno = parsing_anno.copy().astype(np.uint8) + vis_parsing_anno = cv2.resize(vis_parsing_anno, None, fx=stride, fy=stride, interpolation=cv2.INTER_NEAREST) + vis_parsing_anno_color = np.zeros((vis_parsing_anno.shape[0], vis_parsing_anno.shape[1], 3)) + 255 + + num_of_class = np.max(vis_parsing_anno) + + img_path = img_path[:-4] + '.png' + + for pi in range(0, num_of_class + 1): + index = np.where(vis_parsing_anno == pi) + if len(index[0]) > 0: + vis_parsing_anno_color[index[0], index[1], :] = PARSING_COLOR_LIST[pi] + + vis_parsing_anno_color = vis_parsing_anno_color.astype(np.uint8) + vis_im = cv2.addWeighted(vis_im, 0.4, vis_parsing_anno_color, 0.6, 0) + + cv2.imwrite(os.path.join(label_path, img_path), vis_parsing_anno) + write_rgb(os.path.join(vis_path, img_path), vis_im) + + +def evaluate(respth, dspth): + if not os.path.exists(respth): + os.makedirs(respth) + + files = os.listdir(dspth) + files.sort() + for image_path in tqdm.tqdm(files): + parsing, image = FaceParsing.parsing_img(read_rgb(osp.join(dspth, image_path))) + parsing = FaceParsing.swap_parsing_label_to_celeba_mask(parsing) + vis_parsing_maps(image, parsing, stride=1, save_im=True, save_path=respth, img_path=image_path) + + +if __name__ == "__main__": + for dn in data_name: + input_dir = os.path.join(GLOBAL_DATA_ROOT, dn, 'images_256') + target_root_dir = os.path.join(GLOBAL_DATA_ROOT, dn) + evaluate(respth=target_root_dir, dspth=input_dir) diff --git a/models/CtrlHair/dataset_scripts/script_get_rgb_hsv_label.py b/models/CtrlHair/dataset_scripts/script_get_rgb_hsv_label.py new file mode 100644 index 0000000000000000000000000000000000000000..46d0361d3467b0633c2b949af4aa65f8c0980b26 --- /dev/null +++ b/models/CtrlHair/dataset_scripts/script_get_rgb_hsv_label.py @@ -0,0 +1,90 @@ +# -*- coding: utf-8 -*- + +""" +# File name: script_get_rgb_hsv_label.py +# Time : 2021/11/16 21:23 +# Author: xyguoo@163.com +# Description: +""" +import os +import sys + +sys.path.append('.') + +from dataset_scripts.utils import merge_pickle_dir_to_dict +import cv2 +import tqdm +import pickle +from common_dataset import DataFilter +import numpy as np +from global_value_utils import HAIR_IDX, GLOBAL_DATA_ROOT, DATASET_NAME + +data_name = DATASET_NAME + +root_dir = GLOBAL_DATA_ROOT +imgs_sub_dir = 'images_256' +target_dir = os.path.join(root_dir, 'hair_info_all_dataset/rgb_stat') + +ds = DataFilter() +ds.total_list.sort() + +path_list = [] +for d in data_name: + data_dir = os.path.join(root_dir, d, imgs_sub_dir) + path_list += [os.path.join(data_dir, pp) for pp in os.listdir(data_dir)] + +if not os.path.exists(target_dir): + os.makedirs(target_dir) + +print('Building RGB dict...') +res_dict = {} +for img_path in tqdm.tqdm(path_list[:], total=len(path_list)): + for dd in data_name: + if img_path.find(dd) != -1: + dataset_name = dd + break + else: + raise NotImplementedError + base_name = os.path.basename(img_path) + hair_path = os.path.join(root_dir, dataset_name, imgs_sub_dir, base_name) + hair_img = cv2.cvtColor(cv2.imread(hair_path), cv2.COLOR_BGR2RGB) + hair_parsing = cv2.imread(os.path.join(root_dir, dataset_name, 'label', base_name), cv2.IMREAD_GRAYSCALE) + mask_img = cv2.resize(hair_parsing.astype('uint8'), hair_img.shape[:2], interpolation=cv2.INTER_NEAREST) + + hair_mask = (mask_img == HAIR_IDX).astype('uint8') + kernel13 = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, ksize=(19, 19)) + hair_mask = cv2.erode(hair_mask, kernel13, iterations=1) + + points = hair_img[hair_mask.astype('bool')] + + moment1 = points.mean(axis=0) + moment2 = ((points - moment1) ** 2).mean(axis=0) + moment3 = ((points - moment1) ** 3).mean(axis=0) + moment4 = ((points - moment1) ** 4).mean(axis=0) + + target_file = os.path.join(target_dir, '%s___%s.pkl' % (dataset_name, base_name[:-4])) + with open(target_file, 'wb') as f: + pickle.dump([moment1, moment2, moment3, moment4], f) + +rgb_stat_path = os.path.join(root_dir, 'rgb_stat_dict.pkl') +merge_pickle_dir_to_dict(target_dir, os.path.join(root_dir, 'rgb_stat_dict.pkl')) + +### get hsv dict +print('Building ordered HSV dict for assisting distribution fitting...') +data_root = GLOBAL_DATA_ROOT +output_path = os.path.join(GLOBAL_DATA_ROOT, 'hsv_stat_dict_ordered.pkl') + +with open(rgb_stat_path, 'rb') as f: + rgb_stat_dict = pickle.load(f) + +files = list(rgb_stat_dict) + +cols = [rgb_stat_dict[f][0] for f in files] +cols = np.array(cols) +cols_hsv = cv2.cvtColor(cols[None, ...].astype('uint8'), cv2.COLOR_RGB2HSV)[0] + +for dim in range(3): + cols_hsv[:, dim].sort() + +with open(output_path, 'wb') as f: + pickle.dump(cols_hsv, f) diff --git a/models/CtrlHair/dataset_scripts/script_get_sean_code.py b/models/CtrlHair/dataset_scripts/script_get_sean_code.py new file mode 100644 index 0000000000000000000000000000000000000000..95dddb0cc799980f844a20db61cb4d9c0a6a954c --- /dev/null +++ b/models/CtrlHair/dataset_scripts/script_get_sean_code.py @@ -0,0 +1,62 @@ +# -*- coding: utf-8 -*- + +""" +# File name: script_get_sean_code.py +# Time : 2021/11/16 15:56 +# Author: xyguoo@163.com +# Description: +""" + +import os +import sys +sys.path.append('.') + +from dataset_scripts.utils import merge_pickle_dir_to_dict +import cv2 +import tqdm +from global_value_utils import GLOBAL_DATA_ROOT, DATASET_NAME + +from hair_editor import HairEditor +import pickle + +data_name = DATASET_NAME + +root_dir = GLOBAL_DATA_ROOT +imgs_sub_dir = 'images_256' +target_dir = os.path.join(root_dir, 'hair_info_all_dataset/sean_code') + +he = HairEditor(load_mask_model=False) +path_list = [] +for d in data_name: + data_dir = os.path.join(root_dir, d, imgs_sub_dir) + path_list += [os.path.join(data_dir, pp) for pp in os.listdir(data_dir)] + +path_list.sort() +# res_dict = {} + +if not os.path.exists(target_dir): + os.makedirs(target_dir) + +for img_path in tqdm.tqdm(path_list): + for dd in data_name: + if img_path.find(dd) != -1: + dataset_name = dd + break + else: + raise NotImplementedError + base_name = os.path.basename(img_path) + hair_path = os.path.join(root_dir, dataset_name, imgs_sub_dir, base_name) + hair_img = cv2.cvtColor(cv2.imread(hair_path), cv2.COLOR_BGR2RGB) + hair_parsing = cv2.imread(os.path.join(root_dir, dataset_name, 'label', base_name), cv2.IMREAD_GRAYSCALE) + # resize + hair_img = he.preprocess_img(hair_img) + hair_parsing = he.preprocess_mask(hair_parsing) + cur_code = he.get_code(hair_img, hair_parsing) + cur_code = cur_code.cpu().numpy()[0] + # res_dict['%s___%s' % (dataset_name, base_name)] = cur_code + + target_file = os.path.join(target_dir, '%s___%s.pkl' % (dataset_name, base_name[:-4])) + with open(target_file, 'wb') as f: + pickle.dump(cur_code, f) + +merge_pickle_dir_to_dict(target_dir, os.path.join(root_dir, 'sean_code_dict.pkl')) diff --git a/models/CtrlHair/dataset_scripts/script_landmark_detection.py b/models/CtrlHair/dataset_scripts/script_landmark_detection.py new file mode 100644 index 0000000000000000000000000000000000000000..a7978fd3a2943eeed5775d88d8d6c432cbd6c90f --- /dev/null +++ b/models/CtrlHair/dataset_scripts/script_landmark_detection.py @@ -0,0 +1,16 @@ +import sys +sys.path.append('.') + +from external_code.landmarks_util import detect_landmarks, predictor_dict +from global_value_utils import GLOBAL_DATA_ROOT, DATASET_NAME + + +data_name = DATASET_NAME + +root_dir = GLOBAL_DATA_ROOT + +for landmark_num in [81, 68]: + print('detect %d landmarks' % landmark_num) + detect_landmarks(root_dir, data_name, + landmark_output_file_path=root_dir + '/landmark%d.pkl' % landmark_num, + output_dir=None, predictor=predictor_dict[landmark_num]) diff --git a/models/CtrlHair/dataset_scripts/utils.py b/models/CtrlHair/dataset_scripts/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..a1675f905e95495d5f32207dee988bb98df8a6e3 --- /dev/null +++ b/models/CtrlHair/dataset_scripts/utils.py @@ -0,0 +1,21 @@ +# -*- coding: utf-8 -*- + +""" +# File name: utils.py +# Time : 2021/11/29 16:36 +# Author: xyguoo@163.com +# Description: +""" +import os +import pickle +import tqdm + + +def merge_pickle_dir_to_dict(dir_name, target_path): + files = os.listdir(dir_name) + res_dir = {} + for f_name in tqdm.tqdm(files): + with open(os.path.join(dir_name, f_name), 'rb') as f: + res_dir[f_name[:-4]] = pickle.load(f) + with open(target_path, 'wb') as f: + pickle.dump(res_dir, f) diff --git a/models/CtrlHair/external_code/__init__.py b/models/CtrlHair/external_code/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..db11b4ed69d71461922f4d9fc1dae30214df11ab --- /dev/null +++ b/models/CtrlHair/external_code/__init__.py @@ -0,0 +1,8 @@ +# -*- coding: utf-8 -*- + +""" +# File name: __init__.py +# Time : 2021/9/3 21:14 +# Author: xyguoo@163.com +# Description: +""" diff --git a/models/CtrlHair/external_code/crop.py b/models/CtrlHair/external_code/crop.py new file mode 100644 index 0000000000000000000000000000000000000000..dc3f7f4b76d4db5b82cc49c0a7608d075bc26318 --- /dev/null +++ b/models/CtrlHair/external_code/crop.py @@ -0,0 +1,117 @@ +# -*- coding: utf-8 -*- + +""" +# File name: crop.py +# Time : 2021/9/30 21:20 +# Author: xyguoo@163.com +# Description: +""" + +import os +import PIL.Image +import PIL.ImageFile +import numpy as np +import scipy.ndimage +import cv2 + +PIL.ImageFile.LOAD_TRUNCATED_IMAGES = True # avoid "Decompressed Data Too Large" error + + +def recreate_aligned_images(img, lm_68, output_size=1024, transform_size=4096, enable_padding=True): + pil_img = PIL.Image.fromarray(img) + lm_chin = lm_68[0: 17] # left-right + lm_eyebrow_left = lm_68[17: 22] # left-right + lm_eyebrow_right = lm_68[22: 27] # left-right + lm_nose = lm_68[27: 31] # top-down + lm_nostrils = lm_68[31: 36] # top-down + lm_eye_left = lm_68[36: 42] # left-clockwise + lm_eye_right = lm_68[42: 48] # left-clockwise + lm_mouth_outer = lm_68[48: 60] # left-clockwise + lm_mouth_inner = lm_68[60: 68] # left-clockwise + + # Calculate auxiliary vectors. + eye_left = np.mean(lm_eye_left, axis=0) + eye_right = np.mean(lm_eye_right, axis=0) + eye_avg = (eye_left + eye_right) * 0.5 + eye_to_eye = eye_right - eye_left + mouth_left = lm_mouth_outer[0] + mouth_right = lm_mouth_outer[6] + mouth_avg = (mouth_left + mouth_right) * 0.5 + eye_to_mouth = mouth_avg - eye_avg + + # Choose oriented crop rectangle. + x = eye_to_eye - np.flipud(eye_to_mouth) * [-1, 1] + x /= np.hypot(*x) + x *= max(np.hypot(*eye_to_eye) * 2.0, np.hypot(*eye_to_mouth) * 1.8) + y = np.flipud(x) * [-1, 1] + c = eye_avg + eye_to_mouth * 0.1 + quad = np.stack([c - x - y, c - x + y, c + x + y, c + x - y]) + qsize = np.hypot(*x) * 2 + + # Load in-the-wild image. + img = pil_img + + trans_points = lm_68 + + # Shrink. + shrink = int(np.floor(qsize / output_size * 0.5)) + if shrink > 1: + rsize = (int(np.rint(float(img.size[0]) / shrink)), int(np.rint(float(img.size[1]) / shrink))) + img = img.resize(rsize, PIL.Image.ANTIALIAS) + quad /= shrink + qsize /= shrink + trans_points = trans_points / shrink + + # Crop. + border = max(int(np.rint(qsize * 0.1)), 3) + crop = (int(np.floor(min(quad[:, 0]))), int(np.floor(min(quad[:, 1]))), int(np.ceil(max(quad[:, 0]))), + int(np.ceil(max(quad[:, 1])))) + crop = (max(crop[0] - border, 0), max(crop[1] - border, 0), min(crop[2] + border, img.size[0]), + min(crop[3] + border, img.size[1])) + if crop[2] - crop[0] < img.size[0] or crop[3] - crop[1] < img.size[1]: + img = img.crop(crop) + quad -= crop[0:2] + trans_points = trans_points - np.array([crop[0], crop[1]]) + + # Pad. + pad = (int(np.floor(min(quad[:, 0]))), int(np.floor(min(quad[:, 1]))), int(np.ceil(max(quad[:, 0]))), + int(np.ceil(max(quad[:, 1])))) + pad = (max(-pad[0] + border, 0), max(-pad[1] + border, 0), max(pad[2] - img.size[0] + border, 0), + max(pad[3] - img.size[1] + border, 0)) + if enable_padding and max(pad) > border - 4: + pad = np.maximum(pad, int(np.rint(qsize * 0.3))) + img = np.pad(np.float32(img), ((pad[1], pad[3]), (pad[0], pad[2]), (0, 0)), 'reflect') + trans_points = trans_points + np.array([pad[0], pad[1]]) + h, w, _ = img.shape + y, x, _ = np.ogrid[:h, :w, :1] + mask = np.maximum(1.0 - np.minimum(np.float32(x) / pad[0], np.float32(w - 1 - x) / pad[2]), + 1.0 - np.minimum(np.float32(y) / pad[1], np.float32(h - 1 - y) / pad[3])) + blur = qsize * 0.02 + img += (scipy.ndimage.gaussian_filter(img, [blur, blur, 0]) - img) * np.clip(mask * 3.0 + 1.0, 0.0, 1.0) + img += (np.median(img, axis=(0, 1)) - img) * np.clip(mask, 0.0, 1.0) + img = PIL.Image.fromarray(np.uint8(np.clip(np.rint(img), 0, 255)), 'RGB') + quad += pad[:2] + + # Transform. + trans_data = (quad + 0.5) + img = img.transform((transform_size, transform_size), PIL.Image.QUAD, trans_data.flatten(), PIL.Image.BILINEAR) + if output_size < transform_size: + img = img.resize((output_size, output_size), PIL.Image.ANTIALIAS) + + projective_matrix = cv2.getPerspectiveTransform(trans_data.astype('float32'), + np.array([[0, 0], [0, 1], [1, 1], [1, 0]], dtype='float32')) + augmented_lm = projective_matrix @ np.concatenate([trans_points, np.ones([trans_points.shape[0], 1])], axis=1).T + trans_points = augmented_lm[:2, :] / augmented_lm[2] * output_size + trans_points = trans_points.T + trans_points = (trans_points + 0.5).astype('int32') + return img, trans_points[:68] + + +def draw_landmarks(landmarks, img_np, font_size=1.0): + font = cv2.FONT_HERSHEY_SIMPLEX + for idx, point in enumerate(landmarks): + pos = (point[0], point[1]) + cv2.circle(img_np, pos, 2, color=(139, 0, 0)) + cv2.putText(img_np, str(idx + 1), pos, font, font_size, (0, 0, 255), 1, cv2.LINE_AA) + return img_np + diff --git a/models/CtrlHair/external_code/face_parsing/.gitignore b/models/CtrlHair/external_code/face_parsing/.gitignore new file mode 100644 index 0000000000000000000000000000000000000000..84f4b705b91b81ad5235418eea25eb1fa7669c9f --- /dev/null +++ b/models/CtrlHair/external_code/face_parsing/.gitignore @@ -0,0 +1,110 @@ +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*$py.class + +# C extensions +*.so + +# Distribution / packaging +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST + +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +.hypothesis/ +.pytest_cache/ + +# Translations +*.mo +*.pot + +# Django stuff: +*.log +local_settings.py +db.sqlite3 + +# Flask stuff: +instance/ +.webassets-cache + +# Scrapy stuff: +.scrapy + +# Sphinx documentation +docs/_build/ + +# PyBuilder +target/ + +# Jupyter Notebook +.ipynb_checkpoints + +# pyenv +.python-version + +# celery beat schedule file +celerybeat-schedule + +# SageMath parsed files +*.sage.py + +# Environments +.env +.venv +env/ +venv/ +ENV/ +env.bak/ +venv.bak/ + +# Spyder project settings +.spyderproject +.spyproject + +# Rope project settings +.ropeproject + +# mkdocs documentation +/site + +# mypy +.mypy_cache/ + +# resources +res/ + +.idea/ + diff --git a/models/CtrlHair/external_code/face_parsing/LICENSE b/models/CtrlHair/external_code/face_parsing/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..bfae0b0c29f885a118e382b445b6eaeca0d3b3e6 --- /dev/null +++ b/models/CtrlHair/external_code/face_parsing/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2019 zll + +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. diff --git a/models/CtrlHair/external_code/face_parsing/README.md b/models/CtrlHair/external_code/face_parsing/README.md new file mode 100644 index 0000000000000000000000000000000000000000..849d55e2789c8852e01707d1ff755dc74e63a7f5 --- /dev/null +++ b/models/CtrlHair/external_code/face_parsing/README.md @@ -0,0 +1,68 @@ +# face-parsing.PyTorch + +

+ + + +

+ +### Contents +- [Training](#training) +- [Demo](#Demo) +- [References](#references) + +## Training + +1. Prepare training data: + -- download [CelebAMask-HQ dataset](https://github.com/switchablenorms/CelebAMask-HQ) + + -- change file path in the `prepropess_data.py` and run +```Shell +python prepropess_data.py +``` + +2. Train the model using CelebAMask-HQ dataset: +Just run the train script: +``` + $ CUDA_VISIBLE_DEVICES=0,1 python -m torch.distributed.launch --nproc_per_node=2 train.py +``` + +If you do not wish to train the model, you can download [our pre-trained model](https://drive.google.com/open?id=154JgKpzCPW82qINcVieuPH3fZ2e0P812) and save it in `res/cp`. + + +## Demo +1. Evaluate the trained model using: +```Shell +# evaluate using GPU +python test.py +``` + +## Face makeup using parsing maps +[**face-makeup.PyTorch**](https://github.com/zllrunning/face-makeup.PyTorch) + + + + + + + + + + + + + + + + + + + + + + +
 HairLip
Original InputOriginal InputOriginal Input
ColorColorColor
+ + +## References +- [BiSeNet](https://github.com/CoinCheung/BiSeNet) \ No newline at end of file diff --git a/models/CtrlHair/external_code/face_parsing/__init__.py b/models/CtrlHair/external_code/face_parsing/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..6f506a5ea11d8716ff556d6868d67965bcb34bdf --- /dev/null +++ b/models/CtrlHair/external_code/face_parsing/__init__.py @@ -0,0 +1,8 @@ +# -*- coding: utf-8 -*- + +""" +# File name: __init__.py +# Time : 2022/07/15 +# Author: xyguoo@163.com +# Description: +""" diff --git a/models/CtrlHair/external_code/face_parsing/evaluate.py b/models/CtrlHair/external_code/face_parsing/evaluate.py new file mode 100644 index 0000000000000000000000000000000000000000..260d5ba9e52e5dfd108a06ef53e7a49e60ab4ede --- /dev/null +++ b/models/CtrlHair/external_code/face_parsing/evaluate.py @@ -0,0 +1,95 @@ +#!/usr/bin/python +# -*- encoding: utf-8 -*- + +from logger import setup_logger +from model import BiSeNet +from face_dataset import FaceMask + +import torch +import torch.nn as nn +from torch.utils.data import DataLoader +import torch.nn.functional as F +import torch.distributed as dist + +import os +import os.path as osp +import logging +import time +import numpy as np +from tqdm import tqdm +import math +from PIL import Image +import torchvision.transforms as transforms +import cv2 + +def vis_parsing_maps(im, parsing_anno, stride, save_im=False, save_path='vis_results/parsing_map_on_im.jpg'): + # Colors for all 20 parts + part_colors = [[255, 0, 0], [255, 85, 0], [255, 170, 0], + [255, 0, 85], [255, 0, 170], + [0, 255, 0], [85, 255, 0], [170, 255, 0], + [0, 255, 85], [0, 255, 170], + [0, 0, 255], [85, 0, 255], [170, 0, 255], + [0, 85, 255], [0, 170, 255], + [255, 255, 0], [255, 255, 85], [255, 255, 170], + [255, 0, 255], [255, 85, 255], [255, 170, 255], + [0, 255, 255], [85, 255, 255], [170, 255, 255]] + + im = np.array(im) + vis_im = im.copy().astype(np.uint8) + vis_parsing_anno = parsing_anno.copy().astype(np.uint8) + vis_parsing_anno = cv2.resize(vis_parsing_anno, None, fx=stride, fy=stride, interpolation=cv2.INTER_NEAREST) + vis_parsing_anno_color = np.zeros((vis_parsing_anno.shape[0], vis_parsing_anno.shape[1], 3)) + 255 + + num_of_class = np.max(vis_parsing_anno) + + for pi in range(1, num_of_class + 1): + index = np.where(vis_parsing_anno == pi) + vis_parsing_anno_color[index[0], index[1], :] = part_colors[pi] + + vis_parsing_anno_color = vis_parsing_anno_color.astype(np.uint8) + # print(vis_parsing_anno_color.shape, vis_im.shape) + vis_im = cv2.addWeighted(cv2.cvtColor(vis_im, cv2.COLOR_RGB2BGR), 0.4, vis_parsing_anno_color, 0.6, 0) + + # Save result or not + if save_im: + cv2.imwrite(save_path, vis_im, [int(cv2.IMWRITE_JPEG_QUALITY), 100]) + + # return vis_im + +def evaluate(respth='./resources/test_res', dspth='./data', cp='model_final_diss.pth'): + + if not os.path.exists(respth): + os.makedirs(respth) + + n_classes = 19 + net = BiSeNet(n_classes=n_classes) + net.cuda() + save_pth = osp.join('res/cp', cp) + net.load_state_dict(torch.load(save_pth)) + net.eval() + + to_tensor = transforms.Compose([ + transforms.ToTensor(), + transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)), + ]) + with torch.no_grad(): + for image_path in os.listdir(dspth): + img = Image.open(osp.join(dspth, image_path)) + image = img.resize((512, 512), Image.BILINEAR) + img = to_tensor(image) + img = torch.unsqueeze(img, 0) + img = img.cuda() + out = net(img)[0] + parsing = out.squeeze(0).cpu().numpy().argmax(0) + + vis_parsing_maps(image, parsing, stride=1, save_im=True, save_path=osp.join(respth, image_path)) + + + + + + + +if __name__ == "__main__": + setup_logger('./resources') + evaluate() diff --git a/models/CtrlHair/external_code/face_parsing/face_dataset.py b/models/CtrlHair/external_code/face_parsing/face_dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..0721105298e8e59410b570f9c2f4c9177ac51518 --- /dev/null +++ b/models/CtrlHair/external_code/face_parsing/face_dataset.py @@ -0,0 +1,106 @@ +#!/usr/bin/python +# -*- encoding: utf-8 -*- + +import torch +from torch.utils.data import Dataset +import torchvision.transforms as transforms + +import os.path as osp +import os +from PIL import Image +import numpy as np +import json +import cv2 + +from transform import * + + + +class FaceMask(Dataset): + def __init__(self, rootpth, cropsize=(640, 480), mode='scripts', *args, **kwargs): + super(FaceMask, self).__init__(*args, **kwargs) + assert mode in ('scripts', 'val', 'test') + self.mode = mode + self.ignore_lb = 255 + self.rootpth = rootpth + + self.imgs = os.listdir(os.path.join(self.rootpth, 'CelebA-HQ-img')) + + # pre-processing + self.to_tensor = transforms.Compose([ + transforms.ToTensor(), + transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)), + ]) + self.trans_train = Compose([ + ColorJitter( + brightness=0.5, + contrast=0.5, + saturation=0.5), + HorizontalFlip(), + RandomScale((0.75, 1.0, 1.25, 1.5, 1.75, 2.0)), + RandomCrop(cropsize) + ]) + + def __getitem__(self, idx): + impth = self.imgs[idx] + img = Image.open(osp.join(self.rootpth, 'CelebA-HQ-img', impth)) + img = img.resize((512, 512), Image.BILINEAR) + label = Image.open(osp.join(self.rootpth, 'mask', impth[:-3]+'png')).convert('P') + # print(np.unique(np.array(label))) + if self.mode == 'scripts': + im_lb = dict(im=img, lb=label) + im_lb = self.trans_train(im_lb) + img, label = im_lb['im'], im_lb['lb'] + img = self.to_tensor(img) + label = np.array(label).astype(np.int64)[np.newaxis, :] + return img, label + + def __len__(self): + return len(self.imgs) + + +if __name__ == "__main__": + face_data = '/home/zll/data/CelebAMask-HQ/CelebA-HQ-img' + face_sep_mask = '/home/zll/data/CelebAMask-HQ/CelebAMask-HQ-mask-anno' + mask_path = '/home/zll/data/CelebAMask-HQ/mask' + counter = 0 + total = 0 + for i in range(15): + # files = os.listdir(osp.join(face_sep_mask, str(i))) + + atts = ['skin', 'l_brow', 'r_brow', 'l_eye', 'r_eye', 'eye_g', 'l_ear', 'r_ear', 'ear_r', + 'nose', 'mouth', 'u_lip', 'l_lip', 'neck', 'neck_l', 'cloth', 'hair', 'hat'] + + for j in range(i*2000, (i+1)*2000): + + mask = np.zeros((512, 512)) + + for l, att in enumerate(atts, 1): + total += 1 + file_name = ''.join([str(j).rjust(5, '0'), '_', att, '.png']) + path = osp.join(face_sep_mask, str(i), file_name) + + if os.path.exists(path): + counter += 1 + sep_mask = np.array(Image.open(path).convert('P')) + # print(np.unique(sep_mask)) + + mask[sep_mask == 225] = l + cv2.imwrite('{}/{}.png'.format(mask_path, j), mask) + print(j) + + print(counter, total) + + + + + + + + + + + + + + diff --git a/models/CtrlHair/external_code/face_parsing/logger.py b/models/CtrlHair/external_code/face_parsing/logger.py new file mode 100644 index 0000000000000000000000000000000000000000..d3f9ddcc2cae221b4dd881d02404e848b5396f7e --- /dev/null +++ b/models/CtrlHair/external_code/face_parsing/logger.py @@ -0,0 +1,23 @@ +#!/usr/bin/python +# -*- encoding: utf-8 -*- + + +import os.path as osp +import time +import sys +import logging + +import torch.distributed as dist + + +def setup_logger(logpth): + logfile = 'BiSeNet-{}.log'.format(time.strftime('%Y-%m-%d-%H-%M-%S')) + logfile = osp.join(logpth, logfile) + FORMAT = '%(levelname)s %(filename)s(%(lineno)d): %(message)s' + log_level = logging.INFO + if dist.is_initialized() and not dist.get_rank()==0: + log_level = logging.ERROR + logging.basicConfig(level=log_level, format=FORMAT, filename=logfile) + logging.root.addHandler(logging.StreamHandler()) + + diff --git a/models/CtrlHair/external_code/face_parsing/loss.py b/models/CtrlHair/external_code/face_parsing/loss.py new file mode 100644 index 0000000000000000000000000000000000000000..f8f65aa05566853cb87678d97926bd03b911e166 --- /dev/null +++ b/models/CtrlHair/external_code/face_parsing/loss.py @@ -0,0 +1,75 @@ +#!/usr/bin/python +# -*- encoding: utf-8 -*- + + +import torch +import torch.nn as nn +import torch.nn.functional as F + +import numpy as np + + +class OhemCELoss(nn.Module): + def __init__(self, thresh, n_min, ignore_lb=255, *args, **kwargs): + super(OhemCELoss, self).__init__() + self.thresh = -torch.log(torch.tensor(thresh, dtype=torch.float)).cuda() + self.n_min = n_min + self.ignore_lb = ignore_lb + self.criteria = nn.CrossEntropyLoss(ignore_index=ignore_lb, reduction='none') + + def forward(self, logits, labels): + N, C, H, W = logits.size() + loss = self.criteria(logits, labels).view(-1) + loss, _ = torch.sort(loss, descending=True) + if loss[self.n_min] > self.thresh: + loss = loss[loss>self.thresh] + else: + loss = loss[:self.n_min] + return torch.mean(loss) + + +class SoftmaxFocalLoss(nn.Module): + def __init__(self, gamma, ignore_lb=255, *args, **kwargs): + super(SoftmaxFocalLoss, self).__init__() + self.gamma = gamma + self.nll = nn.NLLLoss(ignore_index=ignore_lb) + + def forward(self, logits, labels): + scores = F.softmax(logits, dim=1) + factor = torch.pow(1.-scores, self.gamma) + log_score = F.log_softmax(logits, dim=1) + log_score = factor * log_score + loss = self.nll(log_score, labels) + return loss + + +if __name__ == '__main__': + torch.manual_seed(15) + criteria1 = OhemCELoss(thresh=0.7, n_min=16*20*20//16).cuda() + criteria2 = OhemCELoss(thresh=0.7, n_min=16*20*20//16).cuda() + net1 = nn.Sequential( + nn.Conv2d(3, 19, kernel_size=3, stride=2, padding=1), + ) + net1.cuda() + net1.train() + net2 = nn.Sequential( + nn.Conv2d(3, 19, kernel_size=3, stride=2, padding=1), + ) + net2.cuda() + net2.train() + + with torch.no_grad(): + inten = torch.randn(16, 3, 20, 20).cuda() + lbs = torch.randint(0, 19, [16, 20, 20]).cuda() + lbs[1, :, :] = 255 + + logits1 = net1(inten) + logits1 = F.interpolate(logits1, inten.size()[2:], mode='bilinear') + logits2 = net2(inten) + logits2 = F.interpolate(logits2, inten.size()[2:], mode='bilinear') + + loss1 = criteria1(logits1, lbs) + loss2 = criteria2(logits2, lbs) + loss = loss1 + loss2 + print(loss.detach().cpu()) + loss.backward() diff --git a/models/CtrlHair/external_code/face_parsing/makeup.py b/models/CtrlHair/external_code/face_parsing/makeup.py new file mode 100644 index 0000000000000000000000000000000000000000..13ae9c0167d7a389be2988ea2c486065c95964fb --- /dev/null +++ b/models/CtrlHair/external_code/face_parsing/makeup.py @@ -0,0 +1,130 @@ +import cv2 +import os +import numpy as np +from skimage.filters import gaussian + + +def sharpen(img): + img = img * 1.0 + gauss_out = gaussian(img, sigma=5, multichannel=True) + + alpha = 1.5 + img_out = (img - gauss_out) * alpha + img + + img_out = img_out / 255.0 + + mask_1 = img_out < 0 + mask_2 = img_out > 1 + + img_out = img_out * (1 - mask_1) + img_out = img_out * (1 - mask_2) + mask_2 + img_out = np.clip(img_out, 0, 1) + img_out = img_out * 255 + return np.array(img_out, dtype=np.uint8) + + +def hair(image, parsing, part=17, color=[230, 50, 20]): + b, g, r = color #[10, 50, 250] # [10, 250, 10] + tar_color = np.zeros_like(image) + tar_color[:, :, 0] = b + tar_color[:, :, 1] = g + tar_color[:, :, 2] = r + + image_hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV) + tar_hsv = cv2.cvtColor(tar_color, cv2.COLOR_BGR2HSV) + + if part == 12 or part == 13: + image_hsv[:, :, 0:2] = tar_hsv[:, :, 0:2] + else: + image_hsv[:, :, 0:1] = tar_hsv[:, :, 0:1] + + changed = cv2.cvtColor(image_hsv, cv2.COLOR_HSV2BGR) + + if part == 17: + changed = sharpen(changed) + + changed[parsing != part] = image[parsing != part] + # changed = cv2.resize(changed, (512, 512)) + return changed + +# +# def lip(image, parsing, part=17, color=[230, 50, 20]): +# b, g, r = color #[10, 50, 250] # [10, 250, 10] +# tar_color = np.zeros_like(image) +# tar_color[:, :, 0] = b +# tar_color[:, :, 1] = g +# tar_color[:, :, 2] = r +# +# image_lab = cv2.cvtColor(image, cv2.COLOR_BGR2Lab) +# il, ia, ib = cv2.split(image_lab) +# +# tar_lab = cv2.cvtColor(tar_color, cv2.COLOR_BGR2Lab) +# tl, ta, tb = cv2.split(tar_lab) +# +# image_lab[:, :, 0] = np.clip(il - np.mean(il) + tl, 0, 100) +# image_lab[:, :, 1] = np.clip(ia - np.mean(ia) + ta, -127, 128) +# image_lab[:, :, 2] = np.clip(ib - np.mean(ib) + tb, -127, 128) +# +# +# changed = cv2.cvtColor(image_lab, cv2.COLOR_Lab2BGR) +# +# if part == 17: +# changed = sharpen(changed) +# +# changed[parsing != part] = image[parsing != part] +# # changed = cv2.resize(changed, (512, 512)) +# return changed + + +if __name__ == '__main__': + # 1 face + # 10 nose + # 11 teeth + # 12 upper lip + # 13 lower lip + # 17 hair + num = 116 + table = { + 'hair': 17, + 'upper_lip': 12, + 'lower_lip': 13 + } + image_path = '/home/zll/data/CelebAMask-HQ/test-img/{}.jpg'.format(num) + parsing_path = 'resources/test_res/{}.png'.format(num) + + image = cv2.imread(image_path) + ori = image.copy() + parsing = np.array(cv2.imread(parsing_path, 0)) + parsing = cv2.resize(parsing, image.shape[0:2], interpolation=cv2.INTER_NEAREST) + + parts = [table['hair'], table['upper_lip'], table['lower_lip']] + # colors = [[20, 20, 200], [100, 100, 230], [100, 100, 230]] + colors = [[100, 200, 100]] + for part, color in zip(parts, colors): + image = hair(image, parsing, part, color) + cv2.imwrite('res/makeup/116_ori.png', cv2.resize(ori, (512, 512))) + cv2.imwrite('res/makeup/116_2.png', cv2.resize(image, (512, 512))) + + cv2.imshow('image', cv2.resize(ori, (512, 512))) + cv2.imshow('color', cv2.resize(image, (512, 512))) + + # cv2.imshow('image', ori) + # cv2.imshow('color', image) + + cv2.waitKey(0) + cv2.destroyAllWindows() + + + + + + + + + + + + + + + diff --git a/models/CtrlHair/external_code/face_parsing/model.py b/models/CtrlHair/external_code/face_parsing/model.py new file mode 100644 index 0000000000000000000000000000000000000000..5119e751c3ae18e4dc1eecde7bfcb5bf9c62fb92 --- /dev/null +++ b/models/CtrlHair/external_code/face_parsing/model.py @@ -0,0 +1,283 @@ +#!/usr/bin/python +# -*- encoding: utf-8 -*- + + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torchvision + +from .resnet import Resnet18 +# from modules.bn import InPlaceABNSync as BatchNorm2d + + +class ConvBNReLU(nn.Module): + def __init__(self, in_chan, out_chan, ks=3, stride=1, padding=1, *args, **kwargs): + super(ConvBNReLU, self).__init__() + self.conv = nn.Conv2d(in_chan, + out_chan, + kernel_size = ks, + stride = stride, + padding = padding, + bias = False) + self.bn = nn.BatchNorm2d(out_chan) + self.init_weight() + + def forward(self, x): + x = self.conv(x) + x = F.relu(self.bn(x)) + return x + + def init_weight(self): + for ly in self.children(): + if isinstance(ly, nn.Conv2d): + nn.init.kaiming_normal_(ly.weight, a=1) + if not ly.bias is None: nn.init.constant_(ly.bias, 0) + +class BiSeNetOutput(nn.Module): + def __init__(self, in_chan, mid_chan, n_classes, *args, **kwargs): + super(BiSeNetOutput, self).__init__() + self.conv = ConvBNReLU(in_chan, mid_chan, ks=3, stride=1, padding=1) + self.conv_out = nn.Conv2d(mid_chan, n_classes, kernel_size=1, bias=False) + self.init_weight() + + def forward(self, x): + x = self.conv(x) + x = self.conv_out(x) + return x + + def init_weight(self): + for ly in self.children(): + if isinstance(ly, nn.Conv2d): + nn.init.kaiming_normal_(ly.weight, a=1) + if not ly.bias is None: nn.init.constant_(ly.bias, 0) + + def get_params(self): + wd_params, nowd_params = [], [] + for name, module in self.named_modules(): + if isinstance(module, nn.Linear) or isinstance(module, nn.Conv2d): + wd_params.append(module.weight) + if not module.bias is None: + nowd_params.append(module.bias) + elif isinstance(module, nn.BatchNorm2d): + nowd_params += list(module.parameters()) + return wd_params, nowd_params + + +class AttentionRefinementModule(nn.Module): + def __init__(self, in_chan, out_chan, *args, **kwargs): + super(AttentionRefinementModule, self).__init__() + self.conv = ConvBNReLU(in_chan, out_chan, ks=3, stride=1, padding=1) + self.conv_atten = nn.Conv2d(out_chan, out_chan, kernel_size= 1, bias=False) + self.bn_atten = nn.BatchNorm2d(out_chan) + self.sigmoid_atten = nn.Sigmoid() + self.init_weight() + + def forward(self, x): + feat = self.conv(x) + atten = F.avg_pool2d(feat, feat.size()[2:]) + atten = self.conv_atten(atten) + atten = self.bn_atten(atten) + atten = self.sigmoid_atten(atten) + out = torch.mul(feat, atten) + return out + + def init_weight(self): + for ly in self.children(): + if isinstance(ly, nn.Conv2d): + nn.init.kaiming_normal_(ly.weight, a=1) + if not ly.bias is None: nn.init.constant_(ly.bias, 0) + + +class ContextPath(nn.Module): + def __init__(self, *args, **kwargs): + super(ContextPath, self).__init__() + self.resnet = Resnet18() + self.arm16 = AttentionRefinementModule(256, 128) + self.arm32 = AttentionRefinementModule(512, 128) + self.conv_head32 = ConvBNReLU(128, 128, ks=3, stride=1, padding=1) + self.conv_head16 = ConvBNReLU(128, 128, ks=3, stride=1, padding=1) + self.conv_avg = ConvBNReLU(512, 128, ks=1, stride=1, padding=0) + + self.init_weight() + + def forward(self, x): + H0, W0 = x.size()[2:] + feat8, feat16, feat32 = self.resnet(x) + H8, W8 = feat8.size()[2:] + H16, W16 = feat16.size()[2:] + H32, W32 = feat32.size()[2:] + + avg = F.avg_pool2d(feat32, feat32.size()[2:]) + avg = self.conv_avg(avg) + avg_up = F.interpolate(avg, (H32, W32), mode='nearest') + + feat32_arm = self.arm32(feat32) + feat32_sum = feat32_arm + avg_up + feat32_up = F.interpolate(feat32_sum, (H16, W16), mode='nearest') + feat32_up = self.conv_head32(feat32_up) + + feat16_arm = self.arm16(feat16) + feat16_sum = feat16_arm + feat32_up + feat16_up = F.interpolate(feat16_sum, (H8, W8), mode='nearest') + feat16_up = self.conv_head16(feat16_up) + + return feat8, feat16_up, feat32_up # x8, x8, x16 + + def init_weight(self): + for ly in self.children(): + if isinstance(ly, nn.Conv2d): + nn.init.kaiming_normal_(ly.weight, a=1) + if not ly.bias is None: nn.init.constant_(ly.bias, 0) + + def get_params(self): + wd_params, nowd_params = [], [] + for name, module in self.named_modules(): + if isinstance(module, (nn.Linear, nn.Conv2d)): + wd_params.append(module.weight) + if not module.bias is None: + nowd_params.append(module.bias) + elif isinstance(module, nn.BatchNorm2d): + nowd_params += list(module.parameters()) + return wd_params, nowd_params + + +### This is not used, since I replace this with the resnet feature with the same size +class SpatialPath(nn.Module): + def __init__(self, *args, **kwargs): + super(SpatialPath, self).__init__() + self.conv1 = ConvBNReLU(3, 64, ks=7, stride=2, padding=3) + self.conv2 = ConvBNReLU(64, 64, ks=3, stride=2, padding=1) + self.conv3 = ConvBNReLU(64, 64, ks=3, stride=2, padding=1) + self.conv_out = ConvBNReLU(64, 128, ks=1, stride=1, padding=0) + self.init_weight() + + def forward(self, x): + feat = self.conv1(x) + feat = self.conv2(feat) + feat = self.conv3(feat) + feat = self.conv_out(feat) + return feat + + def init_weight(self): + for ly in self.children(): + if isinstance(ly, nn.Conv2d): + nn.init.kaiming_normal_(ly.weight, a=1) + if not ly.bias is None: nn.init.constant_(ly.bias, 0) + + def get_params(self): + wd_params, nowd_params = [], [] + for name, module in self.named_modules(): + if isinstance(module, nn.Linear) or isinstance(module, nn.Conv2d): + wd_params.append(module.weight) + if not module.bias is None: + nowd_params.append(module.bias) + elif isinstance(module, nn.BatchNorm2d): + nowd_params += list(module.parameters()) + return wd_params, nowd_params + + +class FeatureFusionModule(nn.Module): + def __init__(self, in_chan, out_chan, *args, **kwargs): + super(FeatureFusionModule, self).__init__() + self.convblk = ConvBNReLU(in_chan, out_chan, ks=1, stride=1, padding=0) + self.conv1 = nn.Conv2d(out_chan, + out_chan//4, + kernel_size = 1, + stride = 1, + padding = 0, + bias = False) + self.conv2 = nn.Conv2d(out_chan//4, + out_chan, + kernel_size = 1, + stride = 1, + padding = 0, + bias = False) + self.relu = nn.ReLU(inplace=True) + self.sigmoid = nn.Sigmoid() + self.init_weight() + + def forward(self, fsp, fcp): + fcat = torch.cat([fsp, fcp], dim=1) + feat = self.convblk(fcat) + atten = F.avg_pool2d(feat, feat.size()[2:]) + atten = self.conv1(atten) + atten = self.relu(atten) + atten = self.conv2(atten) + atten = self.sigmoid(atten) + feat_atten = torch.mul(feat, atten) + feat_out = feat_atten + feat + return feat_out + + def init_weight(self): + for ly in self.children(): + if isinstance(ly, nn.Conv2d): + nn.init.kaiming_normal_(ly.weight, a=1) + if not ly.bias is None: nn.init.constant_(ly.bias, 0) + + def get_params(self): + wd_params, nowd_params = [], [] + for name, module in self.named_modules(): + if isinstance(module, nn.Linear) or isinstance(module, nn.Conv2d): + wd_params.append(module.weight) + if not module.bias is None: + nowd_params.append(module.bias) + elif isinstance(module, nn.BatchNorm2d): + nowd_params += list(module.parameters()) + return wd_params, nowd_params + + +class BiSeNet(nn.Module): + def __init__(self, n_classes, *args, **kwargs): + super(BiSeNet, self).__init__() + self.cp = ContextPath() + ## here self.sp is deleted + self.ffm = FeatureFusionModule(256, 256) + self.conv_out = BiSeNetOutput(256, 256, n_classes) + self.conv_out16 = BiSeNetOutput(128, 64, n_classes) + self.conv_out32 = BiSeNetOutput(128, 64, n_classes) + self.init_weight() + + def forward(self, x): + H, W = x.size()[2:] + feat_res8, feat_cp8, feat_cp16 = self.cp(x) # here return res3b1 feature + feat_sp = feat_res8 # use res3b1 feature to replace spatial path feature + feat_fuse = self.ffm(feat_sp, feat_cp8) + + feat_out = self.conv_out(feat_fuse) + feat_out16 = self.conv_out16(feat_cp8) + feat_out32 = self.conv_out32(feat_cp16) + + feat_out = F.interpolate(feat_out, (H, W), mode='bilinear', align_corners=True) + feat_out16 = F.interpolate(feat_out16, (H, W), mode='bilinear', align_corners=True) + feat_out32 = F.interpolate(feat_out32, (H, W), mode='bilinear', align_corners=True) + return feat_out, feat_out16, feat_out32 + + def init_weight(self): + for ly in self.children(): + if isinstance(ly, nn.Conv2d): + nn.init.kaiming_normal_(ly.weight, a=1) + if not ly.bias is None: nn.init.constant_(ly.bias, 0) + + def get_params(self): + wd_params, nowd_params, lr_mul_wd_params, lr_mul_nowd_params = [], [], [], [] + for name, child in self.named_children(): + child_wd_params, child_nowd_params = child.get_params() + if isinstance(child, FeatureFusionModule) or isinstance(child, BiSeNetOutput): + lr_mul_wd_params += child_wd_params + lr_mul_nowd_params += child_nowd_params + else: + wd_params += child_wd_params + nowd_params += child_nowd_params + return wd_params, nowd_params, lr_mul_wd_params, lr_mul_nowd_params + + +if __name__ == "__main__": + net = BiSeNet(19) + net.cuda() + net.eval() + in_ten = torch.randn(16, 3, 640, 480).cuda() + out, out16, out32 = net(in_ten) + print(out.shape) + + net.get_params() diff --git a/models/CtrlHair/external_code/face_parsing/modules/__init__.py b/models/CtrlHair/external_code/face_parsing/modules/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..8a098dee5911f3613d320d23db37bc401cf57fa4 --- /dev/null +++ b/models/CtrlHair/external_code/face_parsing/modules/__init__.py @@ -0,0 +1,5 @@ +from .bn import ABN, InPlaceABN, InPlaceABNSync +from .functions import ACT_RELU, ACT_LEAKY_RELU, ACT_ELU, ACT_NONE +from .misc import GlobalAvgPool2d, SingleGPU +from .residual import IdentityResidualBlock +from .dense import DenseModule diff --git a/models/CtrlHair/external_code/face_parsing/modules/bn.py b/models/CtrlHair/external_code/face_parsing/modules/bn.py new file mode 100644 index 0000000000000000000000000000000000000000..da1f4b8a5df982c8407787af25a7def3809d6de3 --- /dev/null +++ b/models/CtrlHair/external_code/face_parsing/modules/bn.py @@ -0,0 +1,130 @@ +import torch +import torch.nn as nn +import torch.nn.functional as functional + +try: + from queue import Queue +except ImportError: + from Queue import Queue + +from .functions import * + + +class ABN(nn.Module): + """Activated Batch Normalization + + This gathers a `BatchNorm2d` and an activation function in a single module + """ + + def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=True, activation="leaky_relu", slope=0.01): + """Creates an Activated Batch Normalization module + + Parameters + ---------- + num_features : int + Number of feature channels in the input and color_texture. + eps : float + Small constant to prevent numerical issues. + momentum : float + Momentum factor applied to compute running statistics as. + affine : bool + If `True` apply learned scale and shift transformation after normalization. + activation : str + Name of the activation functions, one of: `leaky_relu`, `elu` or `none`. + slope : float + Negative slope for the `leaky_relu` activation. + """ + super(ABN, self).__init__() + self.num_features = num_features + self.affine = affine + self.eps = eps + self.momentum = momentum + self.activation = activation + self.slope = slope + if self.affine: + self.weight = nn.Parameter(torch.ones(num_features)) + self.bias = nn.Parameter(torch.zeros(num_features)) + else: + self.register_parameter('weight', None) + self.register_parameter('bias', None) + self.register_buffer('running_mean', torch.zeros(num_features)) + self.register_buffer('running_var', torch.ones(num_features)) + self.reset_parameters() + + def reset_parameters(self): + nn.init.constant_(self.running_mean, 0) + nn.init.constant_(self.running_var, 1) + if self.affine: + nn.init.constant_(self.weight, 1) + nn.init.constant_(self.bias, 0) + + def forward(self, x): + x = functional.batch_norm(x, self.running_mean, self.running_var, self.weight, self.bias, + self.training, self.momentum, self.eps) + + if self.activation == ACT_RELU: + return functional.relu(x, inplace=True) + elif self.activation == ACT_LEAKY_RELU: + return functional.leaky_relu(x, negative_slope=self.slope, inplace=True) + elif self.activation == ACT_ELU: + return functional.elu(x, inplace=True) + else: + return x + + def __repr__(self): + rep = '{name}({num_features}, eps={eps}, momentum={momentum},' \ + ' affine={affine}, activation={activation}' + if self.activation == "leaky_relu": + rep += ', slope={slope})' + else: + rep += ')' + return rep.format(name=self.__class__.__name__, **self.__dict__) + + +class InPlaceABN(ABN): + """InPlace Activated Batch Normalization""" + + def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=True, activation="leaky_relu", slope=0.01): + """Creates an InPlace Activated Batch Normalization module + + Parameters + ---------- + num_features : int + Number of feature channels in the input and color_texture. + eps : float + Small constant to prevent numerical issues. + momentum : float + Momentum factor applied to compute running statistics as. + affine : bool + If `True` apply learned scale and shift transformation after normalization. + activation : str + Name of the activation functions, one of: `leaky_relu`, `elu` or `none`. + slope : float + Negative slope for the `leaky_relu` activation. + """ + super(InPlaceABN, self).__init__(num_features, eps, momentum, affine, activation, slope) + + def forward(self, x): + return inplace_abn(x, self.weight, self.bias, self.running_mean, self.running_var, + self.training, self.momentum, self.eps, self.activation, self.slope) + + +class InPlaceABNSync(ABN): + """InPlace Activated Batch Normalization with cross-GPU synchronization + This assumes that it will be replicated across GPUs using the same mechanism as in `nn.DistributedDataParallel`. + """ + + def forward(self, x): + return inplace_abn_sync(x, self.weight, self.bias, self.running_mean, self.running_var, + self.training, self.momentum, self.eps, self.activation, self.slope) + + def __repr__(self): + rep = '{name}({num_features}, eps={eps}, momentum={momentum},' \ + ' affine={affine}, activation={activation}' + if self.activation == "leaky_relu": + rep += ', slope={slope})' + else: + rep += ')' + return rep.format(name=self.__class__.__name__, **self.__dict__) + + diff --git a/models/CtrlHair/external_code/face_parsing/modules/deeplab.py b/models/CtrlHair/external_code/face_parsing/modules/deeplab.py new file mode 100644 index 0000000000000000000000000000000000000000..fd25b78369b27ef02c183a0b17b9bf8354c5f7c3 --- /dev/null +++ b/models/CtrlHair/external_code/face_parsing/modules/deeplab.py @@ -0,0 +1,84 @@ +import torch +import torch.nn as nn +import torch.nn.functional as functional + +from models._util import try_index +from .bn import ABN + + +class DeeplabV3(nn.Module): + def __init__(self, + in_channels, + out_channels, + hidden_channels=256, + dilations=(12, 24, 36), + norm_act=ABN, + pooling_size=None): + super(DeeplabV3, self).__init__() + self.pooling_size = pooling_size + + self.map_convs = nn.ModuleList([ + nn.Conv2d(in_channels, hidden_channels, 1, bias=False), + nn.Conv2d(in_channels, hidden_channels, 3, bias=False, dilation=dilations[0], padding=dilations[0]), + nn.Conv2d(in_channels, hidden_channels, 3, bias=False, dilation=dilations[1], padding=dilations[1]), + nn.Conv2d(in_channels, hidden_channels, 3, bias=False, dilation=dilations[2], padding=dilations[2]) + ]) + self.map_bn = norm_act(hidden_channels * 4) + + self.global_pooling_conv = nn.Conv2d(in_channels, hidden_channels, 1, bias=False) + self.global_pooling_bn = norm_act(hidden_channels) + + self.red_conv = nn.Conv2d(hidden_channels * 4, out_channels, 1, bias=False) + self.pool_red_conv = nn.Conv2d(hidden_channels, out_channels, 1, bias=False) + self.red_bn = norm_act(out_channels) + + self.reset_parameters(self.map_bn.activation, self.map_bn.slope) + + def reset_parameters(self, activation, slope): + gain = nn.init.calculate_gain(activation, slope) + for m in self.modules(): + if isinstance(m, nn.Conv2d): + nn.init.xavier_normal_(m.weight.data, gain) + if hasattr(m, "bias") and m.bias is not None: + nn.init.constant_(m.bias, 0) + elif isinstance(m, ABN): + if hasattr(m, "weight") and m.weight is not None: + nn.init.constant_(m.weight, 1) + if hasattr(m, "bias") and m.bias is not None: + nn.init.constant_(m.bias, 0) + + def forward(self, x): + # Map convolutions + out = torch.cat([m(x) for m in self.map_convs], dim=1) + out = self.map_bn(out) + out = self.red_conv(out) + + # Global pooling + pool = self._global_pooling(x) + pool = self.global_pooling_conv(pool) + pool = self.global_pooling_bn(pool) + pool = self.pool_red_conv(pool) + if self.training or self.pooling_size is None: + pool = pool.repeat(1, 1, x.size(2), x.size(3)) + + out += pool + out = self.red_bn(out) + return out + + def _global_pooling(self, x): + if self.training or self.pooling_size is None: + pool = x.view(x.size(0), x.size(1), -1).mean(dim=-1) + pool = pool.view(x.size(0), x.size(1), 1, 1) + else: + pooling_size = (min(try_index(self.pooling_size, 0), x.shape[2]), + min(try_index(self.pooling_size, 1), x.shape[3])) + padding = ( + (pooling_size[1] - 1) // 2, + (pooling_size[1] - 1) // 2 if pooling_size[1] % 2 == 1 else (pooling_size[1] - 1) // 2 + 1, + (pooling_size[0] - 1) // 2, + (pooling_size[0] - 1) // 2 if pooling_size[0] % 2 == 1 else (pooling_size[0] - 1) // 2 + 1 + ) + + pool = functional.avg_pool2d(x, pooling_size, stride=1) + pool = functional.pad(pool, pad=padding, mode="replicate") + return pool diff --git a/models/CtrlHair/external_code/face_parsing/modules/dense.py b/models/CtrlHair/external_code/face_parsing/modules/dense.py new file mode 100644 index 0000000000000000000000000000000000000000..9638d6e86d2ae838550fefa9002a984af52e6cc8 --- /dev/null +++ b/models/CtrlHair/external_code/face_parsing/modules/dense.py @@ -0,0 +1,42 @@ +from collections import OrderedDict + +import torch +import torch.nn as nn + +from .bn import ABN + + +class DenseModule(nn.Module): + def __init__(self, in_channels, growth, layers, bottleneck_factor=4, norm_act=ABN, dilation=1): + super(DenseModule, self).__init__() + self.in_channels = in_channels + self.growth = growth + self.layers = layers + + self.convs1 = nn.ModuleList() + self.convs3 = nn.ModuleList() + for i in range(self.layers): + self.convs1.append(nn.Sequential(OrderedDict([ + ("bn", norm_act(in_channels)), + ("conv", nn.Conv2d(in_channels, self.growth * bottleneck_factor, 1, bias=False)) + ]))) + self.convs3.append(nn.Sequential(OrderedDict([ + ("bn", norm_act(self.growth * bottleneck_factor)), + ("conv", nn.Conv2d(self.growth * bottleneck_factor, self.growth, 3, padding=dilation, bias=False, + dilation=dilation)) + ]))) + in_channels += self.growth + + @property + def out_channels(self): + return self.in_channels + self.growth * self.layers + + def forward(self, x): + inputs = [x] + for i in range(self.layers): + x = torch.cat(inputs, dim=1) + x = self.convs1[i](x) + x = self.convs3[i](x) + inputs += [x] + + return torch.cat(inputs, dim=1) diff --git a/models/CtrlHair/external_code/face_parsing/modules/functions.py b/models/CtrlHair/external_code/face_parsing/modules/functions.py new file mode 100644 index 0000000000000000000000000000000000000000..093615ff4f383e95712c96b57286338ec3b28f3b --- /dev/null +++ b/models/CtrlHair/external_code/face_parsing/modules/functions.py @@ -0,0 +1,234 @@ +from os import path +import torch +import torch.distributed as dist +import torch.autograd as autograd +import torch.cuda.comm as comm +from torch.autograd.function import once_differentiable +from torch.utils.cpp_extension import load + +_src_path = path.join(path.dirname(path.abspath(__file__)), "src") +_backend = load(name="inplace_abn", + extra_cflags=["-O3"], + sources=[path.join(_src_path, f) for f in [ + "inplace_abn.cpp", + "inplace_abn_cpu.cpp", + "inplace_abn_cuda.cu", + "inplace_abn_cuda_half.cu" + ]], + extra_cuda_cflags=["--expt-extended-lambda"]) + +# Activation names +ACT_RELU = "relu" +ACT_LEAKY_RELU = "leaky_relu" +ACT_ELU = "elu" +ACT_NONE = "none" + + +def _check(fn, *args, **kwargs): + success = fn(*args, **kwargs) + if not success: + raise RuntimeError("CUDA Error encountered in {}".format(fn)) + + +def _broadcast_shape(x): + out_size = [] + for i, s in enumerate(x.size()): + if i != 1: + out_size.append(1) + else: + out_size.append(s) + return out_size + + +def _reduce(x): + if len(x.size()) == 2: + return x.sum(dim=0) + else: + n, c = x.size()[0:2] + return x.contiguous().view((n, c, -1)).sum(2).sum(0) + + +def _count_samples(x): + count = 1 + for i, s in enumerate(x.size()): + if i != 1: + count *= s + return count + + +def _act_forward(ctx, x): + if ctx.activation == ACT_LEAKY_RELU: + _backend.leaky_relu_forward(x, ctx.slope) + elif ctx.activation == ACT_ELU: + _backend.elu_forward(x) + elif ctx.activation == ACT_NONE: + pass + + +def _act_backward(ctx, x, dx): + if ctx.activation == ACT_LEAKY_RELU: + _backend.leaky_relu_backward(x, dx, ctx.slope) + elif ctx.activation == ACT_ELU: + _backend.elu_backward(x, dx) + elif ctx.activation == ACT_NONE: + pass + + +class InPlaceABN(autograd.Function): + @staticmethod + def forward(ctx, x, weight, bias, running_mean, running_var, + training=True, momentum=0.1, eps=1e-05, activation=ACT_LEAKY_RELU, slope=0.01): + # Save context + ctx.training = training + ctx.momentum = momentum + ctx.eps = eps + ctx.activation = activation + ctx.slope = slope + ctx.affine = weight is not None and bias is not None + + # Prepare inputs + count = _count_samples(x) + x = x.contiguous() + weight = weight.contiguous() if ctx.affine else x.new_empty(0) + bias = bias.contiguous() if ctx.affine else x.new_empty(0) + + if ctx.training: + mean, var = _backend.mean_var(x) + + # Update running stats + running_mean.mul_((1 - ctx.momentum)).add_(ctx.momentum * mean) + running_var.mul_((1 - ctx.momentum)).add_(ctx.momentum * var * count / (count - 1)) + + # Mark in-place modified tensors + ctx.mark_dirty(x, running_mean, running_var) + else: + mean, var = running_mean.contiguous(), running_var.contiguous() + ctx.mark_dirty(x) + + # BN forward + activation + _backend.forward(x, mean, var, weight, bias, ctx.affine, ctx.eps) + _act_forward(ctx, x) + + # Output + ctx.var = var + ctx.save_for_backward(x, var, weight, bias) + return x + + @staticmethod + @once_differentiable + def backward(ctx, dz): + z, var, weight, bias = ctx.saved_tensors + dz = dz.contiguous() + + # Undo activation + _act_backward(ctx, z, dz) + + if ctx.training: + edz, eydz = _backend.edz_eydz(z, dz, weight, bias, ctx.affine, ctx.eps) + else: + # TODO: implement simplified CUDA backward for inference mode + edz = dz.new_zeros(dz.size(1)) + eydz = dz.new_zeros(dz.size(1)) + + dx = _backend.backward(z, dz, var, weight, bias, edz, eydz, ctx.affine, ctx.eps) + dweight = eydz * weight.sign() if ctx.affine else None + dbias = edz if ctx.affine else None + + return dx, dweight, dbias, None, None, None, None, None, None, None + +class InPlaceABNSync(autograd.Function): + @classmethod + def forward(cls, ctx, x, weight, bias, running_mean, running_var, + training=True, momentum=0.1, eps=1e-05, activation=ACT_LEAKY_RELU, slope=0.01, equal_batches=True): + # Save context + ctx.training = training + ctx.momentum = momentum + ctx.eps = eps + ctx.activation = activation + ctx.slope = slope + ctx.affine = weight is not None and bias is not None + + # Prepare inputs + ctx.world_size = dist.get_world_size() if dist.is_initialized() else 1 + + #count = _count_samples(x) + batch_size = x.new_tensor([x.shape[0]],dtype=torch.long) + + x = x.contiguous() + weight = weight.contiguous() if ctx.affine else x.new_empty(0) + bias = bias.contiguous() if ctx.affine else x.new_empty(0) + + if ctx.training: + mean, var = _backend.mean_var(x) + if ctx.world_size>1: + # get global batch size + if equal_batches: + batch_size *= ctx.world_size + else: + dist.all_reduce(batch_size, dist.ReduceOp.SUM) + + ctx.factor = x.shape[0]/float(batch_size.item()) + + mean_all = mean.clone() * ctx.factor + dist.all_reduce(mean_all, dist.ReduceOp.SUM) + + var_all = (var + (mean - mean_all) ** 2) * ctx.factor + dist.all_reduce(var_all, dist.ReduceOp.SUM) + + mean = mean_all + var = var_all + + # Update running stats + running_mean.mul_((1 - ctx.momentum)).add_(ctx.momentum * mean) + count = batch_size.item() * x.view(x.shape[0],x.shape[1],-1).shape[-1] + running_var.mul_((1 - ctx.momentum)).add_(ctx.momentum * var * (float(count) / (count - 1))) + + # Mark in-place modified tensors + ctx.mark_dirty(x, running_mean, running_var) + else: + mean, var = running_mean.contiguous(), running_var.contiguous() + ctx.mark_dirty(x) + + # BN forward + activation + _backend.forward(x, mean, var, weight, bias, ctx.affine, ctx.eps) + _act_forward(ctx, x) + + # Output + ctx.var = var + ctx.save_for_backward(x, var, weight, bias) + return x + + @staticmethod + @once_differentiable + def backward(ctx, dz): + z, var, weight, bias = ctx.saved_tensors + dz = dz.contiguous() + + # Undo activation + _act_backward(ctx, z, dz) + + if ctx.training: + edz, eydz = _backend.edz_eydz(z, dz, weight, bias, ctx.affine, ctx.eps) + edz_local = edz.clone() + eydz_local = eydz.clone() + + if ctx.world_size>1: + edz *= ctx.factor + dist.all_reduce(edz, dist.ReduceOp.SUM) + + eydz *= ctx.factor + dist.all_reduce(eydz, dist.ReduceOp.SUM) + else: + edz_local = edz = dz.new_zeros(dz.size(1)) + eydz_local = eydz = dz.new_zeros(dz.size(1)) + + dx = _backend.backward(z, dz, var, weight, bias, edz, eydz, ctx.affine, ctx.eps) + dweight = eydz_local * weight.sign() if ctx.affine else None + dbias = edz_local if ctx.affine else None + + return dx, dweight, dbias, None, None, None, None, None, None, None + +inplace_abn = InPlaceABN.apply +inplace_abn_sync = InPlaceABNSync.apply + +__all__ = ["inplace_abn", "inplace_abn_sync", "ACT_RELU", "ACT_LEAKY_RELU", "ACT_ELU", "ACT_NONE"] diff --git a/models/CtrlHair/external_code/face_parsing/modules/misc.py b/models/CtrlHair/external_code/face_parsing/modules/misc.py new file mode 100644 index 0000000000000000000000000000000000000000..3c50b69b38c950801baacba8b3684ffd23aef08b --- /dev/null +++ b/models/CtrlHair/external_code/face_parsing/modules/misc.py @@ -0,0 +1,21 @@ +import torch.nn as nn +import torch +import torch.distributed as dist + +class GlobalAvgPool2d(nn.Module): + def __init__(self): + """Global average pooling over the input's spatial dimensions""" + super(GlobalAvgPool2d, self).__init__() + + def forward(self, inputs): + in_size = inputs.size() + return inputs.view((in_size[0], in_size[1], -1)).mean(dim=2) + +class SingleGPU(nn.Module): + def __init__(self, module): + super(SingleGPU, self).__init__() + self.module=module + + def forward(self, input): + return self.module(input.cuda(non_blocking=True)) + diff --git a/models/CtrlHair/external_code/face_parsing/modules/residual.py b/models/CtrlHair/external_code/face_parsing/modules/residual.py new file mode 100644 index 0000000000000000000000000000000000000000..b7d51ad274f3841813c1584a0ceb60ce58979d94 --- /dev/null +++ b/models/CtrlHair/external_code/face_parsing/modules/residual.py @@ -0,0 +1,88 @@ +from collections import OrderedDict + +import torch.nn as nn + +from .bn import ABN + + +class IdentityResidualBlock(nn.Module): + def __init__(self, + in_channels, + channels, + stride=1, + dilation=1, + groups=1, + norm_act=ABN, + dropout=None): + """Configurable identity-mapping residual block + + Parameters + ---------- + in_channels : int + Number of input channels. + channels : list of int + Number of channels in the internal feature maps. Can either have two or three elements: if three construct + a residual block with two `3 x 3` convolutions, otherwise construct a bottleneck block with `1 x 1`, then + `3 x 3` then `1 x 1` convolutions. + stride : int + Stride of the first `3 x 3` convolution + dilation : int + Dilation to apply to the `3 x 3` convolutions. + groups : int + Number of convolution groups. This is used to create ResNeXt-style blocks and is only compatible with + bottleneck blocks. + norm_act : callable + Function to create normalization / activation Module. + dropout: callable + Function to create Dropout Module. + """ + super(IdentityResidualBlock, self).__init__() + + # Check parameters for inconsistencies + if len(channels) != 2 and len(channels) != 3: + raise ValueError("channels must contain either two or three values") + if len(channels) == 2 and groups != 1: + raise ValueError("groups > 1 are only valid if len(channels) == 3") + + is_bottleneck = len(channels) == 3 + need_proj_conv = stride != 1 or in_channels != channels[-1] + + self.bn1 = norm_act(in_channels) + if not is_bottleneck: + layers = [ + ("conv1", nn.Conv2d(in_channels, channels[0], 3, stride=stride, padding=dilation, bias=False, + dilation=dilation)), + ("bn2", norm_act(channels[0])), + ("conv2", nn.Conv2d(channels[0], channels[1], 3, stride=1, padding=dilation, bias=False, + dilation=dilation)) + ] + if dropout is not None: + layers = layers[0:2] + [("dropout", dropout())] + layers[2:] + else: + layers = [ + ("conv1", nn.Conv2d(in_channels, channels[0], 1, stride=stride, padding=0, bias=False)), + ("bn2", norm_act(channels[0])), + ("conv2", nn.Conv2d(channels[0], channels[1], 3, stride=1, padding=dilation, bias=False, + groups=groups, dilation=dilation)), + ("bn3", norm_act(channels[1])), + ("conv3", nn.Conv2d(channels[1], channels[2], 1, stride=1, padding=0, bias=False)) + ] + if dropout is not None: + layers = layers[0:4] + [("dropout", dropout())] + layers[4:] + self.convs = nn.Sequential(OrderedDict(layers)) + + if need_proj_conv: + self.proj_conv = nn.Conv2d(in_channels, channels[-1], 1, stride=stride, padding=0, bias=False) + + def forward(self, x): + if hasattr(self, "proj_conv"): + bn1 = self.bn1(x) + shortcut = self.proj_conv(bn1) + else: + shortcut = x.clone() + bn1 = self.bn1(x) + + out = self.convs(bn1) + out.add_(shortcut) + + return out diff --git a/models/CtrlHair/external_code/face_parsing/modules/src/checks.h b/models/CtrlHair/external_code/face_parsing/modules/src/checks.h new file mode 100644 index 0000000000000000000000000000000000000000..e761a6fe34d0789815b588eba7e3726026e0e868 --- /dev/null +++ b/models/CtrlHair/external_code/face_parsing/modules/src/checks.h @@ -0,0 +1,15 @@ +#pragma once + +#include + +// Define AT_CHECK for old version of ATen where the same function was called AT_ASSERT +#ifndef AT_CHECK +#define AT_CHECK AT_ASSERT +#endif + +#define CHECK_CUDA(x) AT_CHECK((x).type().is_cuda(), #x " must be a CUDA tensor") +#define CHECK_CPU(x) AT_CHECK(!(x).type().is_cuda(), #x " must be a CPU tensor") +#define CHECK_CONTIGUOUS(x) AT_CHECK((x).is_contiguous(), #x " must be contiguous") + +#define CHECK_CUDA_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x) +#define CHECK_CPU_INPUT(x) CHECK_CPU(x); CHECK_CONTIGUOUS(x) \ No newline at end of file diff --git a/models/CtrlHair/external_code/face_parsing/modules/src/inplace_abn.cpp b/models/CtrlHair/external_code/face_parsing/modules/src/inplace_abn.cpp new file mode 100644 index 0000000000000000000000000000000000000000..0a6b1128cc20cbfc476134154e23e5869a92b856 --- /dev/null +++ b/models/CtrlHair/external_code/face_parsing/modules/src/inplace_abn.cpp @@ -0,0 +1,95 @@ +#include + +#include + +#include "inplace_abn.h" + +std::vector mean_var(at::Tensor x) { + if (x.is_cuda()) { + if (x.type().scalarType() == at::ScalarType::Half) { + return mean_var_cuda_h(x); + } else { + return mean_var_cuda(x); + } + } else { + return mean_var_cpu(x); + } +} + +at::Tensor forward(at::Tensor x, at::Tensor mean, at::Tensor var, at::Tensor weight, at::Tensor bias, + bool affine, float eps) { + if (x.is_cuda()) { + if (x.type().scalarType() == at::ScalarType::Half) { + return forward_cuda_h(x, mean, var, weight, bias, affine, eps); + } else { + return forward_cuda(x, mean, var, weight, bias, affine, eps); + } + } else { + return forward_cpu(x, mean, var, weight, bias, affine, eps); + } +} + +std::vector edz_eydz(at::Tensor z, at::Tensor dz, at::Tensor weight, at::Tensor bias, + bool affine, float eps) { + if (z.is_cuda()) { + if (z.type().scalarType() == at::ScalarType::Half) { + return edz_eydz_cuda_h(z, dz, weight, bias, affine, eps); + } else { + return edz_eydz_cuda(z, dz, weight, bias, affine, eps); + } + } else { + return edz_eydz_cpu(z, dz, weight, bias, affine, eps); + } +} + +at::Tensor backward(at::Tensor z, at::Tensor dz, at::Tensor var, at::Tensor weight, at::Tensor bias, + at::Tensor edz, at::Tensor eydz, bool affine, float eps) { + if (z.is_cuda()) { + if (z.type().scalarType() == at::ScalarType::Half) { + return backward_cuda_h(z, dz, var, weight, bias, edz, eydz, affine, eps); + } else { + return backward_cuda(z, dz, var, weight, bias, edz, eydz, affine, eps); + } + } else { + return backward_cpu(z, dz, var, weight, bias, edz, eydz, affine, eps); + } +} + +void leaky_relu_forward(at::Tensor z, float slope) { + at::leaky_relu_(z, slope); +} + +void leaky_relu_backward(at::Tensor z, at::Tensor dz, float slope) { + if (z.is_cuda()) { + if (z.type().scalarType() == at::ScalarType::Half) { + return leaky_relu_backward_cuda_h(z, dz, slope); + } else { + return leaky_relu_backward_cuda(z, dz, slope); + } + } else { + return leaky_relu_backward_cpu(z, dz, slope); + } +} + +void elu_forward(at::Tensor z) { + at::elu_(z); +} + +void elu_backward(at::Tensor z, at::Tensor dz) { + if (z.is_cuda()) { + return elu_backward_cuda(z, dz); + } else { + return elu_backward_cpu(z, dz); + } +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("mean_var", &mean_var, "Mean and variance computation"); + m.def("forward", &forward, "In-place forward computation"); + m.def("edz_eydz", &edz_eydz, "First part of backward computation"); + m.def("backward", &backward, "Second part of backward computation"); + m.def("leaky_relu_forward", &leaky_relu_forward, "Leaky relu forward computation"); + m.def("leaky_relu_backward", &leaky_relu_backward, "Leaky relu backward computation and inversion"); + m.def("elu_forward", &elu_forward, "Elu forward computation"); + m.def("elu_backward", &elu_backward, "Elu backward computation and inversion"); +} diff --git a/models/CtrlHair/external_code/face_parsing/modules/src/inplace_abn.h b/models/CtrlHair/external_code/face_parsing/modules/src/inplace_abn.h new file mode 100644 index 0000000000000000000000000000000000000000..17afd1196449ecb6376f28961e54b55e1537492f --- /dev/null +++ b/models/CtrlHair/external_code/face_parsing/modules/src/inplace_abn.h @@ -0,0 +1,88 @@ +#pragma once + +#include + +#include + +std::vector mean_var_cpu(at::Tensor x); +std::vector mean_var_cuda(at::Tensor x); +std::vector mean_var_cuda_h(at::Tensor x); + +at::Tensor forward_cpu(at::Tensor x, at::Tensor mean, at::Tensor var, at::Tensor weight, at::Tensor bias, + bool affine, float eps); +at::Tensor forward_cuda(at::Tensor x, at::Tensor mean, at::Tensor var, at::Tensor weight, at::Tensor bias, + bool affine, float eps); +at::Tensor forward_cuda_h(at::Tensor x, at::Tensor mean, at::Tensor var, at::Tensor weight, at::Tensor bias, + bool affine, float eps); + +std::vector edz_eydz_cpu(at::Tensor z, at::Tensor dz, at::Tensor weight, at::Tensor bias, + bool affine, float eps); +std::vector edz_eydz_cuda(at::Tensor z, at::Tensor dz, at::Tensor weight, at::Tensor bias, + bool affine, float eps); +std::vector edz_eydz_cuda_h(at::Tensor z, at::Tensor dz, at::Tensor weight, at::Tensor bias, + bool affine, float eps); + +at::Tensor backward_cpu(at::Tensor z, at::Tensor dz, at::Tensor var, at::Tensor weight, at::Tensor bias, + at::Tensor edz, at::Tensor eydz, bool affine, float eps); +at::Tensor backward_cuda(at::Tensor z, at::Tensor dz, at::Tensor var, at::Tensor weight, at::Tensor bias, + at::Tensor edz, at::Tensor eydz, bool affine, float eps); +at::Tensor backward_cuda_h(at::Tensor z, at::Tensor dz, at::Tensor var, at::Tensor weight, at::Tensor bias, + at::Tensor edz, at::Tensor eydz, bool affine, float eps); + +void leaky_relu_backward_cpu(at::Tensor z, at::Tensor dz, float slope); +void leaky_relu_backward_cuda(at::Tensor z, at::Tensor dz, float slope); +void leaky_relu_backward_cuda_h(at::Tensor z, at::Tensor dz, float slope); + +void elu_backward_cpu(at::Tensor z, at::Tensor dz); +void elu_backward_cuda(at::Tensor z, at::Tensor dz); + +static void get_dims(at::Tensor x, int64_t& num, int64_t& chn, int64_t& sp) { + num = x.size(0); + chn = x.size(1); + sp = 1; + for (int64_t i = 2; i < x.ndimension(); ++i) + sp *= x.size(i); +} + +/* + * Specialized CUDA reduction functions for BN + */ +#ifdef __CUDACC__ + +#include "utils/cuda.cuh" + +template +__device__ T reduce(Op op, int plane, int N, int S) { + T sum = (T)0; + for (int batch = 0; batch < N; ++batch) { + for (int x = threadIdx.x; x < S; x += blockDim.x) { + sum += op(batch, plane, x); + } + } + + // sum over NumThreads within a warp + sum = warpSum(sum); + + // 'transpose', and reduce within warp again + __shared__ T shared[32]; + __syncthreads(); + if (threadIdx.x % WARP_SIZE == 0) { + shared[threadIdx.x / WARP_SIZE] = sum; + } + if (threadIdx.x >= blockDim.x / WARP_SIZE && threadIdx.x < WARP_SIZE) { + // zero out the other entries in shared + shared[threadIdx.x] = (T)0; + } + __syncthreads(); + if (threadIdx.x / WARP_SIZE == 0) { + sum = warpSum(shared[threadIdx.x]); + if (threadIdx.x == 0) { + shared[0] = sum; + } + } + __syncthreads(); + + // Everyone picks it up, should be broadcast into the whole gradInput + return shared[0]; +} +#endif diff --git a/models/CtrlHair/external_code/face_parsing/modules/src/inplace_abn_cpu.cpp b/models/CtrlHair/external_code/face_parsing/modules/src/inplace_abn_cpu.cpp new file mode 100644 index 0000000000000000000000000000000000000000..ffc6d38c52ea31661b8dd438dc3fe1958f50b61e --- /dev/null +++ b/models/CtrlHair/external_code/face_parsing/modules/src/inplace_abn_cpu.cpp @@ -0,0 +1,119 @@ +#include + +#include + +#include "utils/checks.h" +#include "inplace_abn.h" + +at::Tensor reduce_sum(at::Tensor x) { + if (x.ndimension() == 2) { + return x.sum(0); + } else { + auto x_view = x.view({x.size(0), x.size(1), -1}); + return x_view.sum(-1).sum(0); + } +} + +at::Tensor broadcast_to(at::Tensor v, at::Tensor x) { + if (x.ndimension() == 2) { + return v; + } else { + std::vector broadcast_size = {1, -1}; + for (int64_t i = 2; i < x.ndimension(); ++i) + broadcast_size.push_back(1); + + return v.view(broadcast_size); + } +} + +int64_t count(at::Tensor x) { + int64_t count = x.size(0); + for (int64_t i = 2; i < x.ndimension(); ++i) + count *= x.size(i); + + return count; +} + +at::Tensor invert_affine(at::Tensor z, at::Tensor weight, at::Tensor bias, bool affine, float eps) { + if (affine) { + return (z - broadcast_to(bias, z)) / broadcast_to(at::abs(weight) + eps, z); + } else { + return z; + } +} + +std::vector mean_var_cpu(at::Tensor x) { + auto num = count(x); + auto mean = reduce_sum(x) / num; + auto diff = x - broadcast_to(mean, x); + auto var = reduce_sum(diff.pow(2)) / num; + + return {mean, var}; +} + +at::Tensor forward_cpu(at::Tensor x, at::Tensor mean, at::Tensor var, at::Tensor weight, at::Tensor bias, + bool affine, float eps) { + auto gamma = affine ? at::abs(weight) + eps : at::ones_like(var); + auto mul = at::rsqrt(var + eps) * gamma; + + x.sub_(broadcast_to(mean, x)); + x.mul_(broadcast_to(mul, x)); + if (affine) x.add_(broadcast_to(bias, x)); + + return x; +} + +std::vector edz_eydz_cpu(at::Tensor z, at::Tensor dz, at::Tensor weight, at::Tensor bias, + bool affine, float eps) { + auto edz = reduce_sum(dz); + auto y = invert_affine(z, weight, bias, affine, eps); + auto eydz = reduce_sum(y * dz); + + return {edz, eydz}; +} + +at::Tensor backward_cpu(at::Tensor z, at::Tensor dz, at::Tensor var, at::Tensor weight, at::Tensor bias, + at::Tensor edz, at::Tensor eydz, bool affine, float eps) { + auto y = invert_affine(z, weight, bias, affine, eps); + auto mul = affine ? at::rsqrt(var + eps) * (at::abs(weight) + eps) : at::rsqrt(var + eps); + + auto num = count(z); + auto dx = (dz - broadcast_to(edz / num, dz) - y * broadcast_to(eydz / num, dz)) * broadcast_to(mul, dz); + return dx; +} + +void leaky_relu_backward_cpu(at::Tensor z, at::Tensor dz, float slope) { + CHECK_CPU_INPUT(z); + CHECK_CPU_INPUT(dz); + + AT_DISPATCH_FLOATING_TYPES(z.type(), "leaky_relu_backward_cpu", ([&] { + int64_t count = z.numel(); + auto *_z = z.data(); + auto *_dz = dz.data(); + + for (int64_t i = 0; i < count; ++i) { + if (_z[i] < 0) { + _z[i] *= 1 / slope; + _dz[i] *= slope; + } + } + })); +} + +void elu_backward_cpu(at::Tensor z, at::Tensor dz) { + CHECK_CPU_INPUT(z); + CHECK_CPU_INPUT(dz); + + AT_DISPATCH_FLOATING_TYPES(z.type(), "elu_backward_cpu", ([&] { + int64_t count = z.numel(); + auto *_z = z.data(); + auto *_dz = dz.data(); + + for (int64_t i = 0; i < count; ++i) { + if (_z[i] < 0) { + _z[i] = log1p(_z[i]); + _dz[i] *= (_z[i] + 1.f); + } + } + })); +} diff --git a/models/CtrlHair/external_code/face_parsing/modules/src/inplace_abn_cuda.cu b/models/CtrlHair/external_code/face_parsing/modules/src/inplace_abn_cuda.cu new file mode 100644 index 0000000000000000000000000000000000000000..b157b06d47173d1645c6a40c89f564b737e84d43 --- /dev/null +++ b/models/CtrlHair/external_code/face_parsing/modules/src/inplace_abn_cuda.cu @@ -0,0 +1,333 @@ +#include + +#include +#include + +#include + +#include "utils/checks.h" +#include "utils/cuda.cuh" +#include "inplace_abn.h" + +#include + +// Operations for reduce +template +struct SumOp { + __device__ SumOp(const T *t, int c, int s) + : tensor(t), chn(c), sp(s) {} + __device__ __forceinline__ T operator()(int batch, int plane, int n) { + return tensor[(batch * chn + plane) * sp + n]; + } + const T *tensor; + const int chn; + const int sp; +}; + +template +struct VarOp { + __device__ VarOp(T m, const T *t, int c, int s) + : mean(m), tensor(t), chn(c), sp(s) {} + __device__ __forceinline__ T operator()(int batch, int plane, int n) { + T val = tensor[(batch * chn + plane) * sp + n]; + return (val - mean) * (val - mean); + } + const T mean; + const T *tensor; + const int chn; + const int sp; +}; + +template +struct GradOp { + __device__ GradOp(T _weight, T _bias, const T *_z, const T *_dz, int c, int s) + : weight(_weight), bias(_bias), z(_z), dz(_dz), chn(c), sp(s) {} + __device__ __forceinline__ Pair operator()(int batch, int plane, int n) { + T _y = (z[(batch * chn + plane) * sp + n] - bias) / weight; + T _dz = dz[(batch * chn + plane) * sp + n]; + return Pair(_dz, _y * _dz); + } + const T weight; + const T bias; + const T *z; + const T *dz; + const int chn; + const int sp; +}; + +/*********** + * mean_var + ***********/ + +template +__global__ void mean_var_kernel(const T *x, T *mean, T *var, int num, int chn, int sp) { + int plane = blockIdx.x; + T norm = T(1) / T(num * sp); + + T _mean = reduce>(SumOp(x, chn, sp), plane, num, sp) * norm; + __syncthreads(); + T _var = reduce>(VarOp(_mean, x, chn, sp), plane, num, sp) * norm; + + if (threadIdx.x == 0) { + mean[plane] = _mean; + var[plane] = _var; + } +} + +std::vector mean_var_cuda(at::Tensor x) { + CHECK_CUDA_INPUT(x); + + // Extract dimensions + int64_t num, chn, sp; + get_dims(x, num, chn, sp); + + // Prepare output tensors + auto mean = at::empty({chn}, x.options()); + auto var = at::empty({chn}, x.options()); + + // Run kernel + dim3 blocks(chn); + dim3 threads(getNumThreads(sp)); + auto stream = at::cuda::getCurrentCUDAStream(); + AT_DISPATCH_FLOATING_TYPES(x.type(), "mean_var_cuda", ([&] { + mean_var_kernel<<>>( + x.data(), + mean.data(), + var.data(), + num, chn, sp); + })); + + return {mean, var}; +} + +/********** + * forward + **********/ + +template +__global__ void forward_kernel(T *x, const T *mean, const T *var, const T *weight, const T *bias, + bool affine, float eps, int num, int chn, int sp) { + int plane = blockIdx.x; + + T _mean = mean[plane]; + T _var = var[plane]; + T _weight = affine ? abs(weight[plane]) + eps : T(1); + T _bias = affine ? bias[plane] : T(0); + + T mul = rsqrt(_var + eps) * _weight; + + for (int batch = 0; batch < num; ++batch) { + for (int n = threadIdx.x; n < sp; n += blockDim.x) { + T _x = x[(batch * chn + plane) * sp + n]; + T _y = (_x - _mean) * mul + _bias; + + x[(batch * chn + plane) * sp + n] = _y; + } + } +} + +at::Tensor forward_cuda(at::Tensor x, at::Tensor mean, at::Tensor var, at::Tensor weight, at::Tensor bias, + bool affine, float eps) { + CHECK_CUDA_INPUT(x); + CHECK_CUDA_INPUT(mean); + CHECK_CUDA_INPUT(var); + CHECK_CUDA_INPUT(weight); + CHECK_CUDA_INPUT(bias); + + // Extract dimensions + int64_t num, chn, sp; + get_dims(x, num, chn, sp); + + // Run kernel + dim3 blocks(chn); + dim3 threads(getNumThreads(sp)); + auto stream = at::cuda::getCurrentCUDAStream(); + AT_DISPATCH_FLOATING_TYPES(x.type(), "forward_cuda", ([&] { + forward_kernel<<>>( + x.data(), + mean.data(), + var.data(), + weight.data(), + bias.data(), + affine, eps, num, chn, sp); + })); + + return x; +} + +/*********** + * edz_eydz + ***********/ + +template +__global__ void edz_eydz_kernel(const T *z, const T *dz, const T *weight, const T *bias, + T *edz, T *eydz, bool affine, float eps, int num, int chn, int sp) { + int plane = blockIdx.x; + + T _weight = affine ? abs(weight[plane]) + eps : 1.f; + T _bias = affine ? bias[plane] : 0.f; + + Pair res = reduce, GradOp>(GradOp(_weight, _bias, z, dz, chn, sp), plane, num, sp); + __syncthreads(); + + if (threadIdx.x == 0) { + edz[plane] = res.v1; + eydz[plane] = res.v2; + } +} + +std::vector edz_eydz_cuda(at::Tensor z, at::Tensor dz, at::Tensor weight, at::Tensor bias, + bool affine, float eps) { + CHECK_CUDA_INPUT(z); + CHECK_CUDA_INPUT(dz); + CHECK_CUDA_INPUT(weight); + CHECK_CUDA_INPUT(bias); + + // Extract dimensions + int64_t num, chn, sp; + get_dims(z, num, chn, sp); + + auto edz = at::empty({chn}, z.options()); + auto eydz = at::empty({chn}, z.options()); + + // Run kernel + dim3 blocks(chn); + dim3 threads(getNumThreads(sp)); + auto stream = at::cuda::getCurrentCUDAStream(); + AT_DISPATCH_FLOATING_TYPES(z.type(), "edz_eydz_cuda", ([&] { + edz_eydz_kernel<<>>( + z.data(), + dz.data(), + weight.data(), + bias.data(), + edz.data(), + eydz.data(), + affine, eps, num, chn, sp); + })); + + return {edz, eydz}; +} + +/*********** + * backward + ***********/ + +template +__global__ void backward_kernel(const T *z, const T *dz, const T *var, const T *weight, const T *bias, const T *edz, + const T *eydz, T *dx, bool affine, float eps, int num, int chn, int sp) { + int plane = blockIdx.x; + + T _weight = affine ? abs(weight[plane]) + eps : 1.f; + T _bias = affine ? bias[plane] : 0.f; + T _var = var[plane]; + T _edz = edz[plane]; + T _eydz = eydz[plane]; + + T _mul = _weight * rsqrt(_var + eps); + T count = T(num * sp); + + for (int batch = 0; batch < num; ++batch) { + for (int n = threadIdx.x; n < sp; n += blockDim.x) { + T _dz = dz[(batch * chn + plane) * sp + n]; + T _y = (z[(batch * chn + plane) * sp + n] - _bias) / _weight; + + dx[(batch * chn + plane) * sp + n] = (_dz - _edz / count - _y * _eydz / count) * _mul; + } + } +} + +at::Tensor backward_cuda(at::Tensor z, at::Tensor dz, at::Tensor var, at::Tensor weight, at::Tensor bias, + at::Tensor edz, at::Tensor eydz, bool affine, float eps) { + CHECK_CUDA_INPUT(z); + CHECK_CUDA_INPUT(dz); + CHECK_CUDA_INPUT(var); + CHECK_CUDA_INPUT(weight); + CHECK_CUDA_INPUT(bias); + CHECK_CUDA_INPUT(edz); + CHECK_CUDA_INPUT(eydz); + + // Extract dimensions + int64_t num, chn, sp; + get_dims(z, num, chn, sp); + + auto dx = at::zeros_like(z); + + // Run kernel + dim3 blocks(chn); + dim3 threads(getNumThreads(sp)); + auto stream = at::cuda::getCurrentCUDAStream(); + AT_DISPATCH_FLOATING_TYPES(z.type(), "backward_cuda", ([&] { + backward_kernel<<>>( + z.data(), + dz.data(), + var.data(), + weight.data(), + bias.data(), + edz.data(), + eydz.data(), + dx.data(), + affine, eps, num, chn, sp); + })); + + return dx; +} + +/************** + * activations + **************/ + +template +inline void leaky_relu_backward_impl(T *z, T *dz, float slope, int64_t count) { + // Create thrust pointers + thrust::device_ptr th_z = thrust::device_pointer_cast(z); + thrust::device_ptr th_dz = thrust::device_pointer_cast(dz); + + auto stream = at::cuda::getCurrentCUDAStream(); + thrust::transform_if(thrust::cuda::par.on(stream), + th_dz, th_dz + count, th_z, th_dz, + [slope] __device__ (const T& dz) { return dz * slope; }, + [] __device__ (const T& z) { return z < 0; }); + thrust::transform_if(thrust::cuda::par.on(stream), + th_z, th_z + count, th_z, + [slope] __device__ (const T& z) { return z / slope; }, + [] __device__ (const T& z) { return z < 0; }); +} + +void leaky_relu_backward_cuda(at::Tensor z, at::Tensor dz, float slope) { + CHECK_CUDA_INPUT(z); + CHECK_CUDA_INPUT(dz); + + int64_t count = z.numel(); + + AT_DISPATCH_FLOATING_TYPES(z.type(), "leaky_relu_backward_cuda", ([&] { + leaky_relu_backward_impl(z.data(), dz.data(), slope, count); + })); +} + +template +inline void elu_backward_impl(T *z, T *dz, int64_t count) { + // Create thrust pointers + thrust::device_ptr th_z = thrust::device_pointer_cast(z); + thrust::device_ptr th_dz = thrust::device_pointer_cast(dz); + + auto stream = at::cuda::getCurrentCUDAStream(); + thrust::transform_if(thrust::cuda::par.on(stream), + th_dz, th_dz + count, th_z, th_z, th_dz, + [] __device__ (const T& dz, const T& z) { return dz * (z + 1.); }, + [] __device__ (const T& z) { return z < 0; }); + thrust::transform_if(thrust::cuda::par.on(stream), + th_z, th_z + count, th_z, + [] __device__ (const T& z) { return log1p(z); }, + [] __device__ (const T& z) { return z < 0; }); +} + +void elu_backward_cuda(at::Tensor z, at::Tensor dz) { + CHECK_CUDA_INPUT(z); + CHECK_CUDA_INPUT(dz); + + int64_t count = z.numel(); + + AT_DISPATCH_FLOATING_TYPES(z.type(), "leaky_relu_backward_cuda", ([&] { + elu_backward_impl(z.data(), dz.data(), count); + })); +} diff --git a/models/CtrlHair/external_code/face_parsing/modules/src/inplace_abn_cuda_half.cu b/models/CtrlHair/external_code/face_parsing/modules/src/inplace_abn_cuda_half.cu new file mode 100644 index 0000000000000000000000000000000000000000..bb63e73f9d90179e5bd5dae5579c4844da9c25e2 --- /dev/null +++ b/models/CtrlHair/external_code/face_parsing/modules/src/inplace_abn_cuda_half.cu @@ -0,0 +1,275 @@ +#include + +#include + +#include + +#include "utils/checks.h" +#include "utils/cuda.cuh" +#include "inplace_abn.h" + +#include + +// Operations for reduce +struct SumOpH { + __device__ SumOpH(const half *t, int c, int s) + : tensor(t), chn(c), sp(s) {} + __device__ __forceinline__ float operator()(int batch, int plane, int n) { + return __half2float(tensor[(batch * chn + plane) * sp + n]); + } + const half *tensor; + const int chn; + const int sp; +}; + +struct VarOpH { + __device__ VarOpH(float m, const half *t, int c, int s) + : mean(m), tensor(t), chn(c), sp(s) {} + __device__ __forceinline__ float operator()(int batch, int plane, int n) { + const auto t = __half2float(tensor[(batch * chn + plane) * sp + n]); + return (t - mean) * (t - mean); + } + const float mean; + const half *tensor; + const int chn; + const int sp; +}; + +struct GradOpH { + __device__ GradOpH(float _weight, float _bias, const half *_z, const half *_dz, int c, int s) + : weight(_weight), bias(_bias), z(_z), dz(_dz), chn(c), sp(s) {} + __device__ __forceinline__ Pair operator()(int batch, int plane, int n) { + float _y = (__half2float(z[(batch * chn + plane) * sp + n]) - bias) / weight; + float _dz = __half2float(dz[(batch * chn + plane) * sp + n]); + return Pair(_dz, _y * _dz); + } + const float weight; + const float bias; + const half *z; + const half *dz; + const int chn; + const int sp; +}; + +/*********** + * mean_var + ***********/ + +__global__ void mean_var_kernel_h(const half *x, float *mean, float *var, int num, int chn, int sp) { + int plane = blockIdx.x; + float norm = 1.f / static_cast(num * sp); + + float _mean = reduce(SumOpH(x, chn, sp), plane, num, sp) * norm; + __syncthreads(); + float _var = reduce(VarOpH(_mean, x, chn, sp), plane, num, sp) * norm; + + if (threadIdx.x == 0) { + mean[plane] = _mean; + var[plane] = _var; + } +} + +std::vector mean_var_cuda_h(at::Tensor x) { + CHECK_CUDA_INPUT(x); + + // Extract dimensions + int64_t num, chn, sp; + get_dims(x, num, chn, sp); + + // Prepare output tensors + auto mean = at::empty({chn},x.options().dtype(at::kFloat)); + auto var = at::empty({chn},x.options().dtype(at::kFloat)); + + // Run kernel + dim3 blocks(chn); + dim3 threads(getNumThreads(sp)); + auto stream = at::cuda::getCurrentCUDAStream(); + mean_var_kernel_h<<>>( + reinterpret_cast(x.data()), + mean.data(), + var.data(), + num, chn, sp); + + return {mean, var}; +} + +/********** + * forward + **********/ + +__global__ void forward_kernel_h(half *x, const float *mean, const float *var, const float *weight, const float *bias, + bool affine, float eps, int num, int chn, int sp) { + int plane = blockIdx.x; + + const float _mean = mean[plane]; + const float _var = var[plane]; + const float _weight = affine ? abs(weight[plane]) + eps : 1.f; + const float _bias = affine ? bias[plane] : 0.f; + + const float mul = rsqrt(_var + eps) * _weight; + + for (int batch = 0; batch < num; ++batch) { + for (int n = threadIdx.x; n < sp; n += blockDim.x) { + half *x_ptr = x + (batch * chn + plane) * sp + n; + float _x = __half2float(*x_ptr); + float _y = (_x - _mean) * mul + _bias; + + *x_ptr = __float2half(_y); + } + } +} + +at::Tensor forward_cuda_h(at::Tensor x, at::Tensor mean, at::Tensor var, at::Tensor weight, at::Tensor bias, + bool affine, float eps) { + CHECK_CUDA_INPUT(x); + CHECK_CUDA_INPUT(mean); + CHECK_CUDA_INPUT(var); + CHECK_CUDA_INPUT(weight); + CHECK_CUDA_INPUT(bias); + + // Extract dimensions + int64_t num, chn, sp; + get_dims(x, num, chn, sp); + + // Run kernel + dim3 blocks(chn); + dim3 threads(getNumThreads(sp)); + auto stream = at::cuda::getCurrentCUDAStream(); + forward_kernel_h<<>>( + reinterpret_cast(x.data()), + mean.data(), + var.data(), + weight.data(), + bias.data(), + affine, eps, num, chn, sp); + + return x; +} + +__global__ void edz_eydz_kernel_h(const half *z, const half *dz, const float *weight, const float *bias, + float *edz, float *eydz, bool affine, float eps, int num, int chn, int sp) { + int plane = blockIdx.x; + + float _weight = affine ? abs(weight[plane]) + eps : 1.f; + float _bias = affine ? bias[plane] : 0.f; + + Pair res = reduce, GradOpH>(GradOpH(_weight, _bias, z, dz, chn, sp), plane, num, sp); + __syncthreads(); + + if (threadIdx.x == 0) { + edz[plane] = res.v1; + eydz[plane] = res.v2; + } +} + +std::vector edz_eydz_cuda_h(at::Tensor z, at::Tensor dz, at::Tensor weight, at::Tensor bias, + bool affine, float eps) { + CHECK_CUDA_INPUT(z); + CHECK_CUDA_INPUT(dz); + CHECK_CUDA_INPUT(weight); + CHECK_CUDA_INPUT(bias); + + // Extract dimensions + int64_t num, chn, sp; + get_dims(z, num, chn, sp); + + auto edz = at::empty({chn},z.options().dtype(at::kFloat)); + auto eydz = at::empty({chn},z.options().dtype(at::kFloat)); + + // Run kernel + dim3 blocks(chn); + dim3 threads(getNumThreads(sp)); + auto stream = at::cuda::getCurrentCUDAStream(); + edz_eydz_kernel_h<<>>( + reinterpret_cast(z.data()), + reinterpret_cast(dz.data()), + weight.data(), + bias.data(), + edz.data(), + eydz.data(), + affine, eps, num, chn, sp); + + return {edz, eydz}; +} + +__global__ void backward_kernel_h(const half *z, const half *dz, const float *var, const float *weight, const float *bias, const float *edz, + const float *eydz, half *dx, bool affine, float eps, int num, int chn, int sp) { + int plane = blockIdx.x; + + float _weight = affine ? abs(weight[plane]) + eps : 1.f; + float _bias = affine ? bias[plane] : 0.f; + float _var = var[plane]; + float _edz = edz[plane]; + float _eydz = eydz[plane]; + + float _mul = _weight * rsqrt(_var + eps); + float count = float(num * sp); + + for (int batch = 0; batch < num; ++batch) { + for (int n = threadIdx.x; n < sp; n += blockDim.x) { + float _dz = __half2float(dz[(batch * chn + plane) * sp + n]); + float _y = (__half2float(z[(batch * chn + plane) * sp + n]) - _bias) / _weight; + + dx[(batch * chn + plane) * sp + n] = __float2half((_dz - _edz / count - _y * _eydz / count) * _mul); + } + } +} + +at::Tensor backward_cuda_h(at::Tensor z, at::Tensor dz, at::Tensor var, at::Tensor weight, at::Tensor bias, + at::Tensor edz, at::Tensor eydz, bool affine, float eps) { + CHECK_CUDA_INPUT(z); + CHECK_CUDA_INPUT(dz); + CHECK_CUDA_INPUT(var); + CHECK_CUDA_INPUT(weight); + CHECK_CUDA_INPUT(bias); + CHECK_CUDA_INPUT(edz); + CHECK_CUDA_INPUT(eydz); + + // Extract dimensions + int64_t num, chn, sp; + get_dims(z, num, chn, sp); + + auto dx = at::zeros_like(z); + + // Run kernel + dim3 blocks(chn); + dim3 threads(getNumThreads(sp)); + auto stream = at::cuda::getCurrentCUDAStream(); + backward_kernel_h<<>>( + reinterpret_cast(z.data()), + reinterpret_cast(dz.data()), + var.data(), + weight.data(), + bias.data(), + edz.data(), + eydz.data(), + reinterpret_cast(dx.data()), + affine, eps, num, chn, sp); + + return dx; +} + +__global__ void leaky_relu_backward_impl_h(half *z, half *dz, float slope, int64_t count) { + for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < count; i += blockDim.x * gridDim.x){ + float _z = __half2float(z[i]); + if (_z < 0) { + dz[i] = __float2half(__half2float(dz[i]) * slope); + z[i] = __float2half(_z / slope); + } + } +} + +void leaky_relu_backward_cuda_h(at::Tensor z, at::Tensor dz, float slope) { + CHECK_CUDA_INPUT(z); + CHECK_CUDA_INPUT(dz); + + int64_t count = z.numel(); + dim3 threads(getNumThreads(count)); + dim3 blocks = (count + threads.x - 1) / threads.x; + auto stream = at::cuda::getCurrentCUDAStream(); + leaky_relu_backward_impl_h<<>>( + reinterpret_cast(z.data()), + reinterpret_cast(dz.data()), + slope, count); +} + diff --git a/models/CtrlHair/external_code/face_parsing/modules/src/utils/checks.h b/models/CtrlHair/external_code/face_parsing/modules/src/utils/checks.h new file mode 100644 index 0000000000000000000000000000000000000000..e761a6fe34d0789815b588eba7e3726026e0e868 --- /dev/null +++ b/models/CtrlHair/external_code/face_parsing/modules/src/utils/checks.h @@ -0,0 +1,15 @@ +#pragma once + +#include + +// Define AT_CHECK for old version of ATen where the same function was called AT_ASSERT +#ifndef AT_CHECK +#define AT_CHECK AT_ASSERT +#endif + +#define CHECK_CUDA(x) AT_CHECK((x).type().is_cuda(), #x " must be a CUDA tensor") +#define CHECK_CPU(x) AT_CHECK(!(x).type().is_cuda(), #x " must be a CPU tensor") +#define CHECK_CONTIGUOUS(x) AT_CHECK((x).is_contiguous(), #x " must be contiguous") + +#define CHECK_CUDA_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x) +#define CHECK_CPU_INPUT(x) CHECK_CPU(x); CHECK_CONTIGUOUS(x) \ No newline at end of file diff --git a/models/CtrlHair/external_code/face_parsing/modules/src/utils/common.h b/models/CtrlHair/external_code/face_parsing/modules/src/utils/common.h new file mode 100644 index 0000000000000000000000000000000000000000..e8403eef8a233b75dd4bb353c16486fe1be2039a --- /dev/null +++ b/models/CtrlHair/external_code/face_parsing/modules/src/utils/common.h @@ -0,0 +1,49 @@ +#pragma once + +#include + +/* + * Functions to share code between CPU and GPU + */ + +#ifdef __CUDACC__ +// CUDA versions + +#define HOST_DEVICE __host__ __device__ +#define INLINE_HOST_DEVICE __host__ __device__ inline +#define FLOOR(x) floor(x) + +#if __CUDA_ARCH__ >= 600 +// Recent compute capabilities have block-level atomicAdd for all data types, so we use that +#define ACCUM(x,y) atomicAdd_block(&(x),(y)) +#else +// Older architectures don't have block-level atomicAdd, nor atomicAdd for doubles, so we defer to atomicAdd for float +// and use the known atomicCAS-based implementation for double +template +__device__ inline data_t atomic_add(data_t *address, data_t val) { + return atomicAdd(address, val); +} + +template<> +__device__ inline double atomic_add(double *address, double val) { + unsigned long long int* address_as_ull = (unsigned long long int*)address; + unsigned long long int old = *address_as_ull, assumed; + do { + assumed = old; + old = atomicCAS(address_as_ull, assumed, __double_as_longlong(val + __longlong_as_double(assumed))); + } while (assumed != old); + return __longlong_as_double(old); +} + +#define ACCUM(x,y) atomic_add(&(x),(y)) +#endif // #if __CUDA_ARCH__ >= 600 + +#else +// CPU versions + +#define HOST_DEVICE +#define INLINE_HOST_DEVICE inline +#define FLOOR(x) std::floor(x) +#define ACCUM(x,y) (x) += (y) + +#endif // #ifdef __CUDACC__ \ No newline at end of file diff --git a/models/CtrlHair/external_code/face_parsing/modules/src/utils/cuda.cuh b/models/CtrlHair/external_code/face_parsing/modules/src/utils/cuda.cuh new file mode 100644 index 0000000000000000000000000000000000000000..60c0023835e02c5f7c539c28ac07b75b72df394b --- /dev/null +++ b/models/CtrlHair/external_code/face_parsing/modules/src/utils/cuda.cuh @@ -0,0 +1,71 @@ +#pragma once + +/* + * General settings and functions + */ +const int WARP_SIZE = 32; +const int MAX_BLOCK_SIZE = 1024; + +static int getNumThreads(int nElem) { + int threadSizes[6] = {32, 64, 128, 256, 512, MAX_BLOCK_SIZE}; + for (int i = 0; i < 6; ++i) { + if (nElem <= threadSizes[i]) { + return threadSizes[i]; + } + } + return MAX_BLOCK_SIZE; +} + +/* + * Reduction utilities + */ +template +__device__ __forceinline__ T WARP_SHFL_XOR(T value, int laneMask, int width = warpSize, + unsigned int mask = 0xffffffff) { +#if CUDART_VERSION >= 9000 + return __shfl_xor_sync(mask, value, laneMask, width); +#else + return __shfl_xor(value, laneMask, width); +#endif +} + +__device__ __forceinline__ int getMSB(int val) { return 31 - __clz(val); } + +template +struct Pair { + T v1, v2; + __device__ Pair() {} + __device__ Pair(T _v1, T _v2) : v1(_v1), v2(_v2) {} + __device__ Pair(T v) : v1(v), v2(v) {} + __device__ Pair(int v) : v1(v), v2(v) {} + __device__ Pair &operator+=(const Pair &a) { + v1 += a.v1; + v2 += a.v2; + return *this; + } +}; + +template +static __device__ __forceinline__ T warpSum(T val) { +#if __CUDA_ARCH__ >= 300 + for (int i = 0; i < getMSB(WARP_SIZE); ++i) { + val += WARP_SHFL_XOR(val, 1 << i, WARP_SIZE); + } +#else + __shared__ T values[MAX_BLOCK_SIZE]; + values[threadIdx.x] = val; + __threadfence_block(); + const int base = (threadIdx.x / WARP_SIZE) * WARP_SIZE; + for (int i = 1; i < WARP_SIZE; i++) { + val += values[base + ((i + threadIdx.x) % WARP_SIZE)]; + } +#endif + return val; +} + +template +static __device__ __forceinline__ Pair warpSum(Pair value) { + value.v1 = warpSum(value.v1); + value.v2 = warpSum(value.v2); + return value; +} \ No newline at end of file diff --git a/models/CtrlHair/external_code/face_parsing/my_parsing_util.py b/models/CtrlHair/external_code/face_parsing/my_parsing_util.py new file mode 100644 index 0000000000000000000000000000000000000000..88cdff20a4ea6808d36686817df7ff0cf142217d --- /dev/null +++ b/models/CtrlHair/external_code/face_parsing/my_parsing_util.py @@ -0,0 +1,95 @@ +# -*- coding: utf-8 -*- + +""" +# File name: my_parsing_util.py +# Time : 2022/07/15 +# Author: xyguoo@163.com +# Description: +""" +import sys + +import numpy as np +import torch +from PIL import Image + +from models.CtrlHair.external_code.face_parsing.model import BiSeNet +from models.CtrlHair.global_value_utils import PARSING_LABEL_LIST +import torchvision.transforms as transforms + + +class FaceParsing: + label_list = {0: 'background', 1: 'skin_other', 2: 'l_brow', 3: 'r_brow', 4: 'l_eye', + 5: 'r_eye', 6: 'eye_g', 7: 'l_ear', 8: 'r_ear', 9: 'ear_r', + 10: 'nose', 11: 'mouth', 12: 'u_lip', 13: 'l_lip', 14: 'neck', + 15: 'neck_l', 16: 'cloth', 17: 'hair', 18: 'hat'} + skin_area = {1, 2, 3, 4, 5, 7, 8, 10, 11, 12, 13} + bise_net = None + to_tensor = transforms.Compose([ + transforms.ToTensor(), + transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)), + ]) + + @staticmethod + def parsing_img(img): + pil_img = Image.fromarray(img) + with torch.no_grad(): + image = pil_img.resize((512, 512), Image.BILINEAR) + img = FaceParsing.to_tensor(image) + img = torch.unsqueeze(img, 0) + # img = img.cuda() + if FaceParsing.bise_net is None: + n_classes = 19 + FaceParsing.bise_net = BiSeNet(n_classes=n_classes) + # FaceParsing.bise_net.cuda() + save_pth = 'models/CtrlHair/external_model_params/face_parsing_79999_iter.pth' + FaceParsing.bise_net.load_state_dict(torch.load(save_pth)) + FaceParsing.bise_net.eval() + out = FaceParsing.bise_net(img)[0] + parsing = out.squeeze(0).cpu().numpy().argmax(0) + return parsing, image + + @staticmethod + def swap_parsing_label_to_celeba_mask(parsing): + celeba_parsing = np.zeros_like(parsing) + label_lists = list(FaceParsing.label_list.values()) + for label_idx, label_name in enumerate(PARSING_LABEL_LIST): + celeba_parsing[label_lists.index(label_name) == parsing] = label_idx + return celeba_parsing + + +class FaceParsing_tensor: + label_list = {0: 'background', 1: 'skin_other', 2: 'l_brow', 3: 'r_brow', 4: 'l_eye', + 5: 'r_eye', 6: 'eye_g', 7: 'l_ear', 8: 'r_ear', 9: 'ear_r', + 10: 'nose', 11: 'mouth', 12: 'u_lip', 13: 'l_lip', 14: 'neck', + 15: 'neck_l', 16: 'cloth', 17: 'hair', 18: 'hat'} + skin_area = {1, 2, 3, 4, 5, 7, 8, 10, 11, 12, 13} + bise_net = None + to_tensor = transforms.Compose([ + transforms.ToTensor(), + transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)), + ]) + + @staticmethod + def parsing_img(img=None): + with torch.no_grad(): + if FaceParsing.bise_net is None: + n_classes = 19 + FaceParsing.bise_net = BiSeNet(n_classes=n_classes) + save_pth = 'pretrained_models/BiSeNet/face_parsing_79999_iter.pth' + FaceParsing.bise_net.load_state_dict(torch.load(save_pth)) + FaceParsing.bise_net.eval() + FaceParsing.bise_net = FaceParsing.bise_net.cuda() + + if img is None: + return + out = FaceParsing.bise_net(img)[0] + parsing = out.squeeze(0).argmax(0) + return parsing, None + + @staticmethod + def swap_parsing_label_to_celeba_mask(parsing): + celeba_parsing = torch.zeros_like(parsing) + label_lists = list(FaceParsing.label_list.values()) + for label_idx, label_name in enumerate(PARSING_LABEL_LIST): + celeba_parsing[label_lists.index(label_name) == parsing] = label_idx + return celeba_parsing diff --git a/models/CtrlHair/external_code/face_parsing/optimizer.py b/models/CtrlHair/external_code/face_parsing/optimizer.py new file mode 100644 index 0000000000000000000000000000000000000000..0c99e0645164b22f1e743ee99daadadd26a1cd80 --- /dev/null +++ b/models/CtrlHair/external_code/face_parsing/optimizer.py @@ -0,0 +1,69 @@ +#!/usr/bin/python +# -*- encoding: utf-8 -*- + + +import torch +import logging + +logger = logging.getLogger() + +class Optimizer(object): + def __init__(self, + model, + lr0, + momentum, + wd, + warmup_steps, + warmup_start_lr, + max_iter, + power, + *args, **kwargs): + self.warmup_steps = warmup_steps + self.warmup_start_lr = warmup_start_lr + self.lr0 = lr0 + self.lr = self.lr0 + self.max_iter = float(max_iter) + self.power = power + self.it = 0 + wd_params, nowd_params, lr_mul_wd_params, lr_mul_nowd_params = model.get_params() + param_list = [ + {'params': wd_params}, + {'params': nowd_params, 'weight_decay': 0}, + {'params': lr_mul_wd_params, 'lr_mul': True}, + {'params': lr_mul_nowd_params, 'weight_decay': 0, 'lr_mul': True}] + self.optim = torch.optim.SGD( + param_list, + lr = lr0, + momentum = momentum, + weight_decay = wd) + self.warmup_factor = (self.lr0/self.warmup_start_lr)**(1./self.warmup_steps) + + + def get_lr(self): + if self.it <= self.warmup_steps: + lr = self.warmup_start_lr*(self.warmup_factor**self.it) + else: + factor = (1-(self.it-self.warmup_steps)/(self.max_iter-self.warmup_steps))**self.power + lr = self.lr0 * factor + return lr + + + def step(self): + self.lr = self.get_lr() + for pg in self.optim.param_groups: + if pg.get('lr_mul', False): + pg['lr'] = self.lr * 10 + else: + pg['lr'] = self.lr + if self.optim.defaults.get('lr_mul', False): + self.optim.defaults['lr'] = self.lr * 10 + else: + self.optim.defaults['lr'] = self.lr + self.it += 1 + self.optim.step() + if self.it == self.warmup_steps+2: + logger.info('==> warmup done, start to implement poly lr strategy') + + def zero_grad(self): + self.optim.zero_grad() + diff --git a/models/CtrlHair/external_code/face_parsing/prepropess_data.py b/models/CtrlHair/external_code/face_parsing/prepropess_data.py new file mode 100644 index 0000000000000000000000000000000000000000..ee7ed56dd8c0372d482e6a53f323da17043bd521 --- /dev/null +++ b/models/CtrlHair/external_code/face_parsing/prepropess_data.py @@ -0,0 +1,38 @@ +#!/usr/bin/python +# -*- encoding: utf-8 -*- + +import os.path as osp +import os +import cv2 +from transform import * +from PIL import Image + +face_data = '/home/zll/data/CelebAMask-HQ/CelebA-HQ-img' +face_sep_mask = '/home/zll/data/CelebAMask-HQ/CelebAMask-HQ-mask-anno' +mask_path = '/home/zll/data/CelebAMask-HQ/mask' +counter = 0 +total = 0 +for i in range(15): + + atts = ['skin', 'l_brow', 'r_brow', 'l_eye', 'r_eye', 'eye_g', 'l_ear', 'r_ear', 'ear_r', + 'nose', 'mouth', 'u_lip', 'l_lip', 'neck', 'neck_l', 'cloth', 'hair', 'hat'] + + for j in range(i * 2000, (i + 1) * 2000): + + mask = np.zeros((512, 512)) + + for l, att in enumerate(atts, 1): + total += 1 + file_name = ''.join([str(j).rjust(5, '0'), '_', att, '.png']) + path = osp.join(face_sep_mask, str(i), file_name) + + if os.path.exists(path): + counter += 1 + sep_mask = np.array(Image.open(path).convert('P')) + # print(np.unique(sep_mask)) + + mask[sep_mask == 225] = l + cv2.imwrite('{}/{}.png'.format(mask_path, j), mask) + print(j) + +print(counter, total) \ No newline at end of file diff --git a/models/CtrlHair/external_code/face_parsing/resnet.py b/models/CtrlHair/external_code/face_parsing/resnet.py new file mode 100644 index 0000000000000000000000000000000000000000..aa2bf95130e9815ba378cb6f73207068b81a04b9 --- /dev/null +++ b/models/CtrlHair/external_code/face_parsing/resnet.py @@ -0,0 +1,109 @@ +#!/usr/bin/python +# -*- encoding: utf-8 -*- + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.utils.model_zoo as modelzoo + +# from modules.bn import InPlaceABNSync as BatchNorm2d + +resnet18_url = 'https://download.pytorch.org/models/resnet18-5c106cde.pth' + + +def conv3x3(in_planes, out_planes, stride=1): + """3x3 convolution with padding""" + return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, + padding=1, bias=False) + + +class BasicBlock(nn.Module): + def __init__(self, in_chan, out_chan, stride=1): + super(BasicBlock, self).__init__() + self.conv1 = conv3x3(in_chan, out_chan, stride) + self.bn1 = nn.BatchNorm2d(out_chan) + self.conv2 = conv3x3(out_chan, out_chan) + self.bn2 = nn.BatchNorm2d(out_chan) + self.relu = nn.ReLU(inplace=True) + self.downsample = None + if in_chan != out_chan or stride != 1: + self.downsample = nn.Sequential( + nn.Conv2d(in_chan, out_chan, + kernel_size=1, stride=stride, bias=False), + nn.BatchNorm2d(out_chan), + ) + + def forward(self, x): + residual = self.conv1(x) + residual = F.relu(self.bn1(residual)) + residual = self.conv2(residual) + residual = self.bn2(residual) + + shortcut = x + if self.downsample is not None: + shortcut = self.downsample(x) + + out = shortcut + residual + out = self.relu(out) + return out + + +def create_layer_basic(in_chan, out_chan, bnum, stride=1): + layers = [BasicBlock(in_chan, out_chan, stride=stride)] + for i in range(bnum-1): + layers.append(BasicBlock(out_chan, out_chan, stride=1)) + return nn.Sequential(*layers) + + +class Resnet18(nn.Module): + def __init__(self): + super(Resnet18, self).__init__() + self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, + bias=False) + self.bn1 = nn.BatchNorm2d(64) + self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) + self.layer1 = create_layer_basic(64, 64, bnum=2, stride=1) + self.layer2 = create_layer_basic(64, 128, bnum=2, stride=2) + self.layer3 = create_layer_basic(128, 256, bnum=2, stride=2) + self.layer4 = create_layer_basic(256, 512, bnum=2, stride=2) + self.init_weight() + + def forward(self, x): + x = self.conv1(x) + x = F.relu(self.bn1(x)) + x = self.maxpool(x) + + x = self.layer1(x) + feat8 = self.layer2(x) # 1/8 + feat16 = self.layer3(feat8) # 1/16 + feat32 = self.layer4(feat16) # 1/32 + return feat8, feat16, feat32 + + def init_weight(self): + state_dict = modelzoo.load_url(resnet18_url) + self_state_dict = self.state_dict() + for k, v in state_dict.items(): + if 'fc' in k: continue + self_state_dict.update({k: v}) + self.load_state_dict(self_state_dict) + + def get_params(self): + wd_params, nowd_params = [], [] + for name, module in self.named_modules(): + if isinstance(module, (nn.Linear, nn.Conv2d)): + wd_params.append(module.weight) + if not module.bias is None: + nowd_params.append(module.bias) + elif isinstance(module, nn.BatchNorm2d): + nowd_params += list(module.parameters()) + return wd_params, nowd_params + + +if __name__ == "__main__": + net = Resnet18() + x = torch.randn(16, 3, 224, 224) + out = net(x) + print(out[0].size()) + print(out[1].size()) + print(out[2].size()) + net.get_params() diff --git a/models/CtrlHair/external_code/face_parsing/train.py b/models/CtrlHair/external_code/face_parsing/train.py new file mode 100644 index 0000000000000000000000000000000000000000..e550289e66ca8398caf7f993aba76dce422fbc5b --- /dev/null +++ b/models/CtrlHair/external_code/face_parsing/train.py @@ -0,0 +1,177 @@ +#!/usr/bin/python +# -*- encoding: utf-8 -*- + +from logger import setup_logger +from model import BiSeNet +from face_dataset import FaceMask +from loss import OhemCELoss +from evaluate import evaluate +from optimizer import Optimizer + +import torch +import torch.nn as nn +from torch.utils.data import DataLoader +import torch.distributed as dist + +import os +import os.path as osp +import logging +import time +import datetime +import argparse + +import my_torchlib.train_utils + +respth = './resources' +if not osp.exists(respth): + os.makedirs(respth) +logger = logging.getLogger() + + +def parse_args(): + parse = argparse.ArgumentParser() + parse.add_argument( + '--local_rank', + dest = 'local_rank', + type = int, + default = -1, + ) + return parse.parse_args() + + +def train(): + args = parse_args() + torch.cuda.set_device(args.local_rank) + dist.init_process_group( + backend = 'nccl', + init_method = 'tcp://127.0.0.1:33241', + world_size = torch.cuda.device_count(), + rank=args.local_rank + ) + setup_logger(respth) + + # dataset + n_classes = 19 + n_img_per_gpu = 16 + n_workers = 8 + cropsize = [448, 448] + data_root = '/home/zll/data/CelebAMask-HQ/' + + ds = FaceMask(data_root, cropsize=cropsize, mode='scripts') + sampler = torch.utils.data.distributed.DistributedSampler(ds) + dl = DataLoader(ds, + batch_size = n_img_per_gpu, + shuffle = False, + sampler = sampler, + num_workers = n_workers, + pin_memory = True, + drop_last = True) + + # model + ignore_idx = -100 + net = BiSeNet(n_classes=n_classes) + net.cuda() + my_torchlib.train_utils.train() + net = nn.parallel.DistributedDataParallel(net, + device_ids = [args.local_rank, ], + output_device = args.local_rank + ) + score_thres = 0.7 + n_min = n_img_per_gpu * cropsize[0] * cropsize[1]//16 + LossP = OhemCELoss(thresh=score_thres, n_min=n_min, ignore_lb=ignore_idx) + Loss2 = OhemCELoss(thresh=score_thres, n_min=n_min, ignore_lb=ignore_idx) + Loss3 = OhemCELoss(thresh=score_thres, n_min=n_min, ignore_lb=ignore_idx) + + ## optimizer + momentum = 0.9 + weight_decay = 5e-4 + lr_start = 1e-2 + max_iter = 80000 + power = 0.9 + warmup_steps = 1000 + warmup_start_lr = 1e-5 + optim = Optimizer( + model = net.module, + lr0 = lr_start, + momentum = momentum, + wd = weight_decay, + warmup_steps = warmup_steps, + warmup_start_lr = warmup_start_lr, + max_iter = max_iter, + power = power) + + ## scripts loop + msg_iter = 50 + loss_avg = [] + st = glob_st = time.time() + diter = iter(dl) + epoch = 0 + for it in range(max_iter): + try: + im, lb = next(diter) + if not im.size()[0] == n_img_per_gpu: + raise StopIteration + except StopIteration: + epoch += 1 + sampler.set_epoch(epoch) + diter = iter(dl) + im, lb = next(diter) + im = im.cuda() + lb = lb.cuda() + H, W = im.size()[2:] + lb = torch.squeeze(lb, 1) + + optim.zero_grad() + out, out16, out32 = net(im) + lossp = LossP(out, lb) + loss2 = Loss2(out16, lb) + loss3 = Loss3(out32, lb) + loss = lossp + loss2 + loss3 + loss.backward() + optim.step() + + loss_avg.append(loss.item()) + + # print training log message + if (it+1) % msg_iter == 0: + loss_avg = sum(loss_avg) / len(loss_avg) + lr = optim.lr + ed = time.time() + t_intv, glob_t_intv = ed - st, ed - glob_st + eta = int((max_iter - it) * (glob_t_intv / it)) + eta = str(datetime.timedelta(seconds=eta)) + msg = ', '.join([ + 'it: {it}/{max_it}', + 'lr: {lr:4f}', + 'loss: {loss:.4f}', + 'eta: {eta}', + 'time: {time:.4f}', + ]).format( + it = it+1, + max_it = max_iter, + lr = lr, + loss = loss_avg, + time = t_intv, + eta = eta + ) + logger.info(msg) + loss_avg = [] + st = ed + if dist.get_rank() == 0: + if (it+1) % 5000 == 0: + state = net.module.state_dict() if hasattr(net, 'module') else net.state_dict() + if dist.get_rank() == 0: + torch.save(state, './resources/cp/{}_iter.pth'.format(it)) + evaluate(dspth='/home/zll/data/CelebAMask-HQ/test-img', cp='{}_iter.pth'.format(it)) + + # dump the final model + save_pth = osp.join(respth, 'model_final_diss.pth') + # net.cpu() + state = net.module.state_dict() if hasattr(net, 'module') else net.state_dict() + if dist.get_rank() == 0: + torch.save(state, save_pth) + logger.info('training done, model saved to: {}'.format(save_pth)) + + +if __name__ == "__main__": + train() diff --git a/models/CtrlHair/external_code/face_parsing/transform.py b/models/CtrlHair/external_code/face_parsing/transform.py new file mode 100644 index 0000000000000000000000000000000000000000..9479ae356a151f5da8eedf288abeae7458739d24 --- /dev/null +++ b/models/CtrlHair/external_code/face_parsing/transform.py @@ -0,0 +1,129 @@ +#!/usr/bin/python +# -*- encoding: utf-8 -*- + + +from PIL import Image +import PIL.ImageEnhance as ImageEnhance +import random +import numpy as np + +class RandomCrop(object): + def __init__(self, size, *args, **kwargs): + self.size = size + + def __call__(self, im_lb): + im = im_lb['im'] + lb = im_lb['lb'] + assert im.size == lb.size + W, H = self.size + w, h = im.size + + if (W, H) == (w, h): return dict(im=im, lb=lb) + if w < W or h < H: + scale = float(W) / w if w < h else float(H) / h + w, h = int(scale * w + 1), int(scale * h + 1) + im = im.resize((w, h), Image.BILINEAR) + lb = lb.resize((w, h), Image.NEAREST) + sw, sh = random.random() * (w - W), random.random() * (h - H) + crop = int(sw), int(sh), int(sw) + W, int(sh) + H + return dict( + im = im.crop(crop), + lb = lb.crop(crop) + ) + + +class HorizontalFlip(object): + def __init__(self, p=0.5, *args, **kwargs): + self.p = p + + def __call__(self, im_lb): + if random.random() > self.p: + return im_lb + else: + im = im_lb['im'] + lb = im_lb['lb'] + + # atts = [1 'skin', 2 'l_brow', 3 'r_brow', 4 'l_eye', 5 'r_eye', 6 'eye_g', 7 'l_ear', 8 'r_ear', 9 'ear_r', + # 10 'nose', 11 'mouth', 12 'u_lip', 13 'l_lip', 14 'neck', 15 'neck_l', 16 'cloth', 17 'hair', 18 'hat'] + + flip_lb = np.array(lb) + flip_lb[lb == 2] = 3 + flip_lb[lb == 3] = 2 + flip_lb[lb == 4] = 5 + flip_lb[lb == 5] = 4 + flip_lb[lb == 7] = 8 + flip_lb[lb == 8] = 7 + flip_lb = Image.fromarray(flip_lb) + return dict(im = im.transpose(Image.FLIP_LEFT_RIGHT), + lb = flip_lb.transpose(Image.FLIP_LEFT_RIGHT), + ) + + +class RandomScale(object): + def __init__(self, scales=(1, ), *args, **kwargs): + self.scales = scales + + def __call__(self, im_lb): + im = im_lb['im'] + lb = im_lb['lb'] + W, H = im.size + scale = random.choice(self.scales) + w, h = int(W * scale), int(H * scale) + return dict(im = im.resize((w, h), Image.BILINEAR), + lb = lb.resize((w, h), Image.NEAREST), + ) + + +class ColorJitter(object): + def __init__(self, brightness=None, contrast=None, saturation=None, *args, **kwargs): + if not brightness is None and brightness>0: + self.brightness = [max(1-brightness, 0), 1+brightness] + if not contrast is None and contrast>0: + self.contrast = [max(1-contrast, 0), 1+contrast] + if not saturation is None and saturation>0: + self.saturation = [max(1-saturation, 0), 1+saturation] + + def __call__(self, im_lb): + im = im_lb['im'] + lb = im_lb['lb'] + r_brightness = random.uniform(self.brightness[0], self.brightness[1]) + r_contrast = random.uniform(self.contrast[0], self.contrast[1]) + r_saturation = random.uniform(self.saturation[0], self.saturation[1]) + im = ImageEnhance.Brightness(im).enhance(r_brightness) + im = ImageEnhance.Contrast(im).enhance(r_contrast) + im = ImageEnhance.Color(im).enhance(r_saturation) + return dict(im = im, + lb = lb, + ) + + +class MultiScale(object): + def __init__(self, scales): + self.scales = scales + + def __call__(self, img): + W, H = img.size + sizes = [(int(W*ratio), int(H*ratio)) for ratio in self.scales] + imgs = [] + [imgs.append(img.resize(size, Image.BILINEAR)) for size in sizes] + return imgs + + +class Compose(object): + def __init__(self, do_list): + self.do_list = do_list + + def __call__(self, im_lb): + for comp in self.do_list: + im_lb = comp(im_lb) + return im_lb + + + + +if __name__ == '__main__': + flip = HorizontalFlip(p = 1) + crop = RandomCrop((321, 321)) + rscales = RandomScale((0.75, 1.0, 1.5, 1.75, 2.0)) + img = Image.open('data/img.jpg') + lb = Image.open('data/label.png') diff --git a/models/CtrlHair/external_code/landmarks_util.py b/models/CtrlHair/external_code/landmarks_util.py new file mode 100644 index 0000000000000000000000000000000000000000..4d490bde19c8a00343b417b12a1472a33bf0f4e4 --- /dev/null +++ b/models/CtrlHair/external_code/landmarks_util.py @@ -0,0 +1,55 @@ +# -*- coding: utf-8 -*- + +""" +# File name: landmarks_util.py +# Time : 2022/07/15 +# Author: xyguoo@163.com +# Description: +""" +import os +import pickle as pkl + +import dlib as dlib +import numpy as np +import tqdm +import cv2 + +detector = dlib.get_frontal_face_detector() +predictor_dict = {68: dlib.shape_predictor('models/CtrlHair/external_model_params/shape_predictor_68_face_landmarks.dat'), + 81: dlib.shape_predictor('models/CtrlHair/external_model_params/shape_predictor_81_face_landmarks.dat')} + + +def detect_landmarks(root_dir, dataset_name, landmark_output_file_path, output_dir=None, predictor=None): + result_dic = {} + for dn in dataset_name: + img_dir = os.path.join(root_dir, dn, 'images_256') + files = os.listdir(img_dir) + files.sort() + + if output_dir and not os.path.exists(output_dir): + os.makedirs(output_dir) + + for f in tqdm.tqdm(files): + file_path = os.path.join(img_dir, f) + img_rd = cv2.imread(file_path) + img_gray = cv2.cvtColor(img_rd, cv2.COLOR_BGR2RGB) + + faces = detector(img_gray, 0) + font = cv2.FONT_HERSHEY_SIMPLEX + + # annotate landmarks + if len(faces) != 0: + landmarks = np.array([[p.x, p.y] for p in predictor(img_rd, faces[0]).parts()]) + result_dic['%s___%s' % (dn, f[:-4])] = landmarks / img_gray.shape[0] + if output_dir: + for idx, point in enumerate(landmarks): + pos = (point[0], point[1]) + cv2.circle(img_rd, pos, 2, color=(139, 0, 0)) + cv2.putText(img_rd, str(idx + 1), pos, font, 0.5, (0, 0, 255), 2, cv2.LINE_AA) + cv2.imwrite(os.path.join(output_dir, f), img_rd) + else: + # not detect face + print('no face for %s' % file_path) + + with open(landmark_output_file_path, 'wb') as f: + pkl.dump(result_dic, f) \ No newline at end of file diff --git a/models/CtrlHair/global_value_utils.py b/models/CtrlHair/global_value_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..8355d7f7ffe2cac1fdcf720da646039103e5feb5 --- /dev/null +++ b/models/CtrlHair/global_value_utils.py @@ -0,0 +1,56 @@ +# -*- coding: utf-8 -*- + +""" +# File name: global_value_utils.py +# Time : 2021/9/11 19:15 +# Author: xyguoo@163.com +# Description: +""" +import os +import pickle as pkl + + +GLOBAL_DATA_ROOT = './dataset_info_ctrlhair' +######################################################## +# Please change these setting if you scripts a new model` +DATASET_NAME = ['ffhq', 'CelebaMask_HQ'] +DEFAULT_CONFIG_COLOR_TEXTURE_BRANCH = '045' +DEFAULT_CONFIG_SHAPE_BRANCH = '054' +######################################################## + +TEMP_FOLDER = 'models/CtrlHair/temp_folder' + +PARSING_COLOR_LIST = [[0, 0, 0], + [204, 0, 0], + [76, 153, 0], + [204, 204, 0], ## + [51, 51, 255], ## + [204, 0, 204], ## + [0, 255, 255], ## + [51, 255, 255], ## + [102, 51, 0], ## + [255, 0, 0], ## + [102, 204, 0], ## + [255, 255, 0], ## + [0, 0, 153], ## + [0, 0, 204], ## hair + [255, 51, 153], ## + [0, 204, 204], ## + [0, 51, 0], ## + [255, 153, 51], + [0, 204, 0], + [255, 85, 255], + [255, 170, 255], + [0, 170, 255], + [85, 255, 255], + [170, 255, 255], + [255, 255, 255]] + +PARSING_LABEL_LIST = ['background', 'skin_other', 'nose', 'eye_g', 'l_eye', 'r_eye', 'l_brow', 'r_brow', + 'l_ear', 'r_ear', 'mouth', 'u_lip', 'l_lip', 'hair', 'hat', + 'ear_r', 'neck_l', 'neck', 'cloth'] +HAIR_IDX = PARSING_LABEL_LIST.index('hair') +HAT_IDX = PARSING_LABEL_LIST.index('hat') +UNKNOWN_IDX = len(PARSING_COLOR_LIST) - 1 + +WRAP_TEMP_FOLDER = [os.path.join(TEMP_FOLDER, '/wrap_triangle/wrap_temp_result')] diff --git a/models/CtrlHair/hair_editor.py b/models/CtrlHair/hair_editor.py new file mode 100644 index 0000000000000000000000000000000000000000..a2f4dd7964f4ca324ab9d936928c7f4ef979084b --- /dev/null +++ b/models/CtrlHair/hair_editor.py @@ -0,0 +1,335 @@ +# -*- coding: utf-8 -*- + +""" +# File name: hair_editor.py +# Time : 2021/11/18 17:21 +# Author: xyguoo@163.com +# Description: +""" + +import os +import pickle +from glob import glob + +import cv2 +import numpy as np +import torch + +import my_torchlib +from color_texture_branch.solver import Solver as SolveFeature +from external_code.face_parsing.my_parsing_util import FaceParsing +from global_value_utils import HAIR_IDX, PARSING_LABEL_LIST +from poisson_blending import poisson_blending +from sean_codes.models.pix2pix_model import Pix2PixModel +from sean_codes.options.test_options import TestOptions +from shape_branch.solver import Solver as SolverMask +from util.imutil import write_rgb + + +# adaptor_root_dir = '/data1/guoxuyang/myWorkSpace/hair_editing' +# sys.path.append(adaptor_root_dir) +# sys.path.append(os.path.join(adaptor_root_dir, 'external_code/face_3DDFA')) + + +def change_status(model, new_status): + for m in model.modules(): + if hasattr(m, 'status'): + m.status = new_status + + +class HairEditor: + """ + This is the basic module, that could achieve many editing task. ui/hair_swap.py/Backend succeed this class. + """ + + def __init__(self, load_feature_model=True, load_mask_model=True): + self.opt = TestOptions().parse() + self.opt.status = 'test' + self.sean_model = Pix2PixModel(self.opt) + self.sean_model.eval() + self.img_size = 256 + self.device = torch.device('cuda', 0) + + if load_feature_model: + from color_texture_branch.config import cfg as cfg_feature + self.solver_feature = SolveFeature(cfg_feature, device=self.device, local_rank=-1, training=False) + + self.feature_encoder = self.solver_feature.dis + self.feature_generator = self.solver_feature.gen + self.feature_rgb_predictor = self.solver_feature.rgb_model + # self.feature_curliness_predictor = self.solver_feature.curliness_model + + # ckpt_dir = 'external_model_params/disentangle_checkpoints/' + cfg_app.experiment_name + '/checkpoints' + ckpt_dir = 'model_trained/color_texture/' + cfg_feature.experiment_name + '/checkpoints' + ckpt = my_torchlib.load_checkpoint(ckpt_dir) + 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} + + self.feature_generator.load_state_dict(ckpt['Model_G'], strict=True) + self.feature_encoder.load_state_dict(ckpt['Model_D'], strict=True) + + # if 'curliness' in cfg_feature.predictor: + # ckpt = my_torchlib.load_checkpoint(cfg_feature.predictor.curliness.root_dir + '/checkpoints') + # self.feature_curliness_predictor.load_state_dict(ckpt['Predictor'], strict=True) + + if 'rgb' in cfg_feature.predictor: + ckpt = my_torchlib.load_checkpoint(cfg_feature.predictor.rgb.root_dir + '/checkpoints') + self.feature_rgb_predictor.load_state_dict(ckpt['Predictor'], strict=True) + + # load unsupervised direction + existing_dirs_dir = os.path.join('model_trained/color_texture', cfg_feature.experiment_name, + 'texture_dir_used') + if os.path.exists(existing_dirs_dir): + existing_dirs_list = os.listdir(existing_dirs_dir) + existing_dirs_list.sort() + 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).to(self.device)) + self.texture_dirs = existing_dirs + + if load_mask_model: + from shape_branch.config import cfg as cfg_mask + self.solver_mask = SolverMask(cfg_mask, device=self.device, local_rank=-1, training=False) + self.mask_generator = self.solver_mask.gen + + ############################################## + # change to your checkpoints dir # + ############################################## + ckpt_dir = 'model_trained/shape/' + cfg_mask.experiment_name + '/checkpoints' + ckpt = my_torchlib.load_checkpoint(ckpt_dir) + 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} + + self.mask_generator.load_state_dict(ckpt['Model_G'], strict=True) + + # load unsupervised direction + existing_dirs_dir = os.path.join('model_trained/shape', cfg_mask.experiment_name, 'shape_dir_used') + if os.path.exists(existing_dirs_dir): + existing_dirs_list = os.listdir(existing_dirs_dir) + existing_dirs_list.sort() + 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).to(self.device)) + self.shape_dirs = existing_dirs + + def preprocess_img(self, img): + img = cv2.resize(img.astype('uint8'), (self.img_size, self.img_size)) + return (np.transpose(img, [2, 0, 1]) / 127.5 - 1.0)[None, ...] + + def preprocess_mask(self, mask_img): + mask_img = cv2.resize(mask_img.astype('uint8'), (self.img_size, self.img_size), + interpolation=cv2.INTER_NEAREST) + return mask_img[None, None, :, :] + + @staticmethod + def load_average_feature(): + ############### load average features + # average_style_code_folder = 'styles_test/mean_style_code/mean/' + average_style_code_folder = 'sean_codes/styles_test/mean_style_code/median/' + input_style_dic = {} + + ############### hard coding for categories + for i in range(19): + input_style_dic[str(i)] = {} + average_category_folder_list = glob(os.path.join(average_style_code_folder, str(i), '*.npy')) + average_category_list = [os.path.splitext(os.path.basename(name))[0] for name in + average_category_folder_list] + + for style_code_path in average_category_list: + input_style_dic[str(i)][style_code_path] = torch.from_numpy( + np.load(os.path.join(average_style_code_folder, str(i), style_code_path + '.npy'))).cuda() + return input_style_dic + + def get_code(self, hair_img, hair_parsing): + # generate style code + data = {'label': torch.tensor(hair_parsing, dtype=torch.float32), + 'instance': torch.tensor(0), + 'image': torch.tensor(hair_img, dtype=torch.float32), + 'path': ['temp/temp_npy']} + change_status(self.sean_model, 'test') + hair_img_code = self.sean_model(data, mode='style_code') + return hair_img_code + + def gen_img(self, code, parsing): + # load style code + if not isinstance(code, torch.Tensor): + code = torch.tensor(code) + obj_dic = self.load_average_feature() + + for idx in range(19): + cur_code = code[0, idx] + if not torch.all(cur_code == 0): + obj_dic[str(idx)]['ACE'] = cur_code + + temp_face_image = torch.zeros((0, 3, self.img_size, self.img_size)) # place holder + + data = {'label': torch.tensor(parsing, dtype=torch.float32), + 'instance': torch.tensor(0), + 'image': torch.tensor(temp_face_image, dtype=torch.float32), + 'obj_dic': obj_dic} + change_status(self.sean_model, 'UI_mode') + # self.model = self.model.to(code.device) + generated = self.sean_model(data, mode='UI_mode')[0] + return generated + + def generate_by_sean(self, face_img_code, hair_code, target_seg): + """ + :param face_img_code: please input with the shape [19, 512] + :param hair_code: please input with the shape [512] + :param target_seg: + :return: + """ + # load style code + obj_dic = self.load_average_feature() + + for idx in range(19): + if idx == HAIR_IDX: + cur_code = hair_code + # cur_code = face_img_code[0, idx] + else: + cur_code = face_img_code[idx] + if not torch.all(face_img_code == 0): + obj_dic[str(idx)]['ACE'] = cur_code + + data = {'label': torch.tensor(target_seg, dtype=torch.float32), + 'instance': torch.tensor(0), + 'obj_dic': obj_dic, + 'image': None} + change_status(self.sean_model, 'UI_mode') + generated = self.sean_model(data, mode='UI_mode')[0] + return generated + + def generate_instance_transfer_img(self, face_img, face_parsing, hair_img, hair_parsing, target_seg, edit_data=None, + temp_path='temp'): + # generate style code + data = {'label': torch.tensor(face_parsing, dtype=torch.float32), + 'instance': torch.tensor(0), + 'image': torch.tensor(face_img, dtype=torch.float32), + 'path': ['temp/temp_npy']} + face_img_code = self.sean_model(data, mode='style_code') + + if hair_img is None: + hair_img_code = face_img_code + else: + data = {'label': torch.tensor(hair_parsing, dtype=torch.float32), + 'instance': torch.tensor(0), + 'image': torch.tensor(hair_img, dtype=torch.float32), + 'path': ['temp/temp_npy']} + change_status(self.sean_model, 'test') + hair_img_code = self.sean_model(data, mode='style_code') + hair_code = hair_img_code[0, HAIR_IDX] + + if edit_data is not None: + hair_code = self.solver_feature.edit_infer(hair_code[None, ...], edit_data)[0] + + return self.generate_by_sean(face_img_code[0], hair_code, target_seg) + + def get_hair_color(self, img): + parsing, _ = FaceParsing.parsing_img(img) + parsing = FaceParsing.swap_parsing_label_to_celeba_mask(parsing) + parsing = cv2.resize(parsing.astype('uint8'), (1024, 1024), interpolation=cv2.INTER_NEAREST) + img = cv2.resize(img.astype('uint8'), (1024, 1024)) + hair_mask = (parsing == HAIR_IDX).astype('uint8') + + kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, ksize=(19, 19)) + hair_mask = cv2.erode(hair_mask, kernel, iterations=1) + points = img[hair_mask.astype('bool')] + moment1 = points.mean(axis=0) + return moment1 + + @staticmethod + def draw_landmarks(img, lms): + lms = lms / 2 + lms = lms.astype('int') + for idx, point in enumerate(lms): + font = cv2.FONT_HERSHEY_SIMPLEX + pos = (point[0], point[1]) + cv2.circle(img, pos, 2, color=(139, 0, 0)) + cv2.putText(img, str(idx + 1), pos, font, 0.18, (255, 0, 0), 1, cv2.LINE_AA) + return img + + def postprocess_blending(self, face_img, res_img, face_parsing, target_parsing, verbose_print=False, blending=True): + """ + Blend original face img and result image with poisson blending. + If not blend, the result image will look slightly different from original image in some details in + non-hair region, but the image quality will be better. + :param face_img: + :param res_img: + :param face_parsing: + :param target_parsing: + :param verbose_print: + :param blending: If `False`, the result image will do some trivial thing like transferring data type + :return: + """ + if verbose_print: + print("Post process for the result image...") + + def from_tensor_order_to_cv2(tensor_img, is_mask=False): + if isinstance(tensor_img, torch.Tensor): + tensor_img = tensor_img.detach().cpu().numpy() + if len(tensor_img.shape) == 4: + tensor_img = tensor_img[0] + if len(tensor_img.shape) == 2: + tensor_img = tensor_img[None, ...] + if tensor_img.shape[2] <= 3: + return tensor_img + res = np.transpose(tensor_img, [1, 2, 0]) + if not is_mask: + res = res * 127.5 + 127.5 + return res + + res_img = from_tensor_order_to_cv2(res_img) + res_img = res_img.astype('uint8') + if blending: + + target_parsing = from_tensor_order_to_cv2(target_parsing, is_mask=True) + face_img = from_tensor_order_to_cv2(face_img) + face_img = face_img.astype('uint8') + + face_parsing = from_tensor_order_to_cv2(face_parsing, is_mask=True) + + res_mask = np.logical_or(target_parsing == HAIR_IDX, face_parsing == HAIR_IDX).astype('uint8') + kernel13 = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, ksize=(13, 13)) + kernel5 = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, ksize=(5, 5)) + res_mask_dilated = cv2.dilate(res_mask, kernel13, iterations=1)[..., None] + + res_mask_dilated5 = cv2.dilate(res_mask, kernel5, iterations=1)[..., None] + + bg_mask = (target_parsing == PARSING_LABEL_LIST.index('background')) + res_mask_dilated = res_mask_dilated * (1 - bg_mask) + res_mask_dilated5 * bg_mask + + face_to_hair = poisson_blending(face_img, res_img, 1 - res_mask_dilated, with_gamma=True) + return face_to_hair, res_mask_dilated + else: + return res_img, None + + def crop_face(self, img_rgb, save_path=None): + """ + crop the face part in the image to adapt the editing system + :param img_rgb: + :param save_path: + :return: + """ + from external_code.crop import recreate_aligned_images + from external_code.landmarks_util import predictor_dict, detector + + predictor_68 = predictor_dict[68] + bbox = detector(img_rgb, 0)[0] + lm_68 = np.array([[p.x, p.y] for p in predictor_68(img_rgb, bbox).parts()]) + crop_img_pil, lm_68 = recreate_aligned_images(img_rgb, lm_68, output_size=self.img_size) + img_rgb = np.array(crop_img_pil) + if save_path is not None: + write_rgb(save_path, img_rgb) + return img_rgb + + def get_mask(self, img_rgb): + parsing, _ = FaceParsing.parsing_img(img_rgb) + parsing = FaceParsing.swap_parsing_label_to_celeba_mask(parsing) + mask_img = cv2.resize(parsing.astype('uint8'), (self.img_size, self.img_size), interpolation=cv2.INTER_NEAREST) + return mask_img diff --git a/models/CtrlHair/my_pylib/__init__.py b/models/CtrlHair/my_pylib/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3ab90debbcbb649ee8c209509221f6397d25031c --- /dev/null +++ b/models/CtrlHair/my_pylib/__init__.py @@ -0,0 +1,7 @@ +# from __future__ import absolute_import +# from __future__ import division +# from __future__ import print_function + +from .path import * +from .timer import * +from .utils import * diff --git a/models/CtrlHair/my_pylib/path.py b/models/CtrlHair/my_pylib/path.py new file mode 100644 index 0000000000000000000000000000000000000000..c373a2f9fbd7218069bde2ccdd55db4c1ba6997a --- /dev/null +++ b/models/CtrlHair/my_pylib/path.py @@ -0,0 +1,62 @@ +"""Complementary functions for os.path.""" + +import fnmatch +import os +import sys + + +def add_path(paths): + if not isinstance(paths, (list, tuple)): + paths = [paths] + for path in paths: + if path not in sys.path: + sys.path.insert(0, path) + + +def mkdir(paths): + if not isinstance(paths, (list, tuple)): + paths = [paths] + for path in paths: + if not os.path.exists(path): + os.makedirs(path) + + +def split(path): + dir, name_ext = os.path.split(path) + name, ext = os.path.splitext(name_ext) + return dir, name, ext + + +def directory(path): + return split(path)[0] + + +def name(path): + return split(path)[1] + + +def ext(path): + return split(path)[2] + + +def name_ext(path): + return ''.join(split(path)[1:]) + + +abspath = os.path.abspath +join = os.path.join + + +def match(dir, pat, recursive=False): + if recursive: + iterator = os.walk(dir) + else: + try: + iterator = [next(os.walk(dir))] + except: + return [] + matches = [] + for root, _, file_names in iterator: + for file_name in fnmatch.filter(file_names, pat): + matches.append(os.path.join(root, file_name)) + return matches diff --git a/models/CtrlHair/my_pylib/timer.py b/models/CtrlHair/my_pylib/timer.py new file mode 100644 index 0000000000000000000000000000000000000000..fffb1c0cf9614f4a520a4d0cd938cac8463e4b1d --- /dev/null +++ b/models/CtrlHair/my_pylib/timer.py @@ -0,0 +1,124 @@ +import datetime +import timeit + + +class Timer(object): + """A timer as a context manager. + + Modified from https://github.com/brouberol/contexttimer/blob/master/contexttimer/__init__.py. + + Wraps around a timer. A custom timer can be passed + to the constructor. The default timer is timeit.default_timer. + + Note that the latter measures wall clock time, not CPU time! + On Unix systems, it corresponds to time.time. + On Windows systems, it corresponds to time.clock. + + Arguments: + print_at_exit : If True, print when exiting context. + format : 'ms', 's' or 'datetime' + """ + + def __init__(self, fmt='s', print_at_exit=True, timer=timeit.default_timer): + assert fmt in ['ms', 's', 'datetime'], "`fmt` should be 'ms', 's' or 'datetime'!" + self._fmt = fmt + self._print_at_exit = print_at_exit + self._timer = timer + self.start() + + def __enter__(self): + """Start the timer in the context manager scope.""" + self.restart() + return self + + def __exit__(self, exc_type, exc_value, exc_traceback): + """Print the end time.""" + if self._print_at_exit: + print(str(self)) + + def __str__(self): + return self.fmt(self.elapsed)[1] + + def start(self): + self.start_time = self._timer() + + restart = start + + @property + def elapsed(self): + """Return the current elapsed time since last (re)start.""" + return self._timer() - self.start_time + + def fmt(self, second): + if self._fmt == 'ms': + time_fmt = second * 1000 + time_str = '%s %s' % (time_fmt, self._fmt) + elif self._fmt == 's': + time_fmt = second + time_str = '%s %s' % (time_fmt, self._fmt) + elif self._fmt == 'datetime': + time_fmt = datetime.timedelta(seconds=second) + time_str = str(time_fmt) + return time_fmt, time_str + + +def timer(run_times=1, **timer_kwargs): + """Function decorator displaying the function execution time. + + All kwargs are the arguments taken by the Timer class constructor. + """ + + # store Timer kwargs in local variable so the namespace isn't polluted + # by different level args and kwargs + + def wrapped_f(f): + def wrapped(*args, **kwargs): + timer_kwargs.update(print_at_exit=False) + with Timer(**timer_kwargs) as t: + for _ in range(run_times): + out = f(*args, **kwargs) + fmt = '[*] function "%(function_name)s" execution time for %(run_times)d runs: %(execution_time)s = %(' \ + 'execution_time_each)s * %(run_times)d [*] ' + context = {'function_name': f.__name__, 'run_times': run_times, 'execution_time': t, + 'execution_time_each': t.fmt(t.elapsed / run_times)[1]} + print(fmt % context) + return out + + return wrapped + + return wrapped_f + + +if __name__ == "__main__": + import time + + # 1 + print(1) + with Timer() as t: + time.sleep(1) + print(t) + time.sleep(1) + + with Timer(fmt='datetime') as t: + time.sleep(1) + + # 2 + print(2) + t = Timer(fmt='ms') + time.sleep(2) + print(t) + + t = Timer(fmt='datetime') + time.sleep(1) + print(t) + + # 3 + print(3) + + + @timer(run_times=5, fmt='s') + def blah(): + time.sleep(2) + + + blah() diff --git a/models/CtrlHair/my_pylib/utils.py b/models/CtrlHair/my_pylib/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..c7e8580c1f06340729b1dc86a8a059ee96c90c83 --- /dev/null +++ b/models/CtrlHair/my_pylib/utils.py @@ -0,0 +1,19 @@ +import json +import pprint + +import addict + + +def save_json(path, obj, *args, **kw_args): + # wrap json.dumps + with open(path, 'w') as f: + f.write(json.dumps(obj, *args, **kw_args)) + + +def load_json(path, *args, **kw_args): + # wrap json.load + with open(path) as f: + return addict.Dict(json.load(f, *args, **kw_args)) + + +pp = pprint.pprint diff --git a/models/CtrlHair/my_torchlib/__init__.py b/models/CtrlHair/my_torchlib/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..ff2ecaba7d8714a7f2ba626dc27ffb5ee8675404 --- /dev/null +++ b/models/CtrlHair/my_torchlib/__init__.py @@ -0,0 +1,2 @@ +from .layers import * +from .utils import * diff --git a/models/CtrlHair/my_torchlib/layers/__init__.py b/models/CtrlHair/my_torchlib/layers/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..69a388db1dea2d5699b716260dfa0902c27c0ab5 --- /dev/null +++ b/models/CtrlHair/my_torchlib/layers/__init__.py @@ -0,0 +1 @@ +from .layers import * diff --git a/models/CtrlHair/my_torchlib/layers/layers.py b/models/CtrlHair/my_torchlib/layers/layers.py new file mode 100644 index 0000000000000000000000000000000000000000..5b6a9a9ef98969866885281b76eaa1645cb43e5c --- /dev/null +++ b/models/CtrlHair/my_torchlib/layers/layers.py @@ -0,0 +1,33 @@ +from torch import nn + + +# ============================================================================== +# = layer = +# ============================================================================== + +class NoOp(nn.Module): + + def __init__(self, *args, **keyword_args): + super(NoOp, self).__init__() + + def forward(self, x): + return x + + +class Reshape(nn.Module): # 0表示保持当前大小 + + def __init__(self, *new_shape): + super(Reshape, self).__init__() + self._new_shape = new_shape + + def forward(self, x): + new_shape = (x.size(i) if self._new_shape[i] == 0 else self._new_shape[i] for i in range(len(self._new_shape))) + return x.view(*new_shape) + + +# ============================================================================== +# = layer wrapper = +# ============================================================================== + +def identity(x, *args, **keyword_args): + return x diff --git a/models/CtrlHair/my_torchlib/module.py b/models/CtrlHair/my_torchlib/module.py new file mode 100644 index 0000000000000000000000000000000000000000..6a7787631eb1e62e75ec9851f6092149ac06550a --- /dev/null +++ b/models/CtrlHair/my_torchlib/module.py @@ -0,0 +1,283 @@ +# -*- coding: utf-8 -*- + +""" +# File name: module.py +# Time : 2021/11/17 15:38 +# Author: xyguoo@163.com +# Description: +""" + +import torch +import torch.nn.functional as F + +from torch import 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 + + +class Conv2dBlock(nn.Module): + + def __init__(self, input_dim, output_dim, kernel_size, stride, padding=0, norm='none', activation='relu', + pad_type='zero', use_bias=True, norm_affine=False, transpose=False, leaky_slope=0.2): + super(Conv2dBlock, self).__init__() + # initialize padding + self.transpose = transpose + + if not transpose: + if pad_type == 'reflect': + self.pad = nn.ReflectionPad2d(padding) + elif pad_type == 'replicate': + self.pad = nn.ReplicationPad2d(padding) + elif pad_type == 'zero': + self.pad = nn.ZeroPad2d(padding) + else: + assert 0, "Unsupported padding type: {}".format(pad_type) + + # initialize convolution + if norm == 'sn': + self.conv = SpectralNorm(nn.Conv2d(input_dim, output_dim, kernel_size, stride, bias=use_bias)) + elif transpose: + self.conv = nn.ConvTranspose2d(input_dim, output_dim, kernel_size, stride, padding, bias=use_bias) + else: + self.conv = nn.Conv2d(input_dim, output_dim, kernel_size, stride, bias=use_bias) + + # initialize normalization + norm_dim = output_dim + if norm == 'bn': + self.norm = nn.BatchNorm2d(norm_dim) + elif norm == 'in': + # self.norm = nn.InstanceNorm2d(norm_dim, track_running_stats=True) + self.norm = nn.InstanceNorm2d(norm_dim, norm_affine) + elif norm == 'myin': + self.norm = MyInstanceNorm2d(norm_dim) + elif norm == 'ln': + self.norm = LayerNorm(norm_dim) + elif norm == 'adain': + self.norm = AdaptiveInstanceNorm2d(norm_dim) + elif norm == 'none' or norm == 'sn': + 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 == 'sigmoid': + self.activation = nn.Sigmoid() + elif activation == 'none': + self.activation = None + else: + assert 0, "Unsupported activation: {}".format(activation) + + def forward(self, x): + if not self.transpose: + x = self.pad(x) + x = self.conv(x) + if self.norm: + x = self.norm(x) + if self.activation: + x = self.activation(x) + return x + + +# ====================================== +# = normalization layers = +# ====================================== + +class AdaptiveInstanceNorm2d(nn.Module): + + def __init__(self, num_features, eps=1e-5, momentum=0.1): + super(AdaptiveInstanceNorm2d, self).__init__() + self.num_features = num_features + self.eps = eps + self.momentum = momentum + # weight and bias are dynamically assigned + self.weight = None + self.bias = None + # just dummy buffers, not used + self.register_buffer('running_mean', torch.zeros(num_features)) + self.register_buffer('running_var', torch.ones(num_features)) + + def forward(self, x): + assert self.weight is not None and self.bias is not None, "Please assign weight and bias before calling AdaIN!" + b, c = x.size(0), x.size(1) + running_mean = self.running_mean.repeat(b) + running_var = self.running_var.repeat(b) + + # Apply instance norm + x_reshaped = x.contiguous().view(1, b * c, *x.size()[2:]) + + out = F.batch_norm( + x_reshaped, running_mean, running_var, self.weight, self.bias, + True, self.momentum, self.eps) + + return out.view(b, c, *x.size()[2:]) + + def __repr__(self): + return self.__class__.__name__ + '(' + str(self.num_features) + ')' + + +class LayerNorm(nn.Module): + + def __init__(self, num_features, eps=1e-5, affine=True): + super(LayerNorm, self).__init__() + self.num_features = num_features + self.affine = affine + self.eps = eps + + if self.affine: + self.gamma = nn.Parameter(torch.Tensor(num_features).uniform_()) + self.beta = nn.Parameter(torch.zeros(num_features)) + + def forward(self, x): + shape = [-1] + [1] * (x.dim() - 1) + # print(x.size()) + if x.size(0) == 1: + # These two lines run much faster in pytorch 0.4 than the two lines listed below. + mean = x.view(-1).mean().view(*shape) + std = x.view(-1).std().view(*shape) + else: + mean = x.view(x.size(0), -1).mean(1).view(*shape) + std = x.view(x.size(0), -1).std(1).view(*shape) + + x = (x - mean) / (std + self.eps) + + if self.affine: + shape = [1, -1] + [1] * (x.dim() - 2) + x = x * self.gamma.view(*shape) + self.beta.view(*shape) + return x + + +class MyInstanceNorm2d(nn.Module): + + def __init__(self, num_features, eps=1e-5): + super(MyInstanceNorm2d, self).__init__() + self.num_features = num_features + self.eps = eps + + def forward(self, x): + mean = x.view(x.size(0), x.size(1), -1).mean(2).view(x.size(0), x.size(1), 1, 1) + std = x.view(x.size(0), x.size(1), -1).std(2).view(x.size(0), x.size(1), 1, 1) + x = (x - mean) / (std + self.eps) + return x + + +def l2normalize(v, eps=1e-12): + return v / (v.norm() + eps) + + +class SpectralNorm(nn.Module): + """ + Based on the paper "Spectral Normalization for Generative Adversarial Networks" by Takeru Miyato, Toshiki Kataoka, Masanori Koyama, Yuichi Yoshida + and the Pytorch implementation https://github.com/christiancosgrove/pytorch-spectral-normalization-gan + """ + + def __init__(self, module, name='weight', power_iterations=1): + super(SpectralNorm, self).__init__() + self.module = module + self.name = name + self.power_iterations = power_iterations + if not self._made_params(): + self._make_params() + + def _update_u_v(self): + u = getattr(self.module, self.name + "_u") + v = getattr(self.module, self.name + "_v") + w = getattr(self.module, self.name + "_bar") + + height = w.data.shape[0] + for _ in range(self.power_iterations): + v.data = l2normalize(torch.mv(torch.t(w.view(height, -1).data), u.data)) + u.data = l2normalize(torch.mv(w.view(height, -1).data, v.data)) + + # sigma = torch.dot(u.data, torch.mv(w.view(height,-1).data, v.data)) + sigma = u.dot(w.view(height, -1).mv(v)) + setattr(self.module, self.name, w / sigma.expand_as(w)) + + def _made_params(self): + try: + u = getattr(self.module, self.name + "_u") + v = getattr(self.module, self.name + "_v") + w = getattr(self.module, self.name + "_bar") + return True + except AttributeError: + return False + + def _make_params(self): + w = getattr(self.module, self.name) + + height = w.data.shape[0] + width = w.view(height, -1).data.shape[1] + + u = nn.Parameter(w.data.new(height).normal_(0, 1), requires_grad=False) + v = nn.Parameter(w.data.new(width).normal_(0, 1), requires_grad=False) + u.data = l2normalize(u.data) + v.data = l2normalize(v.data) + w_bar = nn.Parameter(w.data) + + del self.module._parameters[self.name] + + self.module.register_parameter(self.name + "_u", u) + self.module.register_parameter(self.name + "_v", v) + self.module.register_parameter(self.name + "_bar", w_bar) + + def forward(self, *args): + self._update_u_v() + return self.module.forward(*args) diff --git a/models/CtrlHair/my_torchlib/train_utils.py b/models/CtrlHair/my_torchlib/train_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..8a9f6ee0dff505b9db942110423d34af2043b0ec --- /dev/null +++ b/models/CtrlHair/my_torchlib/train_utils.py @@ -0,0 +1,89 @@ +# -*- coding: utf-8 -*- + +""" +# File name: train_utils.py +# Time : 2021/12/15 11:21 +# Author: xyguoo@163.com +# Description: +""" +import numpy as np +import torch +from scipy import stats + +class LossUpdater: + def __init__(self, cfg): + self.cfg = cfg + self.register_list = {} + for ke in cfg: + if ke.startswith('lambda_') and isinstance(cfg[ke], dict): + self.register_list[ke] = cfg[ke] + + def update(self, step): + for ke in self.register_list: + loss_dict = self.register_list[ke] + weight = None + for start_step in loss_dict: + if start_step > step: + break + weight = loss_dict[start_step] + if weight is None: + raise Exception() + self.cfg[ke] = weight + + +def tensor2numpy(tensor): + return tensor.cpu().numpy().transpose(1, 2, 0) + + +def to_device(dat, device): + for ke in dat: + if isinstance(dat[ke], torch.Tensor): + dat[ke] = dat[ke].to(device) + + +def generate_noise(bs, dim, label=None): + # trunc = stats.truncnorm(-3, 3) + # noise = trunc.rvs(bs * dim).reshape(bs, dim) + # noise = torch.tensor(noise).float() + noise = torch.randn((bs, dim)) + if label is not None: + noise = (noise.abs() * label).float() + return noise + + +def train(cfg, loss_dict, optimizers, step, writer, flag, retain_graph=False, write_log=False): + """ + :param loss_dict: + :param optimizers: + :param step: + :param writer: + :param flag: + :return: + """ + if len(loss_dict) == 0: + return + loss_total = 0 + for k, v in loss_dict.items(): + if np.isnan(np.array(v.detach().cpu())): + print('!!!!!!!!! %s is nan' % k) + raise Exception() + if np.isinf(np.array(v.detach().cpu())): + print('!!!!!!!!! %s is inf' % k) + raise Exception() + if k not in cfg: # skip rgs_zp + continue + loss_total = loss_total + v * cfg[k] + + for o in optimizers: + o.zero_grad() + loss_total.backward(retain_graph=retain_graph) + for o in optimizers: + o.step() + + # summary + if write_log: + for k, v in loss_dict.items(): + writer.add_scalar('%s/%s' % (flag, k), loss_dict[k].data.mean().cpu().numpy(), + global_step=step) + writer.add_scalar('%s/%s' % (flag, 'total'), loss_total.data.mean().cpu().numpy(), + global_step=step) \ No newline at end of file diff --git a/models/CtrlHair/my_torchlib/utils.py b/models/CtrlHair/my_torchlib/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..abb3cff17ebcc47fdcb68b36c42a7e1b71c60fe2 --- /dev/null +++ b/models/CtrlHair/my_torchlib/utils.py @@ -0,0 +1,85 @@ +import os +import shutil +import torch +import torchvision.transforms as tform +from PIL import Image + + +def cpu(xs): + if not isinstance(xs, (list, tuple)): + return xs.cpu() + else: + return [x.cpu() for x in xs] + + +def cuda(xs): + if torch.cuda.is_available(): + if not isinstance(xs, (list, tuple)): + return xs.cuda() + else: + return [x.cuda() for x in xs] + else: + return xs + + +def load_checkpoint(ckpt_dir_or_file, map_location=None, load_best=False): + if os.path.isdir(ckpt_dir_or_file): + if load_best: + ckpt_path = os.path.join(ckpt_dir_or_file, 'best_model.ckpt') + else: + with open(os.path.join(ckpt_dir_or_file, 'latest_checkpoint')) as f: + ckpt_path = os.path.join(ckpt_dir_or_file, f.readline()[:-1]) # -1去掉换行符 + else: + ckpt_path = ckpt_dir_or_file + ckpt = torch.load(ckpt_path, map_location=map_location) + print(' [*] Loading checkpoint succeeds! Copy variables from % s!' % ckpt_path) + return ckpt + + +def save_checkpoint(obj, save_path, is_best=False, max_keep=None): + # save checkpoint + torch.save(obj, save_path) + + # deal with max_keep + save_dir = os.path.dirname(save_path) + list_path = os.path.join(save_dir, 'latest_checkpoint') + + save_path = os.path.basename(save_path) + if os.path.exists(list_path): + with open(list_path) as f: + ckpt_list = f.readlines() + ckpt_list = [save_path + '\n'] + ckpt_list + else: + ckpt_list = [save_path + '\n'] + + if max_keep is not None: + for ckpt in ckpt_list[max_keep:]: + ckpt = os.path.join(save_dir, ckpt[:-1]) + if os.path.exists(ckpt): + os.remove(ckpt) + ckpt_list[max_keep:] = [] + + with open(list_path, 'w') as f: + f.writelines(ckpt_list) + + # copy best + if is_best: + shutil.copyfile(save_path, os.path.join(save_dir, 'best_model.ckpt')) + + +def get_img_from_file(file_name, target_device, transform=False): + img = tform.ToTensor()(Image.open(file_name)) # [0, 1.0] tensor + if transform: + val_test_img_transform = tform.Compose([ + # crop face area 190 * 178 + # tform.Lambda(lambda x: x[:, 14:204, :]), + # center crop and resize on PLI image + tform.ToPILImage(), + tform.CenterCrop(170), # origin # elegant crop + # tform.CenterCrop(178), + tform.Resize(256, Image.BICUBIC), # elegant crop + # back to tensor + tform.ToTensor(), + ]) + img = val_test_img_transform(img) + return (img * 2 - 1).to(target_device) diff --git a/models/CtrlHair/poisson_blending.py b/models/CtrlHair/poisson_blending.py new file mode 100644 index 0000000000000000000000000000000000000000..d2bb103e218591d10e29776c5cff482af78aa9d3 --- /dev/null +++ b/models/CtrlHair/poisson_blending.py @@ -0,0 +1,87 @@ +# -*- coding: utf-8 -*- + +""" +# File name: poisson_blending.py +# Time : 2021/11/5 15:22 +# Author: xyguoo@163.com +# Description: +""" + +import numpy as np +import scipy.sparse +from scipy.sparse.linalg import spsolve + + +def laplacian_matrix(n, m): + mat_D = scipy.sparse.lil_matrix((m, m)) + mat_D.setdiag(-1, -1) + mat_D.setdiag(4) + mat_D.setdiag(-1, 1) + + mat_A = scipy.sparse.block_diag([mat_D] * n).tolil() + + mat_A.setdiag(-1, 1 * m) + mat_A.setdiag(-1, -1 * m) + + return mat_A + + +def poisson_blending(source, target, mask, with_gamma=True): + """ + source: H * W * 3, cv2 image + target: H * W * 3, cv2 image + mask: H * W * 1 + """ + if with_gamma: + gamma_value = 2.2 + else: + gamma_value = 1 + source = source.astype('float') + target = target.astype('float') + source = np.power(source, 1 / gamma_value) + target = np.power(target, 1 / gamma_value) + + res = target.copy() + y_range, x_range = source.shape[:2] + mat_A = laplacian_matrix(y_range, x_range) + laplacian = mat_A.tocsc() + mask[mask != 0] = 1 + + for y in range(1, y_range - 1): + for x in range(1, x_range - 1): + if mask[y, x] == 0: + k = x + y * x_range + mat_A[k, k] = 1 + mat_A[k, k + 1] = 0 + mat_A[k, k - 1] = 0 + mat_A[k, k + x_range] = 0 + mat_A[k, k - x_range] = 0 + mat_A = mat_A.tocsc() + + y_min, y_max = 0, y_range + x_min, x_max = 0, x_range + + mask_flat = mask.flatten() + for channel in range(source.shape[2]): + source_flat = source[y_min:y_max, x_min:x_max, channel].flatten() + target_flat = target[y_min:y_max, x_min:x_max, channel].flatten() + + # inside the mask: + # \Delta f = div v = \Delta g + alpha = 1 + mat_b = laplacian.dot(source_flat) * alpha + + # outside the mask: + # f = t + mat_b[mask_flat == 0] = target_flat[mask_flat == 0] + + x = spsolve(mat_A, mat_b) + x = x.reshape((y_range, x_range)) + res[:, :, channel] = x + + res = np.power(res, gamma_value) + + res[res > 255] = 255 + res[res < 0] = 0 + res = res.astype('uint8') + return res diff --git a/models/CtrlHair/sean_codes/__init__.py b/models/CtrlHair/sean_codes/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3ee410e3c065678a2877bbb16cb2612c8c8a4c51 --- /dev/null +++ b/models/CtrlHair/sean_codes/__init__.py @@ -0,0 +1,8 @@ +# -*- coding: utf-8 -*- + +""" +# File name: __init__.py.py +# Time : 2021/12/7 17:27 +# Author: xyguoo@163.com +# Description: +""" diff --git a/models/CtrlHair/sean_codes/data/__init__.py b/models/CtrlHair/sean_codes/data/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..7743f99330ba072b3a1773b7ba172c7c6540fb77 --- /dev/null +++ b/models/CtrlHair/sean_codes/data/__init__.py @@ -0,0 +1,55 @@ +""" +Copyright (C) 2019 NVIDIA Corporation. All rights reserved. +Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode). +""" + +import importlib +import torch.utils.data +from .base_dataset import BaseDataset + + +def find_dataset_using_name(dataset_name): + # Given the option --dataset [datasetname], + # the file "datasets/datasetname_dataset.py" + # will be imported. + dataset_filename = "sean_codes.data." + dataset_name + "_dataset" + datasetlib = importlib.import_module(dataset_filename) + + # In the file, the class called DatasetNameDataset() will + # be instantiated. It has to be a subclass of BaseDataset, + # and it is case-insensitive. + dataset = None + target_dataset_name = dataset_name.replace('_', '') + 'dataset' + for name, cls in datasetlib.__dict__.items(): + if name.lower() == target_dataset_name.lower() \ + and issubclass(cls, BaseDataset): + dataset = cls + + if dataset is None: + raise ValueError("In %s.py, there should be a subclass of BaseDataset " + "with class name that matches %s in lowercase." % + (dataset_filename, target_dataset_name)) + + return dataset + + +def get_option_setter(dataset_name): + dataset_class = find_dataset_using_name(dataset_name) + return dataset_class.modify_commandline_options + + +def create_dataloader(opt): + dataset = find_dataset_using_name(opt.dataset_mode) + instance = dataset() + instance.initialize(opt) + print("dataset [%s] of size %d was created" % + (type(instance).__name__, len(instance))) + dataloader = torch.utils.data.DataLoader( + instance, + batch_size=opt.batchSize, + shuffle=not opt.serial_batches, + num_workers=min(opt.batchSize, int(opt.nThreads)), + drop_last=opt.isTrain, + pin_memory=True + ) + return dataloader diff --git a/models/CtrlHair/sean_codes/data/base_dataset.py b/models/CtrlHair/sean_codes/data/base_dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..2bb4b3cbf7d73e89eb53eeb14d8c5389a241fba2 --- /dev/null +++ b/models/CtrlHair/sean_codes/data/base_dataset.py @@ -0,0 +1,128 @@ +""" +Copyright (C) 2019 NVIDIA Corporation. All rights reserved. +Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode). +""" + +import torch.utils.data as data +from PIL import Image +import torchvision.transforms as transforms +import numpy as np +import random + + +class BaseDataset(data.Dataset): + def __init__(self): + super(BaseDataset, self).__init__() + + @staticmethod + def modify_commandline_options(parser, is_train): + return parser + + def initialize(self, opt): + pass + + +def get_params(opt, size): + w, h = size + new_h = h + new_w = w + if opt.preprocess_mode == 'resize_and_crop': + new_h = new_w = opt.load_size + elif opt.preprocess_mode == 'scale_width_and_crop': + new_w = opt.load_size + new_h = opt.load_size * h // w + elif opt.preprocess_mode == 'scale_shortside_and_crop': + ss, ls = min(w, h), max(w, h) # shortside and longside + width_is_shorter = w == ss + ls = int(opt.load_size * ls / ss) + new_w, new_h = (ss, ls) if width_is_shorter else (ls, ss) + + x = random.randint(0, np.maximum(0, new_w - opt.crop_size)) + y = random.randint(0, np.maximum(0, new_h - opt.crop_size)) + + flip = random.random() > 0.5 + return {'crop_pos': (x, y), 'flip': flip} + + +def get_transform(opt, params, method=Image.BICUBIC, normalize=True, toTensor=True): + transform_list = [] + if 'resize' in opt.preprocess_mode: + osize = [opt.load_size, opt.load_size] + transform_list.append(transforms.Resize(osize, interpolation=method)) + elif 'scale_width' in opt.preprocess_mode: + transform_list.append(transforms.Lambda(lambda img: __scale_width(img, opt.load_size, method))) + elif 'scale_shortside' in opt.preprocess_mode: + transform_list.append(transforms.Lambda(lambda img: __scale_shortside(img, opt.load_size, method))) + + if 'crop' in opt.preprocess_mode: + transform_list.append(transforms.Lambda(lambda img: __crop(img, params['crop_pos'], opt.crop_size))) + + if opt.preprocess_mode == 'none': + base = 32 + transform_list.append(transforms.Lambda(lambda img: __make_power_2(img, base, method))) + + if opt.preprocess_mode == 'fixed': + w = opt.crop_size + h = round(opt.crop_size / opt.aspect_ratio) + transform_list.append(transforms.Lambda(lambda img: __resize(img, w, h, method))) + + if opt.isTrain and not opt.no_flip: + transform_list.append(transforms.Lambda(lambda img: __flip(img, params['flip']))) + + if toTensor: + transform_list += [transforms.ToTensor()] + + if normalize: + transform_list += [transforms.Normalize((0.5, 0.5, 0.5), + (0.5, 0.5, 0.5))] + return transforms.Compose(transform_list) + + +def normalize(): + return transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) + + +def __resize(img, w, h, method=Image.BICUBIC): + return img.resize((w, h), method) + + +def __make_power_2(img, base, method=Image.BICUBIC): + ow, oh = img.size + h = int(round(oh / base) * base) + w = int(round(ow / base) * base) + if (h == oh) and (w == ow): + return img + return img.resize((w, h), method) + + +def __scale_width(img, target_width, method=Image.BICUBIC): + ow, oh = img.size + if (ow == target_width): + return img + w = target_width + h = int(target_width * oh / ow) + return img.resize((w, h), method) + + +def __scale_shortside(img, target_width, method=Image.BICUBIC): + ow, oh = img.size + ss, ls = min(ow, oh), max(ow, oh) # shortside and longside + width_is_shorter = ow == ss + if (ss == target_width): + return img + ls = int(target_width * ls / ss) + nw, nh = (ss, ls) if width_is_shorter else (ls, ss) + return img.resize((nw, nh), method) + + +def __crop(img, pos, size): + ow, oh = img.size + x1, y1 = pos + tw = th = size + return img.crop((x1, y1, x1 + tw, y1 + th)) + + +def __flip(img, flip): + if flip: + return img.transpose(Image.FLIP_LEFT_RIGHT) + return img diff --git a/models/CtrlHair/sean_codes/data/custom_dataset.py b/models/CtrlHair/sean_codes/data/custom_dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..4b0b4eca3d1dd647088a7060d1c4fcc1c1d10782 --- /dev/null +++ b/models/CtrlHair/sean_codes/data/custom_dataset.py @@ -0,0 +1,50 @@ +""" +Copyright (C) 2019 NVIDIA Corporation. All rights reserved. +Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode). +""" + +from .pix2pix_dataset import Pix2pixDataset +from .image_folder import make_dataset + + +class CustomDataset(Pix2pixDataset): + """ Dataset that loads images from directories + Use option --label_dir, --image_dir, --instance_dir to specify the directories. + The images in the directories are sorted in alphabetical order and paired in order. + """ + + @staticmethod + def modify_commandline_options(parser, is_train): + parser = Pix2pixDataset.modify_commandline_options(parser, is_train) + parser.set_defaults(preprocess_mode='resize_and_crop') + load_size = 286 if is_train else 256 + parser.set_defaults(load_size=load_size) + parser.set_defaults(crop_size=256) + parser.set_defaults(display_winsize=256) + parser.set_defaults(label_nc=19) + parser.set_defaults(contain_dontcare_label=False) + + parser.add_argument('--label_dir', type=str, default='datasets/CelebA-HQ/test/labels', + help='path to the directory that contains label images') + parser.add_argument('--image_dir', type=str, default='datasets/CelebA-HQ/test/images', + help='path to the directory that contains photo images') + parser.add_argument('--instance_dir', type=str, default='', + help='path to the directory that contains instance maps. Leave black if not exists') + return parser + + def get_paths(self, opt): + label_dir = opt.label_dir + label_paths = make_dataset(label_dir, recursive=False, read_cache=True) + + image_dir = opt.image_dir + image_paths = make_dataset(image_dir, recursive=False, read_cache=True) + + if len(opt.instance_dir) > 0: + instance_dir = opt.instance_dir + instance_paths = make_dataset(instance_dir, recursive=False, read_cache=True) + else: + instance_paths = [] + + assert len(label_paths) == len(image_paths), "The #images in %s and %s do not match. Is there something wrong?" + + return label_paths, image_paths, instance_paths diff --git a/models/CtrlHair/sean_codes/data/image_folder.py b/models/CtrlHair/sean_codes/data/image_folder.py new file mode 100644 index 0000000000000000000000000000000000000000..fd157935070526c823a831d4ad4e0b1970ccd14d --- /dev/null +++ b/models/CtrlHair/sean_codes/data/image_folder.py @@ -0,0 +1,101 @@ +""" +Copyright (C) 2019 NVIDIA Corporation. All rights reserved. +Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode). +""" + +############################################################################### +# Code from +# https://github.com/pytorch/vision/blob/master/torchvision/datasets/folder.py +# Modified the original code so that it also loads images from the current +# directory as well as the subdirectories +############################################################################### +import torch.utils.data as data +from PIL import Image +import os + +IMG_EXTENSIONS = [ + '.jpg', '.JPG', '.jpeg', '.JPEG', + '.png', '.PNG', '.ppm', '.PPM', '.bmp', '.BMP', '.tiff', '.webp' +] + + +def is_image_file(filename): + return any(filename.endswith(extension) for extension in IMG_EXTENSIONS) + + +def make_dataset_rec(dir, images): + assert os.path.isdir(dir), '%s is not a valid directory' % dir + + for root, dnames, fnames in sorted(os.walk(dir, followlinks=True)): + for fname in fnames: + if is_image_file(fname): + path = os.path.join(root, fname) + images.append(path) + + +def make_dataset(dir, recursive=False, read_cache=False, write_cache=False): + """ + xuyang: could get image path list + """ + images = [] + + if read_cache: + possible_filelist = os.path.join(dir, 'files.list') + if os.path.isfile(possible_filelist): + with open(possible_filelist, 'r') as f: + images = f.read().splitlines() + return images + + if recursive: + make_dataset_rec(dir, images) + else: + assert os.path.isdir(dir) or os.path.islink(dir), '%s is not a valid directory' % dir + + for root, dnames, fnames in sorted(os.walk(dir)): + for fname in fnames: + if is_image_file(fname): + path = os.path.join(root, fname) + images.append(path) + + if write_cache: + filelist_cache = os.path.join(dir, 'files.list') + with open(filelist_cache, 'w') as f: + for path in images: + f.write("%s\n" % path) + print('wrote filelist cache at %s' % filelist_cache) + + return images + + +def default_loader(path): + return Image.open(path).convert('RGB') + + +class ImageFolder(data.Dataset): + + def __init__(self, root, transform=None, return_paths=False, + loader=default_loader): + imgs = make_dataset(root) + if len(imgs) == 0: + raise(RuntimeError("Found 0 images in: " + root + "\n" + "Supported image extensions are: " + + ",".join(IMG_EXTENSIONS))) + + self.root = root + self.imgs = imgs + self.transform = transform + self.return_paths = return_paths + self.loader = loader + + def __getitem__(self, index): + path = self.imgs[index] + img = self.loader(path) + if self.transform is not None: + img = self.transform(img) + if self.return_paths: + return img, path + else: + return img + + def __len__(self): + return len(self.imgs) diff --git a/models/CtrlHair/sean_codes/data/pix2pix_dataset.py b/models/CtrlHair/sean_codes/data/pix2pix_dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..7001596886d4f06554d926958f6de6ffe25155af --- /dev/null +++ b/models/CtrlHair/sean_codes/data/pix2pix_dataset.py @@ -0,0 +1,139 @@ +""" +Copyright (C) 2019 NVIDIA Corporation. All rights reserved. +Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode). +""" + +from .base_dataset import BaseDataset, get_params, get_transform +from PIL import Image +import util.util as util +import os +import torch + + +class Pix2pixDataset(BaseDataset): + @staticmethod + def modify_commandline_options(parser, is_train): + parser.add_argument('--no_pairing_check', action='store_true', + help='If specified, skip sanity check of correct label-image file pairing') + return parser + + def initialize(self, opt): + self.opt = opt + + label_paths, image_paths, instance_paths = self.get_paths(opt) + + util.natural_sort(label_paths) + util.natural_sort(image_paths) + if not opt.no_instance: + util.natural_sort(instance_paths) + + label_paths = label_paths[:opt.max_dataset_size] + image_paths = image_paths[:opt.max_dataset_size] + instance_paths = instance_paths[:opt.max_dataset_size] + + if not opt.no_pairing_check: + for path1, path2 in zip(label_paths, image_paths): + assert self.paths_match(path1, path2), \ + "The label-image pair (%s, %s) do not look like the right pair because the filenames are quite different. Are you sure about the pairing? Please see data/pix2pix_dataset.py to see what is going on, and use --no_pairing_check to bypass this." % ( + path1, path2) + + self.label_paths = label_paths + self.image_paths = image_paths + self.instance_paths = instance_paths + + size = len(self.label_paths) + self.dataset_size = size + + def get_paths(self, opt): + label_paths = [] + image_paths = [] + instance_paths = [] + assert False, "A subclass of Pix2pixDataset must override self.get_paths(self, opt)" + return label_paths, image_paths, instance_paths + + def paths_match(self, path1, path2): + filename1_without_ext = os.path.splitext(os.path.basename(path1))[0] + filename2_without_ext = os.path.splitext(os.path.basename(path2))[0] + return filename1_without_ext == filename2_without_ext + + def __getitem__(self, index): + # Label Image + label_path = self.label_paths[index] + label = Image.open(label_path) + params = get_params(self.opt, label.size) + transform_label = get_transform(self.opt, params, method=Image.NEAREST, normalize=False) + label_tensor = transform_label(label) * 255.0 + label_tensor[label_tensor == 255] = self.opt.label_nc # 'unknown' is opt.label_nc + + # input image (real images) + image_path = self.image_paths[index] + assert self.paths_match(label_path, image_path), \ + "The label_path %s and image_path %s don't match." % \ + (label_path, image_path) + image = Image.open(image_path) + image = image.convert('RGB') + + transform_image = get_transform(self.opt, params) + image_tensor = transform_image(image) + + # if using instance maps + if self.opt.no_instance: + instance_tensor = 0 + else: + instance_path = self.instance_paths[index] + instance = Image.open(instance_path) + if instance.mode == 'L': + instance_tensor = transform_label(instance) * 255 + instance_tensor = instance_tensor.long() + else: + instance_tensor = transform_label(instance) + + input_dict = {'label': label_tensor, + 'instance': instance_tensor, + 'image': image_tensor, + 'path': image_path, + } + + # Give subclasses a chance to modify the final color_texture + + self.postprocess(input_dict) + + return input_dict + + def postprocess(self, input_dict): + return input_dict + + def __len__(self): + return self.dataset_size + + # Our codes get input images and labels + def get_input_by_names(self, image_path, image, label_img): + label = Image.fromarray(label_img) + params = get_params(self.opt, label.size) + transform_label = get_transform(self.opt, params, method=Image.NEAREST, normalize=False) + label_tensor = transform_label(label) * 255.0 + label_tensor[label_tensor == 255] = self.opt.label_nc # 'unknown' is opt.label_nc + label_tensor.unsqueeze_(0) + + # input image (real images)] + # image = Image.open(image_path) + # image = image.convert('RGB') + + transform_image = get_transform(self.opt, params) + image_tensor = transform_image(image) + image_tensor.unsqueeze_(0) + + # if using instance maps + if self.opt.no_instance: + instance_tensor = torch.Tensor([0]) + + input_dict = {'label': label_tensor, + 'instance': instance_tensor, + 'image': image_tensor, + 'path': image_path, + } + + # Give subclasses a chance to modify the final color_texture + self.postprocess(input_dict) + + return input_dict diff --git a/models/CtrlHair/sean_codes/get_mean_code.py b/models/CtrlHair/sean_codes/get_mean_code.py new file mode 100644 index 0000000000000000000000000000000000000000..31f5dae9633b5aa20a42722113878f7b16c06990 --- /dev/null +++ b/models/CtrlHair/sean_codes/get_mean_code.py @@ -0,0 +1,45 @@ +# -*- coding: utf-8 -*- + +""" +# File name: get_mean_code.py +# Time : 2022/2/22 17:22 +# Author: xyguoo@163.com +# Description: +""" + +from glob import glob + +import numpy as np +import os + +layers_list = ['ACE.npy'] + +style_list = [] + +for cat_i in range(19): + for layer_j in layers_list: + tmp_list = glob('styles_test/style_codes/*/' + str(cat_i) + '/' + layer_j) + style_list = [] + + for k in tmp_list: + style_list.append(np.load(k)) + + if len(style_list) > 0: + style_list = np.array(style_list) + + style_list_norm2 = np.linalg.norm(style_list, axis=1, keepdims=True) ** 2 + dist_matrix = (style_list_norm2 + style_list_norm2.T -2 * style_list @ style_list.T) + dist_matrix[dist_matrix < 0] = 0 + dist_matrix = dist_matrix ** 0.5 + median_idx = dist_matrix.sum(axis=1).argmin() + feature = style_list[median_idx] + + save_folder = os.path.join('styles_test/mean_style_code/median', str(cat_i)) + + if not os.path.exists(save_folder): + os.makedirs(save_folder) + + save_name = os.path.join(save_folder, layer_j) + np.save(save_name, feature) + +print(100) diff --git a/models/CtrlHair/sean_codes/models/__init__.py b/models/CtrlHair/sean_codes/models/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..810d272b2e64dbeb6c7b941221dba59d1a62df8a --- /dev/null +++ b/models/CtrlHair/sean_codes/models/__init__.py @@ -0,0 +1,44 @@ +""" +Copyright (C) 2019 NVIDIA Corporation. All rights reserved. +Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode). +""" + +import importlib +import torch + + +def find_model_using_name(model_name): + # Given the option --model [modelname], + # the file "models/modelname_model.py" + # will be imported. + model_filename = "sean_codes.models." + model_name + "_model" + modellib = importlib.import_module(model_filename) + + # In the file, the class called ModelNameModel() will + # be instantiated. It has to be a subclass of torch.nn.Module, + # and it is case-insensitive. + model = None + target_model_name = model_name.replace('_', '') + 'model' + for name, cls in modellib.__dict__.items(): + if name.lower() == target_model_name.lower() \ + and issubclass(cls, torch.nn.Module): + model = cls + + if model is None: + print("In %s.py, there should be a subclass of torch.nn.Module with class name that matches %s in lowercase." % (model_filename, target_model_name)) + exit(0) + + return model + + +def get_option_setter(model_name): + model_class = find_model_using_name(model_name) + return model_class.modify_commandline_options + + +def create_model(opt): + model = find_model_using_name(opt.sean_model) + instance = model(opt) + print("model [%s] was created" % (type(instance).__name__)) + + return instance diff --git a/models/CtrlHair/sean_codes/models/networks/__init__.py b/models/CtrlHair/sean_codes/models/networks/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..0073164257fccadb1eaac9a67404df473241b7f1 --- /dev/null +++ b/models/CtrlHair/sean_codes/models/networks/__init__.py @@ -0,0 +1,62 @@ +""" +Copyright (C) 2019 NVIDIA Corporation. All rights reserved. +Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode). +""" + +from .base_network import BaseNetwork +from .loss import * +# from sean_codes.models.networks import * +from . import * +from .encoder import * +import util.util as util + + +def find_network_using_name(target_network_name, filename): + target_class_name = target_network_name + filename + module_name = 'sean_codes.models.networks.' + filename + network = util.find_class_in_module(target_class_name, module_name) + + assert issubclass(network, BaseNetwork), \ + "Class %s should be a subclass of BaseNetwork" % network + + return network + + +def modify_commandline_options(parser, is_train): + opt, _ = parser.parse_known_args() + + netG_cls = find_network_using_name(opt.netG, 'generator') + parser = netG_cls.modify_commandline_options(parser, is_train) + if is_train: + netD_cls = find_network_using_name(opt.netD, 'discriminator') + parser = netD_cls.modify_commandline_options(parser, is_train) + netE_cls = find_network_using_name('conv', 'encoder') + parser = netE_cls.modify_commandline_options(parser, is_train) + + return parser + + +def create_network(cls, opt): + net = cls(opt) + net.print_network() + if len(opt.gpu_ids) > 0: + assert(torch.cuda.is_available()) + net.cuda() + net.init_weights(opt.init_type, opt.init_variance) + return net + + +def define_G(opt): + netG_cls = find_network_using_name(opt.netG, 'generator') + return create_network(netG_cls, opt) + + +def define_D(opt): + netD_cls = find_network_using_name(opt.netD, 'discriminator') + return create_network(netD_cls, opt) + + +def define_E(opt): + # there exists only one encoder type + netE_cls = find_network_using_name('conv', 'encoder') + return create_network(netE_cls, opt) diff --git a/models/CtrlHair/sean_codes/models/networks/architecture.py b/models/CtrlHair/sean_codes/models/networks/architecture.py new file mode 100644 index 0000000000000000000000000000000000000000..75f713905dfd3d7b3e57903ea358382545a9be7c --- /dev/null +++ b/models/CtrlHair/sean_codes/models/networks/architecture.py @@ -0,0 +1,207 @@ +""" +Copyright (C) 2019 NVIDIA Corporation. All rights reserved. +Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode). +""" + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torchvision +import torch.nn.utils.spectral_norm as spectral_norm +from sean_codes.models.networks.normalization import ACE + + +# ResNet block that uses SPADE. +# It differs from the ResNet block of pix2pixHD in that +# it takes in the segmentation map as input, learns the skip connection if necessary, +# and applies normalization first and then convolution. +# This architecture seemed like a standard architecture for unconditional or +# class-conditional GAN architecture using residual block. +# The code was inspired from https://github.com/LMescheder/GAN_stability. +class SPADEResnetBlock(nn.Module): + def __init__(self, fin, fout, opt, Block_Name=None, use_rgb=True): + super().__init__() + + self.use_rgb = use_rgb + + self.Block_Name = Block_Name + self.status = opt.status + + # Attributes + self.learned_shortcut = (fin != fout) + fmiddle = min(fin, fout) + + # create conv layers + self.conv_0 = nn.Conv2d(fin, fmiddle, kernel_size=3, padding=1) + self.conv_1 = nn.Conv2d(fmiddle, fout, kernel_size=3, padding=1) + if self.learned_shortcut: + self.conv_s = nn.Conv2d(fin, fout, kernel_size=1, bias=False) + + # apply spectral norm if specified + if 'spectral' in opt.norm_G: + self.conv_0 = spectral_norm(self.conv_0) + self.conv_1 = spectral_norm(self.conv_1) + if self.learned_shortcut: + self.conv_s = spectral_norm(self.conv_s) + + # define normalization layers + spade_config_str = opt.norm_G.replace('spectral', '') + + ########### Modifications 1 + normtype_list = ['spadeinstance3x3', 'spadesyncbatch3x3', 'spadebatch3x3'] + our_norm_type = 'spadesyncbatch3x3' + + self.ace_0 = ACE(our_norm_type, fin, 3, ACE_Name=Block_Name + '_ACE_0', status=self.status, + spade_params=[spade_config_str, fin, opt.semantic_nc], use_rgb=use_rgb) + ########### Modifications 1 + + ########### Modifications 1 + self.ace_1 = ACE(our_norm_type, fmiddle, 3, ACE_Name=Block_Name + '_ACE_1', status=self.status, + spade_params=[spade_config_str, fmiddle, opt.semantic_nc], use_rgb=use_rgb) + ########### Modifications 1 + + if self.learned_shortcut: + self.ace_s = ACE(our_norm_type, fin, 3, ACE_Name=Block_Name + '_ACE_s', status=self.status, + spade_params=[spade_config_str, fin, opt.semantic_nc], use_rgb=use_rgb) + + # note the resnet block with SPADE also takes in |seg|, + # the semantic segmentation map as input + def forward(self, x, seg, style_codes, obj_dic=None): + + x_s = self.shortcut(x, seg, style_codes, obj_dic) + + ########### Modifications 1 + dx = self.ace_0(x, seg, style_codes, obj_dic) + + dx = self.conv_0(self.actvn(dx)) + + dx = self.ace_1(dx, seg, style_codes, obj_dic) + + dx = self.conv_1(self.actvn(dx)) + ########### Modifications 1 + + out = x_s + dx + return out + + def shortcut(self, x, seg, style_codes, obj_dic): + if self.learned_shortcut: + x_s = self.ace_s(x, seg, style_codes, obj_dic) + x_s = self.conv_s(x_s) + + else: + x_s = x + return x_s + + def actvn(self, x): + return F.leaky_relu(x, 2e-1) + + +# ResNet block used in pix2pixHD +# We keep the same architecture as pix2pixHD. +class ResnetBlock(nn.Module): + def __init__(self, dim, norm_layer, activation=nn.ReLU(False), kernel_size=3): + super().__init__() + + pw = (kernel_size - 1) // 2 + self.conv_block = nn.Sequential( + nn.ReflectionPad2d(pw), + norm_layer(nn.Conv2d(dim, dim, kernel_size=kernel_size)), + activation, + nn.ReflectionPad2d(pw), + norm_layer(nn.Conv2d(dim, dim, kernel_size=kernel_size)) + ) + + def forward(self, x): + y = self.conv_block(x) + out = x + y + return out + + +# VGG architecter, used for the perceptual loss using a pretrained VGG network +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 + + +class Zencoder(torch.nn.Module): + def __init__(self, input_nc, output_nc, ngf=32, n_downsampling=2, norm_layer=nn.InstanceNorm2d): + super(Zencoder, self).__init__() + self.output_nc = output_nc + + model = [nn.ReflectionPad2d(1), nn.Conv2d(input_nc, ngf, kernel_size=3, padding=0), + norm_layer(ngf), nn.LeakyReLU(0.2, False)] + ### downsample + for i in range(n_downsampling): + mult = 2 ** i + model += [nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3, stride=2, padding=1), + norm_layer(ngf * mult * 2), nn.LeakyReLU(0.2, False)] + + ### upsample + for i in range(1): + mult = 2 ** (n_downsampling - i) + model += [ + nn.ConvTranspose2d(ngf * mult, ngf * mult * 2, kernel_size=3, stride=2, padding=1, output_padding=1), + norm_layer(int(ngf * mult / 2)), nn.LeakyReLU(0.2, False)] + + model += [nn.ReflectionPad2d(1), nn.Conv2d(256, output_nc, kernel_size=3, padding=0), nn.Tanh()] + self.model = nn.Sequential(*model) + + def forward(self, input, segmap): + + codes = self.model(input) + # codes shape: [N, 512, 128, 128] + segmap = F.interpolate(segmap, size=codes.size()[2:], mode='nearest') + # segmap shape: [N, 19, 128, 128] + # print(segmap.shape) + # print(codes.shape) + + b_size = codes.shape[0] + # h_size = codes.shape[2] + # w_size = codes.shape[3] + f_size = codes.shape[1] + + s_size = segmap.shape[1] + + codes_vector = torch.zeros((b_size, s_size, f_size), dtype=codes.dtype, device=codes.device) + + for i in range(b_size): + for j in range(s_size): + component_mask_area = torch.sum(segmap.bool()[i, j]) + + if component_mask_area > 0: + codes_component_feature = codes[i].masked_select(segmap.bool()[i, j]).reshape(f_size, + component_mask_area).mean( + 1) + codes_vector[i][j] = codes_component_feature + + # codes_avg[i].masked_scatter_(segmap.bool()[i, j], codes_component_mu) + + return codes_vector diff --git a/models/CtrlHair/sean_codes/models/networks/base_network.py b/models/CtrlHair/sean_codes/models/networks/base_network.py new file mode 100644 index 0000000000000000000000000000000000000000..f79189f3c4c9b0032ea5e8ec28c8915a9112a754 --- /dev/null +++ b/models/CtrlHair/sean_codes/models/networks/base_network.py @@ -0,0 +1,59 @@ +""" +Copyright (C) 2019 NVIDIA Corporation. All rights reserved. +Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode). +""" + +import torch.nn as nn +from torch.nn import init + + +class BaseNetwork(nn.Module): + def __init__(self): + super(BaseNetwork, self).__init__() + + @staticmethod + def modify_commandline_options(parser, is_train): + return parser + + def print_network(self): + if isinstance(self, list): + self = self[0] + num_params = 0 + for param in self.parameters(): + num_params += param.numel() + print('Network [%s] was created. Total number of parameters: %.1f million. ' + 'To see the architecture, do print(network).' + % (type(self).__name__, num_params / 1000000)) + + def init_weights(self, init_type='normal', gain=0.02): + def init_func(m): + classname = m.__class__.__name__ + if classname.find('BatchNorm2d') != -1: + if hasattr(m, 'weight') and m.weight is not None: + init.normal_(m.weight.data, 1.0, gain) + if hasattr(m, 'bias') and m.bias is not None: + init.constant_(m.bias.data, 0.0) + elif hasattr(m, 'weight') and (classname.find('Conv') != -1 or classname.find('Linear') != -1): + if init_type == 'normal': + init.normal_(m.weight.data, 0.0, gain) + elif init_type == 'xavier': + init.xavier_normal_(m.weight.data, gain=gain) + elif init_type == 'xavier_uniform': + init.xavier_uniform_(m.weight.data, gain=1.0) + 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=gain) + elif init_type == 'none': # uses pytorch's default init method + m.reset_parameters() + 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) + + self.apply(init_func) + + # propagate to children + for m in self.children(): + if hasattr(m, 'init_weights'): + m.init_weights(init_type, gain) diff --git a/models/CtrlHair/sean_codes/models/networks/discriminator.py b/models/CtrlHair/sean_codes/models/networks/discriminator.py new file mode 100644 index 0000000000000000000000000000000000000000..215f1bfae07d95d74c9c40799988e3ab5124c0bf --- /dev/null +++ b/models/CtrlHair/sean_codes/models/networks/discriminator.py @@ -0,0 +1,119 @@ +""" +Copyright (C) 2019 NVIDIA Corporation. All rights reserved. +Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode). +""" + +import torch.nn as nn +import numpy as np +import torch.nn.functional as F +from sean_codes.models.networks.base_network import BaseNetwork +from sean_codes.models.networks.normalization import get_nonspade_norm_layer +import util.util as util + + +class MultiscaleDiscriminator(BaseNetwork): + @staticmethod + def modify_commandline_options(parser, is_train): + parser.add_argument('--netD_subarch', type=str, default='n_layer', + help='architecture of each discriminator') + parser.add_argument('--num_D', type=int, default=2, + help='number of discriminators to be used in multiscale') + opt, _ = parser.parse_known_args() + + # define properties of each discriminator of the multiscale discriminator + subnetD = util.find_class_in_module(opt.netD_subarch + 'discriminator', + 'models.networks.discriminator') + subnetD.modify_commandline_options(parser, is_train) + + return parser + + def __init__(self, opt): + super().__init__() + self.opt = opt + + for i in range(opt.num_D): + subnetD = self.create_single_discriminator(opt) + self.add_module('discriminator_%d' % i, subnetD) + + def create_single_discriminator(self, opt): + subarch = opt.netD_subarch + if subarch == 'n_layer': + netD = NLayerDiscriminator(opt) + else: + raise ValueError('unrecognized discriminator subarchitecture %s' % subarch) + return netD + + def downsample(self, input): + return F.avg_pool2d(input, kernel_size=3, + stride=2, padding=[1, 1], + count_include_pad=False) + + # Returns list of lists of discriminator outputs. + # The final result is of size opt.num_D x opt.n_layers_D + def forward(self, input): + result = [] + get_intermediate_features = not self.opt.no_ganFeat_loss + for name, D in self.named_children(): + out = D(input) + if not get_intermediate_features: + out = [out] + result.append(out) + input = self.downsample(input) + + return result + + +# Defines the PatchGAN discriminator with the specified arguments. +class NLayerDiscriminator(BaseNetwork): + @staticmethod + def modify_commandline_options(parser, is_train): + parser.add_argument('--n_layers_D', type=int, default=3, + help='# layers in each discriminator') + return parser + + def __init__(self, opt): + super().__init__() + self.opt = opt + + kw = 4 + padw = int(np.ceil((kw - 1.0) / 2)) + nf = opt.ndf + input_nc = self.compute_D_input_nc(opt) + + norm_layer = get_nonspade_norm_layer(opt, opt.norm_D) + sequence = [[nn.Conv2d(input_nc, nf, kernel_size=kw, stride=2, padding=padw), + nn.LeakyReLU(0.2, False)]] + + for n in range(1, opt.n_layers_D): + nf_prev = nf + nf = min(nf * 2, 512) + sequence += [[norm_layer(nn.Conv2d(nf_prev, nf, kernel_size=kw, + stride=2, padding=padw)), + nn.LeakyReLU(0.2, False) + ]] + + sequence += [[nn.Conv2d(nf, 1, kernel_size=kw, stride=1, padding=padw)]] + + # We divide the layers into groups to extract intermediate layer outputs + for n in range(len(sequence)): + self.add_module('model' + str(n), nn.Sequential(*sequence[n])) + + def compute_D_input_nc(self, opt): + input_nc = opt.label_nc + opt.output_nc + if opt.contain_dontcare_label: + input_nc += 1 + if not opt.no_instance: + input_nc += 1 + return input_nc + + def forward(self, input): + results = [input] + for submodel in self.children(): + intermediate_output = submodel(results[-1]) + results.append(intermediate_output) + + get_intermediate_features = not self.opt.no_ganFeat_loss + if get_intermediate_features: + return results[1:] + else: + return results[-1] diff --git a/models/CtrlHair/sean_codes/models/networks/encoder.py b/models/CtrlHair/sean_codes/models/networks/encoder.py new file mode 100644 index 0000000000000000000000000000000000000000..81d9bd542b2170588649b3fa6b6c76014b1ebbfc --- /dev/null +++ b/models/CtrlHair/sean_codes/models/networks/encoder.py @@ -0,0 +1,55 @@ +""" +Copyright (C) 2019 NVIDIA Corporation. All rights reserved. +Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode). +""" + +import torch.nn as nn +import numpy as np +import torch.nn.functional as F +from sean_codes.models.networks.base_network import BaseNetwork +from sean_codes.models.networks.normalization import get_nonspade_norm_layer + + +class ConvEncoder(BaseNetwork): + """ Same architecture as the image discriminator """ + + def __init__(self, opt): + super().__init__() + + kw = 3 + pw = int(np.ceil((kw - 1.0) / 2)) + ndf = opt.ngf + norm_layer = get_nonspade_norm_layer(opt, opt.norm_E) + self.layer1 = norm_layer(nn.Conv2d(3, ndf, kw, stride=2, padding=pw)) + self.layer2 = norm_layer(nn.Conv2d(ndf * 1, ndf * 2, kw, stride=2, padding=pw)) + self.layer3 = norm_layer(nn.Conv2d(ndf * 2, ndf * 4, kw, stride=2, padding=pw)) + self.layer4 = norm_layer(nn.Conv2d(ndf * 4, ndf * 8, kw, stride=2, padding=pw)) + self.layer5 = norm_layer(nn.Conv2d(ndf * 8, ndf * 8, kw, stride=2, padding=pw)) + if opt.crop_size >= 256: + self.layer6 = norm_layer(nn.Conv2d(ndf * 8, ndf * 8, kw, stride=2, padding=pw)) + + self.so = s0 = 4 + self.fc_mu = nn.Linear(ndf * 8 * s0 * s0, 256) + self.fc_var = nn.Linear(ndf * 8 * s0 * s0, 256) + + self.actvn = nn.LeakyReLU(0.2, False) + self.opt = opt + + def forward(self, x): + if x.size(2) != 256 or x.size(3) != 256: + x = F.interpolate(x, size=(256, 256), mode='bilinear') + + x = self.layer1(x) + x = self.layer2(self.actvn(x)) + x = self.layer3(self.actvn(x)) + x = self.layer4(self.actvn(x)) + x = self.layer5(self.actvn(x)) + if self.opt.crop_size >= 256: + x = self.layer6(self.actvn(x)) + x = self.actvn(x) + + x = x.view(x.size(0), -1) + mu = self.fc_mu(x) + logvar = self.fc_var(x) + + return mu, logvar diff --git a/models/CtrlHair/sean_codes/models/networks/generator.py b/models/CtrlHair/sean_codes/models/networks/generator.py new file mode 100644 index 0000000000000000000000000000000000000000..4600e6e15d352f183537756d737e6aac221de503 --- /dev/null +++ b/models/CtrlHair/sean_codes/models/networks/generator.py @@ -0,0 +1,172 @@ +""" +Copyright (C) 2019 NVIDIA Corporation. All rights reserved. +Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode). +""" + +import torch.nn as nn +import torch.nn.functional as F +from .base_network import BaseNetwork +from .architecture import SPADEResnetBlock as SPADEResnetBlock +from .architecture import Zencoder +import torch + + +class SPADEGenerator(BaseNetwork): + @staticmethod + def modify_commandline_options(parser, is_train): + parser.set_defaults(norm_G='spectralspadesyncbatch3x3') + parser.add_argument('--num_upsampling_layers', + choices=('normal', 'more', 'most'), default='normal', + help="If 'more', adds upsampling layer between the two middle resnet blocks. If 'most', also add one more upsampling + resnet layer at the end of the generator") + + return parser + + def __init__(self, opt): + super().__init__() + self.opt = opt + nf = opt.ngf + + self.sw, self.sh = self.compute_latent_vector_size(opt) + + self.Zencoder = Zencoder(3, 512) + + self.fc = nn.Conv2d(self.opt.semantic_nc, 16 * nf, 3, padding=1) + + self.head_0 = SPADEResnetBlock(16 * nf, 16 * nf, opt, Block_Name='head_0') + + self.G_middle_0 = SPADEResnetBlock(16 * nf, 16 * nf, opt, Block_Name='G_middle_0') + self.G_middle_1 = SPADEResnetBlock(16 * nf, 16 * nf, opt, Block_Name='G_middle_1') + + self.up_0 = SPADEResnetBlock(16 * nf, 8 * nf, opt, Block_Name='up_0') + self.up_1 = SPADEResnetBlock(8 * nf, 4 * nf, opt, Block_Name='up_1') + self.up_2 = SPADEResnetBlock(4 * nf, 2 * nf, opt, Block_Name='up_2') + self.up_3 = SPADEResnetBlock(2 * nf, 1 * nf, opt, Block_Name='up_3', use_rgb=False) + + final_nc = nf + + if opt.num_upsampling_layers == 'most': + self.up_4 = SPADEResnetBlock(1 * nf, nf // 2, opt, Block_Name='up_4') + final_nc = nf // 2 + + self.conv_img = nn.Conv2d(final_nc, 3, 3, padding=1) + + self.up = nn.Upsample(scale_factor=2) + # self.up = nn.Upsample(scale_factor=2, mode='bilinear') + + def compute_latent_vector_size(self, opt): + if opt.num_upsampling_layers == 'normal': + num_up_layers = 5 + elif opt.num_upsampling_layers == 'more': + num_up_layers = 6 + elif opt.num_upsampling_layers == 'most': + num_up_layers = 7 + else: + raise ValueError('opt.num_upsampling_layers [%s] not recognized' % + opt.num_upsampling_layers) + + sw = opt.crop_size // (2 ** num_up_layers) + sh = round(sw / opt.aspect_ratio) + + return sw, sh + + def forward(self, input, rgb_img, obj_dic=None): + seg = input + + x = F.interpolate(seg, size=(self.sh, self.sw)) + x = self.fc(x) + + if rgb_img is None: + style_codes = None + else: + style_codes = self.Zencoder(input=rgb_img, segmap=seg) + + x = self.head_0(x, seg, style_codes, obj_dic=obj_dic) + + x = self.up(x) + x = self.G_middle_0(x, seg, style_codes, obj_dic=obj_dic) + + if self.opt.num_upsampling_layers == 'more' or \ + self.opt.num_upsampling_layers == 'most': + x = self.up(x) + + x = self.G_middle_1(x, seg, style_codes, obj_dic=obj_dic) + + x = self.up(x) + x = self.up_0(x, seg, style_codes, obj_dic=obj_dic) + x = self.up(x) + x = self.up_1(x, seg, style_codes, obj_dic=obj_dic) + x = self.up(x) + x = self.up_2(x, seg, style_codes, obj_dic=obj_dic) + x = self.up(x) + x = self.up_3(x, seg, style_codes, obj_dic=obj_dic) + + # if self.opt.num_upsampling_layers == 'most': + # x = self.up(x) + # x= self.up_4(x, seg, style_codes, obj_dic=obj_dic) + + x = self.conv_img(F.leaky_relu(x, 2e-1)) + x = torch.tanh(x) + return x + +# class Pix2PixHDGenerator(BaseNetwork): +# @staticmethod +# def modify_commandline_options(parser, is_train): +# parser.add_argument('--resnet_n_downsample', type=int, default=4, help='number of downsampling layers in netG') +# parser.add_argument('--resnet_n_blocks', type=int, default=9, help='number of residual blocks in the global generator network') +# parser.add_argument('--resnet_kernel_size', type=int, default=3, +# help='kernel size of the resnet block') +# parser.add_argument('--resnet_initial_kernel_size', type=int, default=7, +# help='kernel size of the first convolution') +# parser.set_defaults(norm_G='instance') +# return parser +# +# def __init__(self, opt): +# super().__init__() +# input_nc = opt.label_nc + (1 if opt.contain_dontcare_label else 0) + (0 if opt.no_instance else 1) +# +# norm_layer = get_nonspade_norm_layer(opt, opt.norm_G) +# activation = nn.ReLU(False) +# +# model = [] +# +# # initial conv +# model += [nn.ReflectionPad2d(opt.resnet_initial_kernel_size // 2), +# norm_layer(nn.Conv2d(input_nc, opt.ngf, +# kernel_size=opt.resnet_initial_kernel_size, +# padding=0)), +# activation] +# +# # downsample +# mult = 1 +# for i in range(opt.resnet_n_downsample): +# model += [norm_layer(nn.Conv2d(opt.ngf * mult, opt.ngf * mult * 2, +# kernel_size=3, stride=2, padding=1)), +# activation] +# mult *= 2 +# +# # resnet blocks +# for i in range(opt.resnet_n_blocks): +# model += [ResnetBlock(opt.ngf * mult, +# norm_layer=norm_layer, +# activation=activation, +# kernel_size=opt.resnet_kernel_size)] +# +# # upsample +# for i in range(opt.resnet_n_downsample): +# nc_in = int(opt.ngf * mult) +# nc_out = int((opt.ngf * mult) / 2) +# model += [norm_layer(nn.ConvTranspose2d(nc_in, nc_out, +# kernel_size=3, stride=2, +# padding=1, output_padding=1)), +# activation] +# mult = mult // 2 +# +# # final color_texture conv +# model += [nn.ReflectionPad2d(3), +# nn.Conv2d(nc_out, opt.output_nc, kernel_size=7, padding=0), +# nn.Tanh()] +# +# self.model = nn.Sequential(*model) +# +# def forward(self, input, z=None): +# return self.model(input) diff --git a/models/CtrlHair/sean_codes/models/networks/loss.py b/models/CtrlHair/sean_codes/models/networks/loss.py new file mode 100644 index 0000000000000000000000000000000000000000..f443cd71ca22a087949e11285996a153307f19cf --- /dev/null +++ b/models/CtrlHair/sean_codes/models/networks/loss.py @@ -0,0 +1,120 @@ +""" +Copyright (C) 2019 NVIDIA Corporation. All rights reserved. +Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode). +""" + +import torch +import torch.nn as nn +import torch.nn.functional as F +from .architecture import VGG19 + + +# Defines the GAN loss which uses either LSGAN or the regular GAN. +# When LSGAN is used, it is basically same as MSELoss, +# but it abstracts away the need to create the target label tensor +# that has the same size as the input +class GANLoss(nn.Module): + def __init__(self, gan_mode, target_real_label=1.0, target_fake_label=0.0, + tensor=torch.FloatTensor, opt=None): + super(GANLoss, self).__init__() + self.real_label = target_real_label + self.fake_label = target_fake_label + self.real_label_tensor = None + self.fake_label_tensor = None + self.zero_tensor = None + self.Tensor = tensor + self.gan_mode = gan_mode + self.opt = opt + if gan_mode == 'ls': + pass + elif gan_mode == 'original': + pass + elif gan_mode == 'w': + pass + elif gan_mode == 'hinge': + pass + else: + raise ValueError('Unexpected gan_mode {}'.format(gan_mode)) + + def get_target_tensor(self, input, target_is_real): + if target_is_real: + if self.real_label_tensor is None: + self.real_label_tensor = self.Tensor(1).fill_(self.real_label) + self.real_label_tensor.requires_grad_(False) + return self.real_label_tensor.expand_as(input) + else: + if self.fake_label_tensor is None: + self.fake_label_tensor = self.Tensor(1).fill_(self.fake_label) + self.fake_label_tensor.requires_grad_(False) + return self.fake_label_tensor.expand_as(input) + + def get_zero_tensor(self, input): + if self.zero_tensor is None: + self.zero_tensor = self.Tensor(1).fill_(0) + self.zero_tensor.requires_grad_(False) + return self.zero_tensor.expand_as(input) + + def loss(self, input, target_is_real, for_discriminator=True): + if self.gan_mode == 'original': # cross entropy loss + target_tensor = self.get_target_tensor(input, target_is_real) + loss = F.binary_cross_entropy_with_logits(input, target_tensor) + return loss + elif self.gan_mode == 'ls': + target_tensor = self.get_target_tensor(input, target_is_real) + return F.mse_loss(input, target_tensor) + elif self.gan_mode == 'hinge': + if for_discriminator: + if target_is_real: + minval = torch.min(input - 1, self.get_zero_tensor(input)) + loss = -torch.mean(minval) + else: + minval = torch.min(-input - 1, self.get_zero_tensor(input)) + loss = -torch.mean(minval) + else: + assert target_is_real, "The generator's hinge loss must be aiming for real" + loss = -torch.mean(input) + return loss + else: + # wgan + if target_is_real: + return -input.mean() + else: + return input.mean() + + def __call__(self, input, target_is_real, for_discriminator=True): + # computing loss is a bit complicated because |input| may not be + # a tensor, but list of tensors in case of multiscale discriminator + if isinstance(input, list): + loss = 0 + for pred_i in input: + if isinstance(pred_i, list): + pred_i = pred_i[-1] + loss_tensor = self.loss(pred_i, target_is_real, for_discriminator) + bs = 1 if len(loss_tensor.size()) == 0 else loss_tensor.size(0) + new_loss = torch.mean(loss_tensor.view(bs, -1), dim=1) + loss += new_loss + return loss / len(input) + else: + return self.loss(input, target_is_real, for_discriminator) + + +# Perceptual loss that uses a pretrained VGG network +class VGGLoss(nn.Module): + def __init__(self, gpu_ids): + 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 + + +# KL Divergence loss used in VAE with an image encoder +# class KLDLoss(nn.Module): +# def forward(self, mu, logvar): +# return -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp()) diff --git a/models/CtrlHair/sean_codes/models/networks/normalization.py b/models/CtrlHair/sean_codes/models/networks/normalization.py new file mode 100644 index 0000000000000000000000000000000000000000..ab8ca20962bef200a32519418a9ce115a670017f --- /dev/null +++ b/models/CtrlHair/sean_codes/models/networks/normalization.py @@ -0,0 +1,257 @@ +""" +Copyright (C) 2019 NVIDIA Corporation. All rights reserved. +Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode). +""" + +import re +import torch +import torch.nn as nn +import torch.nn.functional as F +from sean_codes.models.networks.sync_batchnorm import SynchronizedBatchNorm2d +import torch.nn.utils.spectral_norm as spectral_norm +import os +import numpy as np + + +# Returns a function that creates a normalization function +# that does not condition on semantic map +def get_nonspade_norm_layer(opt, norm_type='instance'): + # helper function to get # color_texture channels of the previous layer + def get_out_channel(layer): + if hasattr(layer, 'out_channels'): + return getattr(layer, 'out_channels') + return layer.weight.size(0) + + # this function will be returned + def add_norm_layer(layer): + nonlocal norm_type + if norm_type.startswith('spectral'): + layer = spectral_norm(layer) + subnorm_type = norm_type[len('spectral'):] + + if subnorm_type == 'none' or len(subnorm_type) == 0: + return layer + + # remove bias in the previous layer, which is meaningless + # since it has no effect after normalization + if getattr(layer, 'bias', None) is not None: + delattr(layer, 'bias') + layer.register_parameter('bias', None) + + if subnorm_type == 'batch': + norm_layer = nn.BatchNorm2d(get_out_channel(layer), affine=True) + elif subnorm_type == 'sync_batch': + norm_layer = SynchronizedBatchNorm2d(get_out_channel(layer), affine=True) + elif subnorm_type == 'instance': + norm_layer = nn.InstanceNorm2d(get_out_channel(layer), affine=False) + else: + raise ValueError('normalization layer %s is not recognized' % subnorm_type) + + return nn.Sequential(layer, norm_layer) + + return add_norm_layer + + +# Creates SPADE normalization layer based on the given configuration +# SPADE consists of two steps. First, it normalizes the activations using +# your favorite normalization method, such as Batch Norm or Instance Norm. +# Second, it applies scale and bias to the normalized color_texture, conditioned on +# the segmentation map. +# The format of |config_text| is spade(norm)(ks), where +# (norm) specifies the type of parameter-free normalization. +# (e.g. syncbatch, batch, instance) +# (ks) specifies the size of kernel in the SPADE module (e.g. 3x3) +# Example |config_text| will be spadesyncbatch3x3, or spadeinstance5x5. +# Also, the other arguments are +# |norm_nc|: the #channels of the normalized activations, hence the color_texture dim of SPADE +# |label_nc|: the #channels of the input semantic map, hence the input dim of SPADE + + +class ACE(nn.Module): + def __init__(self, config_text, norm_nc, label_nc, ACE_Name=None, status='scripts', spade_params=None, use_rgb=True): + super().__init__() + + self.ACE_Name = ACE_Name + self.status = status + self.save_npy = True + self.Spade = SPADE(*spade_params) + self.use_rgb = use_rgb + self.style_length = 512 + self.blending_gamma = nn.Parameter(torch.zeros(1), requires_grad=True) + self.blending_beta = nn.Parameter(torch.zeros(1), requires_grad=True) + self.noise_var = nn.Parameter(torch.zeros(norm_nc), requires_grad=True) + + assert config_text.startswith('spade') + parsed = re.search('spade(\D+)(\d)x\d', config_text) + param_free_norm_type = str(parsed.group(1)) + ks = int(parsed.group(2)) + pw = ks // 2 + + if param_free_norm_type == 'instance': + self.param_free_norm = nn.InstanceNorm2d(norm_nc, affine=False) + elif param_free_norm_type == 'syncbatch': + self.param_free_norm = SynchronizedBatchNorm2d(norm_nc, affine=False) + elif param_free_norm_type == 'batch': + self.param_free_norm = nn.BatchNorm2d(norm_nc, affine=False) + else: + raise ValueError('%s is not a recognized param-free norm type in SPADE' + % param_free_norm_type) + + # The dimension of the intermediate embedding space. Yes, hardcoded. + + if self.use_rgb: + self.create_gamma_beta_fc_layers() + + self.conv_gamma = nn.Conv2d(self.style_length, norm_nc, kernel_size=ks, padding=pw) + self.conv_beta = nn.Conv2d(self.style_length, norm_nc, kernel_size=ks, padding=pw) + + def forward(self, x, segmap, style_codes=None, obj_dic=None): + + # Part 1. generate parameter-free normalized activations + added_noise = (torch.randn(x.shape[0], x.shape[3], x.shape[2], 1).cuda() * self.noise_var).transpose(1, 3) + normalized = self.param_free_norm(x + added_noise) + + # Part 2. produce scaling and bias conditioned on semantic map + segmap = F.interpolate(segmap, size=x.size()[2:], mode='nearest') + + if self.use_rgb: + [b_size, f_size, h_size, w_size] = normalized.shape + middle_avg = torch.zeros((b_size, self.style_length, h_size, w_size), device=normalized.device) + + if self.status == 'UI_mode': + ############## hard coding + + for i in range(1): + for j in range(segmap.shape[1]): + + component_mask_area = torch.sum(segmap.bool()[i, j]) + + if component_mask_area > 0: + if obj_dic is None: + print('wrong even it is the first input') + else: + style_code_tmp = obj_dic[str(j)]['ACE'] + + middle_mu = F.relu(self.__getattr__('fc_mu' + str(j))(style_code_tmp)) + component_mu = middle_mu.reshape(self.style_length, 1).expand(self.style_length, + component_mask_area) + + middle_avg[i].masked_scatter_(segmap.bool()[i, j], component_mu) + + else: + + for i in range(b_size): + for j in range(segmap.shape[1]): + component_mask_area = torch.sum(segmap.bool()[i, j]) + + if component_mask_area > 0: + + middle_mu = F.relu(self.__getattr__('fc_mu' + str(j))(style_codes[i][j])) + component_mu = middle_mu.reshape(self.style_length, 1).expand(self.style_length, + component_mask_area) + + middle_avg[i].masked_scatter_(segmap.bool()[i, j], component_mu) + + if self.status == 'test' and self.save_npy and self.ACE_Name == 'up_2_ACE_0': + tmp = style_codes[i][j].cpu().numpy() + # gil + if obj_dic[i].find('temp_path') != -1: + dir_path = os.path.split(obj_dic[i])[0] + else: + dir_path = 'styles_test' + im_name = os.path.basename(obj_dic[i]) + folder_path = os.path.join(dir_path, 'style_codes', im_name, str(j)) + ############### some problem with obj_dic[i] + + if not os.path.exists(folder_path): + os.makedirs(folder_path) + + style_code_path = os.path.join(folder_path, 'ACE.npy') + np.save(style_code_path, tmp) + + gamma_avg = self.conv_gamma(middle_avg) + beta_avg = self.conv_beta(middle_avg) + + gamma_spade, beta_spade = self.Spade(segmap) + + gamma_alpha = torch.sigmoid(self.blending_gamma) + beta_alpha = torch.sigmoid(self.blending_beta) + + gamma_final = gamma_alpha * gamma_avg + (1 - gamma_alpha) * gamma_spade + beta_final = beta_alpha * beta_avg + (1 - beta_alpha) * beta_spade + out = normalized * (1 + gamma_final) + beta_final + else: + gamma_spade, beta_spade = self.Spade(segmap) + gamma_final = gamma_spade + beta_final = beta_spade + out = normalized * (1 + gamma_final) + beta_final + + return out + + def create_gamma_beta_fc_layers(self): + + ################### These codes should be replaced with torch.nn.ModuleList + + style_length = self.style_length + + self.fc_mu0 = nn.Linear(style_length, style_length) + self.fc_mu1 = nn.Linear(style_length, style_length) + self.fc_mu2 = nn.Linear(style_length, style_length) + self.fc_mu3 = nn.Linear(style_length, style_length) + self.fc_mu4 = nn.Linear(style_length, style_length) + self.fc_mu5 = nn.Linear(style_length, style_length) + self.fc_mu6 = nn.Linear(style_length, style_length) + self.fc_mu7 = nn.Linear(style_length, style_length) + self.fc_mu8 = nn.Linear(style_length, style_length) + self.fc_mu9 = nn.Linear(style_length, style_length) + self.fc_mu10 = nn.Linear(style_length, style_length) + self.fc_mu11 = nn.Linear(style_length, style_length) + self.fc_mu12 = nn.Linear(style_length, style_length) + self.fc_mu13 = nn.Linear(style_length, style_length) + self.fc_mu14 = nn.Linear(style_length, style_length) + self.fc_mu15 = nn.Linear(style_length, style_length) + self.fc_mu16 = nn.Linear(style_length, style_length) + self.fc_mu17 = nn.Linear(style_length, style_length) + self.fc_mu18 = nn.Linear(style_length, style_length) + + +class SPADE(nn.Module): + def __init__(self, config_text, norm_nc, label_nc): + super().__init__() + + assert config_text.startswith('spade') + parsed = re.search('spade(\D+)(\d)x\d', config_text) + param_free_norm_type = str(parsed.group(1)) + ks = int(parsed.group(2)) + + if param_free_norm_type == 'instance': + self.param_free_norm = nn.InstanceNorm2d(norm_nc, affine=False) + elif param_free_norm_type == 'syncbatch': + self.param_free_norm = SynchronizedBatchNorm2d(norm_nc, affine=False) + elif param_free_norm_type == 'batch': + self.param_free_norm = nn.BatchNorm2d(norm_nc, affine=False) + else: + raise ValueError('%s is not a recognized param-free norm type in SPADE' + % param_free_norm_type) + + # The dimension of the intermediate embedding space. Yes, hardcoded. + nhidden = 128 + + pw = ks // 2 + self.mlp_shared = nn.Sequential( + nn.Conv2d(label_nc, nhidden, kernel_size=ks, padding=pw), + nn.ReLU() + ) + + self.mlp_gamma = nn.Conv2d(nhidden, norm_nc, kernel_size=ks, padding=pw) + self.mlp_beta = nn.Conv2d(nhidden, norm_nc, kernel_size=ks, padding=pw) + + def forward(self, segmap): + + inputmap = segmap + + actv = self.mlp_shared(inputmap) + gamma = self.mlp_gamma(actv) + beta = self.mlp_beta(actv) + + return gamma, beta diff --git a/models/CtrlHair/sean_codes/models/networks/sync_batchnorm/__init__.py b/models/CtrlHair/sean_codes/models/networks/sync_batchnorm/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..84ef0a02ec3d1649a62052c65ef1c75e2eaeb5bb --- /dev/null +++ b/models/CtrlHair/sean_codes/models/networks/sync_batchnorm/__init__.py @@ -0,0 +1,13 @@ +# -*- coding: utf-8 -*- +# File : __init__.py +# Author : Jiayuan Mao +# Email : maojiayuan@gmail.com +# Date : 27/01/2018 +# +# This file is part of Synchronized-BatchNorm-PyTorch. +# https://github.com/vacancy/Synchronized-BatchNorm-PyTorch +# Distributed under MIT License. + +from .batchnorm import SynchronizedBatchNorm1d, SynchronizedBatchNorm2d, SynchronizedBatchNorm3d +from .batchnorm import convert_model +from .replicate import DataParallelWithCallback, patch_replication_callback diff --git a/models/CtrlHair/sean_codes/models/networks/sync_batchnorm/batchnorm.py b/models/CtrlHair/sean_codes/models/networks/sync_batchnorm/batchnorm.py new file mode 100644 index 0000000000000000000000000000000000000000..27f4b8bc7a24d19d5ab75e2d2bc2f5b23d2a95d2 --- /dev/null +++ b/models/CtrlHair/sean_codes/models/networks/sync_batchnorm/batchnorm.py @@ -0,0 +1,361 @@ +# -*- coding: utf-8 -*- +# File : batchnorm.py +# Author : Jiayuan Mao +# Email : maojiayuan@gmail.com +# Date : 27/01/2018 +# +# This file is part of Synchronized-BatchNorm-PyTorch. +# https://github.com/vacancy/Synchronized-BatchNorm-PyTorch +# Distributed under MIT License. + +import collections + +import torch +import torch.nn.functional as F + +from torch.nn.modules.batchnorm import _BatchNorm +from torch.nn.parallel._functions import ReduceAddCoalesced, Broadcast + +from .comm import SyncMaster +from .replicate import DataParallelWithCallback + +__all__ = ['SynchronizedBatchNorm1d', 'SynchronizedBatchNorm2d', + 'SynchronizedBatchNorm3d', 'convert_model'] + + +def _sum_ft(tensor): + """sum over the first and last dimention""" + return tensor.sum(dim=0).sum(dim=-1) + + +def _unsqueeze_ft(tensor): + """add new dementions at the front and the tail""" + return tensor.unsqueeze(0).unsqueeze(-1) + + +_ChildMessage = collections.namedtuple('_ChildMessage', ['sum', 'ssum', 'sum_size']) +_MasterMessage = collections.namedtuple('_MasterMessage', ['sum', 'inv_std']) + + +class _SynchronizedBatchNorm(_BatchNorm): + def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=True): + super(_SynchronizedBatchNorm, self).__init__(num_features, eps=eps, momentum=momentum, affine=affine) + + self._sync_master = SyncMaster(self._data_parallel_master) + + self._is_parallel = False + self._parallel_id = None + self._slave_pipe = None + + def forward(self, input): + # If it is not parallel computation or is in evaluation mode, use PyTorch's implementation. + if not (self._is_parallel and self.training): + return F.batch_norm( + input, self.running_mean, self.running_var, self.weight, self.bias, + self.training, self.momentum, self.eps) + + # Resize the input to (B, C, -1). + input_shape = input.size() + input = input.view(input.size(0), self.num_features, -1) + + # Compute the sum and square-sum. + sum_size = input.size(0) * input.size(2) + input_sum = _sum_ft(input) + input_ssum = _sum_ft(input ** 2) + + # Reduce-and-broadcast the statistics. + if self._parallel_id == 0: + mean, inv_std = self._sync_master.run_master(_ChildMessage(input_sum, input_ssum, sum_size)) + else: + mean, inv_std = self._slave_pipe.run_slave(_ChildMessage(input_sum, input_ssum, sum_size)) + + # Compute the color_texture. + if self.affine: + # MJY:: Fuse the multiplication for speed. + output = (input - _unsqueeze_ft(mean)) * _unsqueeze_ft(inv_std * self.weight) + _unsqueeze_ft(self.bias) + else: + output = (input - _unsqueeze_ft(mean)) * _unsqueeze_ft(inv_std) + + # Reshape it. + return output.view(input_shape) + + def __data_parallel_replicate__(self, ctx, copy_id): + self._is_parallel = True + self._parallel_id = copy_id + + # parallel_id == 0 means master device. + if self._parallel_id == 0: + ctx.sync_master = self._sync_master + else: + self._slave_pipe = ctx.sync_master.register_slave(copy_id) + + def _data_parallel_master(self, intermediates): + """Reduce the sum and square-sum, compute the statistics, and broadcast it.""" + + # Always using same "device order" makes the ReduceAdd operation faster. + # Thanks to:: Tete Xiao (http://tetexiao.com/) + intermediates = sorted(intermediates, key=lambda i: i[1].sum.get_device()) + + to_reduce = [i[1][:2] for i in intermediates] + to_reduce = [j for i in to_reduce for j in i] # flatten + target_gpus = [i[1].sum.get_device() for i in intermediates] + + sum_size = sum([i[1].sum_size for i in intermediates]) + sum_, ssum = ReduceAddCoalesced.apply(target_gpus[0], 2, *to_reduce) + mean, inv_std = self._compute_mean_std(sum_, ssum, sum_size) + + broadcasted = Broadcast.apply(target_gpus, mean, inv_std) + + outputs = [] + for i, rec in enumerate(intermediates): + outputs.append((rec[0], _MasterMessage(*broadcasted[i*2:i*2+2]))) + + return outputs + + def _compute_mean_std(self, sum_, ssum, size): + """Compute the mean and standard-deviation with sum and square-sum. This method + also maintains the moving average on the master device.""" + assert size > 1, 'BatchNorm computes unbiased standard-deviation, which requires size > 1.' + mean = sum_ / size + sumvar = ssum - sum_ * mean + unbias_var = sumvar / (size - 1) + bias_var = sumvar / size + + self.running_mean = (1 - self.momentum) * self.running_mean + self.momentum * mean.data + self.running_var = (1 - self.momentum) * self.running_var + self.momentum * unbias_var.data + + return mean, bias_var.clamp(self.eps) ** -0.5 + + +class SynchronizedBatchNorm1d(_SynchronizedBatchNorm): + r"""Applies Synchronized Batch Normalization over a 2d or 3d input that is seen as a + mini-batch. + + .. math:: + + y = \frac{x - mean[x]}{ \sqrt{Var[x] + \epsilon}} * gamma + beta + + This module differs from the built-in PyTorch BatchNorm1d as the mean and + standard-deviation are reduced across all devices during training. + + For example, when one uses `nn.DataParallel` to wrap the network during + training, PyTorch's implementation normalize the tensor on each device using + the statistics only on that device, which accelerated the computation and + is also easy to implement, but the statistics might be inaccurate. + Instead, in this synchronized version, the statistics will be computed + over all training samples distributed on multiple devices. + + Note that, for one-GPU or CPU-only case, this module behaves exactly same + as the built-in PyTorch implementation. + + The mean and standard-deviation are calculated per-dimension over + the mini-batches and gamma and beta are learnable parameter vectors + of size C (where C is the input size). + + During training, this layer keeps a running estimate of its computed mean + and variance. The running sum is kept with a default momentum of 0.1. + + During evaluation, this running mean/variance is used for normalization. + + Because the BatchNorm is done over the `C` dimension, computing statistics + on `(N, L)` slices, it's common terminology to call this Temporal BatchNorm + + Args: + num_features: num_features from an expected input of size + `batch_size x num_features [x width]` + eps: a value added to the denominator for numerical stability. + Default: 1e-5 + momentum: the value used for the running_mean and running_var + computation. Default: 0.1 + affine: a boolean value that when set to ``True``, gives the layer learnable + affine parameters. Default: ``True`` + + Shape: + - Input: :math:`(N, C)` or :math:`(N, C, L)` + - Output: :math:`(N, C)` or :math:`(N, C, L)` (same shape as input) + + Examples: + >>> # With Learnable Parameters + >>> m = SynchronizedBatchNorm1d(100) + >>> # Without Learnable Parameters + >>> m = SynchronizedBatchNorm1d(100, affine=False) + >>> input = torch.autograd.Variable(torch.randn(20, 100)) + >>> color_texture = m(input) + """ + + def _check_input_dim(self, input): + if input.dim() != 2 and input.dim() != 3: + raise ValueError('expected 2D or 3D input (got {}D input)' + .format(input.dim())) + super(SynchronizedBatchNorm1d, self)._check_input_dim(input) + + +class SynchronizedBatchNorm2d(_SynchronizedBatchNorm): + r"""Applies Batch Normalization over a 4d input that is seen as a mini-batch + of 3d inputs + + .. math:: + + y = \frac{x - mean[x]}{ \sqrt{Var[x] + \epsilon}} * gamma + beta + + This module differs from the built-in PyTorch BatchNorm2d as the mean and + standard-deviation are reduced across all devices during training. + + For example, when one uses `nn.DataParallel` to wrap the network during + training, PyTorch's implementation normalize the tensor on each device using + the statistics only on that device, which accelerated the computation and + is also easy to implement, but the statistics might be inaccurate. + Instead, in this synchronized version, the statistics will be computed + over all training samples distributed on multiple devices. + + Note that, for one-GPU or CPU-only case, this module behaves exactly same + as the built-in PyTorch implementation. + + The mean and standard-deviation are calculated per-dimension over + the mini-batches and gamma and beta are learnable parameter vectors + of size C (where C is the input size). + + During training, this layer keeps a running estimate of its computed mean + and variance. The running sum is kept with a default momentum of 0.1. + + During evaluation, this running mean/variance is used for normalization. + + Because the BatchNorm is done over the `C` dimension, computing statistics + on `(N, H, W)` slices, it's common terminology to call this Spatial BatchNorm + + Args: + num_features: num_features from an expected input of + size batch_size x num_features x height x width + eps: a value added to the denominator for numerical stability. + Default: 1e-5 + momentum: the value used for the running_mean and running_var + computation. Default: 0.1 + affine: a boolean value that when set to ``True``, gives the layer learnable + affine parameters. Default: ``True`` + + Shape: + - Input: :math:`(N, C, H, W)` + - Output: :math:`(N, C, H, W)` (same shape as input) + + Examples: + >>> # With Learnable Parameters + >>> m = SynchronizedBatchNorm2d(100) + >>> # Without Learnable Parameters + >>> m = SynchronizedBatchNorm2d(100, affine=False) + >>> input = torch.autograd.Variable(torch.randn(20, 100, 35, 45)) + >>> color_texture = m(input) + """ + + def _check_input_dim(self, input): + if input.dim() != 4: + raise ValueError('expected 4D input (got {}D input)' + .format(input.dim())) + super(SynchronizedBatchNorm2d, self)._check_input_dim(input) + + +class SynchronizedBatchNorm3d(_SynchronizedBatchNorm): + r"""Applies Batch Normalization over a 5d input that is seen as a mini-batch + of 4d inputs + + .. math:: + + y = \frac{x - mean[x]}{ \sqrt{Var[x] + \epsilon}} * gamma + beta + + This module differs from the built-in PyTorch BatchNorm3d as the mean and + standard-deviation are reduced across all devices during training. + + For example, when one uses `nn.DataParallel` to wrap the network during + training, PyTorch's implementation normalize the tensor on each device using + the statistics only on that device, which accelerated the computation and + is also easy to implement, but the statistics might be inaccurate. + Instead, in this synchronized version, the statistics will be computed + over all training samples distributed on multiple devices. + + Note that, for one-GPU or CPU-only case, this module behaves exactly same + as the built-in PyTorch implementation. + + The mean and standard-deviation are calculated per-dimension over + the mini-batches and gamma and beta are learnable parameter vectors + of size C (where C is the input size). + + During training, this layer keeps a running estimate of its computed mean + and variance. The running sum is kept with a default momentum of 0.1. + + During evaluation, this running mean/variance is used for normalization. + + Because the BatchNorm is done over the `C` dimension, computing statistics + on `(N, D, H, W)` slices, it's common terminology to call this Volumetric BatchNorm + or Spatio-temporal BatchNorm + + Args: + num_features: num_features from an expected input of + size batch_size x num_features x depth x height x width + eps: a value added to the denominator for numerical stability. + Default: 1e-5 + momentum: the value used for the running_mean and running_var + computation. Default: 0.1 + affine: a boolean value that when set to ``True``, gives the layer learnable + affine parameters. Default: ``True`` + + Shape: + - Input: :math:`(N, C, D, H, W)` + - Output: :math:`(N, C, D, H, W)` (same shape as input) + + Examples: + >>> # With Learnable Parameters + >>> m = SynchronizedBatchNorm3d(100) + >>> # Without Learnable Parameters + >>> m = SynchronizedBatchNorm3d(100, affine=False) + >>> input = torch.autograd.Variable(torch.randn(20, 100, 35, 45, 10)) + >>> color_texture = m(input) + """ + + def _check_input_dim(self, input): + if input.dim() != 5: + raise ValueError('expected 5D input (got {}D input)' + .format(input.dim())) + super(SynchronizedBatchNorm3d, self)._check_input_dim(input) + + +def convert_model(module): + """Traverse the input module and its child recursively + and replace all instance of torch.nn.modules.batchnorm.BatchNorm*N*d + to SynchronizedBatchNorm*N*d + + Args: + module: the input module needs to be convert to SyncBN model + + Examples: + >>> import torch.nn as nn + >>> import torchvision + >>> # m is a standard pytorch model + >>> m = torchvision.models.resnet18(True) + >>> m = nn.DataParallel(m) + >>> # after convert, m is using SyncBN + >>> m = convert_model(m) + """ + if isinstance(module, torch.nn.DataParallel): + mod = module.module + mod = convert_model(mod) + mod = DataParallelWithCallback(mod) + return mod + + mod = module + for pth_module, sync_module in zip([torch.nn.modules.batchnorm.BatchNorm1d, + torch.nn.modules.batchnorm.BatchNorm2d, + torch.nn.modules.batchnorm.BatchNorm3d], + [SynchronizedBatchNorm1d, + SynchronizedBatchNorm2d, + SynchronizedBatchNorm3d]): + if isinstance(module, pth_module): + mod = sync_module(module.num_features, module.eps, module.momentum, module.affine) + mod.running_mean = module.running_mean + mod.running_var = module.running_var + if module.affine: + mod.weight.data = module.weight.data.clone().detach() + mod.bias.data = module.bias.data.clone().detach() + + for name, child in module.named_children(): + mod.add_module(name, convert_model(child)) + + return mod diff --git a/models/CtrlHair/sean_codes/models/networks/sync_batchnorm/batchnorm_reimpl.py b/models/CtrlHair/sean_codes/models/networks/sync_batchnorm/batchnorm_reimpl.py new file mode 100644 index 0000000000000000000000000000000000000000..7afcdaff9c56d7ac9c487f2dbe61fe6cb9c353a0 --- /dev/null +++ b/models/CtrlHair/sean_codes/models/networks/sync_batchnorm/batchnorm_reimpl.py @@ -0,0 +1,74 @@ +#! /usr/bin/env python3 +# -*- coding: utf-8 -*- +# File : batchnorm_reimpl.py +# Author : acgtyrant +# Date : 11/01/2018 +# +# This file is part of Synchronized-BatchNorm-PyTorch. +# https://github.com/vacancy/Synchronized-BatchNorm-PyTorch +# Distributed under MIT License. + +import torch +import torch.nn as nn +import torch.nn.init as init + +__all__ = ['BatchNormReimpl'] + + +class BatchNorm2dReimpl(nn.Module): + """ + A re-implementation of batch normalization, used for testing the numerical + stability. + + Author: acgtyrant + See also: + https://github.com/vacancy/Synchronized-BatchNorm-PyTorch/issues/14 + """ + def __init__(self, num_features, eps=1e-5, momentum=0.1): + super().__init__() + + self.num_features = num_features + self.eps = eps + self.momentum = momentum + self.weight = nn.Parameter(torch.empty(num_features)) + self.bias = nn.Parameter(torch.empty(num_features)) + self.register_buffer('running_mean', torch.zeros(num_features)) + self.register_buffer('running_var', torch.ones(num_features)) + self.reset_parameters() + + def reset_running_stats(self): + self.running_mean.zero_() + self.running_var.fill_(1) + + def reset_parameters(self): + self.reset_running_stats() + init.uniform_(self.weight) + init.zeros_(self.bias) + + def forward(self, input_): + batchsize, channels, height, width = input_.size() + numel = batchsize * height * width + input_ = input_.permute(1, 0, 2, 3).contiguous().view(channels, numel) + sum_ = input_.sum(1) + sum_of_square = input_.pow(2).sum(1) + mean = sum_ / numel + sumvar = sum_of_square - sum_ * mean + + self.running_mean = ( + (1 - self.momentum) * self.running_mean + + self.momentum * mean.detach() + ) + unbias_var = sumvar / (numel - 1) + self.running_var = ( + (1 - self.momentum) * self.running_var + + self.momentum * unbias_var.detach() + ) + + bias_var = sumvar / numel + inv_std = 1 / (bias_var + self.eps).pow(0.5) + output = ( + (input_ - mean.unsqueeze(1)) * inv_std.unsqueeze(1) * + self.weight.unsqueeze(1) + self.bias.unsqueeze(1)) + + return output.view(channels, batchsize, height, width).permute(1, 0, 2, 3).contiguous() + diff --git a/models/CtrlHair/sean_codes/models/networks/sync_batchnorm/comm.py b/models/CtrlHair/sean_codes/models/networks/sync_batchnorm/comm.py new file mode 100644 index 0000000000000000000000000000000000000000..922f8c4a3adaa9b32fdcaef09583be03b0d7eb2b --- /dev/null +++ b/models/CtrlHair/sean_codes/models/networks/sync_batchnorm/comm.py @@ -0,0 +1,137 @@ +# -*- coding: utf-8 -*- +# File : comm.py +# Author : Jiayuan Mao +# Email : maojiayuan@gmail.com +# Date : 27/01/2018 +# +# This file is part of Synchronized-BatchNorm-PyTorch. +# https://github.com/vacancy/Synchronized-BatchNorm-PyTorch +# Distributed under MIT License. + +import queue +import collections +import threading + +__all__ = ['FutureResult', 'SlavePipe', 'SyncMaster'] + + +class FutureResult(object): + """A thread-safe future implementation. Used only as one-to-one pipe.""" + + def __init__(self): + self._result = None + self._lock = threading.Lock() + self._cond = threading.Condition(self._lock) + + def put(self, result): + with self._lock: + assert self._result is None, 'Previous result has\'t been fetched.' + self._result = result + self._cond.notify() + + def get(self): + with self._lock: + if self._result is None: + self._cond.wait() + + res = self._result + self._result = None + return res + + +_MasterRegistry = collections.namedtuple('MasterRegistry', ['result']) +_SlavePipeBase = collections.namedtuple('_SlavePipeBase', ['identifier', 'queue', 'result']) + + +class SlavePipe(_SlavePipeBase): + """Pipe for master-slave communication.""" + + def run_slave(self, msg): + self.queue.put((self.identifier, msg)) + ret = self.result.get() + self.queue.put(True) + return ret + + +class SyncMaster(object): + """An abstract `SyncMaster` object. + + - During the replication, as the data parallel will trigger an callback of each module, all slave devices should + call `register(id)` and obtain an `SlavePipe` to communicate with the master. + - During the forward pass, master device invokes `run_master`, all messages from slave devices will be collected, + and passed to a registered callback. + - After receiving the messages, the master device should gather the information and determine to message passed + back to each slave devices. + """ + + def __init__(self, master_callback): + """ + + Args: + master_callback: a callback to be invoked after having collected messages from slave devices. + """ + self._master_callback = master_callback + self._queue = queue.Queue() + self._registry = collections.OrderedDict() + self._activated = False + + def __getstate__(self): + return {'master_callback': self._master_callback} + + def __setstate__(self, state): + self.__init__(state['master_callback']) + + def register_slave(self, identifier): + """ + Register an slave device. + + Args: + identifier: an identifier, usually is the device id. + + Returns: a `SlavePipe` object which can be used to communicate with the master device. + + """ + if self._activated: + assert self._queue.empty(), 'Queue is not clean before next initialization.' + self._activated = False + self._registry.clear() + future = FutureResult() + self._registry[identifier] = _MasterRegistry(future) + return SlavePipe(identifier, self._queue, future) + + def run_master(self, master_msg): + """ + Main entry for the master device in each forward pass. + The messages were first collected from each devices (including the master device), and then + an callback will be invoked to compute the message to be sent back to each devices + (including the master device). + + Args: + master_msg: the message that the master want to send to itself. This will be placed as the first + message when calling `master_callback`. For detailed usage, see `_SynchronizedBatchNorm` for an example. + + Returns: the message to be sent back to the master device. + + """ + self._activated = True + + intermediates = [(0, master_msg)] + for i in range(self.nr_slaves): + intermediates.append(self._queue.get()) + + results = self._master_callback(intermediates) + assert results[0][0] == 0, 'The first result should belongs to the master.' + + for i, res in results: + if i == 0: + continue + self._registry[i].result.put(res) + + for i in range(self.nr_slaves): + assert self._queue.get() is True + + return results[0][1] + + @property + def nr_slaves(self): + return len(self._registry) diff --git a/models/CtrlHair/sean_codes/models/networks/sync_batchnorm/replicate.py b/models/CtrlHair/sean_codes/models/networks/sync_batchnorm/replicate.py new file mode 100644 index 0000000000000000000000000000000000000000..b71c7b8ed51a1d6c55b1f753bdd8d90bad79bd06 --- /dev/null +++ b/models/CtrlHair/sean_codes/models/networks/sync_batchnorm/replicate.py @@ -0,0 +1,94 @@ +# -*- coding: utf-8 -*- +# File : replicate.py +# Author : Jiayuan Mao +# Email : maojiayuan@gmail.com +# Date : 27/01/2018 +# +# This file is part of Synchronized-BatchNorm-PyTorch. +# https://github.com/vacancy/Synchronized-BatchNorm-PyTorch +# Distributed under MIT License. + +import functools + +from torch.nn.parallel.data_parallel import DataParallel + +__all__ = [ + 'CallbackContext', + 'execute_replication_callbacks', + 'DataParallelWithCallback', + 'patch_replication_callback' +] + + +class CallbackContext(object): + pass + + +def execute_replication_callbacks(modules): + """ + Execute an replication callback `__data_parallel_replicate__` on each module created by original replication. + + The callback will be invoked with arguments `__data_parallel_replicate__(ctx, copy_id)` + + Note that, as all modules are isomorphism, we assign each sub-module with a context + (shared among multiple copies of this module on different devices). + Through this context, different copies can share some information. + + We guarantee that the callback on the master copy (the first copy) will be called ahead of calling the callback + of any slave copies. + """ + master_copy = modules[0] + nr_modules = len(list(master_copy.modules())) + ctxs = [CallbackContext() for _ in range(nr_modules)] + + for i, module in enumerate(modules): + for j, m in enumerate(module.modules()): + if hasattr(m, '__data_parallel_replicate__'): + m.__data_parallel_replicate__(ctxs[j], i) + + +class DataParallelWithCallback(DataParallel): + """ + Data Parallel with a replication callback. + + An replication callback `__data_parallel_replicate__` of each module will be invoked after being created by + original `replicate` function. + The callback will be invoked with arguments `__data_parallel_replicate__(ctx, copy_id)` + + Examples: + > sync_bn = SynchronizedBatchNorm1d(10, eps=1e-5, affine=False) + > sync_bn = DataParallelWithCallback(sync_bn, device_ids=[0, 1]) + # sync_bn.__data_parallel_replicate__ will be invoked. + """ + + def replicate(self, module, device_ids): + modules = super(DataParallelWithCallback, self).replicate(module, device_ids) + execute_replication_callbacks(modules) + return modules + + +def patch_replication_callback(data_parallel): + """ + Monkey-patch an existing `DataParallel` object. Add the replication callback. + Useful when you have customized `DataParallel` implementation. + + Examples: + > sync_bn = SynchronizedBatchNorm1d(10, eps=1e-5, affine=False) + > sync_bn = DataParallel(sync_bn, device_ids=[0, 1]) + > patch_replication_callback(sync_bn) + # this is equivalent to + > sync_bn = SynchronizedBatchNorm1d(10, eps=1e-5, affine=False) + > sync_bn = DataParallelWithCallback(sync_bn, device_ids=[0, 1]) + """ + + assert isinstance(data_parallel, DataParallel) + + old_replicate = data_parallel.replicate + + @functools.wraps(old_replicate) + def new_replicate(module, device_ids): + modules = old_replicate(module, device_ids) + execute_replication_callbacks(modules) + return modules + + data_parallel.replicate = new_replicate diff --git a/models/CtrlHair/sean_codes/models/networks/sync_batchnorm/unittest.py b/models/CtrlHair/sean_codes/models/networks/sync_batchnorm/unittest.py new file mode 100644 index 0000000000000000000000000000000000000000..bed56f1caa929ac3e9a57c583f8d3e42624f58be --- /dev/null +++ b/models/CtrlHair/sean_codes/models/networks/sync_batchnorm/unittest.py @@ -0,0 +1,29 @@ +# -*- coding: utf-8 -*- +# File : unittest.py +# Author : Jiayuan Mao +# Email : maojiayuan@gmail.com +# Date : 27/01/2018 +# +# This file is part of Synchronized-BatchNorm-PyTorch. +# https://github.com/vacancy/Synchronized-BatchNorm-PyTorch +# Distributed under MIT License. + +import unittest +import torch + + +class TorchTestCase(unittest.TestCase): + def assertTensorClose(self, x, y): + adiff = float((x - y).abs().max()) + if (y == 0).all(): + rdiff = 'NaN' + else: + rdiff = float((adiff / y).abs().max()) + + message = ( + 'Tensor close check failed\n' + 'adiff={}\n' + 'rdiff={}\n' + ).format(adiff, rdiff) + self.assertTrue(torch.allclose(x, y), message) + diff --git a/models/CtrlHair/sean_codes/models/pix2pix_model.py b/models/CtrlHair/sean_codes/models/pix2pix_model.py new file mode 100644 index 0000000000000000000000000000000000000000..afeaf03994e5dab1b5fbe5b31587a8fe4931c99f --- /dev/null +++ b/models/CtrlHair/sean_codes/models/pix2pix_model.py @@ -0,0 +1,266 @@ +""" +Copyright (C) 2019 NVIDIA Corporation. All rights reserved. +Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode). +""" + +import torch +import sean_codes.models.networks as networks +import util.util as util + + +class Pix2PixModel(torch.nn.Module): + @staticmethod + def modify_commandline_options(parser, is_train): + networks.modify_commandline_options(parser, is_train) + return parser + + def __init__(self, opt): + super().__init__() + self.opt = opt + self.FloatTensor = torch.cuda.FloatTensor if self.use_gpu() \ + else torch.FloatTensor + self.ByteTensor = torch.cuda.ByteTensor if self.use_gpu() \ + else torch.ByteTensor + + self.netG, self.netD, self.netE = self.initialize_networks(opt) + + # set loss functions + if opt.isTrain: + self.criterionGAN = networks.GANLoss( + opt.gan_mode, tensor=self.FloatTensor, opt=self.opt) + self.criterionFeat = torch.nn.L1Loss() + if not opt.no_vgg_loss: + self.criterionVGG = networks.VGGLoss(self.opt.gpu_ids) + + # Entry point for all calls involving forward pass + # of deep networks. We used this approach since DataParallel module + # can't parallelize custom functions, we branch to different + # routines based on |mode|. + def forward(self, data, mode): + input_semantics, real_image = self.preprocess_input(data) + + if mode == 'generator': + g_loss, generated = self.compute_generator_loss( + input_semantics, real_image) + return g_loss, generated + elif mode == 'discriminator': + d_loss = self.compute_discriminator_loss( + input_semantics, real_image) + return d_loss + elif mode == 'encode_only': + z, mu, logvar = self.encode_z(real_image) + return mu, logvar + elif mode == 'inference': + with torch.no_grad(): + # fake_image, _ = self.generate_fake(input_semantics, real_image) + obj_dic = data['path'] + fake_image = self.save_style_codes(input_semantics, real_image, obj_dic) + return fake_image + elif mode == 'UI_mode': + with torch.no_grad(): + # fake_image, _ = self.generate_fake(input_semantics, real_image) + + ################### some problems here + obj_dic = data['obj_dic'] + # if isinstance(obj_dic, str): + # obj_dic = [obj_dic] + fake_image = self.use_style_codes(input_semantics, real_image, obj_dic) + return fake_image + elif mode == 'style_code': + with torch.no_grad(): + style_codes = self.netG.Zencoder(input=real_image, segmap=input_semantics) + return style_codes + else: + raise ValueError("|mode| is invalid") + + def create_optimizers(self, opt): + G_params = list(self.netG.parameters()) + if opt.use_vae: + G_params += list(self.netE.parameters()) + if opt.isTrain: + D_params = list(self.netD.parameters()) + + if opt.no_TTUR: + beta1, beta2 = opt.beta1, opt.beta2 + G_lr, D_lr = opt.lr, opt.lr + else: + beta1, beta2 = 0, 0.9 + G_lr, D_lr = opt.lr / 2, opt.lr * 2 + + optimizer_G = torch.optim.Adam(G_params, lr=G_lr, betas=(beta1, beta2)) + optimizer_D = torch.optim.Adam(D_params, lr=D_lr, betas=(beta1, beta2)) + + return optimizer_G, optimizer_D + + def save(self, epoch): + util.save_network(self.netG, 'G', epoch, self.opt) + util.save_network(self.netD, 'D', epoch, self.opt) + + ############################################################################ + # Private helper methods + ############################################################################ + + def initialize_networks(self, opt): + netG = networks.define_G(opt) + netD = networks.define_D(opt) if opt.isTrain else None + netE = networks.define_E(opt) if opt.use_vae else None + + if not opt.isTrain or opt.continue_train: + netG = util.load_network(netG, 'G', opt.which_epoch, opt) + if opt.isTrain: + netD = util.load_network(netD, 'D', opt.which_epoch, opt) + + return netG, netD, netE + + # preprocess the input, such as moving the tensors to GPUs and + # transforming the label map to one-hot encoding + # |data|: dictionary of the input data + + def preprocess_input(self, data): + # move to GPU and change data types + data['label'] = data['label'].long() + if self.use_gpu(): + for param in ['label', 'instance', 'image']: + if param in data and data[param] is not None: + data[param] = data[param].cuda(non_blocking=True) + if 'obj_dic' in data: + for idx in range(19): + if data['obj_dic'][str(idx)]['ACE'].device.type == 'cpu': + data['obj_dic'][str(idx)]['ACE'] = data['obj_dic'][str(idx)]['ACE'].cuda(non_blocking=True) + # create one-hot label map + label_map = data['label'] + bs, _, h, w = label_map.size() + nc = self.opt.label_nc + 1 if self.opt.contain_dontcare_label \ + else self.opt.label_nc + input_label = self.FloatTensor(bs, nc, h, w).zero_() + input_semantics = input_label.scatter_(1, label_map, 1.0) + + # concatenate instance map if it exists + if not self.opt.no_instance: + inst_map = data['instance'] + instance_edge_map = self.get_edges(inst_map) + input_semantics = torch.cat((input_semantics, instance_edge_map), dim=1) + + return input_semantics, data['image'] + + def compute_generator_loss(self, input_semantics, real_image): + G_losses = {} + + fake_image = self.generate_fake( + input_semantics, real_image, compute_kld_loss=self.opt.use_vae) + + pred_fake, pred_real = self.discriminate( + input_semantics, fake_image, real_image) + + G_losses['GAN'] = self.criterionGAN(pred_fake, True, + for_discriminator=False) + + if not self.opt.no_ganFeat_loss: + num_D = len(pred_fake) + GAN_Feat_loss = self.FloatTensor(1).fill_(0) + for i in range(num_D): # for each discriminator + # last color_texture is the final prediction, so we exclude it + num_intermediate_outputs = len(pred_fake[i]) - 1 + for j in range(num_intermediate_outputs): # for each layer color_texture + unweighted_loss = self.criterionFeat( + pred_fake[i][j], pred_real[i][j].detach()) + GAN_Feat_loss += unweighted_loss * self.opt.lambda_feat / num_D + G_losses['GAN_Feat'] = GAN_Feat_loss + + if not self.opt.no_vgg_loss: + G_losses['VGG'] = self.criterionVGG(fake_image, real_image) \ + * self.opt.lambda_vgg + + return G_losses, fake_image + + def compute_discriminator_loss(self, input_semantics, real_image): + D_losses = {} + with torch.no_grad(): + fake_image = self.generate_fake(input_semantics, real_image) + fake_image = fake_image.detach() + fake_image.requires_grad_() + + pred_fake, pred_real = self.discriminate( + input_semantics, fake_image, real_image) + + D_losses['D_Fake'] = self.criterionGAN(pred_fake, False, + for_discriminator=True) + D_losses['D_real'] = self.criterionGAN(pred_real, True, + for_discriminator=True) + + return D_losses + + def encode_z(self, real_image): + mu, logvar = self.netE(real_image) + z = self.reparameterize(mu, logvar) + return z, mu, logvar + + def generate_fake(self, input_semantics, real_image, compute_kld_loss=False): + + fake_image = self.netG(input_semantics, real_image) + + return fake_image + + ############################################################### + + def save_style_codes(self, input_semantics, real_image, obj_dic): + + fake_image = self.netG(input_semantics, real_image, obj_dic=obj_dic) + + return fake_image + + def use_style_codes(self, input_semantics, real_image, obj_dic): + fake_image = self.netG(input_semantics, real_image, obj_dic=obj_dic) + + return fake_image + + # Given fake and real image, return the prediction of discriminator + # for each fake and real image. + + def discriminate(self, input_semantics, fake_image, real_image): + fake_concat = torch.cat([input_semantics, fake_image], dim=1) + real_concat = torch.cat([input_semantics, real_image], dim=1) + + # In Batch Normalization, the fake and real images are + # recommended to be in the same batch to avoid disparate + # statistics in fake and real images. + # So both fake and real images are fed to D all at once. + fake_and_real = torch.cat([fake_concat, real_concat], dim=0) + + discriminator_out = self.netD(fake_and_real) + + pred_fake, pred_real = self.divide_pred(discriminator_out) + + return pred_fake, pred_real + + # Take the prediction of fake and real images from the combined batch + def divide_pred(self, pred): + # the prediction contains the intermediate outputs of multiscale GAN, + # so it's usually a list + if type(pred) == list: + fake = [] + real = [] + for p in pred: + fake.append([tensor[:tensor.size(0) // 2] for tensor in p]) + real.append([tensor[tensor.size(0) // 2:] for tensor in p]) + else: + fake = pred[:pred.size(0) // 2] + real = pred[pred.size(0) // 2:] + + return fake, real + + def get_edges(self, t): + edge = self.ByteTensor(t.size()).zero_() + edge[:, :, :, 1:] = edge[:, :, :, 1:] | (t[:, :, :, 1:] != t[:, :, :, :-1]) + edge[:, :, :, :-1] = edge[:, :, :, :-1] | (t[:, :, :, 1:] != t[:, :, :, :-1]) + edge[:, :, 1:, :] = edge[:, :, 1:, :] | (t[:, :, 1:, :] != t[:, :, :-1, :]) + edge[:, :, :-1, :] = edge[:, :, :-1, :] | (t[:, :, 1:, :] != t[:, :, :-1, :]) + return edge.float() + + def reparameterize(self, mu, logvar): + std = torch.exp(0.5 * logvar) + eps = torch.randn_like(std) + return eps.mul(std) + mu + + def use_gpu(self): + return len(self.opt.gpu_ids) > 0 diff --git a/models/CtrlHair/sean_codes/options/__init__.py b/models/CtrlHair/sean_codes/options/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..ba4ddb2d78be7bbf79bb0671eefd9296d3a168b7 --- /dev/null +++ b/models/CtrlHair/sean_codes/options/__init__.py @@ -0,0 +1,4 @@ +""" +Copyright (C) 2019 NVIDIA Corporation. All rights reserved. +Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode). +""" \ No newline at end of file diff --git a/models/CtrlHair/sean_codes/options/base_options.py b/models/CtrlHair/sean_codes/options/base_options.py new file mode 100644 index 0000000000000000000000000000000000000000..7764e869119caa27e736b9cb6294842bb540286b --- /dev/null +++ b/models/CtrlHair/sean_codes/options/base_options.py @@ -0,0 +1,180 @@ +""" +Copyright (C) 2019 NVIDIA Corporation. All rights reserved. +Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode). +""" + +import sys +import argparse +import os +from util import util +import torch +from sean_codes import models, data +import pickle + + +class BaseOptions(): + def __init__(self): + self.initialized = False + + def initialize(self, parser): + # experiment specifics + parser.add_argument('--name', type=str, default='CelebA-HQ_pretrained', + help='name of the experiment. It decides where to store samples and models') + + parser.add_argument('--gpu_ids', type=str, default='0', help='gpu ids: e.g. 0 0,1,2, 0,2. use -1 for CPU') + parser.add_argument('--checkpoints_dir', type=str, default='external_model_params/sean_checkpoints', + help='models are saved here') + parser.add_argument('--model', type=str, default='pix2pix', help='which model to use') + parser.add_argument('--norm_G', type=str, default='spectralinstance', help='instance normalization or batch normalization') + parser.add_argument('--norm_D', type=str, default='spectralinstance', help='instance normalization or batch normalization') + parser.add_argument('--norm_E', type=str, default='spectralinstance', help='instance normalization or batch normalization') + parser.add_argument('--phase', type=str, default='scripts', help='scripts, val, test, etc') + + # input/color_texture sizes + parser.add_argument('--batchSize', type=int, default=1, help='input batch size') + parser.add_argument('--preprocess_mode', type=str, default='scale_width_and_crop', help='scaling and cropping of images at load time.', choices=("resize_and_crop", "crop", "scale_width", "scale_width_and_crop", "scale_shortside", "scale_shortside_and_crop", "fixed", "none")) + parser.add_argument('--load_size', type=int, default=256, help='Scale images to this size. The final image will be cropped to --crop_size.') + parser.add_argument('--crop_size', type=int, default=256, help='Crop to the width of crop_size (after initially scaling the images to load_size.)') + parser.add_argument('--aspect_ratio', type=float, default=1.0, help='The ratio width/height. The final height of the load image will be crop_size/aspect_ratio') + parser.add_argument('--label_nc', type=int, default=19, help='# of input label classes without unknown class. If you have unknown class as class label, specify --contain_dopntcare_label.') + parser.add_argument('--contain_dontcare_label', action='store_true', help='if the label map contains dontcare label (dontcare=255)') + parser.add_argument('--output_nc', type=int, default=3, help='# of color_texture image channels') + + # for setting inputs + parser.add_argument('--dataroot', type=str, default='./datasets/cityscapes/') + parser.add_argument('--dataset_mode', type=str, default='custom') + parser.add_argument('--serial_batches', action='store_true', help='if true, takes images in order to make batches, otherwise takes them randomly') + parser.add_argument('--no_flip', action='store_true', help='if specified, do not flip the images for data argumentation') + parser.add_argument('--nThreads', default=28, type=int, help='# threads for loading data') + parser.add_argument('--max_dataset_size', type=int, default=sys.maxsize, help='Maximum number of samples allowed per dataset. If the dataset directory contains more than max_dataset_size, only a subset is loaded.') + parser.add_argument('--load_from_opt_file', action='store_true', help='load the options from sean_checkpoints and use that as default') + parser.add_argument('--cache_filelist_write', action='store_true', help='saves the current filelist into a text file, so that it loads faster') + parser.add_argument('--cache_filelist_read', action='store_true', help='reads from the file list cache') + + # for displays + parser.add_argument('--display_winsize', type=int, default=400, help='display window size') + + # for generator + parser.add_argument('--netG', type=str, default='spade', help='selects model to use for netG (pix2pixhd | spade)') + parser.add_argument('--ngf', type=int, default=64, help='# of gen filters in first conv layer') + parser.add_argument('--init_type', type=str, default='xavier', help='network initialization [normal|xavier|kaiming|orthogonal]') + parser.add_argument('--init_variance', type=float, default=0.02, help='variance of the initialization distribution') + parser.add_argument('--z_dim', type=int, default=256, + help="dimension of the latent z vector") + + # for instance-wise features + parser.add_argument('--no_instance', type=bool, default=True, + help='if specified, do *not* add instance map as input') + parser.add_argument('--nef', type=int, default=16, help='# of encoder filters in the first conv layer') + parser.add_argument('--use_vae', action='store_true', help='enable training with an image encoder.') + + self.initialized = True + return parser + + def gather_options(self): + # initialize parser with basic options + if not self.initialized: + parser = argparse.ArgumentParser( + formatter_class=argparse.ArgumentDefaultsHelpFormatter) + parser = self.initialize(parser) + + # get the basic options + opt, unknown = parser.parse_known_args() + + # modify model-related parser options + model_name = opt.model + model_option_setter = models.get_option_setter(model_name) + parser = model_option_setter(parser, self.isTrain) + + # modify dataset-related parser options + dataset_mode = opt.dataset_mode + dataset_option_setter = data.get_option_setter(dataset_mode) + parser = dataset_option_setter(parser, self.isTrain) + + opt, unknown = parser.parse_known_args() + + # if there is opt_file, load it. + # The previous default options will be overwritten + if opt.load_from_opt_file: + parser = self.update_options_from_file(parser, opt) + + opt = parser.parse_args() + self.parser = parser + return opt + + def print_options(self, opt): + message = '' + message += '----------------- Options ---------------\n' + for k, v in sorted(vars(opt).items()): + comment = '' + default = self.parser.get_default(k) + if v != default: + comment = '\t[default: %s]' % str(default) + message += '{:>25}: {:<30}{}\n'.format(str(k), str(v), comment) + message += '----------------- End -------------------' + print(message) + + def option_file_path(self, opt, makedir=False): + expr_dir = os.path.join(opt.checkpoints_dir, opt.name) + if makedir: + util.mkdirs(expr_dir) + file_name = os.path.join(expr_dir, 'opt') + return file_name + + def save_options(self, opt): + file_name = self.option_file_path(opt, makedir=True) + with open(file_name + '.txt', 'wt') as opt_file: + for k, v in sorted(vars(opt).items()): + comment = '' + default = self.parser.get_default(k) + if v != default: + comment = '\t[default: %s]' % str(default) + opt_file.write('{:>25}: {:<30}{}\n'.format(str(k), str(v), comment)) + + with open(file_name + '.pkl', 'wb') as opt_file: + pickle.dump(opt, opt_file) + + def update_options_from_file(self, parser, opt): + new_opt = self.load_options(opt) + for k, v in sorted(vars(opt).items()): + if hasattr(new_opt, k) and v != getattr(new_opt, k): + new_val = getattr(new_opt, k) + parser.set_defaults(**{k: new_val}) + return parser + + def load_options(self, opt): + file_name = self.option_file_path(opt, makedir=False) + new_opt = pickle.load(open(file_name + '.pkl', 'rb')) + return new_opt + + def parse(self, save=False): + + opt = self.gather_options() + opt.isTrain = self.isTrain # scripts or test + + # self.print_options(opt) + if opt.isTrain: + self.save_options(opt) + + # Set semantic_nc based on the option. + # This will be convenient in many places + opt.semantic_nc = opt.label_nc + \ + (1 if opt.contain_dontcare_label else 0) + \ + (0 if opt.no_instance else 1) + + # set gpu ids + str_ids = opt.gpu_ids.split(',') + opt.gpu_ids = [] + for str_id in str_ids: + id = int(str_id) + if id >= 0: + opt.gpu_ids.append(id) + if len(opt.gpu_ids) > 0: + torch.cuda.set_device(opt.gpu_ids[0]) + + assert len(opt.gpu_ids) == 0 or opt.batchSize % len(opt.gpu_ids) == 0, \ + "Batch size %d is wrong. It must be a multiple of # GPUs %d." \ + % (opt.batchSize, len(opt.gpu_ids)) + + self.opt = opt + return self.opt diff --git a/models/CtrlHair/sean_codes/options/test_options.py b/models/CtrlHair/sean_codes/options/test_options.py new file mode 100644 index 0000000000000000000000000000000000000000..d7e06df9ac5e17c875623c9c54cfeb1d6c54616e --- /dev/null +++ b/models/CtrlHair/sean_codes/options/test_options.py @@ -0,0 +1,27 @@ +""" +Copyright (C) 2019 NVIDIA Corporation. All rights reserved. +Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode). +""" +from util.common_options import ctrl_hair_parser_options +from .base_options import BaseOptions + + +class TestOptions(BaseOptions): + def initialize(self, parser): + BaseOptions.initialize(self, parser) + parser.add_argument('--results_dir', type=str, default='./results/', help='saves results here.') + parser.add_argument('--which_epoch', type=str, default='latest', + help='which epoch to load? set to latest to use latest cached model') + parser.add_argument('--how_many', type=int, default=float("inf"), help='how many test images to run') + + parser.set_defaults(preprocess_mode='scale_width_and_crop', crop_size=256, load_size=256, display_winsize=256) + parser.set_defaults(serial_batches=True) + parser.set_defaults(no_flip=True) + parser.set_defaults(phase='test') + + parser.add_argument('--status', type=str, default='test') + + ctrl_hair_parser_options(parser) + + self.isTrain = False + return parser diff --git a/models/CtrlHair/sean_codes/options/train_options.py b/models/CtrlHair/sean_codes/options/train_options.py new file mode 100644 index 0000000000000000000000000000000000000000..1eb1496ab13f5d57eac9460583d82a931fd9a68c --- /dev/null +++ b/models/CtrlHair/sean_codes/options/train_options.py @@ -0,0 +1,46 @@ +""" +Copyright (C) 2019 NVIDIA Corporation. All rights reserved. +Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode). +""" + +from .base_options import BaseOptions + + +class TrainOptions(BaseOptions): + def initialize(self, parser): + BaseOptions.initialize(self, parser) + # for displays + parser.add_argument('--display_freq', type=int, default=100, help='frequency of showing training results on screen') + parser.add_argument('--print_freq', type=int, default=100, help='frequency of showing training results on console') + parser.add_argument('--save_latest_freq', type=int, default=5000, help='frequency of saving the latest results') + parser.add_argument('--save_epoch_freq', type=int, default=10, help='frequency of saving sean_checkpoints at the end of epochs') + parser.add_argument('--no_html', default=False, help='do not save intermediate training results to [opt.checkpoints_dir]/[opt.name]/web/') + parser.add_argument('--debug', action='store_true', help='only do one epoch and displays at each iteration') + parser.add_argument('--tf_log', action='store_true', help='if specified, use tensorboard logging. Requires tensorflow installed') + + # for training + parser.add_argument('--continue_train', action='store_true', help='continue training: load the latest model') + parser.add_argument('--which_epoch', type=str, default='latest', help='which epoch to load? set to latest to use latest cached model') + parser.add_argument('--niter', type=int, default=50, help='# of iter at starting learning rate. This is NOT the total #epochs. Totla #epochs is niter + niter_decay') + parser.add_argument('--niter_decay', type=int, default=50, help='# of iter to linearly decay learning rate to zero') + parser.add_argument('--optimizer', type=str, default='adam') + parser.add_argument('--beta1', type=float, default=0.5, help='momentum term of adam') + parser.add_argument('--beta2', type=float, default=0.999, help='momentum term of adam') + parser.add_argument('--lr', type=float, default=0.0002, help='initial learning rate for adam') + parser.add_argument('--D_steps_per_G', type=int, default=1, help='number of discriminator iterations per generator iterations.') + + # for discriminators + parser.add_argument('--ndf', type=int, default=64, help='# of discrim filters in first conv layer') + parser.add_argument('--lambda_feat', type=float, default=10.0, help='weight for feature matching loss') + parser.add_argument('--lambda_vgg', type=float, default=10.0, help='weight for vgg loss') + parser.add_argument('--no_ganFeat_loss', action='store_true', help='if specified, do *not* use discriminator feature matching loss') + parser.add_argument('--no_vgg_loss', action='store_true', help='if specified, do *not* use VGG feature matching loss') + parser.add_argument('--gan_mode', type=str, default='hinge', help='(ls|original|hinge)') + parser.add_argument('--netD', type=str, default='multiscale', help='(n_layers|multiscale|image)') + parser.add_argument('--no_TTUR', action='store_true', help='Use TTUR training scheme') + parser.add_argument('--lambda_kld', type=float, default=0.005) + + parser.add_argument('--status', type=str, default='scripts') + + self.isTrain = True + return parser diff --git a/models/CtrlHair/sean_codes/trainers/__init__.py b/models/CtrlHair/sean_codes/trainers/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..7c23f2831952e1a35fd8f170760c627f7fb20ed9 --- /dev/null +++ b/models/CtrlHair/sean_codes/trainers/__init__.py @@ -0,0 +1,4 @@ +""" +Copyright (C) 2019 NVIDIA Corporation. All rights reserved. +Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode). +""" diff --git a/models/CtrlHair/sean_codes/trainers/pix2pix_trainer.py b/models/CtrlHair/sean_codes/trainers/pix2pix_trainer.py new file mode 100644 index 0000000000000000000000000000000000000000..988f3244fe6d5152b0e51561635066432dfd0b77 --- /dev/null +++ b/models/CtrlHair/sean_codes/trainers/pix2pix_trainer.py @@ -0,0 +1,86 @@ +""" +Copyright (C) 2019 NVIDIA Corporation. All rights reserved. +Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode). +""" + +from sean_codes.models.networks.sync_batchnorm import DataParallelWithCallback +from sean_codes.models.pix2pix_model import Pix2PixModel + + +class Pix2PixTrainer(): + """ + Trainer creates the model and optimizers, and uses them to + updates the weights of the network while reporting losses + and the latest visuals to visualize the progress in training. + """ + + def __init__(self, opt): + self.opt = opt + self.pix2pix_model = Pix2PixModel(opt) + if len(opt.gpu_ids) > 0: + self.pix2pix_model = DataParallelWithCallback(self.pix2pix_model, + device_ids=opt.gpu_ids) + self.pix2pix_model_on_one_gpu = self.pix2pix_model.module + else: + self.pix2pix_model_on_one_gpu = self.pix2pix_model + + self.generated = None + if opt.isTrain: + self.optimizer_G, self.optimizer_D = \ + self.pix2pix_model_on_one_gpu.create_optimizers(opt) + self.old_lr = opt.lr + + def run_generator_one_step(self, data): + self.optimizer_G.zero_grad() + g_losses, generated = self.pix2pix_model(data, mode='generator') + g_loss = sum(g_losses.values()).mean() + g_loss.backward() + self.optimizer_G.step() + self.g_losses = g_losses + self.generated = generated + + def run_discriminator_one_step(self, data): + self.optimizer_D.zero_grad() + d_losses = self.pix2pix_model(data, mode='discriminator') + d_loss = sum(d_losses.values()).mean() + d_loss.backward() + self.optimizer_D.step() + self.d_losses = d_losses + + def get_latest_losses(self): + return {**self.g_losses, **self.d_losses} + + def get_latest_generated(self): + return self.generated + + def update_learning_rate(self, epoch): + self.update_learning_rate(epoch) + + def save(self, epoch): + self.pix2pix_model_on_one_gpu.save(epoch) + + ################################################################## + # Helper functions + ################################################################## + + def update_learning_rate(self, epoch): + if epoch > self.opt.niter: + lrd = self.opt.lr / self.opt.niter_decay + new_lr = self.old_lr - lrd + else: + new_lr = self.old_lr + + if new_lr != self.old_lr: + if self.opt.no_TTUR: + new_lr_G = new_lr + new_lr_D = new_lr + else: + new_lr_G = new_lr / 2 + new_lr_D = new_lr * 2 + + for param_group in self.optimizer_D.param_groups: + param_group['lr'] = new_lr_D + for param_group in self.optimizer_G.param_groups: + param_group['lr'] = new_lr_G + print('update learning rate: %f -> %f' % (self.old_lr, new_lr)) + self.old_lr = new_lr diff --git a/models/CtrlHair/shape_branch/__init__.py b/models/CtrlHair/shape_branch/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..0f414ceb4ca3fb00e06f3d1bd91786d4f1b0b3dd --- /dev/null +++ b/models/CtrlHair/shape_branch/__init__.py @@ -0,0 +1,8 @@ +# -*- coding: utf-8 -*- + +""" +# File name: __init__.py.py +# Time : 2021/12/31 15:28 +# Author: xyguoo@163.com +# Description: +""" diff --git a/models/CtrlHair/shape_branch/adaptor_generation.py b/models/CtrlHair/shape_branch/adaptor_generation.py new file mode 100644 index 0000000000000000000000000000000000000000..b6fb7208561e7c7d9b338a6eb6c62181ec5befa1 --- /dev/null +++ b/models/CtrlHair/shape_branch/adaptor_generation.py @@ -0,0 +1,114 @@ +# -*- coding: utf-8 -*- + +""" +# File name: adaptor_generation.py +# Time : 2021/12/31 19:35 +# Author: xyguoo@163.com +# Description: +""" + +import sys +import os + +pp = os.path.abspath(os.path.join(os.path.abspath(__file__), '../..')) +sys.path.append(pp) + +import os + +os.environ['CUDA_VISIBLE_DEVICES'] = '0' +from wrap_codes.mask_adaptor import wrap_for_image_with_idx +from common_dataset import DataFilter +from shape_branch.config import cfg +from global_value_utils import HAIR_IDX, TEMP_FOLDER +from util.util import path_join_abs +import cv2 +import threading +import random +import time +import glob +import tqdm + +def generate_mask_to_pool(root_dir, hair_dir, face_dir, hair_base_path, face_base_path, pool_dir, temp_wrap_dir, + thread_id, find_repeat, only_hair): + face_num = face_base_path[-9:-4] + hair_num = hair_base_path[-9:-4] + output_name = os.path.join(pool_dir, '%s___%s___%s___%s___%02d.png') % ( + hair_dir, hair_num, face_dir, face_num, thread_id) + find_res = [] + if find_repeat: + find_res = glob.glob(os.path.join(pool_dir, '%s___%s___%s___%s___*.png') % ( + hair_dir, hair_num, face_dir, face_num)) + if len(find_res) == 0: + print('Generate thread %d: %s %s to %s %s' % (thread_id, hair_dir, hair_num, face_dir, face_num)) + align_mask = wrap_for_image_with_idx(root_dir, hair_dir, face_dir, + hair_base_path, face_base_path, + wrap_temp_folder=temp_wrap_dir)[0] + if only_hair: + output_img = (align_mask == HAIR_IDX) * 255 + else: + output_img = align_mask + cv2.imwrite(output_name, output_img) + else: + print('Hit for %s' % output_name) + + +class AdaptorPoolGeneration: + def __init__(self, only_hair, dir_name, test_dir_name, thread_num=10, max_file=1e7): + self.data_filter = DataFilter(cfg) + self.pool_dir = os.path.join(cfg.data_root, dir_name) + self.pool_test_dir = os.path.join(cfg.data_root, test_dir_name) + self.max_file = max_file + self.only_hair = only_hair + for p in [self.pool_dir, self.pool_test_dir]: + if not os.path.exists(p): + os.makedirs(p) + self.thread_num = thread_num + + def generate_test_set(self, img_num=100): + temp_wrap_dir = os.path.join(TEMP_FOLDER, 'wrap_triangle/temp_wrap_test') + + for hair in tqdm.tqdm(self.data_filter.test_hair_list[:img_num]): + for face in self.data_filter.test_face_list[:img_num]: + hair_dir = hair.split('/')[-3] + face_dir = face.split('/')[-3] + base_hair = os.path.split(hair)[-1] + base_face = os.path.split(face)[-1] + generate_mask_to_pool(cfg.data_root, hair_dir, face_dir, base_hair, base_face, + self.pool_test_dir, temp_wrap_dir, 0, find_repeat=False, only_hair=self.only_hair) + + def run(self): + self.threads = [] + for idx in range(self.thread_num): + t = threading.Thread(target=self.generate_thread, args=[idx]) + self.threads.append(t) + for thread in self.threads: + thread.start() + + def generate_thread(self, thread_idx): + temp_wrap_dir = os.path.join(TEMP_FOLDER, 'wrap_triangle/temp_wrap_%d' % thread_idx) + random.seed(time.time()) + while True: + if len(os.listdir(self.pool_dir)) < self.max_file: + for _ in range(100): + while True: + hair_path = random.choice(self.data_filter.train_list) + hair_num = hair_path[-9:-4] + hair_dir = hair_path.split('/')[-3] + if self.data_filter.valid_hair(path_join_abs(hair_path, '../..'), hair_num): + break + while True: + face_path = random.choice(self.data_filter.train_list) + face_num = face_path[-9:-4] + face_dir = face_path.split('/')[-3] + if self.data_filter.valid_face(path_join_abs(face_path, '../..'), face_num): + break + try: + generate_mask_to_pool(cfg.data_root, hair_dir, face_dir, + os.path.basename(hair_path), os.path.basename(face_path), + self.pool_dir, temp_wrap_dir, + thread_idx, find_repeat=False, only_hair=self.only_hair) + except Exception as e: + print(repr(e)) + else: + print('Full, so sleep in thread %d' % thread_idx) + time.sleep(3.0) diff --git a/models/CtrlHair/shape_branch/config.py b/models/CtrlHair/shape_branch/config.py new file mode 100644 index 0000000000000000000000000000000000000000..ca6f79c7fbe9e71f83b1d4b3cf30d1b934cf58ed --- /dev/null +++ b/models/CtrlHair/shape_branch/config.py @@ -0,0 +1,129 @@ +# -*- 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 models.CtrlHair.global_value_utils import GLOBAL_DATA_ROOT, DEFAULT_CONFIG_SHAPE_BRANCH + +configs = [ + addict.Dict({ + "experiment_name": "054__succeed__049__gan_fake_0.5_from_noise", + 'hair_dim': 16, + 'pos_encoding_order': 10, + 'lambda_hair': 100, + 'lambda_non_hair': 100, + 'lambda_face': 20, + 'lambda_self_rec': 5, + 'lambda_kl': 0.1, + 'regular_method': 'ce', + 'full_dataset': True, + 'only_celeba_as_real': True, + 'g_norm': 'ln', + 'd_norm': 'none', + 'lr_g': 0.0002, + 'lambda_adv_noise': 1, + 'lambda_gp_0_noise': 10, + 'total_batch_size': 4, + 'random_ae_prob': 0.5, + 'lr_dz': 0.00005, + 'adaptor_test_pool_dir': 'shape_testing_wrap_pool', + 'adaptor_pool_dir': 'shape_training_wrap_pool' + }), +] + + +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_d', 0.0001) + add_default_value_to_cfg(cfg, 'lr_g', 0.0002) + add_default_value_to_cfg(cfg, 'lr_dz', 0.0001) + 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', 380002) + add_default_value_to_cfg(cfg, 'log_step', 10) + add_default_value_to_cfg(cfg, 'sample_step', 10000) + add_default_value_to_cfg(cfg, 'model_save_step', 10000) + add_default_value_to_cfg(cfg, 'sample_batch_size', 16) + add_default_value_to_cfg(cfg, 'max_save', 1) + add_default_value_to_cfg(cfg, 'vae_var_output', 'var') + add_default_value_to_cfg(cfg, 'SEAN_code', 512) + add_default_value_to_cfg(cfg, 'd_hidden_in_channel', 16) + + # Model configuration + add_default_value_to_cfg(cfg, 'total_batch_size', 4) + add_default_value_to_cfg(cfg, 'gan_type', 'hinge2') + add_default_value_to_cfg(cfg, 'lambda_gp_0', 10.0) + add_default_value_to_cfg(cfg, 'lambda_adv', 1.0) + + add_default_value_to_cfg(cfg, 'g_norm', 'bn') + add_default_value_to_cfg(cfg, 'd_norm', 'bn') + add_default_value_to_cfg(cfg, 'init_type', 'normal') + add_default_value_to_cfg(cfg, 'G_D_train_num', {'G': 1, 'D': 1}, ) + add_default_value_to_cfg(cfg, 'vae_hair_mode', True) + + output_root_dir = 'model_trained/shape/%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 + '/summaries') + add_default_value_to_cfg(cfg, 'checkpoints_dir', output_root_dir + '/checkpoints') + add_default_value_to_cfg(cfg, 'sample_dir', output_root_dir + '/sample_training') + + 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, 'img_size', 256) + add_default_value_to_cfg(cfg, 'data_root', GLOBAL_DATA_ROOT) + + # dz discriminator + add_default_value_to_cfg(cfg, 'd_hidden_dim', 256) + add_default_value_to_cfg(cfg, 'd_noise_hidden_layer_num', 3) + + +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 + + +def get_basic_arg_parser(): + parser = argparse.ArgumentParser() + parser.add_argument('-c', '--config', type=str, help='Specify config number', default=DEFAULT_CONFIG_SHAPE_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('shape_branch/scripts.py') or sys.argv[0].endswith('shape_branch/script_find_direction.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_SHAPE_BRANCH) + back_process(cfg) diff --git a/models/CtrlHair/shape_branch/dataset.py b/models/CtrlHair/shape_branch/dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..8530fa42615057a9cbfeae43ff153ad6fd24a584 --- /dev/null +++ b/models/CtrlHair/shape_branch/dataset.py @@ -0,0 +1,192 @@ +# -*- coding: utf-8 -*- + +""" +# File name: dataset.py +# Time : 2021/11/16 21:24 +# Author: xyguoo@163.com +# Description: +""" +import os + +import cv2 + +from common_dataset import DataFilter +import random +import pickle +import torch +from torchvision import transforms as tform +from PIL import Image + +from global_value_utils import GLOBAL_DATA_ROOT +from util.util import path_join_abs + + +class Dataset(DataFilter): + + def __init__(self, cfg, rank=0): + super().__init__(cfg) + self.cfg = cfg + + img_size = 256 + self.__setattr__('mask_transform_%d' % img_size, + tform.Compose([tform.Resize(img_size, interpolation=Image.NEAREST), + tform.ToTensor(), + tform.Lambda(lambda x: 255 * x)])) + self.__setattr__('mask_transform_mirror_%d' % img_size, + tform.Compose([tform.Resize(img_size, interpolation=Image.NEAREST), + tform.Lambda(lambda x: tform.functional.hflip(x)), + tform.ToTensor(), + tform.Lambda(lambda x: 255 * x)])) + + self.mask_pool_dir = os.path.join(cfg.data_root, cfg.adaptor_pool_dir) + self.mask_test_pool_dir = os.path.join(cfg.data_root, cfg.adaptor_test_pool_dir) + self.mask_buffer = [] + self.local_rank = rank + random.seed(self.random_seed + self.local_rank + 1) + + if self.cfg.only_celeba_as_real: # CelebA Mask is the manual mask, which has strong realism + self.dis_real_list = [st for st in self.train_list if 'CelebaMask' in st] + + self.data_root = GLOBAL_DATA_ROOT + with open(os.path.join(self.data_root, 'sean_code_dict.pkl'), 'rb') as f: + self.sean_code_dict = pickle.load(f) + + def get_by_file_name(self, size, img_path=None, validate_func=None, mirror=False, data_list=None, need_img=False): + if data_list is None: + data_list = self.train_list + + if img_path is None: + while True: + random_idx = random.randint(0, len(data_list) - 1) + img_path = data_list[random_idx] + if validate_func is None or validate_func(path_join_abs(img_path, '../..'), img_path[-9:-4]): + break + if data_list is self.train_list: + self.train_list = self.train_list[:random_idx] + self.train_list[random_idx + 1:] + else: + data_list = data_list[:random_idx] + data_list[random_idx + 1:] + + if mirror: + mask_transform = self.__getattribute__('mask_transform_mirror_%d' % size) + else: + mask_transform = self.__getattribute__('mask_transform_%d' % size) + + base_num = img_path[-9:-4] + mask_path = path_join_abs(img_path, '../../label', base_num + '.png') + mask = mask_transform(Image.open(mask_path)) + + if need_img: + img = cv2.imread(img_path) + img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) + return mask, img + return mask + + def get_pair_file_randomly_from_buffer(self): + if len(self.mask_buffer) == 0: + self.mask_buffer = os.listdir(self.mask_pool_dir) + random.shuffle(self.mask_buffer) + return self.mask_buffer.pop() + + def get_pair_randomly_from_pool(self, img_size): + mask_file = self.get_pair_file_randomly_from_buffer() + file_parts = mask_file.split('___') + hair_dir, hair, face_dir, face = file_parts[:4] + images_dir = 'label' + mirror = (random.random() > 0.5) + face_img, target_hair_mask, hair_img = self.get_adaptor_pair_mask( + path_join_abs(self.cfg.data_root, face_dir, images_dir, '%s.png' % face), + path_join_abs(self.cfg.data_root, hair_dir, images_dir, '%s.png' % hair), + img_size, self.mask_pool_dir, mirror=mirror, mask_thread_num=file_parts[-1][:2]) + return face_img, target_hair_mask, hair_img + + def get_adaptor_pair_mask(self, face_path, hair_path, img_size, mask_pool_dir, mirror=False, mask_thread_num='00', + need_img=False): + if need_img: + face_mask, face_img = self.get_by_file_name(img_size, face_path, mirror=mirror, need_img=need_img) + hair_mask, hair_img = self.get_by_file_name(img_size, hair_path, mirror=mirror, need_img=need_img) + else: + face_mask = self.get_by_file_name(img_size, face_path, mirror=mirror, need_img=need_img) + hair_mask = self.get_by_file_name(img_size, hair_path, mirror=mirror, need_img=need_img) + face_base_path = os.path.basename(face_path)[:-4] + hair_base_path = os.path.basename(hair_path)[:-4] + face_dir = face_path.split('/')[-3] + hair_dir = hair_path.split('/')[-3] + target_mask_path = os.path.join(mask_pool_dir, '%s___%s___%s___%s___%s.png' % ( + hair_dir, hair_base_path, face_dir, face_base_path, mask_thread_num)) + if mirror: + transform_func = self.__getattribute__('mask_transform_mirror_%d' % img_size) + else: + transform_func = self.__getattribute__('mask_transform_%d' % img_size) + + target_hair_mask = transform_func(Image.open(target_mask_path)) + if need_img: + return face_mask, target_hair_mask, hair_mask, face_img, hair_img + else: + return face_mask, target_hair_mask, hair_mask + + def get_random_pair_batch(self, batch_size, img_size=None): + if not img_size: + img_size = self.cfg.img_size + face_imgs, target_hair_masks, hair_imgs = [], [], [] + while len(face_imgs) < batch_size: + face_img, target_hair_mask, hair_img = \ + self.get_pair_randomly_from_pool(img_size) + face_imgs.append(face_img) + target_hair_masks.append(target_hair_mask) + hair_imgs.append(hair_img) + results = [face_imgs, target_hair_masks, hair_imgs] + for idx in range(len(results)): + results[idx] = torch.stack(results[idx], dim=0) + return {'face': results[0], 'target': results[1], 'hair': results[2]} + + def get_random_single_batch(self, batch_size): + face_imgs = [] + while len(face_imgs) < batch_size: + if self.cfg.only_celeba_as_real: + face_img = self.get_by_file_name(256, validate_func=self.valid_hair, mirror=(random.random() > 0.5), + data_list=self.dis_real_list) + else: + face_img = self.get_by_file_name(256, validate_func=self.valid_hair, mirror=(random.random() > 0.5)) + face_imgs.append(face_img) + face_imgs = torch.stack(face_imgs, dim=0) + return face_imgs + + def get_test_batch(self, batch_size=32, img_size=None): + if not img_size: + img_size = self.cfg.img_size + face_masks, target_hair_masks, hair_masks, sean_codes, face_imgs, hair_imgs = [], [], [], [], [], [] + + idx = 0 + while len(face_masks) < batch_size: + face_maks_path = self.test_face_list[idx] + hair_mask_path = self.test_hair_list[idx] + idx += 1 + face_mask, target_hair_mask, hair_mask, face_img, hair_img = \ + self.get_adaptor_pair_mask(face_maks_path, hair_mask_path, img_size, self.mask_test_pool_dir, + need_img=True) + face_masks.append(face_mask) + target_hair_masks.append(target_hair_mask) + hair_masks.append(hair_mask) + + face_path_parts = face_maks_path.split('/') + sean_codes.append(torch.tensor( + self.sean_code_dict['___'.join([face_path_parts[-3], face_path_parts[-1][:-4]])])) + + face_imgs.append(torch.tensor(face_img)) + hair_imgs.append(torch.tensor(hair_img)) + + face_masks = torch.stack(face_masks, dim=0) + hair_masks = torch.stack(hair_masks, dim=0) + target_hair_masks = torch.stack(target_hair_masks, dim=0) + + return {'face': face_masks, 'target': target_hair_masks, 'hair': hair_masks, 'sean_code': sean_codes, + 'face_imgs': face_imgs, 'hair_imgs': hair_imgs} + + +if __name__ == '__main__': + from shape_branch.config import cfg + + ds = Dataset(cfg) + # resources = ds.get_training_batch(8) + res = ds.get_random_inpainting_batch(9) + pass diff --git a/models/CtrlHair/shape_branch/model.py b/models/CtrlHair/shape_branch/model.py new file mode 100644 index 0000000000000000000000000000000000000000..ddcaac120222e6b27639f8b16e4b3ce65ad5f9fe --- /dev/null +++ b/models/CtrlHair/shape_branch/model.py @@ -0,0 +1,232 @@ +# -*- coding: utf-8 -*- + +""" +# File name: model.py +# Time : 2021/11/17 15:37 +# Author: xyguoo@163.com +# Description: +""" + +import torch.nn as nn +from models.CtrlHair.my_torchlib.module import LinearBlock, Conv2dBlock +import torch +from torch.nn import init +import numpy as np +from models.CtrlHair.global_value_utils import HAIR_IDX + + +def generate_pos_embedding(img_size, order=10): + coordinators = np.linspace(0, 1, img_size, endpoint=False) + bi_coordinators = np.stack(np.meshgrid(coordinators, coordinators), 0) + bi_coordinators = bi_coordinators[None, ...] + nums = np.arange(0, order) + nums = 2 ** nums * np.pi + nums = nums[:, None, None, None] + gamma1 = np.sin(nums * bi_coordinators) + gamma2 = np.cos(nums * bi_coordinators) + gamma = np.concatenate([gamma1, gamma2], axis=0) + gamma = gamma.reshape([-1, img_size, img_size]) + gamma = torch.tensor(gamma, requires_grad=False).float() + return gamma + + +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 + + +class MaskEncoder(nn.Module): + + def __init__(self, input_channel, output_dim, norm, layer_num, input_size=256, vae_mode=False, + pos_encoding_order=10, max_batch_size=1, hidden_in_channel=32): + super(MaskEncoder, self).__init__() + self.vae_mode = vae_mode + layers = [] + + in_channel = input_channel + pos_encoding_order * 4 + for cur_num in range(layer_num): + out_channel = min(2048, 2 ** cur_num * hidden_in_channel) + cur_conv = Conv2dBlock( + in_channel, out_channel, kernel_size=4, stride=2, padding=1, norm=norm, activation='lrelu') + layers.append(cur_conv) + in_channel = out_channel + + out_size = input_size // (2 ** layer_num) + self.layers = nn.Sequential(*layers) + + fc_in_dim = out_size ** 2 * out_channel + self.out_layer = LinearBlock(fc_in_dim, output_dim, norm='none', activation='none') + if self.vae_mode: + self.std_out_layer = LinearBlock(fc_in_dim, output_dim, norm='none', activation='none') + + self.input_embedding = generate_pos_embedding(img_size=input_size, order=pos_encoding_order) + self.input_embedding = self.input_embedding.repeat((max_batch_size, 1, 1, 1)) + + def forward(self, input_mask: torch.Tensor): + batch_size = input_mask.shape[0] + if self.input_embedding.device != input_mask.device: + self.input_embedding = self.input_embedding.to(input_mask.device) + input_with_pos = torch.cat([input_mask, self.input_embedding[:batch_size]], axis=1) + feature = self.layers(input_with_pos) + feature = feature.flatten(1) + out_mean = self.out_layer(feature) + if self.vae_mode: + out_std = self.std_out_layer(feature).abs() + return self.vae_resampling(out_mean, out_std), out_mean, out_std + else: + return out_mean, out_mean, None + + def vae_resampling(self, mean, std): + z = torch.randn(mean.shape).to(mean.device) + res = z * std + mean + return res + + +class MaskDecoder(nn.Module): + def __init__(self, input_dim, output_channel, norm, layer_num, output_size=256): + super(MaskDecoder, self).__init__() + + self.in_channel = min(32 * 2 ** layer_num, 2048) + self.input_size = output_size // (2 ** layer_num) + self.in_layer = LinearBlock(input_dim, self.in_channel * self.input_size ** 2, norm='none', activation='none') + + layers = [] + in_channel = self.in_channel + for cur_num in range(layer_num): + up = nn.Upsample(scale_factor=2, mode='nearest') + out_channel = min(32 * 2 ** (layer_num - 1 - cur_num), 2048) + cur_conv = Conv2dBlock(in_channel, out_channel, kernel_size=3, stride=1, padding=1, norm=norm, + activation='lrelu') + layers.append(up) + layers.append(cur_conv) + in_channel = out_channel + self.layers = nn.Sequential(*layers) + self.out_layer = Conv2dBlock(in_channel, output_channel, kernel_size=3, stride=1, padding=1, norm='none', + activation='none') + + def forward(self, input_vector): + feature = self.in_layer(input_vector) + feature = feature.reshape(-1, self.in_channel, self.input_size, self.input_size) + feature = self.layers(feature) + feature = self.out_layer(feature) + return feature + + +class Generator(nn.Module): + """Generator network.""" + + def __init__(self, cfg): + super(Generator, self).__init__() + self.cfg = cfg + + self.hair_encoder = MaskEncoder(1, cfg.hair_dim, norm=cfg.g_norm, layer_num=7, vae_mode=cfg.vae_hair_mode, + # self.hair_encoder = MaskEncoder(1, cfg.hair_dim, norm='none', layer_num=7, vae_mode=cfg.vae_hair_mode, + pos_encoding_order=cfg.pos_encoding_order, + max_batch_size=max(cfg.total_batch_size, cfg.sample_batch_size)) + self.face_encoder = MaskEncoder(18, 1024, norm=cfg.g_norm, layer_num=7, vae_mode=False, + # self.face_encoder = MaskEncoder(18, 1024, norm='none', layer_num=7, vae_mode=False, + pos_encoding_order=cfg.pos_encoding_order, + max_batch_size=max(cfg.total_batch_size, cfg.sample_batch_size)) + self.hair_decoder = MaskDecoder(1024 + cfg.hair_dim, output_channel=1, norm=cfg.g_norm, layer_num=7) + self.face_decoder = MaskDecoder(1024, output_channel=18, norm=cfg.g_norm, layer_num=7) + + def forward_hair_encoder(self, hair, testing=False): + code, mean, std = self.hair_encoder(hair) + if testing: + return mean + else: + return code, mean, std + + def forward_face_encoder(self, face): + code, _, _ = self.face_encoder(face) + return code + + def forward_hair_decoder(self, hair_code, face_code): + code = torch.cat([face_code, hair_code], dim=1) + hair = self.hair_decoder(code) + return hair + + def forward_face_decoder(self, face_code): + face = self.face_decoder(face_code) + return face + + def forward_decoder(self, hair_logit, face_logit): + logit = torch.cat([face_logit[:, :HAIR_IDX], hair_logit, face_logit[:, HAIR_IDX:]], dim=1) + mask = torch.softmax(logit, dim=1) + return mask + + def forward_edit_directly_in_test(self, hair, face): + _, hair_code, _ = self.forward_hair_encoder(hair) + face_code = self.forward_face_encoder(face) + mask = self.forward_decode_by_code(hair_code, face_code) + return mask + + def forward_decode_by_code(self, hair_code, face_code): + hair_logit = self.forward_hair_decoder(hair_code, face_code) + face_logit = self.forward_face_decoder(face_code) + mask = self.forward_decoder(hair_logit, face_logit) + return mask + + +class Discriminator(nn.Module): + """Discriminator network.""" + + def __init__(self, cfg): + super(Discriminator, self).__init__() + self.cfg = cfg + self.dis = MaskEncoder(19, 1, norm=cfg.d_norm, layer_num=7, vae_mode=False, + pos_encoding_order=cfg.pos_encoding_order, max_batch_size=cfg.total_batch_size, + hidden_in_channel=cfg.d_hidden_in_channel) + + def forward(self, mask): + dis_res, _, _ = self.dis(mask) + return dis_res + + +class DiscriminatorNoise(nn.Module): + """Discriminator network.""" + + def __init__(self, cfg): + super(DiscriminatorNoise, self).__init__() + self.cfg = cfg + input_dim = cfg.hair_dim + layers = [LinearBlock(input_dim, cfg.d_hidden_dim, cfg.d_norm, activation='lrelu')] + 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='lrelu')) + output_dim = 1 + layers.append(LinearBlock(cfg.d_hidden_dim, output_dim, 'none', 'none')) + self.net = nn.Sequential(*layers) + + def forward(self, x): + return self.net(x)[:, [0]] diff --git a/models/CtrlHair/shape_branch/script_adaptor_test_pool.py b/models/CtrlHair/shape_branch/script_adaptor_test_pool.py new file mode 100644 index 0000000000000000000000000000000000000000..87cef04fa7804a68bf9fec7da76b5f0212636317 --- /dev/null +++ b/models/CtrlHair/shape_branch/script_adaptor_test_pool.py @@ -0,0 +1,24 @@ +# -*- coding: utf-8 -*- + +""" +# File name: script_adaptor_test_pool.py +# Time : 2022/07/14 +# Author: xyguoo@163.com +# Description: +""" + +import sys +sys.path.append('.') + +from shape_branch.adaptor_generation import AdaptorPoolGeneration +from shape_branch.config import cfg + +if __name__ == '__main__': + pp = AdaptorPoolGeneration(only_hair=False, + dir_name=cfg.adaptor_pool_dir, + test_dir_name=cfg.adaptor_test_pool_dir, thread_num=10) + + ###################################################### + # Run this for generating wrap pool for testing + ###################################################### + pp.generate_test_set(cfg.sample_batch_size) diff --git a/models/CtrlHair/shape_branch/script_adaptor_train_pool.py b/models/CtrlHair/shape_branch/script_adaptor_train_pool.py new file mode 100644 index 0000000000000000000000000000000000000000..132f498c536dc6e4e13d894e8b3c4e8c1fbaf035 --- /dev/null +++ b/models/CtrlHair/shape_branch/script_adaptor_train_pool.py @@ -0,0 +1,26 @@ +# -*- coding: utf-8 -*- + +""" +# File name: script_adaptor_train_pool.py +# Time : 2022/07/14 +# Author: xyguoo@163.com +# Description: +""" + +import sys +sys.path.append('.') + +from shape_branch.adaptor_generation import AdaptorPoolGeneration +from shape_branch.config import cfg + +if __name__ == '__main__': + pp = AdaptorPoolGeneration(only_hair=False, + dir_name=cfg.adaptor_pool_dir, + test_dir_name=cfg.adaptor_test_pool_dir, thread_num=10) + + ##################################################### + # Run this for generating wrap pool for training dataset, + # which is a multi-thread procedure for speeding + ##################################################### + pp.run() + diff --git a/models/CtrlHair/shape_branch/script_find_direction.py b/models/CtrlHair/shape_branch/script_find_direction.py new file mode 100644 index 0000000000000000000000000000000000000000..f6a7d302497b7ecd333f35e4af4fc355b0112081 --- /dev/null +++ b/models/CtrlHair/shape_branch/script_find_direction.py @@ -0,0 +1,76 @@ +# -*- coding: utf-8 -*- + +""" +# File name: script_find_direction.py +# Time : 2022/2/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 util.mask_color_util import mask_to_rgb +from common_dataset import DataFilter +from util.imutil import read_rgb, write_rgb +from shape_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/shape/%s' % cfg.experiment_name +code_dim = cfg.hair_dim +att_name = 'shape' +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 * 2, 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 * 2, 0) + canvas.process_draw_image(img_parsing, img_idx * 2 + 1, 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 * 2, inter_idx + 1) + canvas.process_draw_image(mask_to_rgb(be.cur_mask), img_idx * 2 + 1, inter_idx + 1) + write_rgb('%s/%d.png' % (img_gen_dir, dir_idx), canvas.canvas) diff --git a/models/CtrlHair/shape_branch/shape_util.py b/models/CtrlHair/shape_branch/shape_util.py new file mode 100644 index 0000000000000000000000000000000000000000..82c7e43733248dc9c1f9402f4625baeb0783e624 --- /dev/null +++ b/models/CtrlHair/shape_branch/shape_util.py @@ -0,0 +1,26 @@ +import torch + +from models.CtrlHair.global_value_utils import HAIR_IDX + + +def mask_label_to_one_hot(img): + # img = img * 255 + img[img == 255] = 19 + bs, _, h, w = img.size() + nc = 19 + input_label = torch.FloatTensor(bs, nc + 1, h, w).zero_().to(img.device) + input_semantics = input_label.scatter_(1, img.long(), 1.0) + input_semantics = input_semantics[:, :-1, :, :] + return input_semantics + + +def mask_one_hot_to_label(one_hot): + mask = torch.argmax(one_hot, dim=1) + mask[one_hot.max(dim=1)[0] == 0] = 255 + return mask + + +def split_hair_face(mask): + hair = mask[:, [HAIR_IDX], :, :] + face = torch.cat([mask[:, :HAIR_IDX, :, :], mask[:, HAIR_IDX + 1:, :, :]], dim=1) + return hair, face diff --git a/models/CtrlHair/shape_branch/solver.py b/models/CtrlHair/shape_branch/solver.py new file mode 100644 index 0000000000000000000000000000000000000000..3b4c929472fc130ccc91578c1abb06470a5037f7 --- /dev/null +++ b/models/CtrlHair/shape_branch/solver.py @@ -0,0 +1,262 @@ +# -*- coding: utf-8 -*- + +""" +# File name: solver.py +# Time : 2021/11/17 16:24 +# Author: xyguoo@163.com +# Description: +""" +import random + +import numpy as np +import torch +import torch.nn.functional as F +import torchvision.transforms as T +from torch.nn.parallel import DistributedDataParallel as DDP + +from models.CtrlHair.my_torchlib.train_utils import generate_noise +from .config import cfg +from .model import Generator, Discriminator, DiscriminatorNoise +# solver +from .shape_util import mask_label_to_one_hot, mask_one_hot_to_label, split_hair_face + + +class Solver: + + def __init__(self, cfg, device, local_rank, training=True): + + self.ce_loss = torch.nn.CrossEntropyLoss() + self.mse_loss = torch.nn.MSELoss() + + self.cfg = cfg + self.gen = Generator(cfg) + self.dis = Discriminator(cfg) + + self.gen.to(device) + self.dis.to(device) + + 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) + self.local_rank = local_rank + self.device = device + + @staticmethod + def kl_loss(mean, std): + var = std ** 2 + var_log = torch.log(var + 1e-4) + kl_loss = 0.5 * (mean ** 2 + 1.0 * (var - 1 - var_log)).mean() + return kl_loss + + def forward(self, data): + self.data_target_mask = data['target'] + self.data_face_mask = data['face'] + self.data_hair_mask = data['hair'] + + if cfg.disturb_real_batch_mask: + data_target_disturb = self.disturb_real(data['target']) + self.ae_in_hair, self.ae_in_target_face = split_hair_face(data_target_disturb) + data_face_disturb = self.disturb_real(data['face']) + _, self.ae_in_face = split_hair_face(data_face_disturb) + else: + self.ae_in_hair, self.ae_in_target_face = split_hair_face(data['target']) + _, self.ae_in_face = split_hair_face(data['face']) + + self.ae_mid_hair_code, self.ae_mid_hair_mean, self.ae_mid_hair_std = self.gen.forward_hair_encoder( + self.ae_in_hair) + self.ae_mid_face_code = self.gen.forward_face_encoder(self.ae_in_face) + self.ae_out_hair_logit = self.gen.forward_hair_decoder(self.ae_mid_hair_code, self.ae_mid_face_code) + self.ae_out_face_logit = self.gen.forward_face_decoder(self.ae_mid_face_code) + + self.ae_out_mask = self.gen.forward_decoder(self.ae_out_hair_logit, self.ae_out_face_logit) + + if 'lambda_adv_noise' in cfg or 'lambda_info' in cfg: + self.real_noise = generate_noise(cfg.batch_size, cfg.hair_dim).to( + self.ae_mid_face_code.device).type_as(self.ae_mid_face_code) + + if 'lambda_info' in cfg: + self.gan_in_hair_code = self.real_noise + self.gan_in_face_code = self.ae_mid_face_code + self.gan_mid_hair_logit = self.gen.forward_hair_decoder(self.gan_in_hair_code, self.gan_in_face_code) + self.gan_mid_face_logit = self.gen.forward_face_decoder(self.gan_in_face_code) + self.gan_mid_mask = self.gen.forward_decoder(self.gan_mid_hair_logit, self.gan_mid_face_logit) + + self.gan_mid_hair, _ = split_hair_face(self.gan_mid_mask) + self.gan_out_hair_code, _, _ = self.gen.forward_hair_encoder(self.gan_mid_hair) + + if random.random() < 0.5: + self.dis_out_fake = self.dis.forward(self.ae_out_mask) + else: + self.dis_out_fake = self.dis.forward(self.gan_mid_mask) + else: + if random.random() < cfg.random_ae_prob: + self.dis_out_fake = self.dis.forward(self.ae_out_mask) + else: + self.gan_in_hair_code = self.real_noise + self.gan_mid_hair_logit = self.gen.forward_hair_decoder(self.gan_in_hair_code, + self.ae_mid_face_code) + self.gan_mid_face_logit = self.ae_out_face_logit + self.gan_mid_mask = self.gen.forward_decoder(self.gan_mid_hair_logit, self.gan_mid_face_logit) + self.dis_out_fake = self.dis.forward(self.gan_mid_mask) + + def forward_g(self, loss_dict): + self.forward_general_gen(self.dis_out_fake, loss_dict) + hair, face = split_hair_face(self.ae_out_mask) + + if cfg.regular_method == 'ce': + loss_dict['lambda_hair'] = -torch.log(hair + 1e-5)[(self.ae_in_hair > 0.5)].mean() + loss_dict['lambda_non_hair'] = -torch.log(1 - hair + 1e-5)[(self.ae_in_hair < 0.5)].mean() + loss_dict['lambda_face'] = -torch.log(face + 1e-5)[(self.ae_in_target_face > 0.5)].mean() + + hair_hair, hair_face = split_hair_face(self.data_hair_mask) + mask = self.gen.forward_edit_directly_in_test(hair_hair, hair_face) + if 'lambda_self_rec' in cfg: + if cfg.regular_method == 'ce': + loss_dict['lambda_self_rec'] = -torch.log(mask + 1e-5)[self.data_hair_mask > 0.5].mean() + + if 'lambda_kl' in cfg and cfg.lambda_kl > 0: + loss_dict['lambda_kl'] = self.kl_loss(self.ae_mid_hair_mean, self.ae_mid_hair_std) + if cfg.lambda_moment_1 or cfg.lambda_moment_2: + noise_mid = self.ae_mid_hair_code + 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) - 0.973) ** 2).mean() + + if 'lambda_info' in cfg: + loss_dict['lambda_info'] = self.mse_loss(self.gan_out_hair_code, self.gan_in_hair_code) + + if cfg.lambda_adv_noise: + self.d_noise_res = self.dis_noise(self.ae_mid_hair_code) + self.forward_general_gen(self.d_noise_res, loss_dict, loss_name_suffix='_noise') + + 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, loss_dict, + dis_model=None, input_real=None, input_fake=None, 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(dis0))) + else: + assert 0, "Unsupported GAN type: {}".format(cfg.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] * (len(input_real.shape) - 1))).type_as(input_real) + x_hat = (alpha_gp * input_real + (1 - alpha_gp) * input_fake).requires_grad_(True) + out_hat = dis_model.forward(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 + + if cfg.lambda_gp_0 and cfg.lambda_gp_0 > 0: + # ii = input_real.requires_grad_(True) + dydx = torch.autograd.grad(outputs=dis1.sum(), inputs=input_real, retain_graph=True, create_graph=True, + only_inputs=True, allow_unused=True)[0] + dydx2 = dydx.pow(2) + dydx_l2norm = dydx2.view(dydx.size(0), -1).sum(1) + loss_gp = dydx_l2norm.mean() + loss_dict['lambda_gp_0' + loss_name_suffix] = loss_gp + + def forward_d(self, loss_dict, real_batch): + if cfg.disturb_real_batch_mask: + real_batch = self.disturb_real(real_batch) + + if 'lambda_gp_0' in cfg and cfg.lambda_gp_0 > 0: + real_batch = real_batch.requires_grad_() + dis_out_real = self.dis.forward(real_batch) + self.forward_general_dis(dis_out_real, self.dis_out_fake, + loss_dict, self.dis, input_real=real_batch, + input_fake=self.ae_out_mask) + + def disturb_real(self, real_batch): + cur = (torch.rand(real_batch.shape).to(real_batch.device) * 0.03 + real_batch) + cur = cur / cur.sum(dim=(1), keepdim=True) + return cur + + def forward_adv_noise(self, loss_dict): + + input_real = self.real_noise + if 'lambda_gp_0' in cfg and cfg.lambda_gp_0 > 0: + input_real = input_real.requires_grad_() + + input_fake = self.ae_mid_hair_code.detach() + + dis1 = self.dis_noise(input_real) + dis0 = self.dis_noise(input_fake) + + self.forward_general_dis(dis1, dis0, loss_dict, self.dis_noise, input_real=input_real, + input_fake=input_fake, loss_name_suffix='_noise') + + +def get_hair_face_code(mask_generator, mask): + mask_batch = mask[None, None, :, :] + mask_batch = T.functional.resize(mask_batch.long(), (256, 256), + interpolation=T.InterpolationMode.NEAREST) + mask_one_hot = mask_label_to_one_hot(mask_batch) + hair, face = split_hair_face(mask_one_hot) + hair_code = mask_generator.forward_hair_encoder(hair, testing=True) + face_code = mask_generator.forward_face_encoder(face) + return face_code, hair_code + + +def get_new_shape(mask_generator, face_code, new_hair_code): + out_mask = mask_generator.forward_decode_by_code(new_hair_code, face_code) + out_mask = mask_one_hot_to_label(out_mask)[0] + return out_mask diff --git a/models/CtrlHair/shape_branch/train.py b/models/CtrlHair/shape_branch/train.py new file mode 100644 index 0000000000000000000000000000000000000000..64423f479a791a7ea5eea61db44ecde3104df6c3 --- /dev/null +++ b/models/CtrlHair/shape_branch/train.py @@ -0,0 +1,135 @@ +# -*- 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 +from shape_branch.config import cfg, args +from shape_branch.dataset import Dataset +import my_pylib +from shape_branch.validation_in_train import print_val_save_model +# distributed training +import torch.distributed as dist +from shape_branch.solver import Solver +from my_torchlib.train_utils import LossUpdater, to_device, generate_noise, train +import my_torchlib +from shape_branch.model import init_weights +from shape_branch.shape_util import mask_label_to_one_hot + + +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 = cfg.checkpoints_dir + 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) + my_pylib.mkdir(cfg.sample_dir) + + try: + ckpt = my_torchlib.load_checkpoint(ckpt_dir) + start_step = ckpt['step'] + 1 + 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) + start_step = 1 + + # writer + if local_rank <= 0: + writer = tensorboardX.SummaryWriter(cfg.log_dir) + else: + writer = None + + # start training + test_batch = ds.get_test_batch(cfg.sample_batch_size) + for ke in ['face', 'target', 'hair']: + test_batch[ke] = mask_label_to_one_hot(test_batch[ke]) + 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_random_pair_batch(cfg.batch_size) + for ke in data: + data[ke] = mask_label_to_one_hot(data[ke]) + to_device(data, device) + loss_dict = {} + solver.forward(data) + if i < cfg.G_D_train_num['D']: + real_batch = ds.get_random_single_batch(cfg.batch_size) + real_batch = mask_label_to_one_hot(real_batch) + real_batch = real_batch.to(device) + solver.forward_d(loss_dict, real_batch) + 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, cfg.sample_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) diff --git a/models/CtrlHair/shape_branch/validation_in_train.py b/models/CtrlHair/shape_branch/validation_in_train.py new file mode 100644 index 0000000000000000000000000000000000000000..b9a46ad6c759ba46b227e64e873a35c854f1811f --- /dev/null +++ b/models/CtrlHair/shape_branch/validation_in_train.py @@ -0,0 +1,165 @@ +# -*- coding: utf-8 -*- + +""" +# File name: validation_in_train.py +# Time : 2021/12/10 12:55 +# Author: xyguoo@163.com +# Description: +""" +import os + +import copy + +import cv2 + +from hair_editor import HairEditor +from my_torchlib.train_utils import generate_noise +from util.canvas_grid import Canvas +from util.imutil import write_rgb +from util.mask_color_util import mask_to_rgb +from .config import cfg +import my_pylib +import torch +import torch.distributed as dist +import my_torchlib +from .shape_util import split_hair_face, mask_one_hot_to_label +import numpy as np + +he = HairEditor(load_feature_model=False, load_mask_model=False) + + +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 + + gen.eval() + # dis.eval() + + with torch.no_grad(): + target_hair_part, _ = split_hair_face(test_batch['target']) + face_hair_part, face_part = split_hair_face(test_batch['face']) + hair_hair_part, hair_face_part = split_hair_face(test_batch['hair']) + + face_face_code = solver.gen.forward_face_encoder(face_part) + _, target_hair_code, _ = solver.gen.forward_hair_encoder(target_hair_part) + hair_face_code = solver.gen.forward_face_encoder(hair_face_part) + _, hair_hair_code, _ = solver.gen.forward_hair_encoder(hair_hair_part) + + hair_mask_label = mask_one_hot_to_label(test_batch['hair']).cpu().numpy() + hair_mask_vis = [mask_to_rgb(ii, draw_type=0) for ii in hair_mask_label] + + face_mask_label = mask_one_hot_to_label(test_batch['face']).cpu().numpy() + face_mask_vis = [mask_to_rgb(ii, draw_type=0) for ii in face_mask_label] + + target_mask_label = mask_one_hot_to_label(test_batch['target']).cpu().numpy() + target_mask_vis = [mask_to_rgb(ii, draw_type=0) for ii in target_mask_label] + + # --------- + # rec and edit random + # --------- + grid_count = 10 + # generate each noise dim + canvas = Canvas(cfg.sample_batch_size, grid_count + 2, margin=3) + + rec_masks = solver.gen.forward_decode_by_code(hair_hair_code, hair_face_code) + rec_masks_label = mask_one_hot_to_label(rec_masks).cpu().numpy() + rec_masks_vis = [mask_to_rgb(ii, draw_type=0) for ii in rec_masks_label] + for draw_idx, idx in enumerate(range(cfg.sample_batch_size)): + canvas.process_draw_image(hair_mask_vis[draw_idx], draw_idx, 0) + canvas.process_draw_image(rec_masks_vis[draw_idx], draw_idx, 1) + + temp_noise = generate_noise(grid_count, cfg.hair_dim).to(device) + for grid_idx in range(grid_count): + cur_hair_code = torch.tile(temp_noise[grid_idx, :], (cfg.sample_batch_size, 1)) + res_masks = solver.gen.forward_decode_by_code(cur_hair_code, hair_face_code) + res_masks_label = mask_one_hot_to_label(res_masks).cpu().numpy() + for draw_idx, idx in enumerate(range(cfg.sample_batch_size)): + canvas.process_draw_image(mask_to_rgb(res_masks_label[draw_idx], draw_type=0), + draw_idx, grid_idx + 2) + + if local_rank <= 0: + canvas.write_(os.path.join(out_dir, '%07d_random.png' % step)) + + # --------- + # transfer + # --------- + res_mask = solver.gen.forward_decode_by_code(target_hair_code, face_face_code) + res_mask_label = mask_one_hot_to_label(res_mask).cpu().numpy() + res_mask_vis = [mask_to_rgb(ii, draw_type=0) for ii in res_mask_label] + vacant_column = np.ones([cfg.sample_batch_size * cfg.img_size, 3, 3], dtype='uint8') * 255 + + face_imgs = np.concatenate(test_batch['face_imgs'], axis=0) + hair_imgs = np.concatenate(test_batch['hair_imgs'], axis=0) + target_imgs = [] + for idx in range(cfg.sample_batch_size): + cur_img = he.gen_img(test_batch['sean_code'][idx][None, ...], + res_mask_label[idx][None, None, ...]).cpu().numpy() * 127.5 + 127.5 + cur_img = np.transpose(cur_img, [1, 2, 0]) + cur_img, _ = he.postprocess_blending(test_batch['face_imgs'][idx], cur_img, + face_mask_label[idx][None, ...], + res_mask_label[idx][None, ...]) + target_imgs.append(cur_img) + target_imgs = np.concatenate(target_imgs, axis=0) + canvas = np.concatenate([np.concatenate(hair_mask_vis, axis=0), vacant_column, + hair_imgs, vacant_column, + np.concatenate(face_mask_vis, axis=0), vacant_column, + face_imgs, vacant_column, + np.concatenate(target_mask_vis, axis=0), vacant_column, + np.concatenate(res_mask_vis, axis=0), vacant_column, + target_imgs], axis=1) + write_rgb(os.path.join(out_dir, '%07d__transfer.png' % step), canvas) + + # --------- + # edit code + # --------- + grid_count = 6 + lin_space = np.linspace(-2.5, 2.5, grid_count) + + for dim_idx in range(cfg.hair_dim): + canvas = Canvas(cfg.sample_batch_size, grid_count + 1, margin=3) + + for draw_idx, idx in enumerate(range(cfg.sample_batch_size)): + canvas.process_draw_image(hair_mask_vis[draw_idx], draw_idx, 0) + + for grid_idx in range(grid_count): + hair_hair_code_copy = copy.deepcopy(hair_hair_code) + hair_hair_code_copy[:, dim_idx] = lin_space[grid_idx] + res_mask = solver.gen.forward_decode_by_code(hair_hair_code_copy, hair_face_code) + res_mask = mask_one_hot_to_label(res_mask).cpu().numpy() + + for draw_idx, idx in enumerate(range(cfg.sample_batch_size)): + draw_mask = mask_to_rgb(res_mask[draw_idx], draw_type=0) + canvas.process_draw_image(draw_mask, draw_idx, grid_idx + 1) + if local_rank <= 0: + canvas.write_(os.path.join(out_dir, '%07d_noise_%02d.png' % (step, dim_idx))) + + gen.train() + # dis.scripts() + 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() diff --git a/models/CtrlHair/ui/__init__.py b/models/CtrlHair/ui/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..a533a6a830fd861255e4c037ede85835b76c6629 --- /dev/null +++ b/models/CtrlHair/ui/__init__.py @@ -0,0 +1,7 @@ +# -*- coding: utf-8 -*- + +""" +# File name: __init__.py.py +# Time : 2022/2/20 15:45 +# Author: xyguoo@163.com +""" diff --git a/models/CtrlHair/ui/backend.py b/models/CtrlHair/ui/backend.py new file mode 100644 index 0000000000000000000000000000000000000000..5996f00d3f050e9d4d1d8efeee16a09b3c5a4dd5 --- /dev/null +++ b/models/CtrlHair/ui/backend.py @@ -0,0 +1,504 @@ +# -*- coding: utf-8 -*- + +""" +# File name: hair_swap.py +# Time : 2022/2/20 17:06 +# Author: xyguoo@163.com +# Description: This is the backend interface for editing. When you want to customize custom editing and + modification functions (such as replacing mask), it is recommended to call these interfaces directly. + In the finally main section of this file, there is an example of using Backend +""" + +import sys +sys.path.append('.') + +import copy +import os + +import cv2 +import torch + +from global_value_utils import HAIR_IDX, TEMP_FOLDER +from my_torchlib.train_utils import generate_noise +from shape_branch.shape_util import mask_label_to_one_hot, mask_one_hot_to_label, split_hair_face +from util.mask_color_util import mask_to_rgb +from wrap_codes.mask_adaptor import wrap_by_imgs + +from hair_editor import HairEditor +from util.color_from_hsv_to_gaussian import DistTranslation + + +class LatentRepresentation: + def __init__(self): + self.color = None + self.curliness = None + self.shape = None + self.texture = None + self.face = None + + +class Backend(HairEditor): + """ + This is the main interface set, please call these interfaces directly to customize custom editing + """ + + def __init__(self, maximum_value_fe, blending=True, temp_path=os.path.join(TEMP_FOLDER, 'demo_output')): + """ + :param maximum_value_fe: The max value in sliding bar of the frontend ui. If you just use the backend, + but not use the frontend, this value is + :param blending: Whether use poisson blending between hair region and others region + :param temp_path: the dir that save temp content, which can be removed after execution + """ + super().__init__(True, True) + self.target_img = None + self.input_img = None + self.target_mask = None + self.input_mask = None + self.cur_latent = None + self.target_latent = None + self.cur_mask = None + self.input_sean_code = None + self.target_size = 256 + self.maximum_value_fe = maximum_value_fe + self.temp_path = temp_path + self.blending = blending + self.dist_translation = DistTranslation() + + def parse_img(self, img_rgb, target_img=False): + """ + mask the input image + :param img_rgb: with channel order rgb + :param target_img: is the target image? + """ + img_ts = cv2.resize(img_rgb, (self.target_size, self.target_size)) + mask = self.get_mask(img_rgb) + lr = LatentRepresentation() + + # infer shape + if target_img: + out_mask = None + else: + mask_batch = self.preprocess_mask(mask) + mask_tensor = torch.tensor(mask_batch, dtype=torch.uint8, device=self.device) + mask_one_hot = mask_label_to_one_hot(mask_tensor) + hair, face = split_hair_face(mask_one_hot) + hair_code = self.mask_generator.forward_hair_encoder(hair, testing=True) + face_code = self.mask_generator.forward_face_encoder(face) + lr.shape = hair_code + lr.face = face_code + out_mask = self.mask_generator.forward_decode_by_code(hair_code, face_code) + out_mask = mask_one_hot_to_label(out_mask).cpu().numpy()[0] + + # infer feature + input_code = self.get_code(self.preprocess_img(img_rgb), mask_batch) + hair_feature = input_code[:, HAIR_IDX] + + out_color = self.feature_rgb_predictor({'code': hair_feature}) + + c = out_color['rgb_mean'].detach().cpu().numpy() + c_hsv = cv2.cvtColor(c[None, ...].astype('uint8'), cv2.COLOR_RGB2HSV) + c_hsv = torch.tensor(c_hsv).to(self.device)[0] + lr.color = {'hsv': c_hsv, 'pca_std': out_color['pca_std']} + + out_enc = self.feature_encoder({'code': hair_feature}) + lr.curliness = out_enc['noise_curliness'] + lr.texture = out_enc['noise'] + return img_ts, out_mask, lr, mask, input_code, hair_feature + + def tensor_hsv_to_rgb(self, hsv): + """ + input tensor with hsv and convert to rgb + """ + c = hsv.detach().cpu().numpy() + c_rgb = cv2.cvtColor(c[None, ...].astype('uint8'), cv2.COLOR_HSV2RGB) + c_rgb = torch.tensor(c_rgb).to(self.device)[0] + return c_rgb + + def tensor_rgb_to_hsv(self, rgb): + """ + input tensor with rgb and convert to hsv + """ + + c = rgb.detach().cpu().numpy() + c_hsv = cv2.cvtColor(c[None, ...].astype('uint8'), cv2.COLOR_RGB2HSV) + c_hsv = torch.tensor(c_hsv).to(self.device)[0] + return c_hsv + + def set_input_img(self, img_rgb): + """ + set and parse the input image + :param img_rgb: + """ + self.input_img, self.cur_mask, self.cur_latent, \ + self.input_mask, self.input_sean_code, self.input_hair_feature = self.parse_img(img_rgb) + input_mask_show = mask_to_rgb(self.cur_mask, draw_type=1) + return self.input_img, input_mask_show + + def set_target_img(self, img_rgb): + """ + set and parse the target image + :param img_rgb: + """ + self.target_img, _, self.target_latent, \ + self.target_mask, _, self.target_hair_feature = self.parse_img(img_rgb) + input_maks_show = mask_to_rgb(self.target_mask, draw_type=1) + return self.target_img, input_maks_show + + def output(self, target_latent=None, feature=None): + """ + generate an color_texture image + :param target_latent: if the latent representation of target is not provided, + `self.cur_latent` and `self.cur_mask` will be used + :param feature: if edited hair feature X is not provided, it will be generated with color and texture branch + :return: color_texture image + """ + if target_latent is None: + target_latent = self.cur_latent + target_mask = self.cur_mask + else: + target_mask = self.refresh_cur_mask(target_latent)[0] + + if 'rgb_mean' in target_latent.color: + target_color_rgb = self.target_latent.color['rgb_mean'] + else: + target_color_rgb = self.tensor_hsv_to_rgb(target_latent.color['hsv']) + + if feature is None: + data = {'noise': target_latent.texture, 'noise_curliness': target_latent.curliness, + 'rgb_mean': target_color_rgb, 'pca_std': target_latent.color['pca_std']} + feature = self.feature_generator(data)['code'] + self.input_sean_code[:, HAIR_IDX] = feature + edit_img = self.gen_img(self.input_sean_code, target_mask[None, None, ...]) + output_img, _ = self.postprocess_blending(self.input_img, edit_img, self.input_mask, target_mask, + blending=self.blending) + + return output_img + + def change_curliness(self, val): + """ + change the latent representation of curliness of texture + """ + self.cur_latent.curliness[0] = val + + # def change_color(self, val, idx): + # val = (val + self.maximum_value_fe) / 2 / self.maximum_value_fe + # if idx == 3: + # self.cur_latent.color['pca_std'][0] = val * 100 + 20 + # else: + # if idx == 0: + # val *= 179 + # if idx == 1: + # val *= 255 + # if idx == 2: + # val *= 255 + # self.cur_latent.color['hsv'][0][idx] = val + + def change_color(self, val, idx): + """ + change the latent representation of color + idx=0 is hue, idx=1 is saturation, idx=2 is brightness, idx=3 is variance + + please note that when idx=3, the variance is not obey gaussian, + the value range is in [-self.maximum_value_fe, self.maximum_value_fe] + """ + if idx == 3: + val = (val + self.maximum_value_fe) / 2 / self.maximum_value_fe + self.cur_latent.color['pca_std'][0] = val * 100 + 20 + else: + val = self.dist_translation.gaussian_to_val(idx, val) + self.cur_latent.color['hsv'][0][idx] = val + + def change_shape(self, val, idx): + """ + change the latent representation of shape + :param val: latent value + :param idx: for current checkpoint, idx=0 is length, idx=1 is volumn, idx=2 is bangs_direction, idx=3 is bangs + """ + self.continue_change_with_direction('shape', self.shape_dirs[idx], val) + self.refresh_cur_mask() + + def change_texture(self, val, idx): + """ + change the latent representation of texture + :param val: latent value + :param idx: for current checkpoint, idx=0 is smoothness, idx=1 is thickness + """ + self.continue_change_with_direction('texture', self.texture_dirs[idx], val) + + def get_curliness_be2fe(self): + """ + convert the latent representation to the value of sliding bar in the frontend of Pyqt ui, and return it + """ + return self.cur_latent.curliness[0] + + def get_color_be2fe(self): + """ + convert the latent representation to the value of sliding bar in the frontend of Pyqt ui, and return it + """ + c_hsv = self.cur_latent.color['hsv'].detach().cpu().numpy()[0] + # color0 = c_hsv[0] / 179 * 2 * self.maximum_value_fe - self.maximum_value_fe + # color1 = c_hsv[1] / 255 * 2 * self.maximum_value_fe - self.maximum_value_fe + # color2 = c_hsv[2] / 255 * 2 * self.maximum_value_fe - self.maximum_value_fe + color0 = self.dist_translation.val_to_gaussian(0, c_hsv[0]) + color1 = self.dist_translation.val_to_gaussian(1, c_hsv[1]) + color2 = self.dist_translation.val_to_gaussian(2, c_hsv[2]) + var_fe = (self.cur_latent.color['pca_std'][0] - 20) / 100 * 2 * self.maximum_value_fe - self.maximum_value_fe + return color0, color1, color2, var_fe + + def get_shape_be2fe(self): + """ + convert the latent representation to the value of sliding bar in the frontend of Pyqt ui, and return it + """ + res = [] + for idx in range(4): + res.append(torch.dot(self.cur_latent.shape[0], self.shape_dirs[idx])) + return res + + def get_texture_be2fe(self): + """ + convert the latent representation to the value of sliding bar in the frontend of Pyqt ui, and return it + """ + res = [] + for idx in range(2): + res.append(torch.dot(self.cur_latent.texture[0], self.texture_dirs[idx])) + return res + + def transfer_latent_representation(self, flag, refresh=True): + """ + transfer the latent representation of target image to input image. + i.e., transfer from self.target_latent to self.cur_latent + :param flag: 'color', 'texture' or 'shape' + :param refresh: whether refresh mask + :return: + """ + if flag == 'shape': + wt, _ = wrap_by_imgs(self.target_img, self.input_img, wrap_temp_folder=self.temp_path, + need_crop=False) + wt = self.preprocess_mask(wt) + self.warp_target = wt[0, 0] + mask_tensor = torch.tensor(wt, dtype=torch.uint8, device=self.device) + mask_one_hot = mask_label_to_one_hot(mask_tensor) + hair, face = split_hair_face(mask_one_hot) + hair_code = self.mask_generator.forward_hair_encoder(hair, testing=True) + face_code = self.mask_generator.forward_face_encoder(face) + self.target_latent.shape = hair_code + self.target_latent.face = face_code + + self.refresh_cur_mask() + + target_att = self.target_latent.__getattribute__(flag) + if isinstance(target_att, torch.Tensor): + self.cur_latent.__setattr__(flag, target_att.clone()) + else: + cp_dict = copy.copy(target_att) + for ke in cp_dict: + cp_dict[ke] = cp_dict[ke].clone() + self.cur_latent.__setattr__(flag, cp_dict) + + if flag == 'shape' and refresh: + self.refresh_cur_mask() + + if flag == 'texture': + self.transfer_latent_representation('curliness') + + def refresh_cur_mask(self, target_latent=None): + """ + refresh and generate current mask + :param target_latent: if no target latent is given, self.cur_mask will be used + :return: + """ + if target_latent is None: + target_latent = self.cur_latent + out_mask = self.mask_generator.forward_decode_by_code(target_latent.shape, target_latent.face) + out_mask = mask_one_hot_to_label(out_mask).cpu().numpy()[0] + self.cur_mask = out_mask + return out_mask, mask_to_rgb(out_mask, draw_type=1) + + def get_cur_mask(self): + """ + get the mask, which can be visited + """ + return mask_to_rgb(self.cur_mask, draw_type=1) + + def interpolate_hsv(self, hsv1, hsv2, alpha): + """ + final hsv = hsv1 * (1 - alpha) + hsv2 * alpha + return: final hsv + """ + rgb1 = self.tensor_hsv_to_rgb(hsv1) + rgb2 = self.tensor_hsv_to_rgb(hsv2) + rgb = rgb1 * (1 - alpha) + rgb2 * alpha + inter_hsv = self.tensor_rgb_to_hsv(rgb) + return inter_hsv + + def interpolate_triple(self, latent1, latent2, latent3, alpha1, alpha2, alpha3): + """ + (latent1 * alpha1 + latent2 * alpha2 + latent3 * alpha3) / (alpha1 + alpha2 + alpha3) + :return: final latent + """ + latent12 = self.interpolate(latent1, latent2, alpha2 / (alpha1 + alpha2)) + latent_res = self.interpolate(latent12, latent3, alpha3) + return latent_res + + def interpolate(self, latent1, latent2, alpha): + """ + final latent = latent1 * (1 - alpha) + latent2 * alpha + :return: final latent + """ + result_latent = LatentRepresentation() + for att in ['curliness', 'shape', 'texture']: + result_latent.__setattr__(att, latent1.__getattribute__(att) * (1 - alpha) + + latent2.__getattribute__(att) * alpha) + color_dic = {} + color_dic['pca_std'] = latent1.color['pca_std'] * (1 - alpha) + latent2.color['pca_std'] * alpha + color_dic['hsv'] = self.interpolate_hsv(latent1.color['hsv'], latent2.color['hsv'], alpha) + + result_latent.color = color_dic + result_latent.face = self.cur_latent.face + return result_latent + + def interpolate_each_att(self, latent1, latent2, alpha, att_name): + """ + interpolate a certain latent + :param att_name: curliness, shape, texture, color + :return: final full latent + """ + result_latent = LatentRepresentation() + for att in ['curliness', 'shape', 'texture']: + result_latent.__setattr__(att, self.cur_latent.__getattribute__(att).clone()) + + if att_name == 'shape': + # keep color + color_dic = {} + for semantic in ['hsv', 'pca_std']: + color_dic[semantic] = self.cur_latent.color[semantic].clone() + result_latent.__setattr__(att_name, latent1.__getattribute__(att_name) * (1 - alpha) + + latent2.__getattribute__(att_name) * alpha) + elif att_name in ['curliness', 'texture']: + # keep color + color_dic = {} + for semantic in ['hsv', 'pca_std']: + color_dic[semantic] = self.cur_latent.color[semantic].clone() + result_latent.__setattr__('curliness', latent1.__getattribute__('curliness') * (1 - alpha) + + latent2.__getattribute__('curliness') * alpha) + result_latent.__setattr__('texture', latent1.__getattribute__('texture') * (1 - alpha) + + latent2.__getattribute__('texture') * alpha) + else: + color_dic = {} + # for semantic in ['hsv', 'pca_std']: + color_dic['pca_std'] = latent1.color['pca_std'] * (1 - alpha) + latent2.color['pca_std'] * alpha + color_dic['hsv'] = self.interpolate_hsv(latent1.color['hsv'], latent2.color['hsv'], alpha) + + result_latent.color = color_dic + result_latent.face = self.cur_latent.face + + return result_latent + + @staticmethod + def show_hair_region(mask, non_hair_value=0): + """ + Get hair region, and full none hair region as `non_hair_value` + :param mask: + :param non_hair_value: + :return: + """ + mask_rgb = mask_to_rgb(mask, draw_type=1) + mask_rgb[mask != HAIR_IDX] = non_hair_value + return mask_rgb + + def directly_change_hair_mask(self, hair_mask): + """ + Directly replace hair region with a hair_mask. This is a recommend method to imply mask transfer. + :param hair_mask: + :return: + """ + hair_mask = hair_mask == HAIR_IDX + face_logit = self.mask_generator.forward_face_decoder(self.cur_latent.face) + hair_logit = torch.tensor(hair_mask)[None, None, ...].type_as(face_logit).to(self.device) + hair_logit = hair_logit * (face_logit.max() - face_logit.min() + 2) + face_logit.min() - 1 + mask = self.mask_generator.forward_decoder(hair_logit, face_logit) + self.cur_mask = mask_one_hot_to_label(mask).cpu().numpy()[0] + + def get_random_texture(self): + """ + sample a texture latent code randomly + :return: + """ + random_latent = generate_noise(1, 8) + random_latent = random_latent.to(self.device) + self.cur_latent.texture = random_latent + + def get_random_shape(self): + """ + sample a shape latent code randomly + :return: + """ + random_latent = generate_noise(1, 16) + random_latent = random_latent.to(self.device) + self.cur_latent.shape = random_latent + self.refresh_cur_mask() + + def get_random_curliness(self): + """ + sample a curliness latent code randomly + :return: + """ + random_latent = generate_noise(1, 1) + random_latent = random_latent.to(self.device) + self.cur_latent.curliness = random_latent + + def continue_change_with_direction(self, att_name, direction, val): + """ + change the latent code value on a projection direction + :param att_name: shape or texture + :param direction: projection direction + :param val: coordinate on this projection direction + :return: + """ + att = self.cur_latent.__getattribute__(att_name) + att = att + (val - torch.dot(att[0], direction)) * direction + self.cur_latent.__setattr__(att_name, att) + if att_name == 'shape': + self.refresh_cur_mask() + + +""" +This is a example of using Backend for costume editing +""" +if __name__ == '__main__': + be = Backend(2.5) + from util.imutil import read_rgb, write_rgb + + input_image = read_rgb('imgs/00079.png') + target_image = read_rgb('imgs/00001.png') + + """ + If the image need crop + """ + # input_image = be.crop_face(input_image) + # target_image = be.crop_face(target_image) + + input_image = cv2.resize(input_image, (256, 256)) + + be.set_input_img(input_image) + be.set_target_img(target_image) + + # transfer all latent code from target image to input image + be.transfer_latent_representation('texture') + be.transfer_latent_representation('color') + be.transfer_latent_representation('shape') + + # change the variance manually + be.change_color(1.0, 2) + + out_mask = be.get_mask(input_image) + output_img = be.output() + write_rgb('temp.png', output_img) + # above is the output image + + im2 = read_rgb('imgs/00037.png') + im2 = cv2.resize(im2, (256, 256)) + be.set_target_img(im2) + be.transfer_latent_representation('shape') + output_img2 = be.output() + # above is the output image 2 diff --git a/models/CtrlHair/ui/frontend_demo.py b/models/CtrlHair/ui/frontend_demo.py new file mode 100644 index 0000000000000000000000000000000000000000..610880ee28302cf605f4a31df1a604606fb856a1 --- /dev/null +++ b/models/CtrlHair/ui/frontend_demo.py @@ -0,0 +1,269 @@ +# -*- coding: utf-8 -*- + +""" +# File name: frontend.py +# Time : 2022/2/20 15:58 +# Author: xyguoo@163.com +# Description: This is the demo frontend +""" + +import sys +sys.path.append('.') + +from global_value_utils import TEMP_FOLDER +import argparse +import os + +from util.common_options import ctrl_hair_parser_options + +parser = argparse.ArgumentParser() +ctrl_hair_parser_options(parser) + +args = parser.parse_args() +os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu +from ui.backend import Backend +from util.imutil import read_rgb, write_rgb + +from PyQt5.QtWidgets import QWidget, QPushButton, QVBoxLayout, QApplication, QLabel, QGridLayout, \ + QSlider, QFileDialog +from PyQt5.QtGui import QPixmap, QFont +from PyQt5.QtCore import Qt + + +class Example(QWidget): + def __init__(self): + super().__init__() + self.temp_path = os.path.join(TEMP_FOLDER, 'demo_output') + self.maximum_value = 2.0 + self.blending = not args.no_blending + self.backend = Backend(self.maximum_value, blending=self.blending) + self.initUI() + self.target_size = 256 + self.need_crop = args.need_crop + if not os.path.exists(self.temp_path): + os.makedirs(self.temp_path) + self.font = QFont() + self.font.setPointSize(15) + self.setFont(self.font) + + self.input_name = None + self.target_name = None + + def initUI(self): + self.lbl_target_img = QLabel(self) + self.lbl_input_img = QLabel(self) + self.lbl_input_seg = QLabel(self) + self.lbl_out_img = QLabel(self) + + self.labels = [self.lbl_target_img, self.lbl_input_img, + self.lbl_input_seg, self.lbl_out_img] + + self.grid1 = QGridLayout() + # tags = ['target image', 'input image', 'hair shape', 'color_texture'] + # for idx in range(len(self.labels)): + # self.grid1.addWidget(QLabel(tags[idx]), 0, idx) + for idx in range(len(self.labels)): + self.grid1.addWidget(self.labels[idx], 1, idx, alignment=Qt.AlignTop) + self.labels[idx].setFixedSize(256, 256) + + self.btn_open_target = QPushButton('Target Image', self) + self.btn_open_target.clicked[bool].connect(self.evt_open_target) + self.grid1.addWidget(self.btn_open_target, 0, 0) + + self.btn_open_input = QPushButton('Input Image', self) + self.btn_open_input.clicked[bool].connect(self.evt_open_input) + self.grid1.addWidget(self.btn_open_input, 0, 1) + + self.grid1.addWidget(QLabel('Hair Shape'), 0, 2, alignment=Qt.AlignCenter) + + self.btn_output = QPushButton('Output', self) + self.btn_output.clicked[bool].connect(self.evt_output) + self.grid1.addWidget(self.btn_output, 0, 3) + self.btn_output.setEnabled(False) + + self.grid2 = QGridLayout() + + self.btn_trans_color = QPushButton('Transfer Color', self) + self.btn_trans_color.clicked[bool].connect(self.evt_trans_color) + self.grid2.addWidget(self.btn_trans_color, 10, 0) + self.btn_trans_color.setEnabled(False) + + self.btn_trans_texture = QPushButton('Transfer Texture', self) + self.btn_trans_texture.clicked[bool].connect(self.evt_trans_texture) + self.grid2.addWidget(self.btn_trans_texture, 10, 1) + self.btn_trans_texture.setEnabled(False) + + self.btn_trans_shape = QPushButton('Transfer Shape', self) + self.btn_trans_shape.clicked[bool].connect(self.evt_trans_shape) + self.grid2.addWidget(self.btn_trans_shape, 10, 2) + self.btn_trans_shape.setEnabled(False) + + self.sld2val = {} + self.val2sld = {} + + self.label_color = ['Color: Hue', 'Color: Saturation', 'Color: Brightness', + 'Color: Variance'] + self.label_shape = ['Shape: Volume', 'Shape: Bangs', 'Shape: Length', 'Shape: Direction'] + self.label_curliness = ['Texture: Curliness'] + self.label_app = ['Texture: Smoothness', 'Texture: Thickness'] + self.label_total = self.label_color + self.label_shape + self.label_curliness + self.label_app + + col_num = 4 + row_num = 3 + for row in range(row_num): + for col in range(col_num): + if col == 3 and row == 2: + continue + num = col_num * row + col + sld = QSlider(Qt.Horizontal, self) + sld.setMinimum(-self.maximum_value * 100) + sld.setMaximum(self.maximum_value * 100) + sld.sliderMoved[int].connect(self.evt_change_value) + self.sld2val[sld] = num + self.val2sld[num] = sld + self.grid2.addWidget(QLabel(self.label_total[num]), row * 2 + 2, col) + self.grid2.addWidget(sld, row * 2 + 2 + 1, col) + sld.setEnabled(False) + + self.grid2.addWidget(QLabel(), 10, 3) + + whole_vbox = QVBoxLayout(self) + whole_vbox.addLayout(self.grid1) + whole_vbox.addLayout(self.grid2) + + self.setLayout(whole_vbox) + self.setGeometry(100, 100, 900, 600) + self.setWindowTitle('CtrlHair') + self.show() + + def evt_open_target(self): + fname = QFileDialog.getOpenFileName(self, 'Open image file') + if fname[0]: + self.target_name = fname[0] + self.load_target_image(fname[0]) + if self.input_name is not None: + self.btn_trans_color.setEnabled(True) + self.btn_trans_shape.setEnabled(True) + self.btn_trans_texture.setEnabled(True) + + def evt_open_input(self): + fname = QFileDialog.getOpenFileName(self, 'Open image file') + if fname[0]: + input_name = fname[0] + self.input_name = input_name + self.load_input_image(input_name) + self.btn_output.setEnabled(True) + if self.target_name is not None: + self.btn_trans_color.setEnabled(True) + self.btn_trans_shape.setEnabled(True) + self.btn_trans_texture.setEnabled(True) + + for kk in self.sld2val: + kk.setEnabled(True) + + def evt_output(self): + output_img = self.backend.output() + img_path = os.path.join(self.temp_path, 'out_img.png') + write_rgb(img_path, output_img) + self.lbl_out_img.setPixmap((QPixmap(img_path))) + + def evt_trans_color(self): + self.backend.transfer_latent_representation('color') + self.refresh_slider() + + def evt_trans_texture(self): + self.backend.transfer_latent_representation('texture') + self.refresh_slider() + + def evt_trans_shape(self): + self.backend.transfer_latent_representation('shape', refresh=True) + self.refresh_slider() + input_parsing_show = self.backend.get_cur_mask() + input_parsing_path = os.path.join(self.temp_path, 'input_parsing.png') + write_rgb(input_parsing_path, input_parsing_show) + self.lbl_input_seg.setPixmap((QPixmap(input_parsing_path))) + + def load_input_image(self, img_path): + img = read_rgb(img_path) + if self.need_crop: + img = self.backend.crop_face(img) + input_img, input_parsing_show = self.backend.set_input_img(img_rgb=img) + input_path = os.path.join(self.temp_path, 'input_img.png') + write_rgb(input_path, input_img) + self.lbl_input_img.setPixmap((QPixmap(input_path))) + + input_parsing_path = os.path.join(self.temp_path, 'input_parsing.png') + write_rgb(input_parsing_path, input_parsing_show) + self.lbl_input_seg.setPixmap((QPixmap(input_parsing_path))) + self.refresh_slider() + + self.lbl_out_img.setPixmap(QPixmap(None)) + + def load_target_image(self, img_path): + img = read_rgb(img_path) + if self.need_crop: + img = self.backend.crop_face(img) + input_img, input_parsing_show = self.backend.set_target_img(img_rgb=img) + input_path = os.path.join(self.temp_path, 'target_img.png') + write_rgb(input_path, input_img) + self.lbl_target_img.setPixmap((QPixmap(input_path))) + + def refresh_slider(self): + idx = 0 + # color + color_val = self.backend.get_color_be2fe() + for ii in range(4): + self.val2sld[idx + ii].setValue(int(color_val[ii] * 100)) + + # shape + idx += len(self.label_color) + shape_val = self.backend.get_shape_be2fe() + for ii in range(4): + self.val2sld[idx + ii].setValue(int(shape_val[ii] * 100)) + + # curliness + idx += len(self.label_shape) + self.val2sld[idx].setValue(self.backend.get_curliness_be2fe() * 100) + # texture + idx += len(self.label_curliness) + app_val = self.backend.get_texture_be2fe() + for ii in range(2): + self.val2sld[idx + ii].setValue(int(app_val[ii] * 100)) + + def evt_change_value(self, sld_v): + """ + change all sliders value + :param v: 0-100 + :return: + """ + v = sld_v / 100.0 + sld_idx = self.sld2val[self.sender()] + if sld_idx < len(self.label_color): + self.backend.change_color(v, sld_idx) + return + sld_idx -= len(self.label_color) + if sld_idx < len(self.label_shape): + self.backend.change_shape(v, sld_idx) + input_parsing_show = self.backend.get_cur_mask() + input_parsing_path = os.path.join(self.temp_path, 'input_parsing.png') + write_rgb(input_parsing_path, input_parsing_show) + self.lbl_input_seg.setPixmap((QPixmap(input_parsing_path))) + return + sld_idx -= len(self.label_shape) + if sld_idx < len(self.label_curliness): + self.backend.change_curliness(v) + return + sld_idx -= len(self.label_curliness) + if sld_idx < len(self.label_app): + self.backend.change_texture(v, sld_idx) + return + + +def main(): + app = QApplication(sys.argv) + ex = Example() + sys.exit(app.exec_()) + + +if __name__ == '__main__': + main() diff --git a/models/CtrlHair/util/__init__.py b/models/CtrlHair/util/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/models/CtrlHair/util/canvas_grid.py b/models/CtrlHair/util/canvas_grid.py new file mode 100644 index 0000000000000000000000000000000000000000..65a200854f842ce8d83876be58bc401c2c24bcc9 --- /dev/null +++ b/models/CtrlHair/util/canvas_grid.py @@ -0,0 +1,34 @@ +# -*- coding: utf-8 -*- + +""" +# File name: canvas_grid.py +# Time : 2021/11/18 16:40 +# Author: xyguoo@163.com +# Description: This is the util that generating image in a big grid picture +""" + +import numpy as np + +from util.imutil import write_rgb + + +class Canvas: + def __init__(self, row, col, img_size=256, margin=0): + self.row = row + self.col = col + self.img_size = img_size + self.margin = margin + self.canvas = np.ones((row * img_size, col * img_size + margin * (col - 1), 3), dtype='uint8') * 255 + + def process_draw_image(self, img, i, j): + if img.dtype in [np.float32, np.float, np.float64]: + if img.min() < 0: + img = img * 127.5 + 127.5 + elif img.max() <= 1: + img = img * 255 + img = img.astype('uint8') + i_start, j_start = int(i * self.img_size), int(j * self.img_size) + int(j * self.margin) + self.canvas[i_start: i_start + img.shape[0], j_start: j_start + img.shape[1], :] = img + + def write_(self, file): + write_rgb(file, self.canvas) diff --git a/models/CtrlHair/util/color_from_hsv_to_gaussian.py b/models/CtrlHair/util/color_from_hsv_to_gaussian.py new file mode 100644 index 0000000000000000000000000000000000000000..5a45b8bacc231d6430eaf971288bd859a5a8d55d --- /dev/null +++ b/models/CtrlHair/util/color_from_hsv_to_gaussian.py @@ -0,0 +1,40 @@ +# -*- coding: utf-8 -*- + +""" +# File name: color_from_hsv_to_gaussian.py +# Time : 2022/3/6 16:00 +# Author: xyguoo@163.com +# Description: transfer HSV value to Gaussian latent value according to the distribution of the dataset +""" +import os +import pickle as pkl +from bisect import bisect_left, bisect_right + +import scipy.stats as st + + +class DistTranslation: + def __init__(self): + hair_root = 'dataset_info_ctrlhair' + with open(os.path.join(hair_root, 'hsv_stat_dict_ordered.pkl'), 'rb') as f: + self.cols_hsv = pkl.load(f) + + def gaussian_to_val(self, dim, val): + # if dim == 0: + # return (val + 2.) / 4 * 179 + return self.cols_hsv[int((st.norm.cdf(val)) * self.cols_hsv.shape[0])][dim] + + def val_to_gaussian(self, dim, val): + # if dim == 0: + # return val / 179 * 2 * 2. - 2. + + left_v = bisect_left(self.cols_hsv[:, dim], val) + right_v = bisect_right(self.cols_hsv[:, dim], val) + return st.norm.ppf((left_v + right_v) / 2 / self.cols_hsv.shape[0]) + +# +# if __name__ == '__main__': +# dt = DistTranslation() +# +# with open('hsv_stat_dict_ordered.pkl', 'wb') as f: +# pkl.dump(dt.cols_hsv, f) diff --git a/models/CtrlHair/util/common_options.py b/models/CtrlHair/util/common_options.py new file mode 100644 index 0000000000000000000000000000000000000000..7f54b7f53ca743e0da96d2a2c05aad5cd234d569 --- /dev/null +++ b/models/CtrlHair/util/common_options.py @@ -0,0 +1,15 @@ +# -*- coding: utf-8 -*- + +""" +# File name: common_options.py +# Time : 2022/7/12 +# Author: xyguoo@163.com +# Description: +""" + +def ctrl_hair_parser_options(parser): + parser.add_argument('-c', '--config', type=str, default='001') + parser.add_argument('-g', '--gpu', type=str, help='Specify GPU number', default='0') + parser.add_argument('-n', '--need_crop', type=bool, help='whether images need crop', default=True) + parser.add_argument('--no_blending', action='store_true', + help='whether using poisson blending as post processing', default=False) \ No newline at end of file diff --git a/models/CtrlHair/util/find_semantic_direction.py b/models/CtrlHair/util/find_semantic_direction.py new file mode 100644 index 0000000000000000000000000000000000000000..e2f8b06abc4cd0906825404cfeb2639a3ac82b31 --- /dev/null +++ b/models/CtrlHair/util/find_semantic_direction.py @@ -0,0 +1,21 @@ +# -*- coding: utf-8 -*- + +""" +# File name: find_semantic_direction.py +# Time : 2022/07/17 +# Author: xyguoo@163.com +# Description: +""" +import torch + + +def get_random_direction(dim, existing_dirs): + dir = torch.randn(dim) + + for dd in existing_dirs: + dir = dir - torch.dot(dir, dd) * dd + + if dir[0] < 0: + dir = -dir + dir = dir / dir.norm() + return dir diff --git a/models/CtrlHair/util/imutil.py b/models/CtrlHair/util/imutil.py new file mode 100644 index 0000000000000000000000000000000000000000..539fa8c449eb192a8c84d5c01588f029a5ecd93d --- /dev/null +++ b/models/CtrlHair/util/imutil.py @@ -0,0 +1,24 @@ +# -*- coding: utf-8 -*- + +""" +# File name: imutil.py +# Time : 2021/12/7 14:55 +# Author: xyguoo@163.com +# Description: +""" +import cv2 +import numpy as np + + +def read_rgb(img_path): + im = cv2.imread(img_path) + im = cv2.cvtColor(im, cv2.COLOR_BGR2RGB) + return im + + +def write_rgb(file_name, img): + if len(img.shape) == 2: + img = np.stack([img] * 3, axis=2) + elif img.shape[2] == 1: + img = np.tile(img, [1, 1, 3]) + cv2.imwrite(file_name, cv2.cvtColor(img, cv2.COLOR_RGB2BGR)) diff --git a/models/CtrlHair/util/mask_color_util.py b/models/CtrlHair/util/mask_color_util.py new file mode 100644 index 0000000000000000000000000000000000000000..8751567c07bd3abad70f7a6bf0257fecda80bf5c --- /dev/null +++ b/models/CtrlHair/util/mask_color_util.py @@ -0,0 +1,64 @@ +# -*- coding: utf-8 -*- + +""" +# File name: mask_color_util.py +# Time : 2021/12/7 11:13 +# Author: xyguoo@163.com +# Description: +""" + +import numpy as np + +from models.CtrlHair.global_value_utils import HAIR_IDX + + +def mask_to_rgb(pred, draw_type=2): + """ + generate visual mask image + :param pred: pred is the label image with pixel vale in {0, 1, 2, ..., 20} + :param draw_type: 0: all part; 1: {bg, face, hair}; 2: {hair, others} + :return: + """ + if len(pred.shape) == 3 and pred.shape[0] == 1: + pred = pred[0] + num_labels = 19 + color = np.array([[0, 128, 64], + [204, 0, 0], + [76, 153, 0], + [204, 204, 0], ## + [51, 51, 255], ## + [204, 0, 204], ## + [0, 255, 255], ## + [51, 255, 255], ## + [102, 51, 0], ## + [255, 0, 0], ## + [102, 204, 0], ## + [255, 255, 0], ## + [0, 0, 153], ## + [0, 0, 204], ## + [255, 51, 153], ## + [0, 204, 204], ## + [0, 51, 0], ## + [255, 153, 51], + [0, 204, 0], + ]) + + for cc in range(len(color)): + if draw_type == 2: + if cc != HAIR_IDX: + color[cc] = [255, 255, 255] + elif draw_type == 1: + if cc != HAIR_IDX and cc != 0: + color[cc] = [237, 28, 36] + + h, w = np.shape(pred) + rgb = np.zeros((h, w, 3), dtype=np.uint8) + # print(color.shape) + for ii in range(num_labels): + # print(ii) + mask = pred == ii + rgb[mask, None] = color[ii, :] + # Correct unk + unk = pred == 255 + rgb[unk, None] = 255 + return rgb diff --git a/models/CtrlHair/util/util.py b/models/CtrlHair/util/util.py new file mode 100644 index 0000000000000000000000000000000000000000..d3eb65e92ea5963f3df4e6ed4c9018e870cd0181 --- /dev/null +++ b/models/CtrlHair/util/util.py @@ -0,0 +1,281 @@ +""" +These interfaces are most based on SEAN. But we modify some of them, and add some new interface for our CtrlHair. +""" + +import argparse +import importlib +import os +import re + +import cv2 +import dill as pickle +import numpy as np +import torch +from PIL import Image + + +def save_obj(obj, name): + with open(name, 'wb') as f: + pickle.dump(obj, f, pickle.HIGHEST_PROTOCOL) + + +def load_obj(name): + with open(name, 'rb') as f: + return pickle.load(f) + +# returns a configuration for creating a generator +# |default_opt| should be the opt of the current experiment +# |**kwargs|: if any configuration should be overriden, it can be specified here + + +def copyconf(default_opt, **kwargs): + conf = argparse.Namespace(**vars(default_opt)) + for key in kwargs: + print(key, kwargs[key]) + setattr(conf, key, kwargs[key]) + return conf + + +def tile_images(imgs, picturesPerRow=4): + """ Code borrowed from + https://stackoverflow.com/questions/26521365/cleanly-tile-numpy-array-of-images-stored-in-a-flattened-1d-format/26521997 + """ + + # Padding + if imgs.shape[0] % picturesPerRow == 0: + rowPadding = 0 + else: + rowPadding = picturesPerRow - imgs.shape[0] % picturesPerRow + if rowPadding > 0: + imgs = np.concatenate([imgs, np.zeros((rowPadding, *imgs.shape[1:]), dtype=imgs.dtype)], axis=0) + + # Tiling Loop (The conditionals are not necessary anymore) + tiled = [] + for i in range(0, imgs.shape[0], picturesPerRow): + tiled.append(np.concatenate([imgs[j] for j in range(i, i + picturesPerRow)], axis=1)) + + tiled = np.concatenate(tiled, axis=0) + return tiled + + +# Converts a Tensor into a Numpy array +# |imtype|: the desired type of the converted numpy array +def tensor2im(image_tensor, imtype=np.uint8, normalize=True, tile=False): + if isinstance(image_tensor, list): + image_numpy = [] + for i in range(len(image_tensor)): + image_numpy.append(tensor2im(image_tensor[i], imtype, normalize)) + return image_numpy + + if image_tensor.dim() == 4: + # transform each image in the batch + images_np = [] + for b in range(image_tensor.size(0)): + one_image = image_tensor[b] + one_image_np = tensor2im(one_image) + images_np.append(one_image_np.reshape(1, *one_image_np.shape)) + images_np = np.concatenate(images_np, axis=0) + if tile: + images_tiled = tile_images(images_np) + return images_tiled + else: + return images_np + + if image_tensor.dim() == 2: + image_tensor = image_tensor.unsqueeze(0) + image_numpy = image_tensor.detach().cpu().float().numpy() + if normalize: + image_numpy = (np.transpose(image_numpy, (1, 2, 0)) + 1) / 2.0 * 255.0 + else: + image_numpy = np.transpose(image_numpy, (1, 2, 0)) * 255.0 + image_numpy = np.clip(image_numpy, 0, 255) + if image_numpy.shape[2] == 1: + image_numpy = image_numpy[:, :, 0] + return image_numpy.astype(imtype) + + +# Converts a one-hot tensor into a colorful label map +def tensor2label(label_tensor, n_label, imtype=np.uint8, tile=False): + if label_tensor.dim() == 4: + # transform each image in the batch + images_np = [] + for b in range(label_tensor.size(0)): + one_image = label_tensor[b] + one_image_np = tensor2label(one_image, n_label, imtype) + images_np.append(one_image_np.reshape(1, *one_image_np.shape)) + images_np = np.concatenate(images_np, axis=0) + if tile: + images_tiled = tile_images(images_np) + return images_tiled + else: + images_np = images_np[0] + return images_np + + if label_tensor.dim() == 1: + return np.zeros((64, 64, 3), dtype=np.uint8) + if n_label == 0: + return tensor2im(label_tensor, imtype) + label_tensor = label_tensor.cpu().float() + if label_tensor.size()[0] > 1: + label_tensor = label_tensor.max(0, keepdim=True)[1] + label_tensor = Colorize(n_label)(label_tensor) + label_numpy = np.transpose(label_tensor.numpy(), (1, 2, 0)) + result = label_numpy.astype(imtype) + return result + + +def save_image(image_numpy, image_path, create_dir=False): + if create_dir: + os.makedirs(os.path.dirname(image_path), exist_ok=True) + if len(image_numpy.shape) == 2: + image_numpy = np.expand_dims(image_numpy, axis=2) + if image_numpy.shape[2] == 1: + image_numpy = np.repeat(image_numpy, 3, 2) + image_pil = Image.fromarray(image_numpy) + + # save to png + image_pil.save(image_path.replace('.jpg', '.png')) + + +def mkdirs(paths): + if isinstance(paths, list) and not isinstance(paths, str): + for path in paths: + mkdir(path) + else: + mkdir(paths) + + +def mkdir(path): + if not os.path.exists(path): + os.makedirs(path) + + +def atoi(text): + return int(text) if text.isdigit() else text + + +def natural_keys(text): + ''' + alist.sort(key=natural_keys) sorts in human order + http://nedbatchelder.com/blog/200712/human_sorting.html + (See Toothy's implementation in the comments) + ''' + return [atoi(c) for c in re.split('(\d+)', text)] + + +def natural_sort(items): + items.sort(key=natural_keys) + + +def str2bool(v): + if v.lower() in ('yes', 'true', 't', 'y', '1'): + return True + elif v.lower() in ('no', 'false', 'f', 'n', '0'): + return False + else: + raise argparse.ArgumentTypeError('Boolean value expected.') + + +def find_class_in_module(target_cls_name, module): + target_cls_name = target_cls_name.replace('_', '').lower() + clslib = importlib.import_module(module) + cls = None + for name, clsobj in clslib.__dict__.items(): + if name.lower() == target_cls_name: + cls = clsobj + + if cls is None: + print("In %s, there should be a class whose name matches %s in lowercase without underscore(_)" % (module, target_cls_name)) + exit(0) + + return cls + + +def save_network(net, label, epoch, opt): + save_filename = '%s_net_%s.pth' % (epoch, label) + save_path = os.path.join(opt.checkpoints_dir, opt.name, save_filename) + torch.save(net.cpu().state_dict(), save_path) + if len(opt.gpu_ids) and torch.cuda.is_available(): + net.cuda() + + +def load_network(net, label, epoch, opt): + save_filename = '%s_net_%s.pth' % (epoch, label) + save_dir = os.path.join(opt.checkpoints_dir, opt.name) + save_path = os.path.join(save_dir, save_filename) + weights = torch.load(save_path) + net.load_state_dict(weights) + return net + + +############################################################################### +# Code from +# https://github.com/ycszen/pytorch-seg/blob/master/transform.py +# Modified so it complies with the Citscape label map colors +############################################################################### +def uint82bin(n, count=8): + """returns the binary of integer n, count refers to amount of bits""" + return ''.join([str((n >> y) & 1) for y in range(count - 1, -1, -1)]) + + +def labelcolormap(N): + if N == 35: # cityscape + cmap = np.array([(0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), (111, 74, 0), (81, 0, 81), + (128, 64, 128), (244, 35, 232), (250, 170, 160), (230, 150, 140), (70, 70, 70), (102, 102, 156), (190, 153, 153), + (180, 165, 180), (150, 100, 100), (150, 120, 90), (153, 153, 153), (153, 153, 153), (250, 170, 30), (220, 220, 0), + (107, 142, 35), (152, 251, 152), (70, 130, 180), (220, 20, 60), (255, 0, 0), (0, 0, 142), (0, 0, 70), + (0, 60, 100), (0, 0, 90), (0, 0, 110), (0, 80, 100), (0, 0, 230), (119, 11, 32), (0, 0, 142)], + dtype=np.uint8) + else: + cmap = np.zeros((N, 3), dtype=np.uint8) + for i in range(N): + r, g, b = 0, 0, 0 + id = i + 1 # let's give 0 a color + for j in range(7): + str_id = uint82bin(id) + r = r ^ (np.uint8(str_id[-1]) << (7 - j)) + g = g ^ (np.uint8(str_id[-2]) << (7 - j)) + b = b ^ (np.uint8(str_id[-3]) << (7 - j)) + id = id >> 3 + cmap[i, 0] = r + cmap[i, 1] = g + cmap[i, 2] = b + + return cmap + + +class Colorize(object): + def __init__(self, n=35): + self.cmap = labelcolormap(n) + self.cmap = torch.from_numpy(self.cmap[:n]) + + def __call__(self, gray_image): + size = gray_image.size() + color_image = torch.ByteTensor(3, size[1], size[2]).fill_(0) + + for label in range(0, len(self.cmap)): + mask = (label == gray_image[0]).cpu() + color_image[0][mask] = self.cmap[label][0] + color_image[1][mask] = self.cmap[label][1] + color_image[2][mask] = self.cmap[label][2] + + return color_image + + +def make_folder(path): + if not os.path.exists(os.path.join(path)): + os.makedirs(os.path.join(path)) + + +def path_join_abs(*paths): + return os.path.abspath(os.path.join(*paths)) + + +def draw_landmark(landmarks, img_rd): + landmarks = landmarks.astype('int') + img_rd = img_rd.copy() + for idx, point in enumerate(landmarks): + pos = (point[0], point[1]) + cv2.circle(img_rd, pos, 2, color=(139, 0, 0)) + cv2.putText(img_rd, str(idx + 1), pos, cv2.FONT_HERSHEY_SIMPLEX, 0.3, (0, 0, 255), 1, cv2.LINE_AA) + return img_rd \ No newline at end of file diff --git a/models/Embedding.py b/models/Embedding.py new file mode 100644 index 0000000000000000000000000000000000000000..4f90b53c66a84bdd8940bf48d7d1e7b4746f4eea --- /dev/null +++ b/models/Embedding.py @@ -0,0 +1,118 @@ +from collections import defaultdict + +import torch +import torch.nn.functional as F +import torchvision.transforms as T +from torch import nn +from torch.utils.data import DataLoader + +from datasets.image_dataset import ImagesDataset, image_collate +from models.FeatureStyleEncoder import FSencoder +from models.Net import Net, get_segmentation +from models.encoder4editing.utils.model_utils import setup_model, get_latents +from utils.bicubic import BicubicDownSample +from utils.save_utils import save_gen_image, save_latents + + +class Embedding(nn.Module): + """ + Module for image embedding + """ + + def __init__(self, opts, net=None): + super().__init__() + self.opts = opts + if net is None: + self.net = Net(self.opts) + else: + self.net = net + + self.encoder = FSencoder.get_trainer(self.opts.device) + self.e4e, _ = setup_model('pretrained_models/encoder4editing/e4e_ffhq_encode.pt', self.opts.device) + + self.normalize = T.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) + self.to_bisenet = T.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)) + + self.downsample_512 = BicubicDownSample(factor=2) + self.downsample_256 = BicubicDownSample(factor=4) + + def setup_dataloader(self, images: dict[torch.Tensor, list[str]] | list[torch.Tensor], batch_size=None): + self.dataset = ImagesDataset(images) + self.dataloader = DataLoader(self.dataset, collate_fn=image_collate, shuffle=False, + batch_size=batch_size or self.opts.batch_size) + + @torch.inference_mode() + def get_e4e_embed(self, images: list[torch.Tensor]) -> dict[str, torch.Tensor]: + device = self.opts.device + self.setup_dataloader(images, batch_size=len(images)) + + for image, _ in self.dataloader: + image = image.to(device) + latent_W = get_latents(self.e4e, image) + latent_F, _ = self.net.generator([latent_W], input_is_latent=True, return_latents=False, + start_layer=0, end_layer=3) + return {"F": latent_F, "W": latent_W} + + @torch.inference_mode() + def embedding_images(self, images_to_name: dict[torch.Tensor, list[str]], **kwargs) -> dict[ + str, dict[str, torch.Tensor]]: + device = self.opts.device + self.setup_dataloader(images_to_name) + + name_to_embed = defaultdict(dict) + for image, names in self.dataloader: + image = image.to(device) + + im_512 = self.downsample_512(image) + im_256 = self.downsample_256(image) + im_256_norm = self.normalize(im_256) + + # E4E + latent_W = get_latents(self.e4e, im_256_norm) + + # FS encoder + output = self.encoder.test(img=self.normalize(image), return_latent=True) + latent = output.pop() # [bs, 512, 16, 16] + latent_S = output.pop() # [bs, 18, 512] + + latent_F, _ = self.net.generator([latent_S], input_is_latent=True, return_latents=False, + start_layer=3, end_layer=3, layer_in=latent) # [bs, 512, 32, 32] + + # BiSeNet + masks = torch.cat([get_segmentation(image.unsqueeze(0)) for image in self.to_bisenet(im_512)]) + + # Mixing if we change the color or shape + if len(images_to_name) > 1: + hair_mask = torch.where(masks == 13, torch.ones_like(masks, device=device), + torch.zeros_like(masks, device=device)) + hair_mask = F.interpolate(hair_mask.float(), size=(32, 32), mode='bicubic') + + latent_F_from_W = self.net.generator([latent_W], input_is_latent=True, return_latents=False, + start_layer=0, end_layer=3)[0] + latent_F = latent_F + self.opts.mixing * hair_mask * (latent_F_from_W - latent_F) + + for k, names in enumerate(names): + for name in names: + name_to_embed[name]['W'] = latent_W[k].unsqueeze(0) + name_to_embed[name]['F'] = latent_F[k].unsqueeze(0) + name_to_embed[name]['S'] = latent_S[k].unsqueeze(0) + name_to_embed[name]['mask'] = masks[k].unsqueeze(0) + name_to_embed[name]['image_256'] = im_256[k].unsqueeze(0) + name_to_embed[name]['image_norm_256'] = im_256_norm[k].unsqueeze(0) + + if self.opts.save_all: + gen_W_im, _ = self.net.generator([latent_W], input_is_latent=True, return_latents=False) + gen_FS_im, _ = self.net.generator([latent_S], input_is_latent=True, return_latents=False, + start_layer=4, end_layer=8, layer_in=latent_F) + + exp_name = exp_name if (exp_name := kwargs.get('exp_name')) is not None else "" + output_dir = self.opts.save_all_dir / exp_name + for name, im_W, lat_W in zip(names, gen_W_im, latent_W): + save_gen_image(output_dir, 'W+', f'{name}.png', im_W) + save_latents(output_dir, 'W+', f'{name}.npz', latent_W=lat_W) + + for name, im_F, lat_S, lat_F in zip(names, gen_FS_im, latent_S, latent_F): + save_gen_image(output_dir, 'FS', f'{name}.png', im_F) + save_latents(output_dir, 'FS', f'{name}.npz', latent_S=lat_S, latent_F=lat_F) + + return name_to_embed diff --git a/models/Encoders.py b/models/Encoders.py new file mode 100644 index 0000000000000000000000000000000000000000..8fe25686776a73cb782b13691efcce37e3373669 --- /dev/null +++ b/models/Encoders.py @@ -0,0 +1,160 @@ +import argparse + +import clip +import torch +import torch.nn as nn +from torch.nn import Linear, LayerNorm, LeakyReLU, Sequential +from torchvision import transforms as T + +from models.Net import FeatureEncoderMult, IBasicBlock, conv1x1 +from models.stylegan2.model import PixelNorm + + +class ModulationModule(nn.Module): + def __init__(self, layernum, last=False, inp=512, middle=512): + super().__init__() + self.layernum = layernum + self.last = last + self.fc = Linear(512, 512) + self.norm = LayerNorm([self.layernum, 512], elementwise_affine=False) + self.gamma_function = Sequential(Linear(inp, middle), LayerNorm([middle]), LeakyReLU(), Linear(middle, 512)) + self.beta_function = Sequential(Linear(inp, middle), LayerNorm([middle]), LeakyReLU(), Linear(middle, 512)) + self.leakyrelu = LeakyReLU() + + def forward(self, x, embedding): + x = self.fc(x) + x = self.norm(x) + gamma = self.gamma_function(embedding) + beta = self.beta_function(embedding) + out = x * (1 + gamma) + beta + if not self.last: + out = self.leakyrelu(out) + return out + + +class FeatureiResnet(nn.Module): + def __init__(self, blocks, inplanes=1024): + super().__init__() + + self.res_blocks = {} + + for n, block in enumerate(blocks, start=1): + planes, num_blocks = block + + for k in range(1, num_blocks + 1): + downsample = None + if inplanes != planes: + downsample = nn.Sequential(conv1x1(inplanes, planes, 1), nn.BatchNorm2d(planes, eps=1e-05, ), ) + + self.res_blocks[f'res_block_{n}_{k}'] = IBasicBlock(inplanes, planes, 1, downsample, 1, 64, 1) + inplanes = planes + + self.res_blocks = nn.ModuleDict(self.res_blocks) + + def forward(self, x): + for module in self.res_blocks.values(): + x = module(x) + return x + + +class RotateModel(nn.Module): + def __init__(self): + super().__init__() + self.pixelnorm = PixelNorm() + self.modulation_module_list = nn.ModuleList([ModulationModule(6, i == 4) for i in range(5)]) + + def forward(self, latent_from, latent_to): + dt_latent = self.pixelnorm(latent_from) + for modulation_module in self.modulation_module_list: + dt_latent = modulation_module(dt_latent, latent_to) + output = latent_from + 0.1 * dt_latent + return output + + +class ClipBlendingModel(nn.Module): + def __init__(self, clip_model="ViT-B/32"): + super().__init__() + self.pixelnorm = PixelNorm() + self.clip_model, _ = clip.load(clip_model, device="cuda") + self.transform = T.Compose( + [T.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711))]) + self.face_pool = torch.nn.AdaptiveAvgPool2d((224, 224)) + self.modulation_module_list = nn.ModuleList( + [ModulationModule(12, i == 4, inp=512 * 3, middle=1024) for i in range(5)] + ) + + for param in self.clip_model.parameters(): + param.requires_grad = False + + def get_image_embed(self, image_tensor): + resized_tensor = self.face_pool(image_tensor) + renormed_tensor = self.transform(resized_tensor * 0.5 + 0.5) + return self.clip_model.encode_image(renormed_tensor) + + def forward(self, latent_face, latent_color, target_face, hair_color): + embed_face = self.get_image_embed(target_face).unsqueeze(1).expand(-1, 12, -1) + embed_color = self.get_image_embed(hair_color).unsqueeze(1).expand(-1, 12, -1) + latent_in = torch.cat((latent_color, embed_face, embed_color), dim=-1) + + dt_latent = self.pixelnorm(latent_face) + for modulation_module in self.modulation_module_list: + dt_latent = modulation_module(dt_latent, latent_in) + output = latent_face + 0.1 * dt_latent + return output + + +class PostProcessModel(nn.Module): + def __init__(self): + super().__init__() + self.encoder_face = FeatureEncoderMult(fs_layers=[9], opts=argparse.Namespace( + **{'arcface_model_path': "pretrained_models/ArcFace/backbone_ir50.pth"})) + + self.latent_avg = torch.load('pretrained_models/PostProcess/latent_avg.pt', map_location=torch.device('cuda')) + self.to_feature = FeatureiResnet([[1024, 2], [768, 2], [512, 2]]) + + self.to_latent_1 = nn.ModuleList([ModulationModule(18, i == 4) for i in range(5)]) + self.to_latent_2 = nn.ModuleList([ModulationModule(18, i == 4) for i in range(5)]) + self.pixelnorm = PixelNorm() + + def forward(self, source, target): + s_face, [f_face] = self.encoder_face(source) + s_hair, [f_hair] = self.encoder_face(target) + + dt_latent_face = self.pixelnorm(s_face) + dt_latent_hair = self.pixelnorm(s_hair) + + for mod_module in self.to_latent_1: + dt_latent_face = mod_module(dt_latent_face, s_hair) + + for mod_module in self.to_latent_2: + dt_latent_hair = mod_module(dt_latent_hair, s_face) + + finall_s = self.latent_avg + 0.1 * (dt_latent_face + dt_latent_hair) + + cat_f = torch.cat((f_face, f_hair), dim=1) + finall_f = self.to_feature(cat_f) + + return finall_s, finall_f + + +class ClipModel(nn.Module): + def __init__(self): + super().__init__() + self.clip_model, _ = clip.load("ViT-B/32", device="cuda") + self.transform = T.Compose( + [T.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711))] + ) + self.face_pool = torch.nn.AdaptiveAvgPool2d((224, 224)) + + for param in self.clip_model.parameters(): + param.requires_grad = False + + def forward(self, image_tensor): + if not image_tensor.is_cuda: + image_tensor = image_tensor.to("cuda") + if image_tensor.dtype == torch.uint8: + image_tensor = image_tensor / 255 + + resized_tensor = self.face_pool(image_tensor) + renormed_tensor = self.transform(resized_tensor) + return self.clip_model.encode_image(renormed_tensor) diff --git a/models/FeatureStyleEncoder/.DS_Store b/models/FeatureStyleEncoder/.DS_Store new file mode 100644 index 0000000000000000000000000000000000000000..24d931501c51af3fb84693dba0fd8b2afa221ab9 Binary files /dev/null and b/models/FeatureStyleEncoder/.DS_Store differ diff --git a/models/FeatureStyleEncoder/.gitignore b/models/FeatureStyleEncoder/.gitignore new file mode 100644 index 0000000000000000000000000000000000000000..e5c5a4e8085394da550fb2da14170ba8c1ba0a49 --- /dev/null +++ b/models/FeatureStyleEncoder/.gitignore @@ -0,0 +1,6 @@ +logs/ +__pycache__/ +matshow/ +.ipynb_checkpoints/ +tmp/ +RAFT/ diff --git a/models/FeatureStyleEncoder/FSencoder.py b/models/FeatureStyleEncoder/FSencoder.py new file mode 100644 index 0000000000000000000000000000000000000000..ef93b00821569a381113695405d09c2adf3d7df1 --- /dev/null +++ b/models/FeatureStyleEncoder/FSencoder.py @@ -0,0 +1,41 @@ +from argparse import Namespace +import glob +import os +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.utils.data as data +import yaml +import sys + +current_dir = os.path.abspath(os.path.dirname(__file__)) +sys.path.insert(0, current_dir) + +from PIL import Image +from tqdm import tqdm +from torchvision import transforms, utils + +from trainer import * + +torch.backends.cudnn.enabled = True +torch.backends.cudnn.deterministic = True +torch.backends.cudnn.benchmark = True +# torch.autograd.set_detect_anomaly(True) +Image.MAX_IMAGE_PIXELS = None + +opts = Namespace(config='001', pretrained_model_path='pretrained_models/FeatureStyleEncoder/143_enc.pth', stylegan_model_path=f'pretrained_models/FeatureStyleEncoder/psp_ffhq_encode.pt', arcface_model_path=f'pretrained_models/FeatureStyleEncoder/backbone.pth', parsing_model_path=f'pretrained_models/FeatureStyleEncoder/79999_iter.pth', log_path='./logs/', resume=False, checkpoint='', checkpoint_noiser='', multigpu=False, input_path='./test/', save_path='./') + +config = yaml.load(open(f'{current_dir}/configs/' + opts.config + '.yaml', 'r'), Loader=yaml.FullLoader) + +def get_trainer(device): + # Initialize trainer + trainer = Trainer(config, opts) + trainer.initialize(opts.stylegan_model_path, opts.arcface_model_path, opts.parsing_model_path) + trainer.to(device) + + # state_dict = torch.load(opts.pretrained_model_path)#os.path.join(opts.log_path, opts.config + '/checkpoint.pth')) + trainer.enc.load_state_dict(torch.load(opts.pretrained_model_path)) + trainer.enc.eval() + + return trainer \ No newline at end of file diff --git a/models/FeatureStyleEncoder/LICENSE.md b/models/FeatureStyleEncoder/LICENSE.md new file mode 100644 index 0000000000000000000000000000000000000000..846612a6045a350bb13c23a278ed25a857d0b494 --- /dev/null +++ b/models/FeatureStyleEncoder/LICENSE.md @@ -0,0 +1,60 @@ +## LIMITED SOFTWARE EVALUATION LICENSE AGREEMENT + +The following Limited Software Evaluation License (the “License”) constitutes an agreement between you (the “Licensee”) and InterDigital Communications, Inc, a company organized and existing under the laws of the State of Delaware, USA, with its registered offices located at 200 Bellevue Parkway, Suite 300, Wilmington, DE 19809, USA (hereinafter “InterDigital”). +This License governs the download and use of the Software (as defined below). Your use of the Software is subject to the terms and conditions set forth in this License. By installing, using, accessing or copying the Software, you hereby irrevocably accept the terms and conditions of this License. If you do not accept all parts of the terms and conditions of this License, you cannot install, use, access nor copy the Software. + +# Article 1. Definitions +“Affiliate” as used herein shall mean any entity that, directly or indirectly, through one or more intermediates, is controlled by, controls, or is under common control with InterDigital or The Licensee, as the case may be. For purposes of this definition only, the term “control” means the possession of the power to direct or cause the direction of the management and policies of an entity, whether by ownership of voting stock or partnership interest, by contract, or otherwise, including direct or indirect ownership of more than fifty percent (50%) of the voting interest in the entity in question. +“Authorized Purpose” means any use of the Software for fundamental research work with the exclusion of any commercial use. A commercial use includes, without limitation, any sublicense granted on the Software against a fee whatever its nature, any use of the Software in a product that is offered (either free or for a price) to any third party, any use of the Software to provide a service to a third party and/or any use of the Software to create a competing product of the Software ("Purpose") +“Documentation” means textual materials delivered by InterDigital to the Licensee pursuant to this License relating to the Software, in written or electronic format, including but not limited to, technical reference manuals, technical notes, user manuals, and application guides. +“Effective Date” means the date Licensee first installs a copy of the Software on any computer. + +“Limited Period” means the life of the copyright owned by InterDigital on the Software in each and every country where such copyright would exist. +“Intellectual Property Rights” means all copyrights, trademarks, trade secrets, patents and any other intellectual property rights recognized in any jurisdiction worldwide, including all applications and registrations with respect thereto. +"Open Source Software" shall mean any software, including where appropriate, any and all modifications, derivative works, enhancements, upgrades, improvements, fixed bugs, and/or statically linked to the source code of such software, released under a free or open source software license that requires, as a condition of usage, copy, modification and/or redistribution of such software, that the party: +• Redistribute the Open Source Software royalty-free; and/or +• Redistribute the Open Source Software under the same license/distribution terms as those contained in the open source or free software license under which it was originally released; and/or +• Release to the public, disclose or otherwise make available the source code of the Open Source Software. +For purposes of this License, by means of example and without limitation, any software that is released or distributed under any of the following licenses shall be qualified as Open Source Software: (i) GNU General Public License (GPL); (ii) GNU Lesser/Library GPL (LGPL); (iii) the Artistic License; (iv) the Mozilla Public License; (v) the Common Public License; (vi) the Sun Community Source License (SCSL); (vii) the Sun Industry Standards Source License (SISSL); (viii) BSD License; (ix) MIT License; (x) Apache Software License; (xi) Open SSL License; (xii) IBM Public License; and (xiii) Open Software License. +“Software” means the Software with which this license was downloaded, namely FeatureStyleEncoder in object code. +# Article 2. License +InterDigital grants Licensee a free, worldwide, non-exclusive, license to InterDigital’s copyright on the Software to download, use and reproduce solely for the Authorized Purpose for the Limited Period. +Licensee shall not pay any royalty, license fee or maintenance fee, or other fee of any nature under this License. +# Article 3. Restrictions on use of the Software +Licensee shall not have the right to correct, adapt, modify, reverse engineer, disassemble, decompile or/and otherwise perform or conduct any action leading to the transformation of the Software. +Licensee shall not remove, obscure or modify any copyright, trademark or other proprietary rights notices, marks or labels contained on or within the Software, falsify or delete any author attributions, legal notices or other labels of the origin or source of the material. +Licensee may reproduce and distribute copies of the Software in any medium provided that Licensee gives any other recipients of the Software a copy of this License. + +# Article 4. Ownership +Title to and ownership of the Software, the Documentation, and/or any Intellectual Property Right protecting the Software and/or the Documentation shall at all times remain with InterDigital. Licensee agrees that except for the limited rights granted to the Software as set forth in Section 2 above, in no event shall anything in this License grant, provide, or convey any other rights, privileges, immunities, or interest in or to any Intellectual Property Rights (including but not limited to patent rights) of InterDigital or any of its Affiliates, whether by implication, estoppel, or otherwise. + +# Article 5. Publication/Communication +Any publication or oral communication resulting from the use of the Software shall be elaborated in good faith and shall not be driven by a deliberate will to denigrate InterDigital or any of its products. In any publication and on any support joined to an oral communication (e.g., a PowerPoint presentation) relating to the Software, the following statement shall be inserted: +“FeatureStyleEncoder” is an InterDigital product” +In any publication, the latest publication about the software shall be properly cited. The latest publication currently is: +A Style-Based GAN Encoder for High Fidelity Reconstruction of Images and Videos, Xu Yao, Alasdair Newson, Yann Gousseau, Pierre Hellier, ECCV European Conference on Computer Vision, 2022 (https://arxiv.org/pdf/2202.02183.pdf). +In any oral communication relating to the Software and/or its use, the Licensee shall orally indicate that the Software is InterDigital’s property. + +# Article 6. No Warranty - Disclaimer +THE SOFTWARE AND DOCUMENTATION ARE PROVIDED TO LICENSEE ON AN “AS IS” BASIS. INTERDIGITAL MAKES NO WARRANTY THAT THE SOFTWARE WILL OPERATE ON ANY PARTICULAR HARDWARE, PLATFORM, OR ENVIRONMENT. THERE IS NO WARRANTY THAT THE OPERATION OF THE SOFTWARE SHALL BE UNINTERRUPTED, WITHOUT BUGS OR ERROR FREE. THE SOFTWARE AND DOCUMENTATION ARE PROVIDED HEREUNDER WITHOUT WARRANTY OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO ANY IMPLIED LIABILITIES AND WARRANTIES OF NONINFRINGEMENT OF INTELLECTUAL PROPERTY, FREEDOM FROM INHERENT DEFECTS, CONFORMITY TO A SAMPLE OR MODEL, MERCHANTABILITY, FITNESS AND/OR SUITABILITY FOR A SPECIFIC OR GENERAL PURPOSE AND THOSE ARISING BY STATUTE OR BY LAW, OR FROM A CAUSE OF DEALING OR USAGE OF TRADE. ANY AND ALL SUCH IMPLIED WARRANTIES ARE FULLY DISCLAIMED BY INTERDIGITAL TO THE MAXIMUM EXTENT ALLOWED BY LAW, AND LICENSEE ACKNOWLEDGES THAT THIS DISCLAIMER OF ALL EXPRESS AND IMPLIED WARRANTIES BY INTERDIGITAL, AS WELL AS LICENSEE’S ACCEPTANCE AND ACKNOWLEDGEMENT OF THE SAME, IS A MATERIAL PART OF THE CONSIDERATION FOR THIS LICENSE. +InterDigital shall not be obligated to perform or provide any modifications, derivative works, enhancements, upgrades, updates or improvements of the Software or Documentation, or to fix any bug that could arise. +Licensee at all times uses the Software at its own cost, risk and responsibility. InterDigital shall not be liable for any damages that could accrue by or to Licensee as a result of its use of the Software, either in accordance with this License or not. +InterDigital shall not be liable for any consequential or indirect losses, including any indirect loss of profits, revenues, business, and/or anticipated savings, whether or not in the contemplation of the Parties at the time of entering into this License unless expressly set out in this License, or arising from gross negligence, willful misconduct or fraud. +Licensee agrees that it will defend, indemnify and hold harmless InterDigital and its Affiliates against any and all losses, damages, costs and expenses arising from a breach by the Licensee of any of its obligations or representations hereunder, including, without limitation, any third party claims, and/or any claims in connection with any such breach and/or any use of the Software, including any claim from third party arising from access, use, or any other activity in relation to this Software. +Licensee shall not make any warranty, representation, or commitment on behalf of InterDigital to any other third party. + +# Article 7. Open Source Software +Licensee hereby represents, warrants, and covenants to InterDigital that Licensee’s use of the Software shall not result in the Contamination of all or any part of the Software, directly or indirectly, or of any Intellectual Property of InterDigital or its Affiliates. +As used herein, “Contamination” shall mean that the licensing terms under which any Open Source Software, distinct from the Software, is released would also apply to the Software herein, by virtue of such Open Source Software being linked to, combined with, or otherwise connected to the Software. +Licensee agree that some Open Source Software are included in the distribution. A list of such is provided in exhibit A with the relevant licenses applicable. For the avoidance of doubt, regarding such open source parts, the relevant license will apply exclusively. + +# Article 8. No Future Contract Obligation +Neither this License nor the furnishing of the Software, nor any other InterDigital information provided to Licensee, shall be construed to obligate either party to: (a) enter into any further agreement or negotiation concerning the deployment of the Software; (b) refrain from entering into any agreement or negotiation with any other third party regarding the same or any other subject matter; or (c) refrain from pursuing its business in whatever manner it elects even if this involves competing with the other party. + +# Article 9. General Provisions +9.1 Severability. If any provision of this License shall be held to be in contravention of applicable law, this License shall be construed as if such provision were not a part thereof, and in all other respects the terms hereof shall remain in full force and effect. +9.2 Governing Law. Regardless of the place of execution, delivery, performance or any other aspect of this License, this License and all of the rights of the parties under this License shall be governed by, construed under and enforced in accordance with the substantive law of the State of Delaware, USA, without regard to conflicts of law principles. In case of a dispute that cannot be settled amicably, the state and federal courts located in New Castle County, Delaware, USA, shall have exclusive jurisdiction over such dispute, and each party hereby irrevocably waives any objection to the jurisdiction of such courts, including but not limited to objections of lack of in personam jurisdiction or based on principles of forum non conveniens. +9.3 Survival. The provisions of articles 1, 3, 4, 6, 7, 8, 9.1, 9.2 and 9.5 shall survive termination of this License. +9.4 Assignment. InterDigital may assign this license to any third Party. Licensee may not assign this agreement to any third party without InterDigital’s prior written approval. +9.5 Entire Agreement. This License constitutes the entire agreement between the parties hereto with respect to the subject matter hereof and supersedes any prior agreements or understanding. +  diff --git a/models/FeatureStyleEncoder/README.md b/models/FeatureStyleEncoder/README.md new file mode 100644 index 0000000000000000000000000000000000000000..625d4cd00adc82e19cebebe0faf093a064e0018d --- /dev/null +++ b/models/FeatureStyleEncoder/README.md @@ -0,0 +1,92 @@ +## A Style-Based GAN Encoder for High Fidelity Reconstruction of Images and Videos + +Official implementation for paper: A Style-Based GAN Encoder for High Fidelity Reconstruction of Images and Videos. + +[[Video Editing Results]](https://drive.google.com/file/d/1ebih6TZxb2eLKxJdbO8GnsInDKSegfYL/view?usp=sharing) + +![teaser](images/teaser.jpg) + +> **Abstract** We propose a novel architecture for GAN inversion, which we call Feature-Style encoder. The style encoder is key for the manipulation of the obtained latent codes, while the feature encoder is crucial for optimal image reconstruction. Our model achieves accurate inversion of real images from the latent space of a pre-trained style-based GAN model, obtaining better perceptual quality and lower reconstruction error than existing methods. Thanks to its encoder structure, the model allows fast and accurate image editing. Additionally, we demonstrate that the proposed encoder is especially well-suited for inversion and editing on videos. We conduct extensive experiments for several style-based generators pre-trained on different data domains. Our proposed method yields state-of-the-art results for style-based GAN inversion, significantly outperforming competing approaches. + + +## Requirements + +### Dependencies + +- Python 3.6 +- PyTorch 1.8 +- Opencv + +You can install a new environment for this repo by running +``` +conda env create -f environment.yml +conda activate feature_style +``` + +### Prepare StyleGAN2 model and other necessary models + +* We adapt the StyleGAN2 model implemented by paper [Encoding in Style: a StyleGAN Encoder for Image-to-Image Translation](https://arxiv.org/pdf/2008.00951.pdf). Here is their [official implementation](https://github.com/eladrich/pixel2style2pixel.git). + +* Download and save the pretrained models running + ``` + sh download_models.sh + ``` + + +## Training + +* Prepare the training data + + To train the encoder for StyleGAN, we use the synthetic images generated by StyleGAN and also the real images [ffhq dataset](https://github.com/NVlabs/ffhq-dataset). + You can generate the synthetic images by running + ``` + python generate_imgs.py + ``` + and download the ffhq dataset (aligned faces) to `data/ffhq-dataset/images/`. + +* Training + + You can modify the training options of the config file in the directory `configs/`. + ``` + python train.py --config 001 + ``` + +## Testing + +* Inversion + + You can test the encoder on the images in `test/`. The output images are saved in `output/image/`. + ``` + python test.py --pretrained_model_path './pretrained_models/143_enc.pth' --input_path './test/' + ``` +* Inversion and editing in notebook + + You can explore the encoder and the attribute editing code in notebook `inference.ipynb`. You can also open it in Google Colab [here](https://colab.research.google.com/github/InterDigitalInc/FeatureStyleEncoder/blob/master/inference.ipynb). + + +## Video Manipulation + +We provide a script to achieve inversion and attribute manipulation for the videos in the test directory `data/video/`. You can upload your own video and modify the options in `run_video_inversion_editing.sh`. + +``` +sh run_video_inversion_editing.sh +``` + +## Citation +``` +@article{xuyao2022, + title={A Style-Based GAN Encoder for High Fidelity Reconstruction of Images and Videos}, + author={Yao, Xu and Newson, Alasdair and Gousseau, Yann and Hellier, Pierre}, + journal={European conference on computer vision}, + year={2022} +} +``` +## License + +Copyright © 2022, InterDigital R&D France. All rights reserved. + +This source code is made available under the license found in the LICENSE.txt in the root directory of this source tree. + + + + diff --git a/models/FeatureStyleEncoder/arcface/__init__.py b/models/FeatureStyleEncoder/arcface/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..4a0bcbf8ecdb946bb6f72e93922006fd6b48dba2 --- /dev/null +++ b/models/FeatureStyleEncoder/arcface/__init__.py @@ -0,0 +1 @@ +from .iresnet import iresnet18, iresnet34, iresnet50, iresnet100, iresnet200 diff --git a/models/FeatureStyleEncoder/arcface/iresnet.py b/models/FeatureStyleEncoder/arcface/iresnet.py new file mode 100644 index 0000000000000000000000000000000000000000..96ec43e32d5fc11871d0b85e34134c7fc6335577 --- /dev/null +++ b/models/FeatureStyleEncoder/arcface/iresnet.py @@ -0,0 +1,195 @@ +import torch +from torch import nn + +__all__ = ['iresnet18', 'iresnet34', 'iresnet50', 'iresnet100', 'iresnet200'] + + +def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1): + """3x3 convolution with padding""" + return nn.Conv2d(in_planes, + out_planes, + kernel_size=3, + stride=stride, + padding=dilation, + groups=groups, + bias=False, + dilation=dilation) + + +def conv1x1(in_planes, out_planes, stride=1): + """1x1 convolution""" + return nn.Conv2d(in_planes, + out_planes, + kernel_size=1, + stride=stride, + bias=False) + + +class IBasicBlock(nn.Module): + expansion = 1 + def __init__(self, inplanes, planes, stride=1, downsample=None, + groups=1, base_width=64, dilation=1): + super(IBasicBlock, self).__init__() + if groups != 1 or base_width != 64: + raise ValueError('BasicBlock only supports groups=1 and base_width=64') + if dilation > 1: + raise NotImplementedError("Dilation > 1 not supported in BasicBlock") + self.bn1 = nn.BatchNorm2d(inplanes, eps=1e-05,) + self.conv1 = conv3x3(inplanes, planes) + self.bn2 = nn.BatchNorm2d(planes, eps=1e-05,) + self.prelu = nn.PReLU(planes) + self.conv2 = conv3x3(planes, planes, stride) + self.bn3 = nn.BatchNorm2d(planes, eps=1e-05,) + self.downsample = downsample + self.stride = stride + + def forward(self, x): + identity = x + out = self.bn1(x) + out = self.conv1(out) + out = self.bn2(out) + out = self.prelu(out) + out = self.conv2(out) + out = self.bn3(out) + if self.downsample is not None: + identity = self.downsample(x) + out += identity + return out + + +class IResNet(nn.Module): + fc_scale = 7 * 7 + def __init__(self, + block, layers, dropout=0, num_features=512, zero_init_residual=False, + groups=1, width_per_group=64, replace_stride_with_dilation=None, fp16=False): + super(IResNet, self).__init__() + self.fp16 = fp16 + self.inplanes = 64 + self.dilation = 1 + if replace_stride_with_dilation is None: + replace_stride_with_dilation = [False, False, False] + if len(replace_stride_with_dilation) != 3: + raise ValueError("replace_stride_with_dilation should be None " + "or a 3-element tuple, got {}".format(replace_stride_with_dilation)) + self.groups = groups + self.base_width = width_per_group + self.conv1 = nn.Conv2d(3, self.inplanes, kernel_size=3, stride=1, padding=1, bias=False) + self.bn1 = nn.BatchNorm2d(self.inplanes, eps=1e-05) + self.prelu = nn.PReLU(self.inplanes) + self.layer1 = self._make_layer(block, 64, layers[0], stride=2) + self.layer2 = self._make_layer(block, + 128, + layers[1], + stride=2, + dilate=replace_stride_with_dilation[0]) + self.layer3 = self._make_layer(block, + 256, + layers[2], + stride=2, + dilate=replace_stride_with_dilation[1]) + self.layer4 = self._make_layer(block, + 512, + layers[3], + stride=2, + dilate=replace_stride_with_dilation[2]) + self.bn2 = nn.BatchNorm2d(512 * block.expansion, eps=1e-05,) + self.dropout = nn.Dropout(p=dropout, inplace=True) + self.fc = nn.Linear(512 * block.expansion * self.fc_scale, num_features) + self.features = nn.BatchNorm1d(num_features, eps=1e-05) + nn.init.constant_(self.features.weight, 1.0) + self.features.weight.requires_grad = False + + for m in self.modules(): + if isinstance(m, nn.Conv2d): + nn.init.normal_(m.weight, 0, 0.1) + elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)): + nn.init.constant_(m.weight, 1) + nn.init.constant_(m.bias, 0) + + if zero_init_residual: + for m in self.modules(): + if isinstance(m, IBasicBlock): + nn.init.constant_(m.bn2.weight, 0) + + def _make_layer(self, block, planes, blocks, stride=1, dilate=False): + downsample = None + previous_dilation = self.dilation + if dilate: + self.dilation *= stride + stride = 1 + if stride != 1 or self.inplanes != planes * block.expansion: + downsample = nn.Sequential( + conv1x1(self.inplanes, planes * block.expansion, stride), + nn.BatchNorm2d(planes * block.expansion, eps=1e-05, ), + ) + layers = [] + layers.append( + block(self.inplanes, planes, stride, downsample, self.groups, + self.base_width, previous_dilation)) + self.inplanes = planes * block.expansion + for _ in range(1, blocks): + layers.append( + block(self.inplanes, + planes, + groups=self.groups, + base_width=self.base_width, + dilation=self.dilation)) + + return nn.Sequential(*layers) + + def forward(self, x, return_features=False): + out = [] + with torch.cuda.amp.autocast(self.fp16): + x = self.conv1(x) + x = self.bn1(x) + x = self.prelu(x) + x = self.layer1(x) + out.append(x) + x = self.layer2(x) + out.append(x) + x = self.layer3(x) + out.append(x) + x = self.layer4(x) + out.append(x) + x = self.bn2(x) + x = torch.flatten(x, 1) + x = self.dropout(x) + x = self.fc(x.float() if self.fp16 else x) + x = self.features(x) + + if return_features: + out.append(x) + return out + return x + + +def _iresnet(arch, block, layers, pretrained, progress, **kwargs): + model = IResNet(block, layers, **kwargs) + if pretrained: + raise ValueError() + return model + + +def iresnet18(pretrained=False, progress=True, **kwargs): + return _iresnet('iresnet18', IBasicBlock, [2, 2, 2, 2], pretrained, + progress, **kwargs) + + +def iresnet34(pretrained=False, progress=True, **kwargs): + return _iresnet('iresnet34', IBasicBlock, [3, 4, 6, 3], pretrained, + progress, **kwargs) + + +def iresnet50(pretrained=False, progress=True, **kwargs): + return _iresnet('iresnet50', IBasicBlock, [3, 4, 14, 3], pretrained, + progress, **kwargs) + + +def iresnet100(pretrained=False, progress=True, **kwargs): + return _iresnet('iresnet100', IBasicBlock, [3, 13, 30, 3], pretrained, + progress, **kwargs) + + +def iresnet200(pretrained=False, progress=True, **kwargs): + return _iresnet('iresnet200', IBasicBlock, [6, 26, 60, 6], pretrained, + progress, **kwargs) diff --git a/models/FeatureStyleEncoder/configs/001.yaml b/models/FeatureStyleEncoder/configs/001.yaml new file mode 100644 index 0000000000000000000000000000000000000000..6292278ccf34bafcdb2efd437fe31b161c5a0916 --- /dev/null +++ b/models/FeatureStyleEncoder/configs/001.yaml @@ -0,0 +1,49 @@ +# Input data +resolution: 1024 +age_min: 20 +age_max: 70 +use_realimg: True +# Training hyperparameters +batch_size: 1 +epochs: 12 +iter_per_epoch: 10000 +device: 'cuda' +# Optimizer parameters +optimizer: 'ranger' +lr: 0.0001 +beta_1: 0.95 +beta_2: 0.999 +weight_decay: 0 +# Learning rate scheduler +step_size: 10 +gamma: 0.1 +# Tensorboard log options +image_save_iter: 100 +log_iter: 10 +# Network setting +use_fs_encoder: True +use_fs_encoder_v2: True +fs_stride: 2 +pretrained_weight_for_fs: False +enc_resolution: 256 +enc_residual: False +truncation_psi: 1 +use_noise: True +randomize_noise: False # If generator use a different random noise at each time of generating a image from z +# Loss setting +use_parsing_net: True +multi_layer_idloss: True +real_image_as_image_loss: False +feature_match_loss: False +feature_match_loss_G: False +use_random_noise: True +optimize_on_z: False +multiscale_lpips: True +# Loss weight +w: + l1: 0 + l2: 1 + lpips: 0.2 + id: 0.1 + landmark: 0.1 + f_recon: 0.01 \ No newline at end of file diff --git a/models/FeatureStyleEncoder/download_models.sh b/models/FeatureStyleEncoder/download_models.sh new file mode 100644 index 0000000000000000000000000000000000000000..a554f924dcdda5a372d14187431f9655331a7e74 --- /dev/null +++ b/models/FeatureStyleEncoder/download_models.sh @@ -0,0 +1,22 @@ +#!/bin/sh +pip install gdown +mkdir pretrained_models +cd pretrained_models + +# download pretrained encoder +gdown --fuzzy https://drive.google.com/file/d/1RnnBL77j_Can0dY1KOiXHvG224MxjvzC/view?usp=sharing + +# download ArcFace pretrained model +gdown --fuzzy https://drive.google.com/file/d/1coFTz-Kkgvoc_gRT8JFzqCgeC3lAFWQp/view?usp=sharing + +# download face parsing model from https://github.com/zllrunning/face-parsing.PyTorch +gdown --fuzzy https://drive.google.com/open?id=154JgKpzCPW82qINcVieuPH3fZ2e0P812 + +# download pSp pretrained model from https://github.com/eladrich/pixel2style2pixel.git +cd ../pixel2style2pixel +mkdir pretrained_models +cd pretrained_models +gdown --fuzzy https://drive.google.com/file/d/1bMTNWkh5LArlaWSc_wa8VKyq2V42T2z0/view?usp=sharing +cd .. +cd .. + diff --git a/models/FeatureStyleEncoder/environment.yml b/models/FeatureStyleEncoder/environment.yml new file mode 100644 index 0000000000000000000000000000000000000000..08e5b22047207f73625a58bbf9ce881c43fa7712 --- /dev/null +++ b/models/FeatureStyleEncoder/environment.yml @@ -0,0 +1,159 @@ +name: feature_style +channels: + - pytorch + - 1adrianb + - conda-forge + - defaults +dependencies: + - _libgcc_mutex=0.1 + - argon2-cffi=20.1.0 + - async_generator=1.10 + - attrs=20.3.0 + - backcall=0.2.0 + - blas=1.0 + - bleach=3.3.0 + - bzip2=1.0.8 + - ca-certificates=2021.10.8 + - certifi=2021.5.30 + - cffi=1.14.5 + - cloudpickle=1.6.0 + - cudatoolkit=10.2.89 + - cycler=0.10.0 + - cytoolz=0.11.0 + - dask-core=1.1.4 + - dataclasses=0.8 + - decorator=5.0.6 + - defusedxml=0.7.1 + - entrypoints=0.3 + - face_alignment=1.3.4 + - ffmpeg=4.3 + - freetype=2.10.4 + - gmp=6.2.1 + - gnutls=3.6.15 + - hdf5=1.10.2 + - imageio=2.9.0 + - importlib-metadata=3.10.0 + - importlib_metadata=3.10.0 + - intel-openmp=2020.2 + - ipykernel=5.3.4 + - ipython=7.16.1 + - ipython_genutils=0.2.0 + - ipywidgets=7.6.3 + - jedi=0.17.0 + - jinja2=2.11.3 + - jpeg=9b + - jsonschema=3.2.0 + - jupyter_client=6.1.12 + - jupyter_core=4.7.1 + - jupyterlab_pygments=0.1.2 + - jupyterlab_widgets=1.0.0 + - kiwisolver=1.3.1 + - lame=3.100 + - lcms2=2.12 + - ld_impl_linux-64=2.33.1 + - libffi=3.3 + - libgcc-ng=9.1.0 + - libgfortran=3.0.0 + - libgfortran-ng=7.3.0 + - libiconv=1.15 + - libidn2=2.3.0 + - libllvm10=10.0.1 + - libpng=1.6.37 + - libsodium=1.0.18 + - libstdcxx-ng=9.1.0 + - libtasn1=4.16.0 + - libtiff=4.1.0 + - libunistring=0.9.10 + - libuv=1.40.0 + - llvmlite=0.36.0 + - lz4-c=1.9.3 + - markupsafe=1.1.1 + - matplotlib-base=3.3.4 + - mistune=0.8.4 + - mkl=2020.2 + - mkl-service=2.3.0 + - mkl_fft=1.3.0 + - mkl_random=1.1.1 + - nbclient=0.5.3 + - nbconvert=6.0.7 + - nbformat=5.1.3 + - ncurses=6.2 + - nest-asyncio=1.5.1 + - nettle=3.7.2 + - networkx=2.2 + - notebook=6.3.0 + - numba=0.53.1 + - numpy=1.19.2 + - numpy-base=1.19.2 + - olefile=0.46 + - opencv=3.4.1 + - openh264=2.1.0 + - openssl=1.1.1l + - packaging=20.9 + - pandoc=2.12 + - pandocfilters=1.4.3 + - parso=0.8.2 + - pexpect=4.8.0 + - pickleshare=0.7.5 + - pillow=8.2.0 + - pip=21.0.1 + - prometheus_client=0.10.1 + - prompt-toolkit=3.0.17 + - ptyprocess=0.7.0 + - pycparser=2.20 + - pygments=2.8.1 + - pyparsing=2.4.7 + - pyrsistent=0.17.3 + - python=3.6.13 + - python-dateutil=2.8.1 + - python_abi=3.6 + - pytorch=1.8.1 + - pywavelets=1.1.1 + - pyzmq=20.0.0 + - readline=8.1 + - scikit-image=0.17.2 + - scipy=1.5.2 + - send2trash=1.5.0 + - setuptools=52.0.0 + - six=1.15.0 + - sqlite=3.35.4 + - tbb=2020.3 + - terminado=0.9.4 + - testpath=0.4.4 + - tifffile=2020.10.1 + - tk=8.6.10 + - toolz=0.11.1 + - torchvision=0.9.1 + - tornado=6.1 + - tqdm=4.59.0 + - traitlets=4.3.3 + - typing_extensions=3.7.4.3 + - wcwidth=0.2.5 + - webencodings=0.5.1 + - wheel=0.36.2 + - widgetsnbextension=3.5.1 + - xz=5.2.5 + - yaml=0.2.5 + - zeromq=4.3.4 + - zipp=3.4.1 + - zlib=1.2.11 + - zstd=1.4.9 + - pip: + - charset-normalizer==2.0.7 + - click==8.0.3 + - cython==0.29.23 + - dlib==19.22.1 + - idna==3.2 + - joblib==1.1.0 + - ninja==1.10.2.2 + - onnxruntime==1.7.0 + - opencv-python==4.5.1.48 + - protobuf==3.15.8 + - pytorch-msssim==0.2.1 + - pyyaml==5.4.1 + - requests==2.26.0 + - scikit-learn==0.22 + - sklearn==0.0 + - tensorboard-logger==0.1.0 + - threadpoolctl==3.0.0 + - urllib3==1.26.7 diff --git a/models/FeatureStyleEncoder/face_parsing/model.py b/models/FeatureStyleEncoder/face_parsing/model.py new file mode 100644 index 0000000000000000000000000000000000000000..f649917a85676ec3ef36fa1e262ddcfaaca12e20 --- /dev/null +++ b/models/FeatureStyleEncoder/face_parsing/model.py @@ -0,0 +1,292 @@ +#!/usr/bin/python +# -*- encoding: utf-8 -*- + + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torchvision + +import sys +sys.path.append('..') +from face_parsing.resnet import Resnet18 +# from modules.bn import InPlaceABNSync as BatchNorm2d + + +class ConvBNReLU(nn.Module): + def __init__(self, in_chan, out_chan, ks=3, stride=1, padding=1, *args, **kwargs): + super(ConvBNReLU, self).__init__() + self.conv = nn.Conv2d(in_chan, + out_chan, + kernel_size = ks, + stride = stride, + padding = padding, + bias = False) + self.bn = nn.BatchNorm2d(out_chan) + self.init_weight() + + def forward(self, x): + x = self.conv(x) + x = F.relu(self.bn(x)) + return x + + def init_weight(self): + for ly in self.children(): + if isinstance(ly, nn.Conv2d): + nn.init.kaiming_normal_(ly.weight, a=1) + if not ly.bias is None: nn.init.constant_(ly.bias, 0) + +class BiSeNetOutput(nn.Module): + def __init__(self, in_chan, mid_chan, n_classes, *args, **kwargs): + super(BiSeNetOutput, self).__init__() + self.conv = ConvBNReLU(in_chan, mid_chan, ks=3, stride=1, padding=1) + self.conv_out = nn.Conv2d(mid_chan, n_classes, kernel_size=1, bias=False) + self.init_weight() + + def forward(self, x): + x = self.conv(x) + x = self.conv_out(x) + return x + + def init_weight(self): + for ly in self.children(): + if isinstance(ly, nn.Conv2d): + nn.init.kaiming_normal_(ly.weight, a=1) + if not ly.bias is None: nn.init.constant_(ly.bias, 0) + + def get_params(self): + wd_params, nowd_params = [], [] + for name, module in self.named_modules(): + if isinstance(module, nn.Linear) or isinstance(module, nn.Conv2d): + wd_params.append(module.weight) + if not module.bias is None: + nowd_params.append(module.bias) + elif isinstance(module, nn.BatchNorm2d): + nowd_params += list(module.parameters()) + return wd_params, nowd_params + + +class AttentionRefinementModule(nn.Module): + def __init__(self, in_chan, out_chan, *args, **kwargs): + super(AttentionRefinementModule, self).__init__() + self.conv = ConvBNReLU(in_chan, out_chan, ks=3, stride=1, padding=1) + self.conv_atten = nn.Conv2d(out_chan, out_chan, kernel_size= 1, bias=False) + self.bn_atten = nn.BatchNorm2d(out_chan) + self.sigmoid_atten = nn.Sigmoid() + self.init_weight() + + def forward(self, x): + feat = self.conv(x) + atten = F.avg_pool2d(feat, feat.size()[2:]) + atten = self.conv_atten(atten) + atten = self.bn_atten(atten) + atten = self.sigmoid_atten(atten) + out = torch.mul(feat, atten) + return out + + def init_weight(self): + for ly in self.children(): + if isinstance(ly, nn.Conv2d): + nn.init.kaiming_normal_(ly.weight, a=1) + if not ly.bias is None: nn.init.constant_(ly.bias, 0) + + +class ContextPath(nn.Module): + def __init__(self, *args, **kwargs): + super(ContextPath, self).__init__() + self.resnet = Resnet18() + self.arm16 = AttentionRefinementModule(256, 128) + self.arm32 = AttentionRefinementModule(512, 128) + self.conv_head32 = ConvBNReLU(128, 128, ks=3, stride=1, padding=1) + self.conv_head16 = ConvBNReLU(128, 128, ks=3, stride=1, padding=1) + self.conv_avg = ConvBNReLU(512, 128, ks=1, stride=1, padding=0) + + self.init_weight() + + def forward(self, x): + H0, W0 = x.size()[2:] + feat8, feat16, feat32 = self.resnet(x) + H8, W8 = feat8.size()[2:] + H16, W16 = feat16.size()[2:] + H32, W32 = feat32.size()[2:] + + avg = F.avg_pool2d(feat32, feat32.size()[2:]) + avg = self.conv_avg(avg) + avg_up = F.interpolate(avg, (H32, W32), mode='nearest') + + feat32_arm = self.arm32(feat32) + feat32_sum = feat32_arm + avg_up + feat32_up = F.interpolate(feat32_sum, (H16, W16), mode='nearest') + feat32_up = self.conv_head32(feat32_up) + + feat16_arm = self.arm16(feat16) + feat16_sum = feat16_arm + feat32_up + feat16_up = F.interpolate(feat16_sum, (H8, W8), mode='nearest') + feat16_up = self.conv_head16(feat16_up) + + return feat8, feat16_up, feat32_up # x8, x8, x16 + + def init_weight(self): + for ly in self.children(): + if isinstance(ly, nn.Conv2d): + nn.init.kaiming_normal_(ly.weight, a=1) + if not ly.bias is None: nn.init.constant_(ly.bias, 0) + + def get_params(self): + wd_params, nowd_params = [], [] + for name, module in self.named_modules(): + if isinstance(module, (nn.Linear, nn.Conv2d)): + wd_params.append(module.weight) + if not module.bias is None: + nowd_params.append(module.bias) + elif isinstance(module, nn.BatchNorm2d): + nowd_params += list(module.parameters()) + return wd_params, nowd_params + + +### This is not used, since I replace this with the resnet feature with the same size +class SpatialPath(nn.Module): + def __init__(self, *args, **kwargs): + super(SpatialPath, self).__init__() + self.conv1 = ConvBNReLU(3, 64, ks=7, stride=2, padding=3) + self.conv2 = ConvBNReLU(64, 64, ks=3, stride=2, padding=1) + self.conv3 = ConvBNReLU(64, 64, ks=3, stride=2, padding=1) + self.conv_out = ConvBNReLU(64, 128, ks=1, stride=1, padding=0) + self.init_weight() + + def forward(self, x): + feat = self.conv1(x) + feat = self.conv2(feat) + feat = self.conv3(feat) + feat = self.conv_out(feat) + return feat + + def init_weight(self): + for ly in self.children(): + if isinstance(ly, nn.Conv2d): + nn.init.kaiming_normal_(ly.weight, a=1) + if not ly.bias is None: nn.init.constant_(ly.bias, 0) + + def get_params(self): + wd_params, nowd_params = [], [] + for name, module in self.named_modules(): + if isinstance(module, nn.Linear) or isinstance(module, nn.Conv2d): + wd_params.append(module.weight) + if not module.bias is None: + nowd_params.append(module.bias) + elif isinstance(module, nn.BatchNorm2d): + nowd_params += list(module.parameters()) + return wd_params, nowd_params + + +class FeatureFusionModule(nn.Module): + def __init__(self, in_chan, out_chan, *args, **kwargs): + super(FeatureFusionModule, self).__init__() + self.convblk = ConvBNReLU(in_chan, out_chan, ks=1, stride=1, padding=0) + self.conv1 = nn.Conv2d(out_chan, + out_chan//4, + kernel_size = 1, + stride = 1, + padding = 0, + bias = False) + self.conv2 = nn.Conv2d(out_chan//4, + out_chan, + kernel_size = 1, + stride = 1, + padding = 0, + bias = False) + self.relu = nn.ReLU(inplace=True) + self.sigmoid = nn.Sigmoid() + self.init_weight() + + def forward(self, fsp, fcp): + fcat = torch.cat([fsp, fcp], dim=1) + feat = self.convblk(fcat) + atten = F.avg_pool2d(feat, feat.size()[2:]) + atten = self.conv1(atten) + atten = self.relu(atten) + atten = self.conv2(atten) + atten = self.sigmoid(atten) + feat_atten = torch.mul(feat, atten) + feat_out = feat_atten + feat + return feat_out + + def init_weight(self): + for ly in self.children(): + if isinstance(ly, nn.Conv2d): + nn.init.kaiming_normal_(ly.weight, a=1) + if not ly.bias is None: nn.init.constant_(ly.bias, 0) + + def get_params(self): + wd_params, nowd_params = [], [] + for name, module in self.named_modules(): + if isinstance(module, nn.Linear) or isinstance(module, nn.Conv2d): + wd_params.append(module.weight) + if not module.bias is None: + nowd_params.append(module.bias) + elif isinstance(module, nn.BatchNorm2d): + nowd_params += list(module.parameters()) + return wd_params, nowd_params + + +class BiSeNet(nn.Module): + def __init__(self, n_classes, *args, **kwargs): + super(BiSeNet, self).__init__() + self.cp = ContextPath() + ## here self.sp is deleted + self.ffm = FeatureFusionModule(256, 256) + self.conv_out = BiSeNetOutput(256, 256, n_classes) + self.conv_out16 = BiSeNetOutput(128, 64, n_classes) + self.conv_out32 = BiSeNetOutput(128, 64, n_classes) + self.init_weight() + + def forward(self, x): + H, W = x.size()[2:] + feat_res8, feat_cp8, feat_cp16 = self.cp(x) # here return res3b1 feature + feat_sp = feat_res8 # use res3b1 feature to replace spatial path feature + feat_fuse = self.ffm(feat_sp, feat_cp8) + + feat_out = self.conv_out(feat_fuse) + feat_out16 = self.conv_out16(feat_cp8) + feat_out32 = self.conv_out32(feat_cp16) + + feat_out = F.interpolate(feat_out, (H, W), mode='bilinear', align_corners=True) + feat_out16 = F.interpolate(feat_out16, (H, W), mode='bilinear', align_corners=True) + feat_out32 = F.interpolate(feat_out32, (H, W), mode='bilinear', align_corners=True) + return feat_out, feat_out16, feat_out32 + + def extract_fuse_layer(self, x): + H, W = x.size()[2:] + feat_res8, feat_cp8, feat_cp16 = self.cp(x) # here return res3b1 feature + feat_sp = feat_res8 # use res3b1 feature to replace spatial path feature + feat_fuse = self.ffm(feat_sp, feat_cp8) + return [feat_fuse] + + def init_weight(self): + for ly in self.children(): + if isinstance(ly, nn.Conv2d): + nn.init.kaiming_normal_(ly.weight, a=1) + if not ly.bias is None: nn.init.constant_(ly.bias, 0) + + def get_params(self): + wd_params, nowd_params, lr_mul_wd_params, lr_mul_nowd_params = [], [], [], [] + for name, child in self.named_children(): + child_wd_params, child_nowd_params = child.get_params() + if isinstance(child, FeatureFusionModule) or isinstance(child, BiSeNetOutput): + lr_mul_wd_params += child_wd_params + lr_mul_nowd_params += child_nowd_params + else: + wd_params += child_wd_params + nowd_params += child_nowd_params + return wd_params, nowd_params, lr_mul_wd_params, lr_mul_nowd_params + + +if __name__ == "__main__": + net = BiSeNet(19) + net.cuda() + net.eval() + in_ten = torch.randn(16, 3, 640, 480).cuda() + out, out16, out32 = net(in_ten) + print(out.shape) + + net.get_params() diff --git a/models/FeatureStyleEncoder/face_parsing/resnet.py b/models/FeatureStyleEncoder/face_parsing/resnet.py new file mode 100644 index 0000000000000000000000000000000000000000..baa4c308aa46c5d592844cb626d9544a769455b0 --- /dev/null +++ b/models/FeatureStyleEncoder/face_parsing/resnet.py @@ -0,0 +1,112 @@ +#!/usr/bin/python +# -*- encoding: utf-8 -*- + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.utils.model_zoo as modelzoo + +import os + + +# from modules.bn import InPlaceABNSync as BatchNorm2d + +resnet18_url = 'https://download.pytorch.org/models/resnet18-5c106cde.pth' + + +def conv3x3(in_planes, out_planes, stride=1): + """3x3 convolution with padding""" + return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, + padding=1, bias=False) + + +class BasicBlock(nn.Module): + def __init__(self, in_chan, out_chan, stride=1): + super(BasicBlock, self).__init__() + self.conv1 = conv3x3(in_chan, out_chan, stride) + self.bn1 = nn.BatchNorm2d(out_chan) + self.conv2 = conv3x3(out_chan, out_chan) + self.bn2 = nn.BatchNorm2d(out_chan) + self.relu = nn.ReLU(inplace=True) + self.downsample = None + if in_chan != out_chan or stride != 1: + self.downsample = nn.Sequential( + nn.Conv2d(in_chan, out_chan, + kernel_size=1, stride=stride, bias=False), + nn.BatchNorm2d(out_chan), + ) + + def forward(self, x): + residual = self.conv1(x) + residual = F.relu(self.bn1(residual)) + residual = self.conv2(residual) + residual = self.bn2(residual) + + shortcut = x + if self.downsample is not None: + shortcut = self.downsample(x) + + out = shortcut + residual + out = self.relu(out) + return out + + +def create_layer_basic(in_chan, out_chan, bnum, stride=1): + layers = [BasicBlock(in_chan, out_chan, stride=stride)] + for i in range(bnum-1): + layers.append(BasicBlock(out_chan, out_chan, stride=1)) + return nn.Sequential(*layers) + + +class Resnet18(nn.Module): + def __init__(self): + super(Resnet18, self).__init__() + self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, + bias=False) + self.bn1 = nn.BatchNorm2d(64) + self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) + self.layer1 = create_layer_basic(64, 64, bnum=2, stride=1) + self.layer2 = create_layer_basic(64, 128, bnum=2, stride=2) + self.layer3 = create_layer_basic(128, 256, bnum=2, stride=2) + self.layer4 = create_layer_basic(256, 512, bnum=2, stride=2) + self.init_weight() + + def forward(self, x): + x = self.conv1(x) + x = F.relu(self.bn1(x)) + x = self.maxpool(x) + + x = self.layer1(x) + feat8 = self.layer2(x) # 1/8 + feat16 = self.layer3(feat8) # 1/16 + feat32 = self.layer4(feat16) # 1/32 + return feat8, feat16, feat32 + + def init_weight(self): + state_dict = modelzoo.load_url(resnet18_url) + self_state_dict = self.state_dict() + for k, v in state_dict.items(): + if 'fc' in k: continue + self_state_dict.update({k: v}) + self.load_state_dict(self_state_dict) + + def get_params(self): + wd_params, nowd_params = [], [] + for name, module in self.named_modules(): + if isinstance(module, (nn.Linear, nn.Conv2d)): + wd_params.append(module.weight) + if not module.bias is None: + nowd_params.append(module.bias) + elif isinstance(module, nn.BatchNorm2d): + nowd_params += list(module.parameters()) + return wd_params, nowd_params + + +if __name__ == "__main__": + net = Resnet18() + x = torch.randn(16, 3, 224, 224) + out = net(x) + print(out[0].size()) + print(out[1].size()) + print(out[2].size()) + net.get_params() diff --git a/models/FeatureStyleEncoder/face_parsing/test.py b/models/FeatureStyleEncoder/face_parsing/test.py new file mode 100644 index 0000000000000000000000000000000000000000..ad94fe9b62f724b2d3957984840699b317657c44 --- /dev/null +++ b/models/FeatureStyleEncoder/face_parsing/test.py @@ -0,0 +1,87 @@ +#!/usr/bin/python +# -*- encoding: utf-8 -*- + +from model import BiSeNet + +import torch + +import os +import os.path as osp +import numpy as np +from PIL import Image +import torchvision.transforms as transforms +import cv2 + +def vis_parsing_maps(im, parsing_anno, stride, save_im=False, save_path='vis_results/parsing_map_on_im.jpg'): + # Colors for all 20 parts + part_colors = [[255, 0, 0], [255, 85, 0], [255, 170, 0], + [255, 0, 85], [255, 0, 170], + [0, 255, 0], [85, 255, 0], [170, 255, 0], + [0, 255, 85], [0, 255, 170], + [0, 0, 255], [85, 0, 255], [170, 0, 255], + [0, 85, 255], [0, 170, 255], + [255, 255, 0], [255, 255, 85], [255, 255, 170], + [255, 0, 255], [255, 85, 255], [255, 170, 255], + [0, 255, 255], [85, 255, 255], [170, 255, 255]] + + im = np.array(im) + vis_im = im.copy().astype(np.uint8) + vis_parsing_anno = parsing_anno.copy().astype(np.uint8) + vis_parsing_anno = cv2.resize(vis_parsing_anno, None, fx=stride, fy=stride, interpolation=cv2.INTER_NEAREST) + vis_parsing_anno_color = np.zeros((vis_parsing_anno.shape[0], vis_parsing_anno.shape[1], 3)) + 255 + + num_of_class = np.max(vis_parsing_anno) + + for pi in range(1, num_of_class + 1): + index = np.where(vis_parsing_anno == pi) + vis_parsing_anno_color[index[0], index[1], :] = part_colors[pi] + + vis_parsing_anno_color = vis_parsing_anno_color.astype(np.uint8) + # print(vis_parsing_anno_color.shape, vis_im.shape) + vis_im = cv2.addWeighted(cv2.cvtColor(vis_im, cv2.COLOR_RGB2BGR), 0.4, vis_parsing_anno_color, 0.6, 0) + + # Save result or not + if save_im: + cv2.imwrite(save_path[:-4] +'.png', vis_parsing_anno) + cv2.imwrite(save_path, vis_im, [int(cv2.IMWRITE_JPEG_QUALITY), 100]) + + # return vis_im + +def evaluate(respth='./res/test_res', dspth='./data', cp='model_final_diss.pth'): + + if not os.path.exists(respth): + os.makedirs(respth) + + n_classes = 19 + net = BiSeNet(n_classes=n_classes) + net.cuda() + save_pth = osp.join('', cp) + net.load_state_dict(torch.load(save_pth)) + net.eval() + + to_tensor = transforms.Compose([ + transforms.ToTensor(), + transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)), + ]) + with torch.no_grad(): + for image_path in os.listdir(dspth): + img = Image.open(osp.join(dspth, image_path)) + image = img.resize((512, 512), Image.BILINEAR) + img = to_tensor(image) + img = torch.unsqueeze(img, 0) + img = img.cuda() + out = net(img) + out = out[0] + parsing = out.squeeze(0).cpu().numpy().argmax(0) + # print(parsing) + print(np.unique(parsing)) + + vis_parsing_maps(image, parsing, stride=1, save_im=True, save_path=osp.join(respth, image_path)) + + + + +if __name__ == "__main__": + evaluate(dspth='../face-parsing/imgs/', cp='./pretrained_models/79999_iter.pth') + + diff --git a/models/FeatureStyleEncoder/generate_imgs.py b/models/FeatureStyleEncoder/generate_imgs.py new file mode 100644 index 0000000000000000000000000000000000000000..f1e0975401a3ecaa4c536f1fbcdcb8bde031700e --- /dev/null +++ b/models/FeatureStyleEncoder/generate_imgs.py @@ -0,0 +1,52 @@ +import argparse +import os +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.utils.data as data +import yaml + +from PIL import Image +from torchvision import transforms, utils +from tensorboard_logger import Logger +from tqdm import tqdm +from utils.functions import * + +import sys +sys.path.append('pixel2style2pixel/') +from pixel2style2pixel.models.stylegan2.model import Generator, get_keys + +torch.backends.cudnn.enabled = True +torch.backends.cudnn.deterministic = True +torch.backends.cudnn.benchmark = True +torch.autograd.set_detect_anomaly(True) +Image.MAX_IMAGE_PIXELS = None +device = torch.device('cuda') + +parser = argparse.ArgumentParser() +parser.add_argument('--config', type=str, default='002', help='Path to the config file.') +parser.add_argument('--dataset_path', type=str, default='./data/stylegan2-generate-images/', help='dataset path') +parser.add_argument('--stylegan_model_path', type=str, default='./pixel2style2pixel/pretrained_models/psp_ffhq_encode.pt', help='pretrained stylegan model') +opts = parser.parse_args() + + +StyleGAN = Generator(1024, 512, 8) +state_dict = torch.load(opts.stylegan_model_path, map_location='cpu') +StyleGAN.load_state_dict(get_keys(state_dict, 'decoder'), strict=True) +StyleGAN.to(device) + +#seeds = np.array([torch.random.seed() for i in range(100000)]) +seeds = np.load(opts.dataset_path + 'seeds_pytorch_1.8.1.npy') + +with torch.no_grad(): + os.makedirs(opts.dataset_path + 'ims/', exist_ok=True) + + for i, seed in enumerate(tqdm(seeds)): + + torch.manual_seed(seed) + z = torch.randn(1, 512).to(device) + n = StyleGAN.make_noise() + w = StyleGAN.get_latent(z) + x, _ = StyleGAN([w], input_is_latent=True, noise=n) + utils.save_image(clip_img(x), opts.dataset_path + 'ims/%06d.jpg'%i) diff --git a/models/FeatureStyleEncoder/lpips/__init__.py b/models/FeatureStyleEncoder/lpips/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..98a780ff8576cfe4d0486388bd0028a980c1aa81 --- /dev/null +++ b/models/FeatureStyleEncoder/lpips/__init__.py @@ -0,0 +1,178 @@ + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import numpy as np +import torch +# from torch.autograd import Variable + +from lpips.trainer import * +from lpips.lpips import * + +# class PerceptualLoss(torch.nn.Module): +# def __init__(self, model='lpips', net='alex', spatial=False, use_gpu=False, gpu_ids=[0], version='0.1'): # 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, spatial=self.spatial, gpu_ids=gpu_ids, version=version) +# 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+1e-8,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.): + from skimage.measure import compare_ssim + return (1 - compare_ssim(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 load_image(path): + if(path[-3:] == 'dng'): + import rawpy + with rawpy.imread(path) as raw: + img = raw.postprocess() + elif(path[-3:]=='bmp' or path[-3:]=='jpg' or path[-3:]=='png' or path[-4:]=='jpeg'): + import cv2 + return cv2.imread(path)[:,:,::-1] + else: + img = (255*plt.imread(path)[:,:,:3]).astype('uint8') + + return img + +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 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))) + + + +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 + diff --git a/models/FeatureStyleEncoder/lpips/lpips.py b/models/FeatureStyleEncoder/lpips/lpips.py new file mode 100644 index 0000000000000000000000000000000000000000..9b979c082fdc45092e4669c52179339e16fa9d1f --- /dev/null +++ b/models/FeatureStyleEncoder/lpips/lpips.py @@ -0,0 +1,219 @@ + +from __future__ import absolute_import + +import torch +import torch.nn as nn +import torch.nn.init as init +from torch.autograd import Variable +import numpy as np +from . import pretrained_networks as pn +import torch.nn + +import lpips + +def spatial_average(in_tens, keepdim=True): + return in_tens.mean([2,3],keepdim=keepdim) + +def upsample(in_tens, out_HW=(64,64)): # assumes scale factor is same for H and W + in_H, in_W = in_tens.shape[2], in_tens.shape[3] + return nn.Upsample(size=out_HW, mode='bilinear', align_corners=False)(in_tens) + +# Learned perceptual metric +class LPIPS(nn.Module): + def __init__(self, pretrained=True, net='alex', version='0.1', lpips=True, spatial=False, + pnet_rand=False, pnet_tune=False, use_dropout=True, model_path=None, eval_mode=True, verbose=True): + # lpips - [True] means with linear calibration on top of base network + # pretrained - [True] means load linear weights + + super(LPIPS, self).__init__() + if(verbose): + print('Setting up [%s] perceptual loss: trunk [%s], v[%s], spatial [%s]'% + ('LPIPS' if lpips else 'baseline', net, version, 'on' if spatial else 'off')) + + self.pnet_type = net + self.pnet_tune = pnet_tune + self.pnet_rand = pnet_rand + self.spatial = spatial + self.lpips = lpips # false means baseline of just averaging all layers + 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] + self.lins = nn.ModuleList(self.lins) + + if(pretrained): + if(model_path is None): + import inspect + import os + model_path = os.path.abspath(os.path.join(inspect.getfile(self.__init__), '..', 'weights/v%s/%s.pth'%(version,net))) + + if(verbose): + print('Loading model from: %s'%model_path) + self.load_state_dict(torch.load(model_path, map_location='cpu'), strict=False) + + if(eval_mode): + self.eval() + + def forward(self, in0, in1, retPerLayer=False, normalize=False): + if normalize: # turn on this flag if input is [0,1] so it can be adjusted to [-1, +1] + in0 = 2 * in0 - 1 + in1 = 2 * in1 - 1 + + # 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] = lpips.normalize_tensor(outs0[kk]), lpips.normalize_tensor(outs1[kk]) + diffs[kk] = (feats0[kk]-feats1[kk])**2 + + if(self.lpips): + if(self.spatial): + res = [upsample(self.lins[kk](diffs[kk]), out_HW=in0.shape[2:]) for kk in range(self.L)] + else: + res = [spatial_average(self.lins[kk](diffs[kk]), keepdim=True) for kk in range(self.L)] + else: + if(self.spatial): + res = [upsample(diffs[kk].sum(dim=1,keepdim=True), out_HW=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] + + # a = spatial_average(self.lins[kk](diffs[kk]), keepdim=True) + # b = torch.max(self.lins[kk](feats0[kk]**2)) + # for kk in range(self.L): + # a += spatial_average(self.lins[kk](diffs[kk]), keepdim=True) + # b = torch.max(b,torch.max(self.lins[kk](feats0[kk]**2))) + # a = a/self.L + # from IPython import embed + # embed() + # return 10*torch.log10(b/a) + + 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) + + def forward(self, x): + return self.model(x) + +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 = lpips.l2(lpips.tensor2np(lpips.tensor2tensorlab(in0.data,to_norm=False)), + lpips.tensor2np(lpips.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 = lpips.dssim(1.*lpips.tensor2im(in0.data), 1.*lpips.tensor2im(in1.data), range=255.).astype('float') + elif(self.colorspace=='Lab'): + value = lpips.dssim(lpips.tensor2np(lpips.tensor2tensorlab(in0.data,to_norm=False)), + lpips.tensor2np(lpips.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) diff --git a/models/FeatureStyleEncoder/lpips/pretrained_networks.py b/models/FeatureStyleEncoder/lpips/pretrained_networks.py new file mode 100644 index 0000000000000000000000000000000000000000..054dba7f4f7fa94bcd41df0aaaba52b4dd852f93 --- /dev/null +++ b/models/FeatureStyleEncoder/lpips/pretrained_networks.py @@ -0,0 +1,182 @@ +from collections import namedtuple +import torch +from torchvision import models as tv + +import os + +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 diff --git a/models/FeatureStyleEncoder/lpips/trainer.py b/models/FeatureStyleEncoder/lpips/trainer.py new file mode 100644 index 0000000000000000000000000000000000000000..52b6112cdc79db7a429ec52e60fcefdb756f776b --- /dev/null +++ b/models/FeatureStyleEncoder/lpips/trainer.py @@ -0,0 +1,280 @@ + +from __future__ import absolute_import + +import numpy as np +import torch +from torch import nn +from collections import OrderedDict +from torch.autograd import Variable +from scipy.ndimage import zoom +from tqdm import tqdm +import lpips +import os + + +class Trainer(): + def name(self): + return self.model_name + + def initialize(self, model='lpips', 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 - ['lpips'] for linearly calibrated network + ['baseline'] 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 + 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 + ''' + self.use_gpu = use_gpu + self.gpu_ids = gpu_ids + self.model = model + self.net = net + self.is_train = is_train + self.spatial = spatial + self.model_name = '%s [%s]'%(model,net) + + if(self.model == 'lpips'): # pretrained net + linear layer + self.net = lpips.LPIPS(pretrained=not is_train, net=net, version=version, lpips=True, spatial=spatial, + pnet_rand=pnet_rand, pnet_tune=pnet_tune, + use_dropout=True, model_path=model_path, eval_mode=False) + elif(self.model=='baseline'): # pretrained network + self.net = lpips.LPIPS(pnet_rand=pnet_rand, net=net, lpips=False) + elif(self.model in ['L2','l2']): + self.net = lpips.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 = lpips.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 = lpips.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 + 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 %f' % (type,self.old_lr, lr)) + self.old_lr = lr + + + 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') + + +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 1: + kernel = kernel * (upsample_factor ** 2) + + self.register_buffer('kernel', kernel) + + self.pad = pad + + def forward(self, input): + out = upfirdn2d(input, self.kernel, pad=self.pad) + + return out + + +class EqualConv2d(nn.Module): + def __init__( + self, in_channel, out_channel, kernel_size, stride=1, padding=0, bias=True + ): + super().__init__() + + self.weight = nn.Parameter( + torch.randn(out_channel, in_channel, kernel_size, kernel_size) + ) + self.scale = 1 / math.sqrt(in_channel * kernel_size ** 2) + + self.stride = stride + self.padding = padding + + if bias: + self.bias = nn.Parameter(torch.zeros(out_channel)) + + else: + self.bias = None + + def forward(self, input): + out = F.conv2d( + input, + self.weight * self.scale, + bias=self.bias, + stride=self.stride, + padding=self.padding, + ) + + return out + + def __repr__(self): + return ( + f'{self.__class__.__name__}({self.weight.shape[1]}, {self.weight.shape[0]},' + f' {self.weight.shape[2]}, stride={self.stride}, padding={self.padding})' + ) + + +class EqualLinear(nn.Module): + def __init__( + self, in_dim, out_dim, bias=True, bias_init=0, lr_mul=1, activation=None + ): + super().__init__() + + self.weight = nn.Parameter(torch.randn(out_dim, in_dim).div_(lr_mul)) + + if bias: + self.bias = nn.Parameter(torch.zeros(out_dim).fill_(bias_init)) + + else: + self.bias = None + + self.activation = activation + + self.scale = (1 / math.sqrt(in_dim)) * lr_mul + self.lr_mul = lr_mul + + def forward(self, input): + if self.activation: + out = F.linear(input, self.weight * self.scale) + out = fused_leaky_relu(out, self.bias * self.lr_mul) + + else: + out = F.linear( + input, self.weight * self.scale, bias=self.bias * self.lr_mul + ) + + return out + + def __repr__(self): + return ( + f'{self.__class__.__name__}({self.weight.shape[1]}, {self.weight.shape[0]})' + ) + + +class ScaledLeakyReLU(nn.Module): + def __init__(self, negative_slope=0.2): + super().__init__() + + self.negative_slope = negative_slope + + def forward(self, input): + out = F.leaky_relu(input, negative_slope=self.negative_slope) + + return out * math.sqrt(2) + + +class ModulatedConv2d(nn.Module): + def __init__( + self, + in_channel, + out_channel, + kernel_size, + style_dim, + demodulate=True, + upsample=False, + downsample=False, + blur_kernel=[1, 3, 3, 1], + ): + super().__init__() + + self.eps = 1e-8 + self.kernel_size = kernel_size + self.in_channel = in_channel + self.out_channel = out_channel + self.upsample = upsample + self.downsample = downsample + + if upsample: + factor = 2 + p = (len(blur_kernel) - factor) - (kernel_size - 1) + pad0 = (p + 1) // 2 + factor - 1 + pad1 = p // 2 + 1 + + self.blur = Blur(blur_kernel, pad=(pad0, pad1), upsample_factor=factor) + + if downsample: + factor = 2 + p = (len(blur_kernel) - factor) + (kernel_size - 1) + pad0 = (p + 1) // 2 + pad1 = p // 2 + + self.blur = Blur(blur_kernel, pad=(pad0, pad1)) + + fan_in = in_channel * kernel_size ** 2 + self.scale = 1 / math.sqrt(fan_in) + self.padding = kernel_size // 2 + + self.weight = nn.Parameter( + torch.randn(1, out_channel, in_channel, kernel_size, kernel_size) + ) + + self.modulation = EqualLinear(style_dim, in_channel, bias_init=1) + + self.demodulate = demodulate + + def __repr__(self): + return ( + f'{self.__class__.__name__}({self.in_channel}, {self.out_channel}, {self.kernel_size}, ' + f'upsample={self.upsample}, downsample={self.downsample})' + ) + + def forward(self, input, style): + batch, in_channel, height, width = input.shape + + style = self.modulation(style).view(batch, 1, in_channel, 1, 1) + weight = self.scale * self.weight * style + + if self.demodulate: + demod = torch.rsqrt(weight.pow(2).sum([2, 3, 4]) + 1e-8) + weight = weight * demod.view(batch, self.out_channel, 1, 1, 1) + + weight = weight.view( + batch * self.out_channel, in_channel, self.kernel_size, self.kernel_size + ) + + if self.upsample: + input = input.view(1, batch * in_channel, height, width) + weight = weight.view( + batch, self.out_channel, in_channel, self.kernel_size, self.kernel_size + ) + weight = weight.transpose(1, 2).reshape( + batch * in_channel, self.out_channel, self.kernel_size, self.kernel_size + ) + out = F.conv_transpose2d(input, weight, padding=0, stride=2, groups=batch) + _, _, height, width = out.shape + out = out.view(batch, self.out_channel, height, width) + out = self.blur(out) + + elif self.downsample: + input = self.blur(input) + _, _, height, width = input.shape + input = input.view(1, batch * in_channel, height, width) + out = F.conv2d(input, weight, padding=0, stride=2, groups=batch) + _, _, height, width = out.shape + out = out.view(batch, self.out_channel, height, width) + + else: + input = input.view(1, batch * in_channel, height, width) + out = F.conv2d(input, weight, padding=self.padding, groups=batch) + _, _, height, width = out.shape + out = out.view(batch, self.out_channel, height, width) + + return out + + +class NoiseInjection(nn.Module): + def __init__(self): + super().__init__() + + self.weight = nn.Parameter(torch.zeros(1)) + + def forward(self, image, noise=None): + if noise is None: + batch, _, height, width = image.shape + noise = image.new_empty(batch, 1, height, width).normal_() + + return image + self.weight * noise + + +class ConstantInput(nn.Module): + def __init__(self, channel, size=4): + super().__init__() + + self.input = nn.Parameter(torch.randn(1, channel, size, size)) + + def forward(self, input): + batch = input.shape[0] + out = self.input.repeat(batch, 1, 1, 1) + + return out + + +class StyledConv(nn.Module): + def __init__( + self, + in_channel, + out_channel, + kernel_size, + style_dim, + upsample=False, + blur_kernel=[1, 3, 3, 1], + demodulate=True, + ): + super().__init__() + + self.conv = ModulatedConv2d( + in_channel, + out_channel, + kernel_size, + style_dim, + upsample=upsample, + blur_kernel=blur_kernel, + demodulate=demodulate, + ) + + self.noise = NoiseInjection() + # self.bias = nn.Parameter(torch.zeros(1, out_channel, 1, 1)) + # self.activate = ScaledLeakyReLU(0.2) + self.activate = FusedLeakyReLU(out_channel) + + def forward(self, input, style, noise=None): + out = self.conv(input, style) + out = self.noise(out, noise=noise) + # out = out + self.bias + out = self.activate(out) + + return out + + +class ToRGB(nn.Module): + def __init__(self, in_channel, style_dim, upsample=True, blur_kernel=[1, 3, 3, 1]): + super().__init__() + + if upsample: + self.upsample = Upsample(blur_kernel) + + self.conv = ModulatedConv2d(in_channel, 3, 1, style_dim, demodulate=False) + self.bias = nn.Parameter(torch.zeros(1, 3, 1, 1)) + + def forward(self, input, style, skip=None): + out = self.conv(input, style) + out = out + self.bias + + if skip is not None: + skip = self.upsample(skip) + + out = out + skip + + return out + + +class Generator(nn.Module): + def __init__( + self, + size, + style_dim, + n_mlp, + channel_multiplier=2, + blur_kernel=[1, 3, 3, 1], + lr_mlp=0.01, + ): + super().__init__() + + self.size = size + + self.style_dim = style_dim + + layers = [PixelNorm()] + + for i in range(n_mlp): + layers.append( + EqualLinear( + style_dim, style_dim, lr_mul=lr_mlp, activation='fused_lrelu' + ) + ) + + self.style = nn.Sequential(*layers) + + self.channels = { + 4: 512, + 8: 512, + 16: 512, + 32: 512, + 64: 256 * channel_multiplier, + 128: 128 * channel_multiplier, + 256: 64 * channel_multiplier, + 512: 32 * channel_multiplier, + 1024: 16 * channel_multiplier, + } + + self.input = ConstantInput(self.channels[4]) + self.conv1 = StyledConv( + self.channels[4], self.channels[4], 3, style_dim, blur_kernel=blur_kernel + ) + self.to_rgb1 = ToRGB(self.channels[4], style_dim, upsample=False) + + self.log_size = int(math.log(size, 2)) + self.num_layers = (self.log_size - 2) * 2 + 1 + + self.convs = nn.ModuleList() + self.upsamples = nn.ModuleList() + self.to_rgbs = nn.ModuleList() + self.noises = nn.Module() + + in_channel = self.channels[4] + + for layer_idx in range(self.num_layers): + res = (layer_idx + 5) // 2 + shape = [1, 1, 2 ** res, 2 ** res] + self.noises.register_buffer(f'noise_{layer_idx}', torch.randn(*shape)) + + for i in range(3, self.log_size + 1): + out_channel = self.channels[2 ** i] + + self.convs.append( + StyledConv( + in_channel, + out_channel, + 3, + style_dim, + upsample=True, + blur_kernel=blur_kernel, + ) + ) + + self.convs.append( + StyledConv( + out_channel, out_channel, 3, style_dim, blur_kernel=blur_kernel + ) + ) + + self.to_rgbs.append(ToRGB(out_channel, style_dim)) + + in_channel = out_channel + + self.n_latent = self.log_size * 2 - 2 + + def make_noise(self): + device = self.input.input.device + + noises = [torch.randn(1, 1, 2 ** 2, 2 ** 2, device=device)] + + for i in range(3, self.log_size + 1): + for _ in range(2): + noises.append(torch.randn(1, 1, 2 ** i, 2 ** i, device=device)) + + return noises + + def mean_latent(self, n_latent): + latent_in = torch.randn( + n_latent, self.style_dim, device=self.input.input.device + ) + latent = self.style(latent_in).mean(0, keepdim=True) + + return latent + + def get_latent(self, input): + return self.style(input) + + def forward( + self, + styles, + return_latents=False, + return_features=False, + inject_index=None, + truncation=1, + truncation_latent=None, + input_is_latent=False, + noise=None, + randomize_noise=True, + features_in=None, + feature_scale=1.0 + ): + if not input_is_latent: + styles = [self.style(s) for s in styles] + + if noise is None: + if randomize_noise: + noise = [None] * self.num_layers + else: + noise = [ + getattr(self.noises, f'noise_{i}') for i in range(self.num_layers) + ] + + if truncation < 1: + style_t = [] + + for style in styles: + style_t.append( + truncation_latent + truncation * (style - truncation_latent) + ) + + styles = style_t + + if len(styles) < 2: + inject_index = self.n_latent + + if styles[0].ndim < 3: + latent = styles[0].unsqueeze(1).repeat(1, inject_index, 1) + else: + latent = styles[0] + + else: + if inject_index is None: + inject_index = random.randint(1, self.n_latent - 1) + + latent = styles[0].unsqueeze(1).repeat(1, inject_index, 1) + latent2 = styles[1].unsqueeze(1).repeat(1, self.n_latent - inject_index, 1) + + latent = torch.cat([latent, latent2], 1) + + def insert_feature(x, layer_idx): + if features_in is not None and features_in[layer_idx] is not None: + x = (1 - feature_scale) * x + feature_scale * features_in[layer_idx].type_as(x) + return x + + outs = [] + out = self.input(latent) + outs.append(out) + out = self.conv1(out, latent[:, 0], noise=noise[0]) + outs.append(out) + + skip = self.to_rgb1(out, latent[:, 1]) + + i = 1 + for conv1, conv2, noise1, noise2, to_rgb in zip( + self.convs[::2], self.convs[1::2], noise[1::2], noise[2::2], self.to_rgbs + ): + out = insert_feature(out, i) + out = conv1(out, latent[:, i], noise=noise1) + outs.append(out) + out = insert_feature(out, i + 1) + out = conv2(out, latent[:, i + 1], noise=noise2) + outs.append(out) + skip = to_rgb(out, latent[:, i + 2], skip) + + i += 2 + + image = skip + + if return_latents: + return image, latent + elif return_features: + return image, outs + else: + return image, None + + +class ConvLayer(nn.Sequential): + def __init__( + self, + in_channel, + out_channel, + kernel_size, + downsample=False, + blur_kernel=[1, 3, 3, 1], + bias=True, + activate=True, + ): + layers = [] + + if downsample: + factor = 2 + p = (len(blur_kernel) - factor) + (kernel_size - 1) + pad0 = (p + 1) // 2 + pad1 = p // 2 + + layers.append(Blur(blur_kernel, pad=(pad0, pad1))) + + stride = 2 + self.padding = 0 + + else: + stride = 1 + self.padding = kernel_size // 2 + + layers.append( + EqualConv2d( + in_channel, + out_channel, + kernel_size, + padding=self.padding, + stride=stride, + bias=bias and not activate, + ) + ) + + if activate: + if bias: + layers.append(FusedLeakyReLU(out_channel)) + + else: + layers.append(ScaledLeakyReLU(0.2)) + + super().__init__(*layers) + + +class ResBlock(nn.Module): + def __init__(self, in_channel, out_channel, blur_kernel=[1, 3, 3, 1]): + super().__init__() + + self.conv1 = ConvLayer(in_channel, in_channel, 3) + self.conv2 = ConvLayer(in_channel, out_channel, 3, downsample=True) + + self.skip = ConvLayer( + in_channel, out_channel, 1, downsample=True, activate=False, bias=False + ) + + def forward(self, input): + out = self.conv1(input) + out = self.conv2(out) + + skip = self.skip(input) + out = (out + skip) / math.sqrt(2) + + return out + + +class Discriminator(nn.Module): + def __init__(self, size, channel_multiplier=2, blur_kernel=[1, 3, 3, 1]): + super().__init__() + + channels = { + 4: 512, + 8: 512, + 16: 512, + 32: 512, + 64: 256 * channel_multiplier, + 128: 128 * channel_multiplier, + 256: 64 * channel_multiplier, + 512: 32 * channel_multiplier, + 1024: 16 * channel_multiplier, + } + + convs = [ConvLayer(3, channels[size], 1)] + + log_size = int(math.log(size, 2)) + + in_channel = channels[size] + + for i in range(log_size, 2, -1): + out_channel = channels[2 ** (i - 1)] + + convs.append(ResBlock(in_channel, out_channel, blur_kernel)) + + in_channel = out_channel + + self.convs = nn.Sequential(*convs) + + self.stddev_group = 4 + self.stddev_feat = 1 + + self.final_conv = ConvLayer(in_channel + 1, channels[4], 3) + self.final_linear = nn.Sequential( + EqualLinear(channels[4] * 4 * 4, channels[4], activation='fused_lrelu'), + EqualLinear(channels[4], 1), + ) + + def forward(self, input): + out = self.convs(input) + + batch, channel, height, width = out.shape + group = min(batch, self.stddev_group) + stddev = out.view( + group, -1, self.stddev_feat, channel // self.stddev_feat, height, width + ) + stddev = torch.sqrt(stddev.var(0, unbiased=False) + 1e-8) + stddev = stddev.mean([2, 3, 4], keepdims=True).squeeze(2) + stddev = stddev.repeat(group, 1, height, width) + out = torch.cat([out, stddev], 1) + + out = self.final_conv(out) + + out = out.view(batch, -1) + out = self.final_linear(out) + + return out diff --git a/models/FeatureStyleEncoder/pixel2style2pixel/models/stylegan2/op/__init__.py b/models/FeatureStyleEncoder/pixel2style2pixel/models/stylegan2/op/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d0918d92285955855be89f00096b888ee5597ce3 --- /dev/null +++ b/models/FeatureStyleEncoder/pixel2style2pixel/models/stylegan2/op/__init__.py @@ -0,0 +1,2 @@ +from .fused_act import FusedLeakyReLU, fused_leaky_relu +from .upfirdn2d import upfirdn2d diff --git a/models/FeatureStyleEncoder/pixel2style2pixel/models/stylegan2/op/fused_act.py b/models/FeatureStyleEncoder/pixel2style2pixel/models/stylegan2/op/fused_act.py new file mode 100644 index 0000000000000000000000000000000000000000..973a84fffde53668d31397da5fb993bbc95f7be0 --- /dev/null +++ b/models/FeatureStyleEncoder/pixel2style2pixel/models/stylegan2/op/fused_act.py @@ -0,0 +1,85 @@ +import os + +import torch +from torch import nn +from torch.autograd import Function +from torch.utils.cpp_extension import load + +module_path = os.path.dirname(__file__) +fused = load( + 'fused', + sources=[ + os.path.join(module_path, 'fused_bias_act.cpp'), + os.path.join(module_path, 'fused_bias_act_kernel.cu'), + ], +) + + +class FusedLeakyReLUFunctionBackward(Function): + @staticmethod + def forward(ctx, grad_output, out, negative_slope, scale): + ctx.save_for_backward(out) + ctx.negative_slope = negative_slope + ctx.scale = scale + + empty = grad_output.new_empty(0) + + grad_input = fused.fused_bias_act( + grad_output, empty, out, 3, 1, negative_slope, scale + ) + + dim = [0] + + if grad_input.ndim > 2: + dim += list(range(2, grad_input.ndim)) + + grad_bias = grad_input.sum(dim).detach() + + return grad_input, grad_bias + + @staticmethod + def backward(ctx, gradgrad_input, gradgrad_bias): + out, = ctx.saved_tensors + gradgrad_out = fused.fused_bias_act( + gradgrad_input, gradgrad_bias, out, 3, 1, ctx.negative_slope, ctx.scale + ) + + return gradgrad_out, None, None, None + + +class FusedLeakyReLUFunction(Function): + @staticmethod + def forward(ctx, input, bias, negative_slope, scale): + empty = input.new_empty(0) + out = fused.fused_bias_act(input, bias, empty, 3, 0, negative_slope, scale) + ctx.save_for_backward(out) + ctx.negative_slope = negative_slope + ctx.scale = scale + + return out + + @staticmethod + def backward(ctx, grad_output): + out, = ctx.saved_tensors + + grad_input, grad_bias = FusedLeakyReLUFunctionBackward.apply( + grad_output, out, ctx.negative_slope, ctx.scale + ) + + return grad_input, grad_bias, None, None + + +class FusedLeakyReLU(nn.Module): + def __init__(self, channel, negative_slope=0.2, scale=2 ** 0.5): + super().__init__() + + self.bias = nn.Parameter(torch.zeros(channel)) + self.negative_slope = negative_slope + self.scale = scale + + def forward(self, input): + return fused_leaky_relu(input, self.bias, self.negative_slope, self.scale) + + +def fused_leaky_relu(input, bias, negative_slope=0.2, scale=2 ** 0.5): + return FusedLeakyReLUFunction.apply(input, bias, negative_slope, scale) diff --git a/models/FeatureStyleEncoder/pixel2style2pixel/models/stylegan2/op/fused_bias_act.cpp b/models/FeatureStyleEncoder/pixel2style2pixel/models/stylegan2/op/fused_bias_act.cpp new file mode 100644 index 0000000000000000000000000000000000000000..02be898f970bcc8ea297867fcaa4e71b24b3d949 --- /dev/null +++ b/models/FeatureStyleEncoder/pixel2style2pixel/models/stylegan2/op/fused_bias_act.cpp @@ -0,0 +1,21 @@ +#include + + +torch::Tensor fused_bias_act_op(const torch::Tensor& input, const torch::Tensor& bias, const torch::Tensor& refer, + int act, int grad, float alpha, float scale); + +#define CHECK_CUDA(x) TORCH_CHECK(x.type().is_cuda(), #x " must be a CUDA tensor") +#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous") +#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x) + +torch::Tensor fused_bias_act(const torch::Tensor& input, const torch::Tensor& bias, const torch::Tensor& refer, + int act, int grad, float alpha, float scale) { + CHECK_CUDA(input); + CHECK_CUDA(bias); + + return fused_bias_act_op(input, bias, refer, act, grad, alpha, scale); +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("fused_bias_act", &fused_bias_act, "fused bias act (CUDA)"); +} \ No newline at end of file diff --git a/models/FeatureStyleEncoder/pixel2style2pixel/models/stylegan2/op/fused_bias_act_kernel.cu b/models/FeatureStyleEncoder/pixel2style2pixel/models/stylegan2/op/fused_bias_act_kernel.cu new file mode 100644 index 0000000000000000000000000000000000000000..c9fa56fea7ede7072dc8925cfb0148f136eb85b8 --- /dev/null +++ b/models/FeatureStyleEncoder/pixel2style2pixel/models/stylegan2/op/fused_bias_act_kernel.cu @@ -0,0 +1,99 @@ +// Copyright (c) 2019, NVIDIA Corporation. All rights reserved. +// +// This work is made available under the Nvidia Source Code License-NC. +// To view a copy of this license, visit +// https://nvlabs.github.io/stylegan2/license.html + +#include + +#include +#include +#include +#include + +#include +#include + + +template +static __global__ void fused_bias_act_kernel(scalar_t* out, const scalar_t* p_x, const scalar_t* p_b, const scalar_t* p_ref, + int act, int grad, scalar_t alpha, scalar_t scale, int loop_x, int size_x, int step_b, int size_b, int use_bias, int use_ref) { + int xi = blockIdx.x * loop_x * blockDim.x + threadIdx.x; + + scalar_t zero = 0.0; + + for (int loop_idx = 0; loop_idx < loop_x && xi < size_x; loop_idx++, xi += blockDim.x) { + scalar_t x = p_x[xi]; + + if (use_bias) { + x += p_b[(xi / step_b) % size_b]; + } + + scalar_t ref = use_ref ? p_ref[xi] : zero; + + scalar_t y; + + switch (act * 10 + grad) { + default: + case 10: y = x; break; + case 11: y = x; break; + case 12: y = 0.0; break; + + case 30: y = (x > 0.0) ? x : x * alpha; break; + case 31: y = (ref > 0.0) ? x : x * alpha; break; + case 32: y = 0.0; break; + } + + out[xi] = y * scale; + } +} + + +torch::Tensor fused_bias_act_op(const torch::Tensor& input, const torch::Tensor& bias, const torch::Tensor& refer, + int act, int grad, float alpha, float scale) { + int curDevice = -1; + cudaGetDevice(&curDevice); + cudaStream_t stream = at::cuda::getCurrentCUDAStream(curDevice); + + auto x = input.contiguous(); + auto b = bias.contiguous(); + auto ref = refer.contiguous(); + + int use_bias = b.numel() ? 1 : 0; + int use_ref = ref.numel() ? 1 : 0; + + int size_x = x.numel(); + int size_b = b.numel(); + int step_b = 1; + + for (int i = 1 + 1; i < x.dim(); i++) { + step_b *= x.size(i); + } + + int loop_x = 4; + int block_size = 4 * 32; + int grid_size = (size_x - 1) / (loop_x * block_size) + 1; + + auto y = torch::empty_like(x); + + AT_DISPATCH_FLOATING_TYPES_AND_HALF(x.scalar_type(), "fused_bias_act_kernel", [&] { + fused_bias_act_kernel<<>>( + y.data_ptr(), + x.data_ptr(), + b.data_ptr(), + ref.data_ptr(), + act, + grad, + alpha, + scale, + loop_x, + size_x, + step_b, + size_b, + use_bias, + use_ref + ); + }); + + return y; +} \ No newline at end of file diff --git a/models/FeatureStyleEncoder/pixel2style2pixel/models/stylegan2/op/upfirdn2d.cpp b/models/FeatureStyleEncoder/pixel2style2pixel/models/stylegan2/op/upfirdn2d.cpp new file mode 100644 index 0000000000000000000000000000000000000000..d2e633dc896433c205e18bc3e455539192ff968e --- /dev/null +++ b/models/FeatureStyleEncoder/pixel2style2pixel/models/stylegan2/op/upfirdn2d.cpp @@ -0,0 +1,23 @@ +#include + + +torch::Tensor upfirdn2d_op(const torch::Tensor& input, const torch::Tensor& kernel, + int up_x, int up_y, int down_x, int down_y, + int pad_x0, int pad_x1, int pad_y0, int pad_y1); + +#define CHECK_CUDA(x) TORCH_CHECK(x.type().is_cuda(), #x " must be a CUDA tensor") +#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous") +#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x) + +torch::Tensor upfirdn2d(const torch::Tensor& input, const torch::Tensor& kernel, + int up_x, int up_y, int down_x, int down_y, + int pad_x0, int pad_x1, int pad_y0, int pad_y1) { + CHECK_CUDA(input); + CHECK_CUDA(kernel); + + return upfirdn2d_op(input, kernel, up_x, up_y, down_x, down_y, pad_x0, pad_x1, pad_y0, pad_y1); +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("upfirdn2d", &upfirdn2d, "upfirdn2d (CUDA)"); +} \ No newline at end of file diff --git a/models/FeatureStyleEncoder/pixel2style2pixel/models/stylegan2/op/upfirdn2d.py b/models/FeatureStyleEncoder/pixel2style2pixel/models/stylegan2/op/upfirdn2d.py new file mode 100644 index 0000000000000000000000000000000000000000..e43c7eeddbc143ea313454e5b64d083c27dafc5b --- /dev/null +++ b/models/FeatureStyleEncoder/pixel2style2pixel/models/stylegan2/op/upfirdn2d.py @@ -0,0 +1,185 @@ +import os + +import torch +from torch.autograd import Function +from torch.utils.cpp_extension import load + + +module_path = os.path.dirname(__file__) +upfirdn2d_op = load( + 'upfirdn2d', + sources=[ + os.path.join(module_path, 'upfirdn2d.cpp'), + os.path.join(module_path, 'upfirdn2d_kernel.cu'), + ], +) + + +class UpFirDn2dBackward(Function): + @staticmethod + def forward( + ctx, grad_output, kernel, grad_kernel, up, down, pad, g_pad, in_size, out_size + ): + up_x, up_y = up + down_x, down_y = down + g_pad_x0, g_pad_x1, g_pad_y0, g_pad_y1 = g_pad + + grad_output = grad_output.reshape(-1, out_size[0], out_size[1], 1) + + grad_input = upfirdn2d_op.upfirdn2d( + grad_output, + grad_kernel, + down_x, + down_y, + up_x, + up_y, + g_pad_x0, + g_pad_x1, + g_pad_y0, + g_pad_y1, + ) + grad_input = grad_input.view(in_size[0], in_size[1], in_size[2], in_size[3]) + + ctx.save_for_backward(kernel) + + pad_x0, pad_x1, pad_y0, pad_y1 = pad + + ctx.up_x = up_x + ctx.up_y = up_y + ctx.down_x = down_x + ctx.down_y = down_y + ctx.pad_x0 = pad_x0 + ctx.pad_x1 = pad_x1 + ctx.pad_y0 = pad_y0 + ctx.pad_y1 = pad_y1 + ctx.in_size = in_size + ctx.out_size = out_size + + return grad_input + + @staticmethod + def backward(ctx, gradgrad_input): + kernel, = ctx.saved_tensors + + gradgrad_input = gradgrad_input.reshape(-1, ctx.in_size[2], ctx.in_size[3], 1) + + gradgrad_out = upfirdn2d_op.upfirdn2d( + gradgrad_input, + kernel, + ctx.up_x, + ctx.up_y, + ctx.down_x, + ctx.down_y, + ctx.pad_x0, + ctx.pad_x1, + ctx.pad_y0, + ctx.pad_y1, + ) + # gradgrad_out = gradgrad_out.view(ctx.in_size[0], ctx.out_size[0], ctx.out_size[1], ctx.in_size[3]) + gradgrad_out = gradgrad_out.view( + ctx.in_size[0], ctx.in_size[1], ctx.out_size[0], ctx.out_size[1] + ) + + return gradgrad_out, None, None, None, None, None, None, None, None + + +class UpFirDn2d(Function): + @staticmethod + def forward(ctx, input, kernel, up, down, pad): + up_x, up_y = up + down_x, down_y = down + pad_x0, pad_x1, pad_y0, pad_y1 = pad + + kernel_h, kernel_w = kernel.shape + batch, channel, in_h, in_w = input.shape + ctx.in_size = input.shape + + input = input.reshape(-1, in_h, in_w, 1) + + ctx.save_for_backward(kernel, torch.flip(kernel, [0, 1])) + + out_h = (in_h * up_y + pad_y0 + pad_y1 - kernel_h) // down_y + 1 + out_w = (in_w * up_x + pad_x0 + pad_x1 - kernel_w) // down_x + 1 + ctx.out_size = (out_h, out_w) + + ctx.up = (up_x, up_y) + ctx.down = (down_x, down_y) + ctx.pad = (pad_x0, pad_x1, pad_y0, pad_y1) + + g_pad_x0 = kernel_w - pad_x0 - 1 + g_pad_y0 = kernel_h - pad_y0 - 1 + g_pad_x1 = in_w * up_x - out_w * down_x + pad_x0 - up_x + 1 + g_pad_y1 = in_h * up_y - out_h * down_y + pad_y0 - up_y + 1 + + ctx.g_pad = (g_pad_x0, g_pad_x1, g_pad_y0, g_pad_y1) + + out = upfirdn2d_op.upfirdn2d( + input, kernel, up_x, up_y, down_x, down_y, pad_x0, pad_x1, pad_y0, pad_y1 + ) + # out = out.view(major, out_h, out_w, minor) + out = out.view(-1, channel, out_h, out_w) + + return out + + @staticmethod + def backward(ctx, grad_output): + kernel, grad_kernel = ctx.saved_tensors + + grad_input = UpFirDn2dBackward.apply( + grad_output, + kernel, + grad_kernel, + ctx.up, + ctx.down, + ctx.pad, + ctx.g_pad, + ctx.in_size, + ctx.out_size, + ) + + return grad_input, None, None, None, None + + +def upfirdn2d(input, kernel, up=1, down=1, pad=(0, 0)): + out = UpFirDn2d.apply( + input, kernel, (up, up), (down, down), (pad[0], pad[1], pad[0], pad[1]) + ) + + return out + + +def upfirdn2d_native( + input, kernel, up_x, up_y, down_x, down_y, pad_x0, pad_x1, pad_y0, pad_y1 +): + _, in_h, in_w, minor = input.shape + kernel_h, kernel_w = kernel.shape + + out = input.view(-1, in_h, 1, in_w, 1, minor) + out = F.pad(out, [0, 0, 0, up_x - 1, 0, 0, 0, up_y - 1]) + out = out.view(-1, in_h * up_y, in_w * up_x, minor) + + out = F.pad( + out, [0, 0, max(pad_x0, 0), max(pad_x1, 0), max(pad_y0, 0), max(pad_y1, 0)] + ) + out = out[ + :, + max(-pad_y0, 0): out.shape[1] - max(-pad_y1, 0), + max(-pad_x0, 0): out.shape[2] - max(-pad_x1, 0), + :, + ] + + out = out.permute(0, 3, 1, 2) + out = out.reshape( + [-1, 1, in_h * up_y + pad_y0 + pad_y1, in_w * up_x + pad_x0 + pad_x1] + ) + w = torch.flip(kernel, [0, 1]).view(1, 1, kernel_h, kernel_w) + out = F.conv2d(out, w) + out = out.reshape( + -1, + minor, + in_h * up_y + pad_y0 + pad_y1 - kernel_h + 1, + in_w * up_x + pad_x0 + pad_x1 - kernel_w + 1, + ) + out = out.permute(0, 2, 3, 1) + + return out[:, ::down_y, ::down_x, :] diff --git a/models/FeatureStyleEncoder/pixel2style2pixel/models/stylegan2/op/upfirdn2d_kernel.cu b/models/FeatureStyleEncoder/pixel2style2pixel/models/stylegan2/op/upfirdn2d_kernel.cu new file mode 100644 index 0000000000000000000000000000000000000000..2a710aa6adc3d43ac93136a1814e3c39970e1c7e --- /dev/null +++ b/models/FeatureStyleEncoder/pixel2style2pixel/models/stylegan2/op/upfirdn2d_kernel.cu @@ -0,0 +1,272 @@ +// Copyright (c) 2019, NVIDIA Corporation. All rights reserved. +// +// This work is made available under the Nvidia Source Code License-NC. +// To view a copy of this license, visit +// https://nvlabs.github.io/stylegan2/license.html + +#include + +#include +#include +#include +#include + +#include +#include + + +static __host__ __device__ __forceinline__ int floor_div(int a, int b) { + int c = a / b; + + if (c * b > a) { + c--; + } + + return c; +} + + +struct UpFirDn2DKernelParams { + int up_x; + int up_y; + int down_x; + int down_y; + int pad_x0; + int pad_x1; + int pad_y0; + int pad_y1; + + int major_dim; + int in_h; + int in_w; + int minor_dim; + int kernel_h; + int kernel_w; + int out_h; + int out_w; + int loop_major; + int loop_x; +}; + + +template +__global__ void upfirdn2d_kernel(scalar_t* out, const scalar_t* input, const scalar_t* kernel, const UpFirDn2DKernelParams p) { + const int tile_in_h = ((tile_out_h - 1) * down_y + kernel_h - 1) / up_y + 1; + const int tile_in_w = ((tile_out_w - 1) * down_x + kernel_w - 1) / up_x + 1; + + __shared__ volatile float sk[kernel_h][kernel_w]; + __shared__ volatile float sx[tile_in_h][tile_in_w]; + + int minor_idx = blockIdx.x; + int tile_out_y = minor_idx / p.minor_dim; + minor_idx -= tile_out_y * p.minor_dim; + tile_out_y *= tile_out_h; + int tile_out_x_base = blockIdx.y * p.loop_x * tile_out_w; + int major_idx_base = blockIdx.z * p.loop_major; + + if (tile_out_x_base >= p.out_w | tile_out_y >= p.out_h | major_idx_base >= p.major_dim) { + return; + } + + for (int tap_idx = threadIdx.x; tap_idx < kernel_h * kernel_w; tap_idx += blockDim.x) { + int ky = tap_idx / kernel_w; + int kx = tap_idx - ky * kernel_w; + scalar_t v = 0.0; + + if (kx < p.kernel_w & ky < p.kernel_h) { + v = kernel[(p.kernel_h - 1 - ky) * p.kernel_w + (p.kernel_w - 1 - kx)]; + } + + sk[ky][kx] = v; + } + + for (int loop_major = 0, major_idx = major_idx_base; loop_major < p.loop_major & major_idx < p.major_dim; loop_major++, major_idx++) { + for (int loop_x = 0, tile_out_x = tile_out_x_base; loop_x < p.loop_x & tile_out_x < p.out_w; loop_x++, tile_out_x += tile_out_w) { + int tile_mid_x = tile_out_x * down_x + up_x - 1 - p.pad_x0; + int tile_mid_y = tile_out_y * down_y + up_y - 1 - p.pad_y0; + int tile_in_x = floor_div(tile_mid_x, up_x); + int tile_in_y = floor_div(tile_mid_y, up_y); + + __syncthreads(); + + for (int in_idx = threadIdx.x; in_idx < tile_in_h * tile_in_w; in_idx += blockDim.x) { + int rel_in_y = in_idx / tile_in_w; + int rel_in_x = in_idx - rel_in_y * tile_in_w; + int in_x = rel_in_x + tile_in_x; + int in_y = rel_in_y + tile_in_y; + + scalar_t v = 0.0; + + if (in_x >= 0 & in_y >= 0 & in_x < p.in_w & in_y < p.in_h) { + v = input[((major_idx * p.in_h + in_y) * p.in_w + in_x) * p.minor_dim + minor_idx]; + } + + sx[rel_in_y][rel_in_x] = v; + } + + __syncthreads(); + for (int out_idx = threadIdx.x; out_idx < tile_out_h * tile_out_w; out_idx += blockDim.x) { + int rel_out_y = out_idx / tile_out_w; + int rel_out_x = out_idx - rel_out_y * tile_out_w; + int out_x = rel_out_x + tile_out_x; + int out_y = rel_out_y + tile_out_y; + + int mid_x = tile_mid_x + rel_out_x * down_x; + int mid_y = tile_mid_y + rel_out_y * down_y; + int in_x = floor_div(mid_x, up_x); + int in_y = floor_div(mid_y, up_y); + int rel_in_x = in_x - tile_in_x; + int rel_in_y = in_y - tile_in_y; + int kernel_x = (in_x + 1) * up_x - mid_x - 1; + int kernel_y = (in_y + 1) * up_y - mid_y - 1; + + scalar_t v = 0.0; + + #pragma unroll + for (int y = 0; y < kernel_h / up_y; y++) + #pragma unroll + for (int x = 0; x < kernel_w / up_x; x++) + v += sx[rel_in_y + y][rel_in_x + x] * sk[kernel_y + y * up_y][kernel_x + x * up_x]; + + if (out_x < p.out_w & out_y < p.out_h) { + out[((major_idx * p.out_h + out_y) * p.out_w + out_x) * p.minor_dim + minor_idx] = v; + } + } + } + } +} + + +torch::Tensor upfirdn2d_op(const torch::Tensor& input, const torch::Tensor& kernel, + int up_x, int up_y, int down_x, int down_y, + int pad_x0, int pad_x1, int pad_y0, int pad_y1) { + int curDevice = -1; + cudaGetDevice(&curDevice); + cudaStream_t stream = at::cuda::getCurrentCUDAStream(curDevice); + + UpFirDn2DKernelParams p; + + auto x = input.contiguous(); + auto k = kernel.contiguous(); + + p.major_dim = x.size(0); + p.in_h = x.size(1); + p.in_w = x.size(2); + p.minor_dim = x.size(3); + p.kernel_h = k.size(0); + p.kernel_w = k.size(1); + p.up_x = up_x; + p.up_y = up_y; + p.down_x = down_x; + p.down_y = down_y; + p.pad_x0 = pad_x0; + p.pad_x1 = pad_x1; + p.pad_y0 = pad_y0; + p.pad_y1 = pad_y1; + + p.out_h = (p.in_h * p.up_y + p.pad_y0 + p.pad_y1 - p.kernel_h + p.down_y) / p.down_y; + p.out_w = (p.in_w * p.up_x + p.pad_x0 + p.pad_x1 - p.kernel_w + p.down_x) / p.down_x; + + auto out = at::empty({p.major_dim, p.out_h, p.out_w, p.minor_dim}, x.options()); + + int mode = -1; + + int tile_out_h; + int tile_out_w; + + if (p.up_x == 1 && p.up_y == 1 && p.down_x == 1 && p.down_y == 1 && p.kernel_h <= 4 && p.kernel_w <= 4) { + mode = 1; + tile_out_h = 16; + tile_out_w = 64; + } + + if (p.up_x == 1 && p.up_y == 1 && p.down_x == 1 && p.down_y == 1 && p.kernel_h <= 3 && p.kernel_w <= 3) { + mode = 2; + tile_out_h = 16; + tile_out_w = 64; + } + + if (p.up_x == 2 && p.up_y == 2 && p.down_x == 1 && p.down_y == 1 && p.kernel_h <= 4 && p.kernel_w <= 4) { + mode = 3; + tile_out_h = 16; + tile_out_w = 64; + } + + if (p.up_x == 2 && p.up_y == 2 && p.down_x == 1 && p.down_y == 1 && p.kernel_h <= 2 && p.kernel_w <= 2) { + mode = 4; + tile_out_h = 16; + tile_out_w = 64; + } + + if (p.up_x == 1 && p.up_y == 1 && p.down_x == 2 && p.down_y == 2 && p.kernel_h <= 4 && p.kernel_w <= 4) { + mode = 5; + tile_out_h = 8; + tile_out_w = 32; + } + + if (p.up_x == 1 && p.up_y == 1 && p.down_x == 2 && p.down_y == 2 && p.kernel_h <= 2 && p.kernel_w <= 2) { + mode = 6; + tile_out_h = 8; + tile_out_w = 32; + } + + dim3 block_size; + dim3 grid_size; + + if (tile_out_h > 0 && tile_out_w) { + p.loop_major = (p.major_dim - 1) / 16384 + 1; + p.loop_x = 1; + block_size = dim3(32 * 8, 1, 1); + grid_size = dim3(((p.out_h - 1) / tile_out_h + 1) * p.minor_dim, + (p.out_w - 1) / (p.loop_x * tile_out_w) + 1, + (p.major_dim - 1) / p.loop_major + 1); + } + + AT_DISPATCH_FLOATING_TYPES_AND_HALF(x.scalar_type(), "upfirdn2d_cuda", [&] { + switch (mode) { + case 1: + upfirdn2d_kernel<<>>( + out.data_ptr(), x.data_ptr(), k.data_ptr(), p + ); + + break; + + case 2: + upfirdn2d_kernel<<>>( + out.data_ptr(), x.data_ptr(), k.data_ptr(), p + ); + + break; + + case 3: + upfirdn2d_kernel<<>>( + out.data_ptr(), x.data_ptr(), k.data_ptr(), p + ); + + break; + + case 4: + upfirdn2d_kernel<<>>( + out.data_ptr(), x.data_ptr(), k.data_ptr(), p + ); + + break; + + case 5: + upfirdn2d_kernel<<>>( + out.data_ptr(), x.data_ptr(), k.data_ptr(), p + ); + + break; + + case 6: + upfirdn2d_kernel<<>>( + out.data_ptr(), x.data_ptr(), k.data_ptr(), p + ); + + break; + } + }); + + return out; +} \ No newline at end of file diff --git a/models/FeatureStyleEncoder/ranger.py b/models/FeatureStyleEncoder/ranger.py new file mode 100644 index 0000000000000000000000000000000000000000..3d63264dda6df0ee40cac143440f0b5f8977a9ad --- /dev/null +++ b/models/FeatureStyleEncoder/ranger.py @@ -0,0 +1,164 @@ +# Ranger deep learning optimizer - RAdam + Lookahead + Gradient Centralization, combined into one optimizer. + +# https://github.com/lessw2020/Ranger-Deep-Learning-Optimizer +# and/or +# https://github.com/lessw2020/Best-Deep-Learning-Optimizers + +# Ranger has now been used to capture 12 records on the FastAI leaderboard. + +# This version = 20.4.11 + +# Credits: +# Gradient Centralization --> https://arxiv.org/abs/2004.01461v2 (a new optimization technique for DNNs), github: https://github.com/Yonghongwei/Gradient-Centralization +# RAdam --> https://github.com/LiyuanLucasLiu/RAdam +# Lookahead --> rewritten by lessw2020, but big thanks to Github @LonePatient and @RWightman for ideas from their code. +# Lookahead paper --> MZhang,G Hinton https://arxiv.org/abs/1907.08610 + +# summary of changes: +# 4/11/20 - add gradient centralization option. Set new testing benchmark for accuracy with it, toggle with use_gc flag at init. +# full code integration with all updates at param level instead of group, moves slow weights into state dict (from generic weights), +# supports group learning rates (thanks @SHolderbach), fixes sporadic load from saved model issues. +# changes 8/31/19 - fix references to *self*.N_sma_threshold; +# changed eps to 1e-5 as better default than 1e-8. + +import math +import torch +from torch.optim.optimizer import Optimizer + + +class Ranger(Optimizer): + + def __init__(self, params, lr=1e-3, # lr + alpha=0.5, k=6, N_sma_threshhold=5, # Ranger options + betas=(.95, 0.999), eps=1e-5, weight_decay=0, # Adam options + use_gc=True, gc_conv_only=False + # Gradient centralization on or off, applied to conv layers only or conv + fc layers + ): + + # parameter checks + if not 0.0 <= alpha <= 1.0: + raise ValueError(f'Invalid slow update rate: {alpha}') + if not 1 <= k: + raise ValueError(f'Invalid lookahead steps: {k}') + if not lr > 0: + raise ValueError(f'Invalid Learning Rate: {lr}') + if not eps > 0: + raise ValueError(f'Invalid eps: {eps}') + + # parameter comments: + # beta1 (momentum) of .95 seems to work better than .90... + # N_sma_threshold of 5 seems better in testing than 4. + # In both cases, worth testing on your dataset (.90 vs .95, 4 vs 5) to make sure which works best for you. + + # prep defaults and init torch.optim base + defaults = dict(lr=lr, alpha=alpha, k=k, step_counter=0, betas=betas, N_sma_threshhold=N_sma_threshhold, + eps=eps, weight_decay=weight_decay) + super().__init__(params, defaults) + + # adjustable threshold + self.N_sma_threshhold = N_sma_threshhold + + # look ahead params + + self.alpha = alpha + self.k = k + + # radam buffer for state + self.radam_buffer = [[None, None, None] for ind in range(10)] + + # gc on or off + self.use_gc = use_gc + + # level of gradient centralization + self.gc_gradient_threshold = 3 if gc_conv_only else 1 + + def __setstate__(self, state): + super(Ranger, self).__setstate__(state) + + def step(self, closure=None): + loss = None + + # Evaluate averages and grad, update param tensors + for group in self.param_groups: + + for p in group['params']: + if p.grad is None: + continue + grad = p.grad.data.float() + + if grad.is_sparse: + raise RuntimeError('Ranger optimizer does not support sparse gradients') + + p_data_fp32 = p.data.float() + + state = self.state[p] # get state dict for this param + + if len(state) == 0: # if first time to run...init dictionary with our desired entries + # if self.first_run_check==0: + # self.first_run_check=1 + # print("Initializing slow buffer...should not see this at load from saved model!") + state['step'] = 0 + state['exp_avg'] = torch.zeros_like(p_data_fp32) + state['exp_avg_sq'] = torch.zeros_like(p_data_fp32) + + # look ahead weight storage now in state dict + state['slow_buffer'] = torch.empty_like(p.data) + state['slow_buffer'].copy_(p.data) + + else: + state['exp_avg'] = state['exp_avg'].type_as(p_data_fp32) + state['exp_avg_sq'] = state['exp_avg_sq'].type_as(p_data_fp32) + + # begin computations + exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq'] + beta1, beta2 = group['betas'] + + # GC operation for Conv layers and FC layers + if grad.dim() > self.gc_gradient_threshold: + grad.add_(-grad.mean(dim=tuple(range(1, grad.dim())), keepdim=True)) + + state['step'] += 1 + + # compute variance mov avg + exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad) + # compute mean moving avg + exp_avg.mul_(beta1).add_(1 - beta1, grad) + + buffered = self.radam_buffer[int(state['step'] % 10)] + + if state['step'] == buffered[0]: + N_sma, step_size = buffered[1], buffered[2] + else: + buffered[0] = state['step'] + beta2_t = beta2 ** state['step'] + N_sma_max = 2 / (1 - beta2) - 1 + N_sma = N_sma_max - 2 * state['step'] * beta2_t / (1 - beta2_t) + buffered[1] = N_sma + if N_sma > self.N_sma_threshhold: + step_size = math.sqrt( + (1 - beta2_t) * (N_sma - 4) / (N_sma_max - 4) * (N_sma - 2) / N_sma * N_sma_max / ( + N_sma_max - 2)) / (1 - beta1 ** state['step']) + else: + step_size = 1.0 / (1 - beta1 ** state['step']) + buffered[2] = step_size + + if group['weight_decay'] != 0: + p_data_fp32.add_(-group['weight_decay'] * group['lr'], p_data_fp32) + + # apply lr + if N_sma > self.N_sma_threshhold: + denom = exp_avg_sq.sqrt().add_(group['eps']) + p_data_fp32.addcdiv_(-step_size * group['lr'], exp_avg, denom) + else: + p_data_fp32.add_(-step_size * group['lr'], exp_avg) + + p.data.copy_(p_data_fp32) + + # integrated look ahead... + # we do it at the param level instead of group level + if state['step'] % group['k'] == 0: + slow_p = state['slow_buffer'] # get access to slow param tensor + slow_p.add_(self.alpha, p.data - slow_p) # (fast weights - slow weights) * alpha + p.data.copy_(slow_p) # copy interpolated weights to RAdam param tensor + + return loss \ No newline at end of file diff --git a/models/FeatureStyleEncoder/run_video_inversion_editing.sh b/models/FeatureStyleEncoder/run_video_inversion_editing.sh new file mode 100644 index 0000000000000000000000000000000000000000..8a221a3b8d44e10b0622bec175818ca8c173b901 --- /dev/null +++ b/models/FeatureStyleEncoder/run_video_inversion_editing.sh @@ -0,0 +1,24 @@ +VideoName='FP010363HD03' +Attribute='Heavy_Makeup' +Scale='1' +Sigma='3' # Choose appropriate gaussian filter size +VideoDir='./data/video/' +OutputDir='./output/video/' + + +# Cut video to frames +python video_processing.py --function 'video_to_frames' --video_path ${VideoDir}/${VideoName}.mp4 --output_path ${OutputDir} #--resize + +# Crop and align the faces in each frame +python video_processing.py --function 'align_frames' --video_path ${VideoDir}/${VideoName}.mp4 --output_path ${OutputDir} --filter_size=${Sigma} --optical_flow + +# Inversion +python test.py --config 143 --input_path ${OutputDir}/${VideoName}/${VideoName}_crop_align/ --save_path ${OutputDir}/${VideoName}/${VideoName}_inversion/ + +# Achieve latent manipulation +python video_processing.py --function 'latent_manipulation' --video_path ${VideoDir}/${VideoName}.mp4 --attr ${Attribute} --alpha=${Scale} + +# Reproject the manipulated frames to the original video +python video_processing.py --function 'reproject_origin' --video_path ${VideoDir}/${VideoName}.mp4 --seamless +python video_processing.py --function 'reproject_manipulate' --video_path ${VideoDir}/${VideoName}.mp4 --attr ${Attribute} --seamless +python video_processing.py --function 'compare_frames' --video_path ${VideoDir}/${VideoName}.mp4 --attr ${Attribute} --strs 'Original,Projected,Manipulated' diff --git a/models/FeatureStyleEncoder/test.py b/models/FeatureStyleEncoder/test.py new file mode 100644 index 0000000000000000000000000000000000000000..618bb00474c297b47700906755d26dec46b395d2 --- /dev/null +++ b/models/FeatureStyleEncoder/test.py @@ -0,0 +1,78 @@ +import argparse +import glob +import os +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.utils.data as data +import yaml + +from PIL import Image +from tqdm import tqdm +from torchvision import transforms, utils + +from utils.datasets import * +from utils.functions import * +from trainer import * + +torch.backends.cudnn.enabled = True +torch.backends.cudnn.deterministic = True +torch.backends.cudnn.benchmark = True +torch.autograd.set_detect_anomaly(True) +Image.MAX_IMAGE_PIXELS = None +device = torch.device('cuda') + +parser = argparse.ArgumentParser() +parser.add_argument('--config', type=str, default='001', help='Path to the config file.') +parser.add_argument('--pretrained_model_path', type=str, default='./pretrained_models/143_enc.pth', help='pretrained stylegan2 model') +parser.add_argument('--stylegan_model_path', type=str, default='./pixel2style2pixel/pretrained_models/psp_ffhq_encode.pt', help='pretrained stylegan2 model') +parser.add_argument('--arcface_model_path', type=str, default='./pretrained_models/backbone.pth', help='pretrained ArcFace model') +parser.add_argument('--parsing_model_path', type=str, default='./pretrained_models/79999_iter.pth', help='pretrained parsing model') +parser.add_argument('--log_path', type=str, default='./logs/', help='log file path') +parser.add_argument('--resume', action='store_true', help='resume from checkpoint') +parser.add_argument('--checkpoint', type=str, default='', help='checkpoint file path') +parser.add_argument('--checkpoint_noiser', type=str, default='', help='checkpoint file path') +parser.add_argument('--multigpu', type=bool, default=False, help='use multiple gpus') +parser.add_argument('--input_path', type=str, default='./test/', help='evaluation data file path') +parser.add_argument('--save_path', type=str, default='./output/image/', help='output data save path') + +opts = parser.parse_args() + +log_dir = os.path.join(opts.log_path, opts.config) + '/' +config = yaml.load(open('./configs/' + opts.config + '.yaml', 'r'), Loader=yaml.FullLoader) + +# Initialize trainer +trainer = Trainer(config, opts) +trainer.initialize(opts.stylegan_model_path, opts.arcface_model_path, opts.parsing_model_path) +trainer.to(device) + +state_dict = torch.load(opts.pretrained_model_path)#os.path.join(opts.log_path, opts.config + '/checkpoint.pth')) +trainer.enc.load_state_dict(torch.load(opts.pretrained_model_path)) +trainer.enc.eval() + +img_to_tensor = transforms.Compose([ + transforms.Resize((1024, 1024)), + transforms.ToTensor(), + transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) +]) + +# simple inference +image_dir = opts.input_path +save_dir = opts.save_path +os.makedirs(save_dir, exist_ok=True) + +with torch.no_grad(): + img_list = [glob.glob1(image_dir, ext) for ext in ['*jpg','*png']] + img_list = [item for sublist in img_list for item in sublist] + img_list.sort() + for i, img_name in enumerate(img_list): + #print(i, img_name) + image_A = img_to_tensor(Image.open(image_dir + img_name)).unsqueeze(0).to(device) + output = trainer.test(img=image_A, return_latent=True) + feature = output.pop() + latent = output.pop() + #np.save(save_dir + 'latent_code_%d.npy'%i, latent.cpu().numpy()) + utils.save_image(clip_img(output[1]), save_dir + img_name) + if i > 1000: + break diff --git a/models/FeatureStyleEncoder/train.py b/models/FeatureStyleEncoder/train.py new file mode 100644 index 0000000000000000000000000000000000000000..ba582aca6615a981e8b0a6b1a51bf4f3c0a3f913 --- /dev/null +++ b/models/FeatureStyleEncoder/train.py @@ -0,0 +1,149 @@ +import argparse +import os +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.utils.data as data +import yaml + +from PIL import Image +from tqdm import tqdm +from torchvision import transforms, utils +from tensorboard_logger import Logger + +from utils.datasets import * +from utils.functions import * +from trainer import * + +torch.backends.cudnn.enabled = True +torch.backends.cudnn.deterministic = True +torch.backends.cudnn.benchmark = True +torch.autograd.set_detect_anomaly(True) +Image.MAX_IMAGE_PIXELS = None +device = torch.device('cuda') + +parser = argparse.ArgumentParser() +parser.add_argument('--config', type=str, default='001', help='Path to the config file.') +parser.add_argument('--real_dataset_path', type=str, default='./data/ffhq-dataset/images/', help='dataset path') +parser.add_argument('--dataset_path', type=str, default='./data/stylegan2-generate-images/ims/', help='dataset path') +parser.add_argument('--label_path', type=str, default='./data/stylegan2-generate-images/seeds_pytorch_1.8.1.npy', help='laebl path') +parser.add_argument('--stylegan_model_path', type=str, default='./pixel2style2pixel/pretrained_models/psp_ffhq_encode.pt', help='pretrained stylegan2 model') +parser.add_argument('--arcface_model_path', type=str, default='./pretrained_models/backbone.pth', help='pretrained ArcFace model') +parser.add_argument('--parsing_model_path', type=str, default='./pretrained_models/79999_iter.pth', help='pretrained parsing model') +parser.add_argument('--log_path', type=str, default='./logs/', help='log file path') +parser.add_argument('--resume', action='store_true', help='resume from checkpoint') +parser.add_argument('--checkpoint', type=str, default='', help='checkpoint file path') +opts = parser.parse_args() + +log_dir = os.path.join(opts.log_path, opts.config) + '/' +os.makedirs(log_dir, exist_ok=True) +logger = Logger(log_dir) + +config = yaml.load(open('./configs/' + opts.config + '.yaml', 'r'), Loader=yaml.FullLoader) + +batch_size = config['batch_size'] +epochs = config['epochs'] +iter_per_epoch = config['iter_per_epoch'] +img_size = (config['resolution'], config['resolution']) +video_data_input = False + + +img_to_tensor = transforms.Compose([ + transforms.ToTensor(), + transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) +]) +img_to_tensor_car = transforms.Compose([ + transforms.Resize((384, 512)), + transforms.Pad(padding=(0, 64, 0, 64)), + transforms.ToTensor(), + transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) +]) + +# Initialize trainer +trainer = Trainer(config, opts) +trainer.initialize(opts.stylegan_model_path, opts.arcface_model_path, opts.parsing_model_path) +trainer.to(device) + +noise_exemple = trainer.noise_inputs +train_data_split = 0.9 if 'train_split' not in config else config['train_split'] + +# Load synthetic dataset +dataset_A = MyDataSet(image_dir=opts.dataset_path, label_dir=opts.label_path, output_size=img_size, noise_in=noise_exemple, training_set=True, train_split=train_data_split) +loader_A = data.DataLoader(dataset_A, batch_size=batch_size, shuffle=True, num_workers=2, pin_memory=True) +# Load real dataset +dataset_B = MyDataSet(image_dir=opts.real_dataset_path, label_dir=None, output_size=img_size, noise_in=noise_exemple, training_set=True, train_split=train_data_split) +loader_B = data.DataLoader(dataset_B, batch_size=batch_size, shuffle=True, num_workers=2, pin_memory=True) + +# Start Training +epoch_0 = 0 + +# check if checkpoint exist +if 'checkpoint.pth' in os.listdir(log_dir): + epoch_0 = trainer.load_checkpoint(os.path.join(log_dir, 'checkpoint.pth')) + +if opts.resume: + epoch_0 = trainer.load_checkpoint(os.path.join(opts.log_path, opts.checkpoint)) + +torch.manual_seed(0) +os.makedirs(log_dir + 'validation/', exist_ok=True) + +print("Start!") + +for n_epoch in tqdm(range(epoch_0, epochs)): + + iter_A = iter(loader_A) + iter_B = iter(loader_B) + iter_0 = n_epoch*iter_per_epoch + + trainer.enc_opt.zero_grad() + + for n_iter in range(iter_0, iter_0 + iter_per_epoch): + + if opts.dataset_path is None: + z, noise = next(iter_A) + img_A = None + else: + z, img_A, noise = next(iter_A) + img_A = img_A.to(device) + + z = z.to(device) + noise = [ee.to(device) for ee in noise] + w = trainer.mapping(z) + if 'fixed_noise' in config and config['fixed_noise']: + img_A, noise = None, None + + img_B = None + if 'use_realimg' in config and config['use_realimg']: + try: + img_B = next(iter_B) + if img_B.size(0) != batch_size: + iter_B = iter(loader_B) + img_B = next(iter_B) + except StopIteration: + iter_B = iter(loader_B) + img_B = next(iter_B) + img_B = img_B.to(device) + + trainer.update(w=w, img=img_A, noise=noise, real_img=img_B, n_iter=n_iter) + if (n_iter+1) % config['log_iter'] == 0: + trainer.log_loss(logger, n_iter, prefix='scripts') + if (n_iter+1) % config['image_save_iter'] == 0: + trainer.save_image(log_dir, n_epoch, n_iter, prefix='/scripts/', w=w, img=img_A, noise=noise) + trainer.save_image(log_dir, n_epoch, n_iter+1, prefix='/scripts/', w=w, img=img_B, noise=noise, training_mode=False) + + trainer.enc_scheduler.step() + trainer.save_checkpoint(n_epoch, log_dir) + + # Test the model on celeba hq dataset + with torch.no_grad(): + trainer.enc.eval() + for i in range(10): + image_A = img_to_tensor(Image.open('./data/celeba_hq/%d.jpg' % i)).unsqueeze(0).to(device) + output = trainer.test(img=image_A) + out_img = torch.cat(output, 3) + utils.save_image(clip_img(out_img[:1]), log_dir + 'validation/' + 'epoch_' +str(n_epoch+1) + '_' + str(i) + '.jpg') + trainer.compute_loss(w=w, img=img_A, noise=noise, real_img=img_B) + trainer.log_loss(logger, n_iter, prefix='validation') + +trainer.save_model(log_dir) \ No newline at end of file diff --git a/models/FeatureStyleEncoder/trainer.py b/models/FeatureStyleEncoder/trainer.py new file mode 100644 index 0000000000000000000000000000000000000000..e69cacf142a968a5be010c5ca4286c7348c3a2c4 --- /dev/null +++ b/models/FeatureStyleEncoder/trainer.py @@ -0,0 +1,431 @@ +import sys +import os +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.utils.data as data + +from PIL import Image +from torch.autograd import grad +from torchvision import transforms, utils + +import face_alignment +import lpips + +current_dir = os.path.abspath(os.path.dirname(__file__)) +sys.path.insert(0, current_dir) +from pixel2style2pixel.models.stylegan2.model import Generator, get_keys + +from nets.feature_style_encoder import * +from arcface.iresnet import * +from face_parsing.model import BiSeNet +from ranger import Ranger + +import os +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.utils.data as data + +from PIL import Image +from torch.autograd import grad + + +def clip_img(x): + """Clip stylegan generated image to range(0,1)""" + img_tmp = x.clone()[0] + img_tmp = (img_tmp + 1) / 2 + img_tmp = torch.clamp(img_tmp, 0, 1) + return [img_tmp.detach().cpu()] + +def tensor_byte(x): + return x.element_size()*x.nelement() + +def count_parameters(net): + s = sum([np.prod(list(mm.size())) for mm in net.parameters()]) + print(s) + +def stylegan_to_classifier(x, out_size=(224, 224)): + """Clip image to range(0,1)""" + img_tmp = x.clone() + img_tmp = torch.clamp((0.5*img_tmp + 0.5), 0, 1) + img_tmp = F.interpolate(img_tmp, size=out_size, mode='bilinear') + img_tmp[:,0] = (img_tmp[:,0] - 0.485)/0.229 + img_tmp[:,1] = (img_tmp[:,1] - 0.456)/0.224 + img_tmp[:,2] = (img_tmp[:,2] - 0.406)/0.225 + #img_tmp = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])(img_tmp) + return img_tmp + +def downscale(x, scale_times=1, mode='bilinear'): + for i in range(scale_times): + x = F.interpolate(x, scale_factor=0.5, mode=mode) + return x + +def upscale(x, scale_times=1, mode='bilinear'): + for i in range(scale_times): + x = F.interpolate(x, scale_factor=2, mode=mode) + return x + +def hist_transform(source_tensor, target_tensor): + """Histogram transformation""" + c, h, w = source_tensor.size() + s_t = source_tensor.view(c, -1) + t_t = target_tensor.view(c, -1) + s_t_sorted, s_t_indices = torch.sort(s_t) + t_t_sorted, t_t_indices = torch.sort(t_t) + for i in range(c): + s_t[i, s_t_indices[i]] = t_t_sorted[i] + return s_t.view(c, h, w) + +def init_weights(m): + """Initialize layers with Xavier uniform distribution""" + if type(m) == nn.Conv2d: + nn.init.xavier_uniform_(m.weight) + elif type(m) == nn.Linear: + nn.init.uniform_(m.weight, 0.0, 1.0) + if m.bias is not None: + nn.init.constant_(m.bias, 0.01) + +def total_variation(x, delta=1): + """Total variation, x: tensor of size (B, C, H, W)""" + out = torch.mean(torch.abs(x[:, :, :, :-delta] - x[:, :, :, delta:]))\ + + torch.mean(torch.abs(x[:, :, :-delta, :] - x[:, :, delta:, :])) + return out + +def vgg_transform(x): + """Adapt image for vgg network, x: image of range(0,1) subtracting ImageNet mean""" + r, g, b = torch.split(x, 1, 1) + out = torch.cat((b, g, r), dim = 1) + out = F.interpolate(out, size=(224, 224), mode='bilinear') + out = out*255. + return out + +# warp image with flow +def normalize_axis(x,L): + return (x-1-(L-1)/2)*2/(L-1) + +def unnormalize_axis(x,L): + return x*(L-1)/2+1+(L-1)/2 + +def torch_flow_to_th_sampling_grid(flow,h_src,w_src,use_cuda=False): + b,c,h_tgt,w_tgt=flow.size() + grid_y, grid_x = torch.meshgrid(torch.tensor(range(1,w_tgt+1)),torch.tensor(range(1,h_tgt+1))) + disp_x=flow[:,0,:,:] + disp_y=flow[:,1,:,:] + source_x=grid_x.unsqueeze(0).repeat(b,1,1).type_as(flow)+disp_x + source_y=grid_y.unsqueeze(0).repeat(b,1,1).type_as(flow)+disp_y + source_x_norm=normalize_axis(source_x,w_src) + source_y_norm=normalize_axis(source_y,h_src) + sampling_grid=torch.cat((source_x_norm.unsqueeze(3), source_y_norm.unsqueeze(3)), dim=3) + if use_cuda: + sampling_grid = sampling_grid.cuda() + return sampling_grid + +def warp_image_torch(image, flow): + """ + Warp image (tensor, shape=[b, 3, h_src, w_src]) with flow (tensor, shape=[b, h_tgt, w_tgt, 2]) + """ + b,c,h_src,w_src=image.size() + sampling_grid_torch = torch_flow_to_th_sampling_grid(flow, h_src, w_src) + warped_image_torch = F.grid_sample(image, sampling_grid_torch) + return warped_image_torch + +class Trainer(nn.Module): + def __init__(self, config, opts): + super(Trainer, self).__init__() + # Load Hyperparameters + self.config = config + self.device = torch.device(self.config['device']) + self.scale = int(np.log2(config['resolution']/config['enc_resolution'])) + self.scale_mode = 'bilinear' + self.opts = opts + self.n_styles = 2 * int(np.log2(config['resolution'])) - 2 + self.idx_k = 5 + if 'idx_k' in self.config: + self.idx_k = self.config['idx_k'] + if 'stylegan_version' in self.config and self.config['stylegan_version'] == 3: + self.n_styles = 16 + # Networks + in_channels = 256 + if 'in_c' in self.config: + in_channels = config['in_c'] + enc_residual = False + if 'enc_residual' in self.config: + enc_residual = self.config['enc_residual'] + enc_residual_coeff = False + if 'enc_residual_coeff' in self.config: + enc_residual_coeff = self.config['enc_residual_coeff'] + resnet_layers = [4,5,6] + if 'enc_start_layer' in self.config: + st_l = self.config['enc_start_layer'] + resnet_layers = [st_l, st_l+1, st_l+2] + if 'scale_mode' in self.config: + self.scale_mode = self.config['scale_mode'] + # Load encoder + self.stride = (self.config['fs_stride'], self.config['fs_stride']) + self.enc = fs_encoder_v2(n_styles=self.n_styles, opts=opts, residual=enc_residual, use_coeff=enc_residual_coeff, resnet_layer=resnet_layers, stride=self.stride) + + ########################## + # Other nets + self.StyleGAN = self.init_stylegan(config) + self.Arcface = iresnet50() + self.parsing_net = BiSeNet(n_classes=19) + # Optimizers + # Latent encoder + self.enc_params = list(self.enc.parameters()) + if 'freeze_iresnet' in self.config and self.config['freeze_iresnet']: + self.enc_params = list(self.enc.styles.parameters()) + if 'optimizer' in self.config and self.config['optimizer'] == 'ranger': + self.enc_opt = Ranger(self.enc_params, lr=config['lr'], betas=(config['beta_1'], config['beta_2']), weight_decay=config['weight_decay']) + else: + self.enc_opt = torch.optim.Adam(self.enc_params, lr=config['lr'], betas=(config['beta_1'], config['beta_2']), weight_decay=config['weight_decay']) + self.enc_scheduler = torch.optim.lr_scheduler.StepLR(self.enc_opt, step_size=config['step_size'], gamma=config['gamma']) + + self.fea_avg = None + + def initialize(self, stylegan_model_path, arcface_model_path, parsing_model_path): + # load StyleGAN model + stylegan_state_dict = torch.load(stylegan_model_path, map_location='cpu') + self.StyleGAN.load_state_dict(get_keys(stylegan_state_dict, 'decoder'), strict=True) + self.StyleGAN.to(self.device) + # get StyleGAN average latent in w space and the noise inputs + self.dlatent_avg = stylegan_state_dict['latent_avg'].to(self.device) + self.noise_inputs = [getattr(self.StyleGAN.noises, f'noise_{i}').to(self.device) for i in range(self.StyleGAN.num_layers)] + # load Arcface weight + self.Arcface.load_state_dict(torch.load(self.opts.arcface_model_path)) + self.Arcface.eval() + # load face parsing net weight + self.parsing_net.load_state_dict(torch.load(self.opts.parsing_model_path)) + self.parsing_net.eval() + # load lpips net weight + # self.loss_fn = lpips.LPIPS(net='alex', spatial=False) + # self.loss_fn.to(self.device) + + def init_stylegan(self, config): + """StyleGAN = G_main( + truncation_psi=config['truncation_psi'], + resolution=config['resolution'], + use_noise=config['use_noise'], + randomize_noise=config['randomize_noise'] + )""" + StyleGAN = Generator(1024, 512, 8) + return StyleGAN + + def mapping(self, z): + return self.StyleGAN.get_latent(z).detach() + + def L1loss(self, input, target): + return nn.L1Loss()(input,target) + + def L2loss(self, input, target): + return nn.MSELoss()(input,target) + + def CEloss(self, x, target_age): + return nn.CrossEntropyLoss()(x, target_age) + + def LPIPS(self, input, target, multi_scale=False): + if multi_scale: + out = 0 + for k in range(3): + out += self.loss_fn.forward(downscale(input, k, self.scale_mode), downscale(target, k, self.scale_mode)).mean() + else: + out = self.loss_fn.forward(downscale(input, self.scale, self.scale_mode), downscale(target, self.scale, self.scale_mode)).mean() + return out + + def IDloss(self, input, target): + x_1 = F.interpolate(input, (112,112)) + x_2 = F.interpolate(target, (112,112)) + cos = nn.CosineSimilarity(dim=1, eps=1e-6) + if 'multi_layer_idloss' in self.config and self.config['multi_layer_idloss']: + id_1 = self.Arcface(x_1, return_features=True) + id_2 = self.Arcface(x_2, return_features=True) + return sum([1 - cos(id_1[i].flatten(start_dim=1), id_2[i].flatten(start_dim=1)) for i in range(len(id_1))]) + else: + id_1 = self.Arcface(x_1) + id_2 = self.Arcface(x_2) + return 1 - cos(id_1, id_2) + + def landmarkloss(self, input, target): + cos = nn.CosineSimilarity(dim=1, eps=1e-6) + x_1 = stylegan_to_classifier(input, out_size=(512, 512)) + x_2 = stylegan_to_classifier(target, out_size=(512,512)) + out_1 = self.parsing_net(x_1) + out_2 = self.parsing_net(x_2) + parsing_loss = sum([1 - cos(out_1[i].flatten(start_dim=1), out_2[i].flatten(start_dim=1)) for i in range(len(out_1))]) + return parsing_loss.mean() + + + def feature_match(self, enc_feat, dec_feat, layer_idx=None): + loss = [] + if layer_idx is None: + layer_idx = [i for i in range(len(enc_feat))] + for i in layer_idx: + loss.append(self.L1loss(enc_feat[i], dec_feat[i])) + return loss + + def encode(self, img): + w_recon, fea = self.enc(downscale(img, self.scale, self.scale_mode)) + w_recon = w_recon + self.dlatent_avg + return w_recon, fea + + def get_image(self, w=None, img=None, noise=None, zero_noise_input=True, training_mode=True): + + x_1, n_1 = img, noise + if x_1 is None: + x_1, _ = self.StyleGAN([w], input_is_latent=True, noise = n_1) + + w_delta = None + fea = None + features = None + return_features = False + # Reconstruction + k = 0 + if 'use_fs_encoder' in self.config and self.config['use_fs_encoder']: + return_features = True + k = self.idx_k + w_recon, fea = self.enc(downscale(x_1, self.scale, self.scale_mode)) + w_recon = w_recon + self.dlatent_avg + features = [None]*k + [fea] + [None]*(17-k) + else: + w_recon = self.enc(downscale(x_1, self.scale, self.scale_mode)) + self.dlatent_avg + + # generate image + x_1_recon, fea_recon = self.StyleGAN([w_recon], input_is_latent=True, return_features=True, features_in=features, feature_scale=min(1.0, 0.0001*self.n_iter)) + fea_recon = fea_recon[k].detach() + return [x_1_recon, x_1[:,:3,:,:], w_recon, w_delta, n_1, fea, fea_recon] + + def compute_loss(self, w=None, img=None, noise=None, real_img=None): + return self.compute_loss_stylegan2(w=w, img=img, noise=noise, real_img=real_img) + + def compute_loss_stylegan2(self, w=None, img=None, noise=None, real_img=None): + + if img is None: + # generate synthetic images + if noise is None: + noise = [torch.randn(w.size()[:1] + ee.size()[1:]).to(self.device) for ee in self.noise_inputs] + img, _ = self.StyleGAN([w], input_is_latent=True, noise = noise) + img = img.detach() + + if img is not None and real_img is not None: + # concat synthetic and real data + img = torch.cat([img, real_img], dim=0) + noise = [torch.cat([ee, ee], dim=0) for ee in noise] + + out = self.get_image(w=w, img=img, noise=noise) + x_1_recon, x_1, w_recon, w_delta, n_1, fea_1, fea_recon = out + + # Loss setting + w_l2, w_lpips, w_id = self.config['w']['l2'], self.config['w']['lpips'], self.config['w']['id'] + b = x_1.size(0)//2 + if 'l2loss_on_real_image' in self.config and self.config['l2loss_on_real_image']: + b = x_1.size(0) + self.l2_loss = self.L2loss(x_1_recon[:b], x_1[:b]) if w_l2 > 0 else torch.tensor(0) # l2 loss only on synthetic data + # LPIPS + multiscale_lpips=False if 'multiscale_lpips' not in self.config else self.config['multiscale_lpips'] + self.lpips_loss = self.LPIPS(x_1_recon, x_1, multi_scale=multiscale_lpips).mean() if w_lpips > 0 else torch.tensor(0) + self.id_loss = self.IDloss(x_1_recon, x_1).mean() if w_id > 0 else torch.tensor(0) + self.landmark_loss = self.landmarkloss(x_1_recon, x_1) if self.config['w']['landmark'] > 0 else torch.tensor(0) + + if 'use_fs_encoder' in self.config and self.config['use_fs_encoder']: + k = self.idx_k + features = [None]*k + [fea_1] + [None]*(17-k) + x_1_recon_2, _ = self.StyleGAN([w_recon], noise=n_1, input_is_latent=True, features_in=features, feature_scale=min(1.0, 0.0001*self.n_iter)) + self.lpips_loss += self.LPIPS(x_1_recon_2, x_1, multi_scale=multiscale_lpips).mean() if w_lpips > 0 else torch.tensor(0) + self.id_loss += self.IDloss(x_1_recon_2, x_1).mean() if w_id > 0 else torch.tensor(0) + self.landmark_loss += self.landmarkloss(x_1_recon_2, x_1) if self.config['w']['landmark'] > 0 else torch.tensor(0) + + # downscale image + x_1 = downscale(x_1, self.scale, self.scale_mode) + x_1_recon = downscale(x_1_recon, self.scale, self.scale_mode) + + # Total loss + w_l2, w_lpips, w_id = self.config['w']['l2'], self.config['w']['lpips'], self.config['w']['id'] + self.loss = w_l2*self.l2_loss + w_lpips*self.lpips_loss + w_id*self.id_loss + + if 'f_recon' in self.config['w']: + self.feature_recon_loss = self.L2loss(fea_1, fea_recon) + self.loss += self.config['w']['f_recon']*self.feature_recon_loss + if 'l1' in self.config['w'] and self.config['w']['l1']>0: + self.l1_loss = self.L1loss(x_1_recon, x_1) + self.loss += self.config['w']['l1']*self.l1_loss + if 'landmark' in self.config['w']: + self.loss += self.config['w']['landmark']*self.landmark_loss + return self.loss + + def test(self, w=None, img=None, noise=None, zero_noise_input=True, return_latent=False, training_mode=False): + if 'n_iter' not in self.__dict__.keys(): + self.n_iter = 1e5 + out = self.get_image(w=w, img=img, noise=noise, training_mode=training_mode) + x_1_recon, x_1, w_recon, w_delta, n_1, fea_1 = out[:6] + output = [x_1, x_1_recon] + if return_latent: + output += [w_recon, fea_1] + return output + + def log_loss(self, logger, n_iter, prefix='scripts'): + logger.log_value(prefix + '/l2_loss', self.l2_loss.item(), n_iter + 1) + logger.log_value(prefix + '/lpips_loss', self.lpips_loss.item(), n_iter + 1) + logger.log_value(prefix + '/id_loss', self.id_loss.item(), n_iter + 1) + logger.log_value(prefix + '/total_loss', self.loss.item(), n_iter + 1) + if 'f_recon' in self.config['w']: + logger.log_value(prefix + '/feature_recon_loss', self.feature_recon_loss.item(), n_iter + 1) + if 'l1' in self.config['w'] and self.config['w']['l1']>0: + logger.log_value(prefix + '/l1_loss', self.l1_loss.item(), n_iter + 1) + if 'landmark' in self.config['w']: + logger.log_value(prefix + '/landmark_loss', self.landmark_loss.item(), n_iter + 1) + + def save_image(self, log_dir, n_epoch, n_iter, prefix='/scripts/', w=None, img=None, noise=None, training_mode=True): + return self.save_image_stylegan2(log_dir=log_dir, n_epoch=n_epoch, n_iter=n_iter, prefix=prefix, w=w, img=img, noise=noise, training_mode=training_mode) + + def save_image_stylegan2(self, log_dir, n_epoch, n_iter, prefix='/scripts/', w=None, img=None, noise=None, training_mode=True): + os.makedirs(log_dir + prefix, exist_ok=True) + with torch.no_grad(): + out = self.get_image(w=w, img=img, noise=noise, training_mode=training_mode) + x_1_recon, x_1, w_recon, w_delta, n_1, fea_1 = out[:6] + x_1 = downscale(x_1, self.scale, self.scale_mode) + x_1_recon = downscale(x_1_recon, self.scale, self.scale_mode) + out_img = torch.cat((x_1, x_1_recon), dim=3) + #fs + if 'use_fs_encoder' in self.config and self.config['use_fs_encoder']: + k = self.idx_k + features = [None]*k + [fea_1] + [None]*(17-k) + x_1_recon_2, _ = self.StyleGAN([w_recon], noise=n_1, input_is_latent=True, features_in=features, feature_scale=min(1.0, 0.0001*self.n_iter)) + x_1_recon_2 = downscale(x_1_recon_2, self.scale, self.scale_mode) + out_img = torch.cat((x_1, x_1_recon, x_1_recon_2), dim=3) + utils.save_image(clip_img(out_img[:1]), log_dir + prefix + 'epoch_' +str(n_epoch+1) + '_iter_' + str(n_iter+1) + '_0.jpg') + if out_img.size(0)>1: + utils.save_image(clip_img(out_img[1:]), log_dir + prefix + 'epoch_' +str(n_epoch+1) + '_iter_' + str(n_iter+1) + '_1.jpg') + + def save_model(self, log_dir): + torch.save(self.enc.state_dict(),'{:s}/enc.pth.tar'.format(log_dir)) + + def save_checkpoint(self, n_epoch, log_dir): + checkpoint_state = { + 'n_epoch': n_epoch, + 'enc_state_dict': self.enc.state_dict(), + 'enc_opt_state_dict': self.enc_opt.state_dict(), + 'enc_scheduler_state_dict': self.enc_scheduler.state_dict() + } + torch.save(checkpoint_state, '{:s}/checkpoint.pth'.format(log_dir)) + if (n_epoch+1)%10 == 0 : + torch.save(checkpoint_state, '{:s}/checkpoint'.format(log_dir)+'_'+str(n_epoch+1)+'.pth') + + def load_model(self, log_dir): + self.enc.load_state_dict(torch.load('{:s}/enc.pth.tar'.format(log_dir))) + + def load_checkpoint(self, checkpoint_path): + state_dict = torch.load(checkpoint_path) + self.enc.load_state_dict(state_dict['enc_state_dict']) + self.enc_opt.load_state_dict(state_dict['enc_opt_state_dict']) + self.enc_scheduler.load_state_dict(state_dict['enc_scheduler_state_dict']) + return state_dict['n_epoch'] + 1 + + def update(self, w=None, img=None, noise=None, real_img=None, n_iter=0): + self.n_iter = n_iter + self.enc_opt.zero_grad() + self.compute_loss(w=w, img=img, noise=noise, real_img=real_img).backward() + self.enc_opt.step() + + diff --git a/models/FeatureStyleEncoder/utils/.DS_Store b/models/FeatureStyleEncoder/utils/.DS_Store new file mode 100644 index 0000000000000000000000000000000000000000..f9853bccefe6d4cf619b94db2b287e29ebd0dc06 Binary files /dev/null and b/models/FeatureStyleEncoder/utils/.DS_Store differ diff --git a/models/FeatureStyleEncoder/utils/datasets.py b/models/FeatureStyleEncoder/utils/datasets.py new file mode 100644 index 0000000000000000000000000000000000000000..d1cbb377efe9eea031123431ff31f1c6a7aab497 --- /dev/null +++ b/models/FeatureStyleEncoder/utils/datasets.py @@ -0,0 +1,158 @@ +import os +import glob +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.utils.data as data + +from PIL import Image +from torchvision import transforms, utils + +class MyDataSet(data.Dataset): + def __init__(self, image_dir=None, label_dir=None, output_size=(256, 256), noise_in=None, training_set=True, video_data=False, train_split=0.9): + self.image_dir = image_dir + self.normalize = transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) + self.resize = transforms.Compose([ + transforms.Resize(output_size), + transforms.ToTensor() + ]) + self.noise_in = noise_in + self.video_data = video_data + self.random_rotation = transforms.Compose([ + transforms.Resize(output_size), + transforms.RandomPerspective(distortion_scale=0.05, p=1.0), + transforms.ToTensor() + ]) + + # load image file + train_len = None + self.length = 0 + self.image_dir = image_dir + if image_dir is not None: + img_list = [glob.glob1(self.image_dir, ext) for ext in ['*jpg','*png']] + image_list = [item for sublist in img_list for item in sublist] + image_list.sort() + train_len = int(train_split*len(image_list)) + if training_set: + self.image_list = image_list[:train_len] + else: + self.image_list = image_list[train_len:] + self.length = len(self.image_list) + + # load label file + self.label_dir = label_dir + if label_dir is not None: + self.seeds = np.load(label_dir) + if train_len is None: + train_len = int(train_split*len(self.seeds)) + if training_set: + self.seeds = self.seeds[:train_len] + else: + self.seeds = self.seeds[train_len:] + if self.length == 0: + self.length = len(self.seeds) + + def __len__(self): + return self.length + + def __getitem__(self, idx): + img = None + if self.image_dir is not None: + img_name = os.path.join(self.image_dir, self.image_list[idx]) + image = Image.open(img_name) + img = self.resize(image) + if img.size(0) == 1: + img = torch.cat((img, img, img), dim=0) + img = self.normalize(img) + + # generate image + if self.label_dir is not None: + torch.manual_seed(self.seeds[idx]) + z = torch.randn(1, 512)[0] + if self.noise_in is None: + n = [torch.randn(1, 1)] + else: + n = [torch.randn(noise.size())[0] for noise in self.noise_in] + if img is None: + return z, n + else: + return z, img, n + else: + return img + +class Car_DataSet(data.Dataset): + def __init__(self, image_dir=None, label_dir=None, output_size=(512, 512), noise_in=None, training_set=True, video_data=False, train_split=0.9): + self.image_dir = image_dir + self.normalize = transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) + self.resize = transforms.Compose([ + transforms.Resize((384, 512)), + transforms.Pad(padding=(0, 64, 0, 64)), + transforms.ToTensor() + ]) + self.noise_in = noise_in + self.video_data = video_data + self.random_rotation = transforms.Compose([ + transforms.Resize(output_size), + transforms.RandomPerspective(distortion_scale=0.05, p=1.0), + transforms.ToTensor() + ]) + + # load image file + train_len = None + self.length = 0 + self.image_dir = image_dir + if image_dir is not None: + img_list = [glob.glob1(self.image_dir, ext) for ext in ['*jpg','*png']] + image_list = [item for sublist in img_list for item in sublist] + image_list.sort() + train_len = int(train_split*len(image_list)) + if training_set: + self.image_list = image_list[:train_len] + else: + self.image_list = image_list[train_len:] + self.length = len(self.image_list) + + # load label file + self.label_dir = label_dir + if label_dir is not None: + self.seeds = np.load(label_dir) + if train_len is None: + train_len = int(train_split*len(self.seeds)) + if training_set: + self.seeds = self.seeds[:train_len] + else: + self.seeds = self.seeds[train_len:] + if self.length == 0: + self.length = len(self.seeds) + + def __len__(self): + return self.length + + def __getitem__(self, idx): + img = None + if self.image_dir is not None: + img_name = os.path.join(self.image_dir, self.image_list[idx]) + image = Image.open(img_name) + img = self.resize(image) + if img.size(0) == 1: + img = torch.cat((img, img, img), dim=0) + img = self.normalize(img) + if self.video_data: + img_2 = self.random_rotation(image) + img_2 = self.normalize(img_2) + img_2 = torch.where(img_2 > -1, img_2, img) + img = torch.cat([img, img_2], dim=0) + + # generate image + if self.label_dir is not None: + torch.manual_seed(self.seeds[idx]) + z = torch.randn(1, 512)[0] + n = [torch.randn_like(noise[0]) for noise in self.noise_in] + if img is None: + return z, n + else: + return z, img, n + else: + return img + diff --git a/models/FeatureStyleEncoder/utils/functions.py b/models/FeatureStyleEncoder/utils/functions.py new file mode 100644 index 0000000000000000000000000000000000000000..2e3fa0ee305302762763830fe7f70d175ad84d59 --- /dev/null +++ b/models/FeatureStyleEncoder/utils/functions.py @@ -0,0 +1,109 @@ +import os +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.utils.data as data + +from PIL import Image +from torch.autograd import grad + + +def clip_img(x): + """Clip stylegan generated image to range(0,1)""" + img_tmp = x.clone()[0] + img_tmp = (img_tmp + 1) / 2 + img_tmp = torch.clamp(img_tmp, 0, 1) + return [img_tmp.detach().cpu()] + +def tensor_byte(x): + return x.element_size()*x.nelement() + +def count_parameters(net): + s = sum([np.prod(list(mm.size())) for mm in net.parameters()]) + print(s) + +def stylegan_to_classifier(x, out_size=(224, 224)): + """Clip image to range(0,1)""" + img_tmp = x.clone() + img_tmp = torch.clamp((0.5*img_tmp + 0.5), 0, 1) + img_tmp = F.interpolate(img_tmp, size=out_size, mode='bilinear') + img_tmp[:,0] = (img_tmp[:,0] - 0.485)/0.229 + img_tmp[:,1] = (img_tmp[:,1] - 0.456)/0.224 + img_tmp[:,2] = (img_tmp[:,2] - 0.406)/0.225 + #img_tmp = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])(img_tmp) + return img_tmp + +def downscale(x, scale_times=1, mode='bilinear'): + for i in range(scale_times): + x = F.interpolate(x, scale_factor=0.5, mode=mode) + return x + +def upscale(x, scale_times=1, mode='bilinear'): + for i in range(scale_times): + x = F.interpolate(x, scale_factor=2, mode=mode) + return x + +def hist_transform(source_tensor, target_tensor): + """Histogram transformation""" + c, h, w = source_tensor.size() + s_t = source_tensor.view(c, -1) + t_t = target_tensor.view(c, -1) + s_t_sorted, s_t_indices = torch.sort(s_t) + t_t_sorted, t_t_indices = torch.sort(t_t) + for i in range(c): + s_t[i, s_t_indices[i]] = t_t_sorted[i] + return s_t.view(c, h, w) + +def init_weights(m): + """Initialize layers with Xavier uniform distribution""" + if type(m) == nn.Conv2d: + nn.init.xavier_uniform_(m.weight) + elif type(m) == nn.Linear: + nn.init.uniform_(m.weight, 0.0, 1.0) + if m.bias is not None: + nn.init.constant_(m.bias, 0.01) + +def total_variation(x, delta=1): + """Total variation, x: tensor of size (B, C, H, W)""" + out = torch.mean(torch.abs(x[:, :, :, :-delta] - x[:, :, :, delta:]))\ + + torch.mean(torch.abs(x[:, :, :-delta, :] - x[:, :, delta:, :])) + return out + +def vgg_transform(x): + """Adapt image for vgg network, x: image of range(0,1) subtracting ImageNet mean""" + r, g, b = torch.split(x, 1, 1) + out = torch.cat((b, g, r), dim = 1) + out = F.interpolate(out, size=(224, 224), mode='bilinear') + out = out*255. + return out + +# warp image with flow +def normalize_axis(x,L): + return (x-1-(L-1)/2)*2/(L-1) + +def unnormalize_axis(x,L): + return x*(L-1)/2+1+(L-1)/2 + +def torch_flow_to_th_sampling_grid(flow,h_src,w_src,use_cuda=False): + b,c,h_tgt,w_tgt=flow.size() + grid_y, grid_x = torch.meshgrid(torch.tensor(range(1,w_tgt+1)),torch.tensor(range(1,h_tgt+1))) + disp_x=flow[:,0,:,:] + disp_y=flow[:,1,:,:] + source_x=grid_x.unsqueeze(0).repeat(b,1,1).type_as(flow)+disp_x + source_y=grid_y.unsqueeze(0).repeat(b,1,1).type_as(flow)+disp_y + source_x_norm=normalize_axis(source_x,w_src) + source_y_norm=normalize_axis(source_y,h_src) + sampling_grid=torch.cat((source_x_norm.unsqueeze(3), source_y_norm.unsqueeze(3)), dim=3) + if use_cuda: + sampling_grid = sampling_grid.cuda() + return sampling_grid + +def warp_image_torch(image, flow): + """ + Warp image (tensor, shape=[b, 3, h_src, w_src]) with flow (tensor, shape=[b, h_tgt, w_tgt, 2]) + """ + b,c,h_src,w_src=image.size() + sampling_grid_torch = torch_flow_to_th_sampling_grid(flow, h_src, w_src) + warped_image_torch = F.grid_sample(image, sampling_grid_torch) + return warped_image_torch \ No newline at end of file diff --git a/models/FeatureStyleEncoder/utils/video_utils.py b/models/FeatureStyleEncoder/utils/video_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..01f0874a0e88dec0c8404dfdaab92a761f2323bf --- /dev/null +++ b/models/FeatureStyleEncoder/utils/video_utils.py @@ -0,0 +1,454 @@ +# Copyright (c) 2021, InterDigital R&D France. All rights reserved. +# +# This source code is made available under the license found in the +# LICENSE.txt in the root directory of this source tree. + +import cv2 +import glob +import numpy as np +import os +import face_alignment +import torch + +from PIL import Image, ImageFilter +from scipy import ndimage +from scipy.ndimage import gaussian_filter1d +from skimage import io +from torchvision import transforms, utils + + +def pil_to_cv2(pil_image): + open_cv_image = np.array(pil_image) + return open_cv_image[:, :, ::-1].copy() + + +def cv2_to_pil(open_cv_image): + return Image.fromarray(open_cv_image[:, :, ::-1].copy()) + + +def put_text(img, text): + font = cv2.FONT_HERSHEY_SIMPLEX + bottomLeftCornerOfText = (10,50) + fontScale = 1.5 + fontColor = (255,255,0) + lineType = 2 + return cv2.putText(img, text, + bottomLeftCornerOfText, + font, + fontScale, + fontColor, + lineType) + + +# Compare frames in two directory +def compare_frames(save_dir, origin_dir, target_dir, strs='Original,Projected,Manipulated', dim=None): + + os.makedirs(save_dir, exist_ok=True) + try: + if not isinstance(target_dir, list): + target_dir = [target_dir] + image_list = glob.glob1(origin_dir,'frame*') + image_list.sort() + for name in image_list: + img_l = [] + for idx, dir_path in enumerate([origin_dir] + list(target_dir)): + img_1 = cv2.imread(dir_path + name) + img_1 = put_text(img_1, strs.split(',')[idx]) + img_l.append(img_1) + img = np.concatenate(img_l, axis=1) + cv2.imwrite(save_dir + name, img) + except FileNotFoundError: + pass + + +# Save frames into video +def create_video(image_folder, fps=24, video_format='.mp4', resize_ratio=1): + + video_name = os.path.dirname(image_folder) + video_format + img_list = glob.glob1(image_folder,'frame*') + img_list.sort() + frame = cv2.imread(os.path.join(image_folder, img_list[0])) + frame = cv2.resize(frame, (0,0), fx=resize_ratio, fy=resize_ratio) + height, width, layers = frame.shape + if video_format == '.mp4': + fourcc = cv2.VideoWriter_fourcc(*'mp4v') + elif video_format == '.avi': + fourcc = cv2.VideoWriter_fourcc(*'XVID') + video = cv2.VideoWriter(video_name, fourcc, fps, (width,height)) + for image_name in img_list: + frame = cv2.imread(os.path.join(image_folder, image_name)) + frame = cv2.resize(frame, (0,0), fx=resize_ratio, fy=resize_ratio) + video.write(frame) + + +# Split video into frames +def video_to_frames(video_path, frame_path, img_format='.jpg', count_num=1000, resize=False): + + os.makedirs(frame_path, exist_ok=True) + vidcap = cv2.VideoCapture(video_path) + success,image = vidcap.read() + count = 0 + while success: + if resize: + image = cv2.resize(image, (0,0), fx=0.5, fy=0.5) + cv2.imwrite(frame_path + '/frame%04d' % count + img_format, image) + success,image = vidcap.read() + count += 1 + if count >= count_num: + break + +# Align faces +def align_frames(img_dir, save_dir, output_size=1024, transform_size=1024, optical_flow=True, gaussian=True, filter_size=3): + + os.makedirs(save_dir, exist_ok=True) + + # load face landmark detector + fa = face_alignment.FaceAlignment(face_alignment.LandmarksType._2D, flip_input=False, device='cuda') + + # list images in the directory + img_list = glob.glob1(img_dir, 'frame*') + img_list.sort() + + # save align statistics + stat_dict = {'quad':[], 'qsize':[], 'coord':[], 'crop':[]} + lms = [] + for idx, img_name in enumerate(img_list): + + img_path = os.path.join(img_dir, img_name) + img = io.imread(img_path) + lm = [] + + preds = fa.get_landmarks(img) + for kk in range(68): + lm.append((preds[0][kk][0], preds[0][kk][1])) + + # Eye distance + lm_eye_left = lm[36 : 42] # left-clockwise + lm_eye_right = lm[42 : 48] # left-clockwise + eye_left = np.mean(lm_eye_left, axis=0) + eye_right = np.mean(lm_eye_right, axis=0) + eye_to_eye = eye_right - eye_left + + if optical_flow: + if idx > 0: + s = int(np.hypot(*eye_to_eye)/4) + lk_params = dict(winSize=(s, s), maxLevel=5, criteria = (cv2.TERM_CRITERIA_COUNT | cv2.TERM_CRITERIA_EPS, 10, 0.03)) + points_arr = np.array(lm, np.float32) + points_prevarr = np.array(prev_lm, np.float32) + points_arr,status, err = cv2.calcOpticalFlowPyrLK(prev_img, img, points_prevarr, points_arr, **lk_params) + sigma =100 + points_arr_float = np.array(points_arr,np.float32) + points = points_arr_float.tolist() + for k in range(0, len(lm)): + d = cv2.norm(np.array(prev_lm[k]) - np.array(lm[k])) + alpha = np.exp(-d*d/sigma) + lm[k] = (1 - alpha) * np.array(lm[k]) + alpha * np.array(points[k]) + prev_img = img + prev_lm = lm + + lms.append(lm) + + # Apply gaussian filter on landmarks + if gaussian: + lm_filtered = np.array(lms) + for kk in range(68): + lm_filtered[:, kk, 0] = gaussian_filter1d(lm_filtered[:, kk, 0], filter_size) + lm_filtered[:, kk, 1] = gaussian_filter1d(lm_filtered[:, kk, 1], filter_size) + lms = lm_filtered.tolist() + + # save landmarks + landmark_out_dir = os.path.dirname(img_dir) + '_landmark/' + os.makedirs(landmark_out_dir, exist_ok=True) + + for idx, img_name in enumerate(img_list): + + img_path = os.path.join(img_dir, img_name) + img = io.imread(img_path) + + lm = lms[idx] + img_lm = img.copy() + for kk in range(68): + img_lm = cv2.circle(img_lm, (int(lm[kk][0]),int(lm[kk][1])), radius=3, color=(255, 0, 255), thickness=-1) + # Save landmark images + cv2.imwrite(landmark_out_dir + img_name, img_lm[:,:,::-1]) + + # Save mask images + """ + seg_mask = np.zeros(img.shape, img.dtype) + poly = np.array(lm[0:17] + lm[17:27][::-1], np.int32) + cv2.fillPoly(seg_mask, [poly], (255, 255, 255)) + cv2.imwrite(img_dir + "mask%04d.jpg"%idx, seg_mask); + """ + + # Parse landmarks. + lm_eye_left = lm[36 : 42] # left-clockwise + lm_eye_right = lm[42 : 48] # left-clockwise + lm_mouth_outer = lm[48 : 60] # left-clockwise + + # Calculate auxiliary vectors. + eye_left = np.mean([lm_eye_left[0], lm_eye_left[3]], axis=0) + eye_right = np.mean([lm_eye_right[0], lm_eye_right[3]], axis=0) + eye_avg = (eye_left + eye_right) * 0.5 + eye_to_eye = eye_right - eye_left + mouth_left = np.array(lm_mouth_outer[0]) + mouth_right = np.array(lm_mouth_outer[6]) + mouth_avg = (mouth_left + mouth_right) * 0.5 + eye_to_mouth = mouth_avg - eye_avg + + # Choose oriented crop rectangle. + x = eye_to_eye - np.flipud(eye_to_mouth) * [-1, 1] + x /= np.hypot(*x) + x *= max(np.hypot(*eye_to_eye) * 2.0, np.hypot(*eye_to_mouth) * 1.8) + y = np.flipud(x) * [-1, 1] + c = eye_avg + eye_to_mouth * 0.1 + quad = np.stack([c - x - y, c - x + y, c + x + y, c + x - y]) + qsize = np.hypot(*x) * 2 + + stat_dict['coord'].append(quad) + stat_dict['qsize'].append(qsize) + + # Apply gaussian filter on crops + if gaussian: + quads = np.array(stat_dict['coord']) + quads = gaussian_filter1d(quads, 2*filter_size, axis=0) + stat_dict['coord'] = quads.tolist() + qsize = np.array(stat_dict['qsize']) + qsize = gaussian_filter1d(qsize, 2*filter_size, axis=0) + stat_dict['qsize'] = qsize.tolist() + + for idx, img_name in enumerate(img_list): + + img_path = os.path.join(img_dir, img_name) + img = Image.open(img_path) + + qsize = stat_dict['qsize'][idx] + quad = np.array(stat_dict['coord'][idx]) + + # Crop. + border = max(int(np.rint(qsize * 0.1)), 3) + crop = (int(np.floor(min(quad[:,0]))), int(np.floor(min(quad[:,1]))), int(np.ceil(max(quad[:,0]))), int(np.ceil(max(quad[:,1])))) + crop = (max(crop[0] - border, 0), max(crop[1] - border, 0), min(crop[2] + border, img.size[0]), min(crop[3] + border, img.size[1])) + if crop[2] - crop[0] < img.size[0] or crop[3] - crop[1] < img.size[1]: + img = img.crop(crop) + quad -= crop[0:2] + + stat_dict['crop'].append(crop) + stat_dict['quad'].append((quad + 0.5).flatten()) + + # Pad. + pad = (int(np.floor(min(quad[:,0]))), int(np.floor(min(quad[:,1]))), int(np.ceil(max(quad[:,0]))), int(np.ceil(max(quad[:,1])))) + pad = (max(-pad[0] + border, 0), max(-pad[1] + border, 0), max(pad[2] - img.size[0] + border, 0), max(pad[3] - img.size[1] + border, 0)) + if max(pad) > border - 4: + pad = np.maximum(pad, int(np.rint(qsize * 0.3))) + img = np.pad(np.float32(img), ((pad[1], pad[3]), (pad[0], pad[2]), (0, 0)), 'reflect') + h, w, _ = img.shape + y, x, _ = np.ogrid[:h, :w, :1] + img = Image.fromarray(np.uint8(np.clip(np.rint(img), 0, 255)), 'RGB') + quad += pad[:2] + # Transform. + img = img.transform((transform_size, transform_size), Image.QUAD, (quad + 0.5).flatten(), Image.BILINEAR) + + # resizing + img_pil = img.resize((output_size, output_size), Image.LANCZOS) + img_pil.save(save_dir+img_name) + + create_video(landmark_out_dir) + np.save(save_dir+'stat_dict.npy', stat_dict) + + +def find_coeffs(pa, pb): + + matrix = [] + for p1, p2 in zip(pa, pb): + matrix.append([p1[0], p1[1], 1, 0, 0, 0, -p2[0]*p1[0], -p2[0]*p1[1]]) + matrix.append([0, 0, 0, p1[0], p1[1], 1, -p2[1]*p1[0], -p2[1]*p1[1]]) + A = np.matrix(matrix, dtype=np.float) + B = np.array(pb).reshape(8) + res = np.dot(np.linalg.inv(A.T * A) * A.T, B) + return np.array(res).reshape(8) + +# reproject aligned frames to the original video +def video_reproject(orig_dir_path, recon_dir_path, save_dir_path, state_dir_path, mask_dir_path, seamless=False): + + if not os.path.exists(save_dir_path): + os.makedirs(save_dir_path) + + img_list_0 = glob.glob1(orig_dir_path,'frame*') + img_list_2 = glob.glob1(recon_dir_path,'frame*') + img_list_0.sort() + img_list_2.sort() + stat_dict = np.load(state_dir_path + 'stat_dict.npy', allow_pickle=True).item() + counter = len(img_list_2) + + for idx in range(counter): + + img_0 = Image.open(orig_dir_path + img_list_0[idx]) + img_2 = Image.open(recon_dir_path + img_list_2[idx]) + + quad_f = stat_dict['quad'][idx] + quad_0 = stat_dict['crop'][idx] + + coeffs = find_coeffs( + [(quad_f[0], quad_f[1]), (quad_f[2] , quad_f[3]), (quad_f[4], quad_f[5]), (quad_f[6], quad_f[7])], + [(0, 0), (0, 1024), (1024, 1024), (1024, 0)]) + crop_size = (quad_0[2] - quad_0[0], quad_0[3] - quad_0[1]) + img_2 = img_2.transform(crop_size, Image.PERSPECTIVE, coeffs, Image.BICUBIC) + output = img_0.copy() + output.paste(img_2, (int(quad_0[0]), int(quad_0[1]))) + + """ + mask = cv2.imread(orig_dir_path + 'mask%04d.jpg'%idx) + kernel = np.ones((10,10), np.uint8) + mask = cv2.dilate(mask, kernel, iterations=5) + """ + crop_mask = Image.open(mask_dir_path + img_list_0[idx]) + crop_mask = crop_mask.transform(crop_size, Image.PERSPECTIVE, coeffs, Image.BICUBIC) + mask = Image.fromarray(np.zeros(np.array(img_0).shape, np.array(img_0).dtype)) + mask.paste(crop_mask, (int(quad_0[0]), int(quad_0[1]))) + mask = pil_to_cv2(mask) + # Apply mask + if not seamless: + mask = cv2_to_pil(mask).filter(ImageFilter.GaussianBlur(radius=10)).convert('L') + mask = np.array(mask)[:, :, np.newaxis]/255. + output = np.array(img_0)*(1-mask) + np.array(output)*mask + output = Image.fromarray(output.astype(np.uint8)) + output.save(save_dir_path + img_list_2[idx]) + else: + src = pil_to_cv2(output) + dst = pil_to_cv2(img_0) + # clone + br = cv2.boundingRect(cv2.split(mask)[0]) # bounding rect (x,y,width,height) + center = (br[0] + br[2] // 2, br[1] + br[3] // 2) + output = cv2.seamlessClone(src, dst, mask, center, cv2.NORMAL_CLONE) + cv2.imwrite(save_dir_path + img_list_2[idx], output) + + + + +# Align faces +def align_image(img_dir, save_dir, output_size=1024, transform_size=1024, format='*.png'): + os.makedirs(save_dir, exist_ok=True) + + # load face landmark detector + fa = face_alignment.FaceAlignment(face_alignment.LandmarksType._2D, flip_input=False, device='cuda') + # list images in the directory + img_list = glob.glob1(img_dir, format) + #img_list = os.listdir(img_dir) + img_list.sort() + + # save align statistics + stat_dict = {'quad':[], 'qsize':[], 'coord':[], 'crop':[]} + + for idx, img_name in enumerate(img_list): + + img_path = os.path.join(img_dir, img_name) + img = Image.open(img_path).convert('RGB') + img_np = np.array(img) + lm = [] + + preds = fa.get_landmarks(img_np) + for kk in range(68): + lm.append((preds[0][kk][0], preds[0][kk][1])) + if len(lm)==0: + continue + + # Parse landmarks. Code extracted from ffhq-dataset + # pylint: disable=unused-variable + lm_chin = lm[0 : 17] # left-right + lm_eyebrow_left = lm[17 : 22] # left-right + lm_eyebrow_right = lm[22 : 27] # left-right + lm_nose = lm[27 : 31] # top-down + lm_nostrils = lm[31 : 36] # top-down + lm_eye_left = lm[36 : 42] # left-clockwise + lm_eye_right = lm[42 : 48] # left-clockwise + lm_mouth_outer = lm[48 : 60] # left-clockwise + lm_mouth_inner = lm[60 : 68] # left-clockwise + + # Calculate auxiliary vectors. + eye_left = np.mean([lm_eye_left[0], lm_eye_left[3]], axis=0) + eye_right = np.mean([lm_eye_right[0], lm_eye_right[3]], axis=0) + eye_avg = (eye_left + eye_right) * 0.5 + eye_to_eye = eye_right - eye_left + mouth_left = np.array(lm_mouth_outer[0]) + mouth_right = np.array(lm_mouth_outer[6]) + mouth_avg = (mouth_left + mouth_right) * 0.5 + eye_to_mouth = mouth_avg - eye_avg + + # Choose oriented crop rectangle. + x = eye_to_eye - np.flipud(eye_to_mouth) * [-1, 1] + x /= np.hypot(*x) + x *= np.hypot(*eye_to_eye) * 2.0#max(np.hypot(*eye_to_eye) * 2.0, np.hypot(*eye_to_mouth) * 1.8) + + y = np.flipud(x) * [-1, 1] + c = eye_avg + eye_to_mouth * 0.1 + quad = np.stack([c - x - y, c - x + y, c + x + y, c + x - y]) + qsize = np.hypot(*x) * 2 + + stat_dict['coord'].append(quad) + stat_dict['qsize'].append(qsize) + + qsize = stat_dict['qsize'][idx] + quad = np.array(stat_dict['coord'][idx]) + """ + # Shrink. + shrink = int(np.floor(qsize / output_size * 0.5)) + if shrink > 1: + print('shrink!') + rsize = (int(np.rint(float(img.size[0]) / shrink)), int(np.rint(float(img.size[1]) / shrink))) + img = img.resize(rsize, Image.ANTIALIAS) + quad /= shrink + qsize /= shrink + """ + # Crop. + border = max(int(np.rint(qsize * 0.1)), 3) + crop = (int(np.floor(min(quad[:,0]))), int(np.floor(min(quad[:,1]))), int(np.ceil(max(quad[:,0]))), int(np.ceil(max(quad[:,1])))) + crop = (max(crop[0] - border, 0), max(crop[1] - border, 0), min(crop[2] + border, img.size[0]), min(crop[3] + border, img.size[1])) + if crop[2] - crop[0] < img.size[0] or crop[3] - crop[1] < img.size[1]: + img = img.crop(crop) + quad -= crop[0:2] + + stat_dict['crop'].append(crop) + stat_dict['quad'].append((quad + 0.5).flatten()) + #img = img.crop(crop) + # Pad. + pad = (int(np.floor(min(quad[:,0]))), int(np.floor(min(quad[:,1]))), int(np.ceil(max(quad[:,0]))), int(np.ceil(max(quad[:,1])))) + pad = (max(-pad[0] + border, 0), max(-pad[1] + border, 0), max(pad[2] - img.size[0] + border, 0), max(pad[3] - img.size[1] + border, 0)) + if max(pad) > border - 4: + pad = np.maximum(pad, int(np.rint(qsize * 0.3))) + img = np.pad(np.float32(img), ((pad[1], pad[3]), (pad[0], pad[2]), (0, 0)), 'edge') + h, w, _ = img.shape + y, x, _ = np.ogrid[:h, :w, :1] + img = Image.fromarray(np.uint8(np.clip(np.rint(img), 0, 255)), 'RGB') + quad += pad[:2] + # Transform. + img = img.transform((transform_size, transform_size), Image.QUAD, (quad + 0.5).flatten(), Image.BILINEAR) + img_pil = img.resize((output_size, output_size), Image.LANCZOS) + + # resizing + img_pil.save(save_dir+img_name) + + np.save(save_dir+'stat_dict.npy', stat_dict) + +img_to_tensor = transforms.Compose([ + transforms.ToTensor(), + transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) +]) + +def generate_mask(img_dir, save_dir, parsing_net, labels=[1,2,3,4,5,6,9,10,11,12,13], output_size=(1024, 1024), device=torch.device('cuda')): + os.makedirs(save_dir, exist_ok=True) + img_list = glob.glob1(img_dir, 'frame*') + img_list.sort() + + for img_name in img_list: + img_path = os.path.join(img_dir, img_name) + img = Image.open(img_path).resize((512, 512), Image.LANCZOS) + x_1 = img_to_tensor(img).unsqueeze(0).to(device) + out_1 = parsing_net(x_1) + parsing = out_1[0].squeeze(0).detach().cpu().numpy().argmax(0) + mask = np.uint8(parsing) + for j in labels: + mask = np.where(mask==j, 255, mask) + mask = np.where(mask==255, 255, 0) + mask_pil = Image.fromarray(np.uint8(mask)).resize(output_size, Image.LANCZOS) + save_path = os.path.join(save_dir, img_name) + mask_pil.save(save_path) \ No newline at end of file diff --git a/models/FeatureStyleEncoder/video_processing.py b/models/FeatureStyleEncoder/video_processing.py new file mode 100644 index 0000000000000000000000000000000000000000..e0fd231e87501d37ba9ba74899a1790731fff4aa --- /dev/null +++ b/models/FeatureStyleEncoder/video_processing.py @@ -0,0 +1,174 @@ +# Copyright (c) 2021, InterDigital R&D France. All rights reserved. +# +# This source code is made available under the license found in the +# LICENSE.txt in the root directory of this source tree. + +import argparse +import copy +import glob +import numpy as np +import os +import torch +import yaml +import time + +from PIL import Image +from torchvision import transforms, utils, models + +from utils.video_utils import * +from face_parsing.model import BiSeNet +from trainer import * + + +torch.backends.cudnn.enabled = True +torch.backends.cudnn.deterministic = True +torch.backends.cudnn.benchmark = True +torch.autograd.set_detect_anomaly(True) +Image.MAX_IMAGE_PIXELS = None +device = torch.device('cuda') + +parser = argparse.ArgumentParser() +parser.add_argument('--config', type=str, default='001', help='Path to the config file.') +parser.add_argument('--attr', type=str, default='Eyeglasses', help='attribute for manipulation.') +parser.add_argument('--alpha', type=str, default='1.', help='scale for manipulation.') +parser.add_argument('--label_file', type=str, default='./data/celebahq_anno.npy', help='label file path') +parser.add_argument('--pretrained_model_path', type=str, default='./pretrained_models/143_enc.pth', help='pretrained stylegan2 model') +parser.add_argument('--stylegan_model_path', type=str, default='./pixel2style2pixel/pretrained_models/psp_ffhq_encode.pt', help='pretrained stylegan2 model') +parser.add_argument('--arcface_model_path', type=str, default='./pretrained_models/backbone.pth', help='pretrained ArcFace model') +parser.add_argument('--parsing_model_path', type=str, default='./pretrained_models/79999_iter.pth', help='pretrained parsing model') +parser.add_argument('--log_path', type=str, default='./logs/', help='log file path') +parser.add_argument('--function', type=str, default='', help='Calling function by name.') +parser.add_argument('--video_path', type=str, default='./data/video/FP006911MD02.mp4', help='video file path') +parser.add_argument('--output_path', type=str, default='./output/video/', help='output video file path') +parser.add_argument('--boundary_path', type=str, default='./boundaries_ours/', help='output video file path') +parser.add_argument('--optical_flow', action='store_true', help='use optical flow') +parser.add_argument('--resize', action='store_true', help='downscale image size') +parser.add_argument('--seamless', action='store_true', help='seamless cloning') +parser.add_argument('--filter_size', type=float, default=3, help='filter size') +parser.add_argument('--strs', type=str, default='Original,Projected,Manipulated', help='strs to be added on video') +opts = parser.parse_args() + +# Celeba attribute list +attr_dict = {'5_o_Clock_Shadow': 0, 'Arched_Eyebrows': 1, 'Attractive': 2, 'Bags_Under_Eyes': 3, \ + 'Bald': 4, 'Bangs': 5, 'Big_Lips': 6, 'Big_Nose': 7, 'Black_Hair': 8, 'Blond_Hair': 9, \ + 'Brown_Hair': 11, 'Bushy_Eyebrows': 12, 'Chubby': 13, 'Double_Chin': 14, \ + 'Eyeglasses': 15, 'Goatee': 16, 'Gray_Hair': 17, 'Heavy_Makeup': 18, 'High_Cheekbones': 19, \ + 'Male': 20, 'Mouth_Slightly_Open': 21, 'Mustache': 22, 'Narrow_Eyes': 23, 'No_Beard': 24, \ + 'Oval_Face': 25, 'Pale_Skin': 26, 'Pointy_Nose': 27, 'Receding_Hairline': 28, 'Rosy_Cheeks': 29, \ + 'Sideburns': 30, 'Smiling': 31, 'Straight_Hair': 32, 'Wavy_Hair': 33, 'Wearing_Earrings': 34, \ + 'Wearing_Hat': 35, 'Wearing_Lipstick': 36, 'Wearing_Necklace': 37, 'Wearing_Necktie': 38, 'Young': 39} + +img_to_tensor = transforms.Compose([ + transforms.Resize((1024, 1024)), + transforms.ToTensor(), + transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) +]) + +# linear interpolation +def linear_interpolate(latent_code, + boundary, + start_distance=-3.0, + end_distance=3.0, + steps=10): + assert (latent_code.shape[0] == 1 and boundary.shape[0] == 1 and + len(boundary.shape) == 2 and + boundary.shape[1] == latent_code.shape[-1]) + + linspace = np.linspace(start_distance, end_distance, steps) + if len(latent_code.shape) == 2: + linspace = linspace.reshape(-1, 1).astype(np.float32) + return latent_code + linspace * boundary + if len(latent_code.shape) == 3: + linspace = linspace.reshape(-1, 1, 1).astype(np.float32) + return latent_code + linspace * boundary.reshape(1, 1, -1) + +# Latent code manipulation +def latent_manipulation(opts, align_dir_path, process_dir_path): + + os.makedirs(process_dir_path, exist_ok=True) + #attrs = opts.attr.split(',') + #alphas = opts.alpha.split(',') + step_scale = 15 * int(opts.alpha) + n_steps = 5 + + boundary = np.load(opts.boundary_path +'%s_boundary.npy'%opts.attr) + + # Initialize trainer + config = yaml.load(open('./configs/' + opts.config + '.yaml', 'r'), Loader=yaml.FullLoader) + trainer = Trainer(config, opts) + trainer.initialize(opts.stylegan_model_path, opts.arcface_model_path, opts.parsing_model_path) + trainer.to(device) + + state_dict = torch.load(opts.pretrained_model_path)#os.path.join(opts.log_path, opts.config + '/checkpoint.pth')) + trainer.enc.load_state_dict(torch.load(opts.pretrained_model_path)) + trainer.enc.eval() + + with torch.no_grad(): + img_list = [glob.glob1(align_dir_path, ext) for ext in ['*jpg','*png']] + img_list = [item for sublist in img_list for item in sublist] + img_list.sort() + n_1 = trainer.StyleGAN.make_noise() + + for i, img_name in enumerate(img_list): + #print(i, img_name) + image_A = img_to_tensor(Image.open(align_dir_path + img_name)).unsqueeze(0).to(device) + w_0, f_0 = trainer.encode(image_A) + + w_0_np = w_0.cpu().numpy().reshape(1, -1) + out = linear_interpolate(w_0_np, boundary, start_distance=-step_scale, end_distance=step_scale, steps=n_steps) + w_1 = torch.tensor(out[-1]).view(1, -1, 512).to(device) + + _, fea_0 = trainer.StyleGAN([w_0], noise=n_1, input_is_latent=True, return_features=True) + _, fea_1 = trainer.StyleGAN([w_1], noise=n_1, input_is_latent=True, return_features=True) + + features = [None]*5 + [f_0 + fea_1[5] - fea_0[5]] + [None]*(17-5) + x_1, _ = trainer.StyleGAN([w_1], noise=n_1, input_is_latent=True, features_in=features, feature_scale=1.0) + utils.save_image(clip_img(x_1), process_dir_path + 'frame%04d'%i+'.jpg') + + +video_path = opts.video_path +video_name = video_path.split('/')[-1] +orig_dir_path = opts.output_path + video_name.split('.')[0] + '/' + video_name.split('.')[0] + '/' +align_dir_path = os.path.dirname(orig_dir_path) + '_crop_align/' +mask_dir_path = os.path.dirname(orig_dir_path) + '_crop_align_mask/' +latent_dir_path = os.path.dirname(orig_dir_path) + '_crop_align_latent/' +process_dir_path = os.path.dirname(orig_dir_path) + '_crop_align_' + opts.attr.replace(',','_') + '/' +reproject_dir_path = os.path.dirname(orig_dir_path) + '_crop_align_' + opts.attr.replace(',','_') + '_reproject/' + + +print(opts.function) +start_time = time.perf_counter() + +if opts.function == 'video_to_frames': + video_to_frames(video_path, orig_dir_path, count_num=120, resize=opts.resize) + create_video(orig_dir_path) +elif opts.function == 'align_frames': + align_frames(orig_dir_path, align_dir_path, output_size=1024, optical_flow=opts.optical_flow, filter_size=opts.filter_size) + # parsing mask + parsing_net = BiSeNet(n_classes=19) + parsing_net.load_state_dict(torch.load(opts.parsing_model_path)) + parsing_net.eval() + parsing_net.to(device) + generate_mask(align_dir_path, mask_dir_path, parsing_net) +elif opts.function == 'latent_manipulation': + latent_manipulation(opts, align_dir_path, process_dir_path) +elif opts.function == 'reproject_origin': + process_dir_path = os.path.dirname(orig_dir_path) + '_inversion/' + reproject_dir_path = os.path.dirname(orig_dir_path) + '_inversion_reproject/' + video_reproject(orig_dir_path, process_dir_path, reproject_dir_path, align_dir_path, mask_dir_path, seamless=opts.seamless) + create_video(reproject_dir_path) +elif opts.function == 'reproject_manipulate': + video_reproject(orig_dir_path, process_dir_path, reproject_dir_path, align_dir_path, mask_dir_path, seamless=opts.seamless) + create_video(reproject_dir_path) +elif opts.function == 'compare_frames': + process_dir_paths = [] + process_dir_paths.append(os.path.dirname(orig_dir_path) + '_inversion_reproject/') + if len(opts.attr.split(','))>0: + process_dir_paths.append(reproject_dir_path) + save_dir = os.path.dirname(orig_dir_path) + '_crop_align_' + opts.attr.replace(',','_') + '_compare/' + compare_frames(save_dir, orig_dir_path, process_dir_paths, strs=opts.strs, dim=1) + create_video(save_dir, video_format='.avi', resize_ratio=1) + + +count_time = time.perf_counter() - start_time +print("Elapsed time: %0.4f seconds"%count_time) \ No newline at end of file diff --git a/models/Net.py b/models/Net.py new file mode 100644 index 0000000000000000000000000000000000000000..6bd1684d5ea1f751735817cccbc78129823c8eed --- /dev/null +++ b/models/Net.py @@ -0,0 +1,484 @@ +import os + +import numpy as np +import torch +from torch import nn +from torchvision import transforms + +from models.CtrlHair.external_code.face_parsing.my_parsing_util import FaceParsing_tensor +from models.stylegan2.model import Generator +from utils.drive import download_weight + +transform_to_256 = transforms.Compose([ + transforms.Resize((256, 256)), +]) + +__all__ = ['Net', 'iresnet18', 'iresnet34', 'iresnet50', 'iresnet100', 'iresnet200', 'FeatureEncoderMult', + 'IBasicBlock', 'conv1x1', 'get_segmentation'] + + +class Net(nn.Module): + + def __init__(self, opts): + super(Net, self).__init__() + self.opts = opts + self.generator = Generator(opts.size, opts.latent, opts.n_mlp, channel_multiplier=opts.channel_multiplier) + self.cal_layer_num() + self.load_weights() + self.load_PCA_model() + FaceParsing_tensor.parsing_img() + + def load_weights(self): + if not os.path.exists(self.opts.ckpt): + print('Downloading StyleGAN2 checkpoint: {}'.format(self.opts.ckpt)) + download_weight(self.opts.ckpt) + + print('Loading StyleGAN2 from checkpoint: {}'.format(self.opts.ckpt)) + checkpoint = torch.load(self.opts.ckpt) + device = self.opts.device + self.generator.load_state_dict(checkpoint['g_ema']) + self.latent_avg = checkpoint['latent_avg'] + self.generator.to(device) + self.latent_avg = self.latent_avg.to(device) + + for param in self.generator.parameters(): + param.requires_grad = False + self.generator.eval() + + def build_PCA_model(self, PCA_path): + + with torch.no_grad(): + latent = torch.randn((1000000, 512), dtype=torch.float32) + # latent = torch.randn((10000, 512), dtype=torch.float32) + self.generator.style.cpu() + pulse_space = torch.nn.LeakyReLU(5)(self.generator.style(latent)).numpy() + self.generator.style.to(self.opts.device) + + from utils.PCA_utils import IPCAEstimator + + transformer = IPCAEstimator(512) + X_mean = pulse_space.mean(0) + transformer.fit(pulse_space - X_mean) + X_comp, X_stdev, X_var_ratio = transformer.get_components() + np.savez(PCA_path, X_mean=X_mean, X_comp=X_comp, X_stdev=X_stdev, X_var_ratio=X_var_ratio) + + def load_PCA_model(self): + device = self.opts.device + + PCA_path = self.opts.ckpt[:-3] + '_PCA.npz' + + if not os.path.isfile(PCA_path): + self.build_PCA_model(PCA_path) + + PCA_model = np.load(PCA_path) + self.X_mean = torch.from_numpy(PCA_model['X_mean']).float().to(device) + self.X_comp = torch.from_numpy(PCA_model['X_comp']).float().to(device) + self.X_stdev = torch.from_numpy(PCA_model['X_stdev']).float().to(device) + + # def make_noise(self): + # noises_single = self.generator.make_noise() + # noises = [] + # for noise in noises_single: + # noises.append(noise.repeat(1, 1, 1, 1).normal_()) + # + # return noises + + def cal_layer_num(self): + if self.opts.size == 1024: + self.layer_num = 18 + elif self.opts.size == 512: + self.layer_num = 16 + elif self.opts.size == 256: + self.layer_num = 14 + + self.S_index = self.layer_num - 11 + + return + + def cal_p_norm_loss(self, latent_in): + latent_p_norm = (torch.nn.LeakyReLU(negative_slope=5)(latent_in) - self.X_mean).bmm( + self.X_comp.T.unsqueeze(0)) / self.X_stdev + p_norm_loss = self.opts.p_norm_lambda * (latent_p_norm.pow(2).mean()) + return p_norm_loss + + def cal_l_F(self, latent_F, F_init): + return self.opts.l_F_lambda * (latent_F - F_init).pow(2).mean() + + +def get_segmentation(img_rgb, resize=True): + parsing, _ = FaceParsing_tensor.parsing_img(img_rgb) + parsing = FaceParsing_tensor.swap_parsing_label_to_celeba_mask(parsing) + mask_img = parsing.long()[None, None, ...] + if resize: + mask_img = transforms.functional.resize(mask_img, (256, 256), + interpolation=transforms.InterpolationMode.NEAREST) + return mask_img + + +fs_kernals = { + 0: (12, 12), + 1: (12, 12), + 2: (6, 6), + 3: (6, 6), + 4: (3, 3), + 5: (3, 3), + 6: (3, 3), + 7: (3, 3), +} + +fs_strides = { + 0: (7, 7), + 1: (7, 7), + 2: (4, 4), + 3: (4, 4), + 4: (2, 2), + 5: (2, 2), + 6: (1, 1), + 7: (1, 1), +} + + +def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1): + """3x3 convolution with padding""" + return nn.Conv2d(in_planes, + out_planes, + kernel_size=3, + stride=stride, + padding=dilation, + groups=groups, + bias=False, + dilation=dilation) + + +def conv1x1(in_planes, out_planes, stride=1): + """1x1 convolution""" + return nn.Conv2d(in_planes, + out_planes, + kernel_size=1, + stride=stride, + bias=False) + + +class IBasicBlock(nn.Module): + expansion = 1 + + def __init__(self, inplanes, planes, stride=1, downsample=None, + groups=1, base_width=64, dilation=1): + super(IBasicBlock, self).__init__() + if groups != 1 or base_width != 64: + raise ValueError('BasicBlock only supports groups=1 and base_width=64') + if dilation > 1: + raise NotImplementedError("Dilation > 1 not supported in BasicBlock") + self.bn1 = nn.BatchNorm2d(inplanes, eps=1e-05, ) + self.conv1 = conv3x3(inplanes, planes) + self.bn2 = nn.BatchNorm2d(planes, eps=1e-05, ) + self.prelu = nn.PReLU(planes) + self.conv2 = conv3x3(planes, planes, stride) + self.bn3 = nn.BatchNorm2d(planes, eps=1e-05, ) + self.downsample = downsample + self.stride = stride + + def forward(self, x): + identity = x + out = self.bn1(x) + out = self.conv1(out) + out = self.bn2(out) + out = self.prelu(out) + out = self.conv2(out) + out = self.bn3(out) + if self.downsample is not None: + identity = self.downsample(x) + out += identity + return out + + +class IResNet(nn.Module): + fc_scale = 7 * 7 + + def __init__(self, + block, layers, dropout=0, num_features=512, zero_init_residual=False, + groups=1, width_per_group=64, replace_stride_with_dilation=None, fp16=False): + super(IResNet, self).__init__() + self.fp16 = fp16 + self.inplanes = 64 + self.dilation = 1 + if replace_stride_with_dilation is None: + replace_stride_with_dilation = [False, False, False] + if len(replace_stride_with_dilation) != 3: + raise ValueError("replace_stride_with_dilation should be None " + "or a 3-element tuple, got {}".format(replace_stride_with_dilation)) + self.groups = groups + self.base_width = width_per_group + self.conv1 = nn.Conv2d(3, self.inplanes, kernel_size=3, stride=1, padding=1, bias=False) + self.bn1 = nn.BatchNorm2d(self.inplanes, eps=1e-05) + self.prelu = nn.PReLU(self.inplanes) + self.layer1 = self._make_layer(block, 64, layers[0], stride=2) + self.layer2 = self._make_layer(block, + 128, + layers[1], + stride=2, + dilate=replace_stride_with_dilation[0]) + self.layer3 = self._make_layer(block, + 256, + layers[2], + stride=2, + dilate=replace_stride_with_dilation[1]) + self.layer4 = self._make_layer(block, + 512, + layers[3], + stride=2, + dilate=replace_stride_with_dilation[2]) + self.bn2 = nn.BatchNorm2d(512 * block.expansion, eps=1e-05, ) + self.dropout = nn.Dropout(p=dropout, inplace=True) + self.fc = nn.Linear(512 * block.expansion * self.fc_scale, num_features) + self.features = nn.BatchNorm1d(num_features, eps=1e-05) + nn.init.constant_(self.features.weight, 1.0) + self.features.weight.requires_grad = False + + for m in self.modules(): + if isinstance(m, nn.Conv2d): + nn.init.normal_(m.weight, 0, 0.1) + elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)): + nn.init.constant_(m.weight, 1) + nn.init.constant_(m.bias, 0) + + if zero_init_residual: + for m in self.modules(): + if isinstance(m, IBasicBlock): + nn.init.constant_(m.bn2.weight, 0) + + def _make_layer(self, block, planes, blocks, stride=1, dilate=False): + downsample = None + previous_dilation = self.dilation + if dilate: + self.dilation *= stride + stride = 1 + if stride != 1 or self.inplanes != planes * block.expansion: + downsample = nn.Sequential( + conv1x1(self.inplanes, planes * block.expansion, stride), + nn.BatchNorm2d(planes * block.expansion, eps=1e-05, ), + ) + layers = [] + layers.append( + block(self.inplanes, planes, stride, downsample, self.groups, + self.base_width, previous_dilation)) + self.inplanes = planes * block.expansion + for _ in range(1, blocks): + layers.append( + block(self.inplanes, + planes, + groups=self.groups, + base_width=self.base_width, + dilation=self.dilation)) + + return nn.Sequential(*layers) + + def forward(self, x, return_features=False): + out = [] + with torch.cuda.amp.autocast(self.fp16): + x = self.conv1(x) + x = self.bn1(x) + x = self.prelu(x) + x = self.layer1(x) + out.append(x) + x = self.layer2(x) + out.append(x) + x = self.layer3(x) + out.append(x) + x = self.layer4(x) + out.append(x) + x = self.bn2(x) + x = torch.flatten(x, 1) + x = self.dropout(x) + x = self.fc(x.float() if self.fp16 else x) + x = self.features(x) + + if return_features: + out.append(x) + return out + return x + + +def _iresnet(arch, block, layers, pretrained, progress, **kwargs): + model = IResNet(block, layers, **kwargs) + if pretrained: + raise ValueError() + return model + + +def iresnet18(pretrained=False, progress=True, **kwargs): + return _iresnet('iresnet18', IBasicBlock, [2, 2, 2, 2], pretrained, + progress, **kwargs) + + +def iresnet34(pretrained=False, progress=True, **kwargs): + return _iresnet('iresnet34', IBasicBlock, [3, 4, 6, 3], pretrained, + progress, **kwargs) + + +def iresnet50(pretrained=False, progress=True, **kwargs): + return _iresnet('iresnet50', IBasicBlock, [3, 4, 14, 3], pretrained, + progress, **kwargs) + + +def iresnet100(pretrained=False, progress=True, **kwargs): + return _iresnet('iresnet100', IBasicBlock, [3, 13, 30, 3], pretrained, + progress, **kwargs) + + +def iresnet200(pretrained=False, progress=True, **kwargs): + return _iresnet('iresnet200', IBasicBlock, [6, 26, 60, 6], pretrained, + progress, **kwargs) + + +class FeatureEncoder(nn.Module): + def __init__(self, n_styles=18, opts=None, residual=False, + use_coeff=False, resnet_layer=None, + video_input=False, f_maps=512, stride=(1, 1)): + super(FeatureEncoder, self).__init__() + + resnet50 = iresnet50() + resnet50.load_state_dict(torch.load(opts.arcface_model_path)) + + # input conv layer + if video_input: + self.conv = nn.Sequential( + nn.Conv2d(6, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False), + *list(resnet50.children())[1:3] + ) + else: + self.conv = nn.Sequential(*list(resnet50.children())[:3]) + + # define layers + self.block_1 = list(resnet50.children())[3] # 15-18 + self.block_2 = list(resnet50.children())[4] # 10-14 + self.block_3 = list(resnet50.children())[5] # 5-9 + self.block_4 = list(resnet50.children())[6] # 1-4 + self.content_layer = nn.Sequential( + nn.BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True), + nn.Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False), + nn.BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True), + nn.PReLU(num_parameters=512), + nn.Conv2d(512, 512, kernel_size=(3, 3), stride=stride, padding=(1, 1), bias=False), + nn.BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True) + ) + + self.avg_pool = nn.AdaptiveAvgPool2d((3, 3)) + self.styles = nn.ModuleList() + for i in range(n_styles): + self.styles.append(nn.Linear(960 * 9, 512)) + + def apply_head(self, x): + latents = [] + for i in range(len(self.styles)): + latents.append(self.styles[i](x)) + out = torch.stack(latents, dim=1) + return out + + def forward(self, x): + latents = [] + features = [] + x = self.conv(x) + x = self.block_1(x) + features.append(self.avg_pool(x)) + x = self.block_2(x) + features.append(self.avg_pool(x)) + x = self.block_3(x) + content = self.content_layer(x) + features.append(self.avg_pool(x)) + x = self.block_4(x) + features.append(self.avg_pool(x)) + x = torch.cat(features, dim=1) + x = x.view(x.size(0), -1) + return self.apply_head(x), content + + +class FeatureEncoderMult(FeatureEncoder): + def __init__(self, fs_layers=(5,), ranks=None, **kwargs): + super().__init__(**kwargs) + + self.fs_layers = fs_layers + self.content_layer = nn.ModuleList() + self.ranks = ranks + shift = 0 if max(fs_layers) <= 7 else 2 + scale = 1 if max(fs_layers) <= 7 else 2 + for i in range(len(fs_layers)): + if ranks is not None: + stride, kern = ranks_data[ranks[i] - shift] + layer1 = nn.Sequential( + nn.BatchNorm2d(256 // scale, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True), + nn.Conv2d(256 // scale, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False), + nn.BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True), + nn.PReLU(num_parameters=512), + nn.Conv2d(512, 512, kernel_size=(fs_kernals[fs_layers[i] - shift][0], kern), + stride=(fs_strides[fs_layers[i] - shift][0], stride), + padding=(1, 1), bias=False), + nn.BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True) + ) + layer2 = nn.Sequential( + nn.BatchNorm2d(256 // scale, eps=1e-05, momentum=0.1, affine=True, + track_running_stats=True), + nn.Conv2d(256 // scale, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), + bias=False), + nn.BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, + track_running_stats=True), + nn.PReLU(num_parameters=512), + nn.Conv2d(512, 512, kernel_size=(kern, fs_kernals[fs_layers[i] - shift][1]), + stride=(stride, fs_strides[fs_layers[i] - shift][1]), + padding=(1, 1), bias=False), + nn.BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, + track_running_stats=True) + ) + layer = nn.ModuleList([layer1, layer2]) + else: + layer = nn.Sequential( + nn.BatchNorm2d(256 // scale, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True), + nn.Conv2d(256 // scale, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False), + nn.BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True), + nn.PReLU(num_parameters=512), + nn.Conv2d(512, 512, kernel_size=fs_kernals[fs_layers[i] - shift], + stride=fs_strides[fs_layers[i] - shift], + padding=(1, 1), bias=False), + nn.BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True) + ) + self.content_layer.append(layer) + + def forward(self, x): + x = transform_to_256(x) + features = [] + content = [] + x = self.conv(x) + x = self.block_1(x) + features.append(self.avg_pool(x)) + x = self.block_2(x) + if max(self.fs_layers) > 7: + for layer in self.content_layer: + if self.ranks is not None: + mat1 = layer[0](x) + mat2 = layer[1](x) + content.append(torch.matmul(mat1, mat2)) + else: + content.append(layer(x)) + features.append(self.avg_pool(x)) + x = self.block_3(x) + if len(content) == 0: + for layer in self.content_layer: + if self.ranks is not None: + mat1 = layer[0](x) + mat2 = layer[1](x) + content.append(torch.matmul(mat1, mat2)) + else: + content.append(layer(x)) + features.append(self.avg_pool(x)) + x = self.block_4(x) + features.append(self.avg_pool(x)) + x = torch.cat(features, dim=1) + x = x.view(x.size(0), -1) + return self.apply_head(x), content + + +def get_keys(d, name, key="state_dict"): + if key in d: + d = d[key] + d_filt = {k[len(name) + 1:]: v for k, v in d.items() if k[: len(name) + 1] == name + '.'} + return d_filt diff --git a/models/STAR/.DS_Store b/models/STAR/.DS_Store new file mode 100644 index 0000000000000000000000000000000000000000..4f0e9920a4ea9e5ab371e5d28a21d90745a36d19 Binary files /dev/null and b/models/STAR/.DS_Store differ diff --git a/models/STAR/README.md b/models/STAR/README.md new file mode 100644 index 0000000000000000000000000000000000000000..68abf0f118dd6641b30e820f1e25b21be3c538f1 --- /dev/null +++ b/models/STAR/README.md @@ -0,0 +1,110 @@ +# STAR Loss: Reducing Semantic Ambiguity in Facial Landmark Detection. + +Paper Link: [arxiv](https://arxiv.org/abs/2306.02763) | [CVPR 2023](https://openaccess.thecvf.com/content/CVPR2023/papers/Zhou_STAR_Loss_Reducing_Semantic_Ambiguity_in_Facial_Landmark_Detection_CVPR_2023_paper.pdf) + + +- Pytorch implementation of **S**elf-adap**T**ive **A**mbiguity **R**eduction (**STAR**) loss. +- STAR loss is a self-adaptive anisotropic direction loss, which can be used in heatmap regression-based methods for facial landmark detection. +- Specifically, we find that semantic ambiguity results in the anisotropic predicted distribution, which inspires us to use predicted distribution to represent semantic ambiguity. So, we use PCA to indicate the character of the predicted distribution and indirectly formulate the direction and intensity of semantic ambiguity. Based on this, STAR loss adaptively suppresses the prediction error in the ambiguity direction to mitigate the impact of ambiguity annotation in training. More details can be found in our paper. +

+ +

+ + + + +## Dependencies + +* python==3.7.3 +* PyTorch=1.6.0 +* requirements.txt + +## Dataset Preparation + + - Step1: Download the raw images from [COFW](http://www.vision.caltech.edu/xpburgos/ICCV13/#dataset), [300W](https://ibug.doc.ic.ac.uk/resources/300-W/), and [WFLW](https://wywu.github.io/projects/LAB/WFLW.html). + - Step2: We follow the data preprocess in [ADNet](https://openaccess.thecvf.com/content/ICCV2021/papers/Huang_ADNet_Leveraging_Error-Bias_Towards_Normal_Direction_in_Face_Alignment_ICCV_2021_paper.pdf), and the metadata can be download from [the corresponding repository](https://github.com/huangyangyu/ADNet). + - Step3: Make them look like this: +```script +# the dataset directory: +|-- ${image_dir} + |-- WFLW + | -- WFLW_images + |-- 300W + | -- afw + | -- helen + | -- ibug + | -- lfpw + |-- COFW + | -- train + | -- test +|-- ${annot_dir} + |-- WFLW + |-- train.tsv, test.tsv + |-- 300W + |-- train.tsv, test.tsv + |--COFW + |-- train.tsv, test.tsv +``` + +## Usage +* Work directory: set the ${ckpt_dir} in ./conf/alignment.py. +* Pretrained model: + +| Dataset | Model | +|:-----------------------------------------------------------------|:--------------------------------------------------------------------------------------------------------------------------------------------------------------------| +| WFLW | [google](https://drive.google.com/file/d/1aOx0wYEZUfBndYy_8IYszLPG_D2fhxrT/view?usp=sharing) / [baidu](https://pan.baidu.com/s/10vvI-ovs3x9NrdmpnXK6sg?pwd=u0yu) | +| 300W | [google](https://drive.google.com/file/d/1Fiu3hjjkQRdKsWE9IgyNPdiJSz9_MzA5/view?usp=sharing) / [baidu](https://pan.baidu.com/s/1bjUhLq1zS1XSl1nX78fU7A?pwd=yb2s) | +| COFW | [google](https://drive.google.com/file/d/1NFcZ9jzql_jnn3ulaSzUlyhS05HWB9n_/view?usp=drive_link) / [baidu](https://pan.baidu.com/s/1XO6hDZ8siJLTgFcpyu1Tzw?pwd=m57n) | + + +### Training +```shell +python main.py --mode=train --device_ids=0,1,2,3 \ + --image_dir=${image_dir} --annot_dir=${annot_dir} \ + --data_definition={WFLW, 300W, COFW} +``` + +### Testing +```shell +python main.py --mode=test --device_ids=0 \ + --image_dir=${image_dir} --annot_dir=${annot_dir} \ + --data_definition={WFLW, 300W, COFW} \ + --pretrained_weight=${model_path} \ +``` + +### Evaluation +```shell +python evaluate.py --device_ids=0 \ + --model_path=${model_path} --metadata_path=${metadata_path} \ + --image_dir=${image_dir} --data_definition={WFLW, 300W, COFW} \ +``` + +To test on your own image, the following code could be considered: +```shell +python demo.py +``` + + +## Results +The models trained by STAR Loss achieved **SOTA** performance in all of COFW, 300W and WFLW datasets. + +

+ +

+ +## BibTeX Citation +Please consider citing our papers in your publications if the project helps your research. BibTeX reference is as follows. +``` +@inproceedings{Zhou_2023_CVPR, + author = {Zhou, Zhenglin and Li, Huaxia and Liu, Hong and Wang, Nanyang and Yu, Gang and Ji, Rongrong}, + title = {STAR Loss: Reducing Semantic Ambiguity in Facial Landmark Detection}, + booktitle = {Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)}, + month = {June}, + year = {2023}, + pages = {15475-15484} +} +``` + +## Acknowledgments +This repository is built on top of [ADNet](https://github.com/huangyangyu/ADNet). +Thanks for this strong baseline. diff --git a/models/STAR/conf/__init__.py b/models/STAR/conf/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..2f92d0e82f402d1599c16deb5d1f0c3bb568bfb3 --- /dev/null +++ b/models/STAR/conf/__init__.py @@ -0,0 +1 @@ +from .alignment import Alignment \ No newline at end of file diff --git a/models/STAR/conf/alignment.py b/models/STAR/conf/alignment.py new file mode 100644 index 0000000000000000000000000000000000000000..34adf12a2ab7d14f189cac4b842cf81e19b65938 --- /dev/null +++ b/models/STAR/conf/alignment.py @@ -0,0 +1,228 @@ +import os.path as osp +from .base import Base + + +class Alignment(Base): + """ + Alignment configure file, which contains training parameters of alignment. + """ + + def __init__(self, args): + super(Alignment, self).__init__('alignment') + self.ckpt_dir = '/apdcephfs_cq3/share_1134483/charlinzhou/ckpts/STAR/' + self.net = "stackedHGnet_v1" + self.nstack = 4 + self.loader_type = "alignment" + self.data_definition = "WFLW" # COFW, 300W, WFLW + self.test_file = "test.tsv" + + # image + self.channels = 3 + self.width = 256 + self.height = 256 + self.means = (127.5, 127.5, 127.5) + self.scale = 1 / 127.5 + self.aug_prob = 1.0 + + self.display_iteration = 10 + self.val_epoch = 1 + self.valset = "test.tsv" + self.norm_type = 'default' + self.encoder_type = 'default' + self.decoder_type = 'default' + + # scheduler & optimizer + self.milestones = [200, 350, 450] + self.max_epoch = 500 + self.optimizer = "adam" + self.learn_rate = 0.001 + self.weight_decay = 0.00001 + self.betas = [0.9, 0.999] + self.gamma = 0.1 + + # batch_size & workers + self.batch_size = 32 + self.train_num_workers = 16 + self.val_batch_size = 32 + self.val_num_workers = 16 + self.test_batch_size = 16 + self.test_num_workers = 0 + + # tricks + self.ema = True + self.add_coord = True + self.use_AAM = True + + # loss + self.loss_func = "STARLoss_v2" + + # STAR Loss paras + self.star_w = 1 + self.star_dist = 'smoothl1' + + self.init_from_args(args) + + # COFW + if self.data_definition == "COFW": + self.edge_info = ( + (True, (0, 4, 2, 5)), # RightEyebrow + (True, (1, 6, 3, 7)), # LeftEyebrow + (True, (8, 12, 10, 13)), # RightEye + (False, (9, 14, 11, 15)), # LeftEye + (True, (18, 20, 19, 21)), # Nose + (True, (22, 26, 23, 27)), # LowerLip + (True, (22, 24, 23, 25)), # UpperLip + ) + if self.norm_type == 'ocular': + self.nme_left_index = 8 # ocular + self.nme_right_index = 9 # ocular + elif self.norm_type in ['pupil', 'default']: + self.nme_left_index = 16 # pupil + self.nme_right_index = 17 # pupil + else: + raise NotImplementedError + self.classes_num = [29, 7, 29] + self.crop_op = True + self.flip_mapping = ( + [0, 1], [4, 6], [2, 3], [5, 7], [8, 9], [10, 11], [12, 14], [16, 17], [13, 15], [18, 19], [22, 23], + ) + self.image_dir = osp.join(self.image_dir, 'COFW') + # 300W + elif self.data_definition == "300W": + self.edge_info = ( + (False, (0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16)), # FaceContour + (False, (17, 18, 19, 20, 21)), # RightEyebrow + (False, (22, 23, 24, 25, 26)), # LeftEyebrow + (False, (27, 28, 29, 30)), # NoseLine + (False, (31, 32, 33, 34, 35)), # Nose + (True, (36, 37, 38, 39, 40, 41)), # RightEye + (True, (42, 43, 44, 45, 46, 47)), # LeftEye + (True, (48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59)), # OuterLip + (True, (60, 61, 62, 63, 64, 65, 66, 67)), # InnerLip + ) + if self.norm_type in ['ocular', 'default']: + self.nme_left_index = 36 # ocular + self.nme_right_index = 45 # ocular + elif self.norm_type == 'pupil': + self.nme_left_index = [36, 37, 38, 39, 40, 41] # pupil + self.nme_right_index = [42, 43, 44, 45, 46, 47] # pupil + else: + raise NotImplementedError + self.classes_num = [68, 9, 68] + self.crop_op = True + self.flip_mapping = ( + [0, 16], [1, 15], [2, 14], [3, 13], [4, 12], [5, 11], [6, 10], [7, 9], + [17, 26], [18, 25], [19, 24], [20, 23], [21, 22], + [31, 35], [32, 34], + [36, 45], [37, 44], [38, 43], [39, 42], [40, 47], [41, 46], + [48, 54], [49, 53], [50, 52], [61, 63], [60, 64], [67, 65], [58, 56], [59, 55], + ) + self.image_dir = osp.join(self.image_dir, '300W') + # 300VW + elif self.data_definition == "300VW": + self.edge_info = ( + (False, (0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16)), # FaceContour + (False, (17, 18, 19, 20, 21)), # RightEyebrow + (False, (22, 23, 24, 25, 26)), # LeftEyebrow + (False, (27, 28, 29, 30)), # NoseLine + (False, (31, 32, 33, 34, 35)), # Nose + (True, (36, 37, 38, 39, 40, 41)), # RightEye + (True, (42, 43, 44, 45, 46, 47)), # LeftEye + (True, (48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59)), # OuterLip + (True, (60, 61, 62, 63, 64, 65, 66, 67)), # InnerLip + ) + if self.norm_type in ['ocular', 'default']: + self.nme_left_index = 36 # ocular + self.nme_right_index = 45 # ocular + elif self.norm_type == 'pupil': + self.nme_left_index = [36, 37, 38, 39, 40, 41] # pupil + self.nme_right_index = [42, 43, 44, 45, 46, 47] # pupil + else: + raise NotImplementedError + self.classes_num = [68, 9, 68] + self.crop_op = True + self.flip_mapping = ( + [0, 16], [1, 15], [2, 14], [3, 13], [4, 12], [5, 11], [6, 10], [7, 9], + [17, 26], [18, 25], [19, 24], [20, 23], [21, 22], + [31, 35], [32, 34], + [36, 45], [37, 44], [38, 43], [39, 42], [40, 47], [41, 46], + [48, 54], [49, 53], [50, 52], [61, 63], [60, 64], [67, 65], [58, 56], [59, 55], + ) + self.image_dir = osp.join(self.image_dir, '300VW_Dataset_2015_12_14') + # WFLW + elif self.data_definition == "WFLW": + self.edge_info = ( + (False, ( + 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, + 27, + 28, 29, 30, 31, 32)), # FaceContour + (True, (33, 34, 35, 36, 37, 38, 39, 40, 41)), # RightEyebrow + (True, (42, 43, 44, 45, 46, 47, 48, 49, 50)), # LeftEyebrow + (False, (51, 52, 53, 54)), # NoseLine + (False, (55, 56, 57, 58, 59)), # Nose + (True, (60, 61, 62, 63, 64, 65, 66, 67)), # RightEye + (True, (68, 69, 70, 71, 72, 73, 74, 75)), # LeftEye + (True, (76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87)), # OuterLip + (True, (88, 89, 90, 91, 92, 93, 94, 95)), # InnerLip + ) + if self.norm_type in ['ocular', 'default']: + self.nme_left_index = 60 # ocular + self.nme_right_index = 72 # ocular + elif self.norm_type == 'pupil': + self.nme_left_index = 96 # pupils + self.nme_right_index = 97 # pupils + else: + raise NotImplementedError + self.classes_num = [98, 9, 98] + self.crop_op = True + self.flip_mapping = ( + [0, 32], [1, 31], [2, 30], [3, 29], [4, 28], [5, 27], [6, 26], [7, 25], [8, 24], [9, 23], [10, 22], + [11, 21], [12, 20], [13, 19], [14, 18], [15, 17], # cheek + [33, 46], [34, 45], [35, 44], [36, 43], [37, 42], [38, 50], [39, 49], [40, 48], [41, 47], # elbrow + [60, 72], [61, 71], [62, 70], [63, 69], [64, 68], [65, 75], [66, 74], [67, 73], + [55, 59], [56, 58], + [76, 82], [77, 81], [78, 80], [87, 83], [86, 84], + [88, 92], [89, 91], [95, 93], [96, 97] + ) + self.image_dir = osp.join(self.image_dir, 'WFLW', 'WFLW_images') + + self.label_num = self.nstack * 3 if self.use_AAM else self.nstack + self.loss_weights, self.criterions, self.metrics = [], [], [] + for i in range(self.nstack): + factor = (2 ** i) / (2 ** (self.nstack - 1)) + if self.use_AAM: + self.loss_weights += [factor * weight for weight in [1.0, 10.0, 10.0]] + self.criterions += [self.loss_func, "AWingLoss", "AWingLoss"] + self.metrics += ["NME", None, None] + else: + self.loss_weights += [factor * weight for weight in [1.0]] + self.criterions += [self.loss_func, ] + self.metrics += ["NME", ] + + self.key_metric_index = (self.nstack - 1) * 3 if self.use_AAM else (self.nstack - 1) + + # data + self.folder = self.get_foldername() + self.work_dir = osp.join(self.ckpt_dir, self.data_definition, self.folder) + self.model_dir = osp.join(self.work_dir, 'model') + self.log_dir = osp.join(self.work_dir, 'log') + + self.train_tsv_file = osp.join(self.annot_dir, self.data_definition, "train.tsv") + self.train_pic_dir = self.image_dir + + self.val_tsv_file = osp.join(self.annot_dir, self.data_definition, self.valset) + self.val_pic_dir = self.image_dir + + self.test_tsv_file = osp.join(self.annot_dir, self.data_definition, self.test_file) + self.test_pic_dir = self.image_dir + + def get_foldername(self): + str = '' + str += '{}_{}x{}_{}_ep{}_lr{}_bs{}'.format(self.data_definition, self.height, self.width, + self.optimizer, self.max_epoch, self.learn_rate, self.batch_size) + str += '_{}'.format(self.loss_func) + str += '_{}_{}'.format(self.star_dist, self.star_w) if self.loss_func == 'STARLoss' else '' + str += '_AAM' if self.use_AAM else '' + str += '_{}'.format(self.valset[:-4]) if self.valset != 'test.tsv' else '' + str += '_{}'.format(self.id) + return str diff --git a/models/STAR/conf/base.py b/models/STAR/conf/base.py new file mode 100644 index 0000000000000000000000000000000000000000..b2545649666ba3dbc60ad50a542374b3bff2bbbe --- /dev/null +++ b/models/STAR/conf/base.py @@ -0,0 +1,95 @@ +import uuid +import logging +import os.path as osp +from argparse import Namespace +# from tensorboardX import SummaryWriter + + +class Base: + """ + Base configure file, which contains the basic training parameters and should be inherited by other attribute configure file. + """ + + def __init__(self, config_name, ckpt_dir='./', image_dir='./', annot_dir='./'): + self.type = config_name + self.id = str(uuid.uuid4()) + self.note = "" + + self.ckpt_dir = ckpt_dir + self.image_dir = image_dir + self.annot_dir = annot_dir + + self.loader_type = "alignment" + self.loss_func = "STARLoss" + + # train + self.batch_size = 128 + self.val_batch_size = 1 + self.test_batch_size = 32 + self.channels = 3 + self.width = 256 + self.height = 256 + + # mean values in r, g, b channel. + self.means = (127, 127, 127) + self.scale = 0.0078125 + + self.display_iteration = 100 + self.milestones = [50, 80] + self.max_epoch = 100 + + self.net = "stackedHGnet_v1" + self.nstack = 4 + + # ["adam", "sgd"] + self.optimizer = "adam" + self.learn_rate = 0.1 + self.momentum = 0.01 # caffe: 0.99 + self.weight_decay = 0.0 + self.nesterov = False + self.scheduler = "MultiStepLR" + self.gamma = 0.1 + + self.loss_weights = [1.0] + self.criterions = ["SoftmaxWithLoss"] + self.metrics = ["Accuracy"] + self.key_metric_index = 0 + self.classes_num = [1000] + self.label_num = len(self.classes_num) + + # model + self.ema = False + self.use_AAM = True + + # visualization + self.writer = None + + # log file + self.logger = None + + def init_instance(self): + self.writer = SummaryWriter(logdir=self.log_dir, comment=self.type) + log_formatter = logging.Formatter("%(asctime)s %(levelname)-8s: %(message)s") + root_logger = logging.getLogger() + file_handler = logging.FileHandler(osp.join(self.log_dir, "log.txt")) + file_handler.setFormatter(log_formatter) + file_handler.setLevel(logging.NOTSET) + root_logger.addHandler(file_handler) + console_handler = logging.StreamHandler() + console_handler.setFormatter(log_formatter) + console_handler.setLevel(logging.NOTSET) + root_logger.addHandler(console_handler) + root_logger.setLevel(logging.NOTSET) + self.logger = root_logger + + def __del__(self): + # tensorboard --logdir self.log_dir + if self.writer is not None: + # self.writer.export_scalars_to_json(self.log_dir + "visual.json") + self.writer.close() + + def init_from_args(self, args: Namespace): + args_vars = vars(args) + for key, value in args_vars.items(): + if hasattr(self, key) and value is not None: + setattr(self, key, value) diff --git a/models/STAR/config.json b/models/STAR/config.json new file mode 100644 index 0000000000000000000000000000000000000000..35831f0d94065ba9b748dfeab3a3bf1aa25f1de3 --- /dev/null +++ b/models/STAR/config.json @@ -0,0 +1,15 @@ +{ + "Token":"bpt4JPotFA6bpdknR9ZDCw", + "business_flag": "shadow_cv_face", + "model_local_file_path": "/apdcephfs_cq3/share_1134483/charlinzhou/Documents/awesome-tools/jizhi/", + "host_num": 1, + "host_gpu_num": 1, + "GPUName": "V100", + "is_elasticity": true, + "enable_evicted_pulled_up": true, + "task_name": "20230312_slpt_star_bb_init_eigen_box_align_smoothl1-1", + "task_flag": "20230312_slpt_star_bb_init_eigen_box_align_smoothl1-1", + "model_name": "20230312_slpt_star_bb_init_eigen_box_align_smoothl1-1", + "image_full_name": "mirrors.tencent.com/haroldzcli/py36-pytorch1.7.1-torchvision0.8.2-cuda10.1-cudnn7.6", + "start_cmd": "./start_slpt.sh /apdcephfs_cq3/share_1134483/charlinzhou/Documents/SLPT_Training train.py --loss_func=star --bb_init --eigen_box --dist_func=align_smoothl1" +} diff --git a/models/STAR/demo.py b/models/STAR/demo.py new file mode 100644 index 0000000000000000000000000000000000000000..e0025275bf8030b646b0fdd7a6a282f14535d14c --- /dev/null +++ b/models/STAR/demo.py @@ -0,0 +1,267 @@ +import os +import cv2 +import copy +import dlib +import math +import argparse +import numpy as np +import gradio as gr +from matplotlib import pyplot as plt +import torch +# private package +from lib import utility + + +class GetCropMatrix(): + """ + from_shape -> transform_matrix + """ + + def __init__(self, image_size, target_face_scale, align_corners=False): + self.image_size = image_size + self.target_face_scale = target_face_scale + self.align_corners = align_corners + + def _compose_rotate_and_scale(self, angle, scale, shift_xy, from_center, to_center): + cosv = math.cos(angle) + sinv = math.sin(angle) + + fx, fy = from_center + tx, ty = to_center + + acos = scale * cosv + asin = scale * sinv + + a0 = acos + a1 = -asin + a2 = tx - acos * fx + asin * fy + shift_xy[0] + + b0 = asin + b1 = acos + b2 = ty - asin * fx - acos * fy + shift_xy[1] + + rot_scale_m = np.array([ + [a0, a1, a2], + [b0, b1, b2], + [0.0, 0.0, 1.0] + ], np.float32) + return rot_scale_m + + def process(self, scale, center_w, center_h): + if self.align_corners: + to_w, to_h = self.image_size - 1, self.image_size - 1 + else: + to_w, to_h = self.image_size, self.image_size + + rot_mu = 0 + scale_mu = self.image_size / (scale * self.target_face_scale * 200.0) + shift_xy_mu = (0, 0) + matrix = self._compose_rotate_and_scale( + rot_mu, scale_mu, shift_xy_mu, + from_center=[center_w, center_h], + to_center=[to_w / 2.0, to_h / 2.0]) + return matrix + + +class TransformPerspective(): + """ + image, matrix3x3 -> transformed_image + """ + + def __init__(self, image_size): + self.image_size = image_size + + def process(self, image, matrix): + return cv2.warpPerspective( + image, matrix, dsize=(self.image_size, self.image_size), + flags=cv2.INTER_LINEAR, borderValue=0) + + +class TransformPoints2D(): + """ + points (nx2), matrix (3x3) -> points (nx2) + """ + + def process(self, srcPoints, matrix): + # nx3 + desPoints = np.concatenate([srcPoints, np.ones_like(srcPoints[:, [0]])], axis=1) + desPoints = desPoints @ np.transpose(matrix) # nx3 + desPoints = desPoints[:, :2] / desPoints[:, [2, 2]] + return desPoints.astype(srcPoints.dtype) + + +class Alignment: + def __init__(self, args, model_path, dl_framework, device_ids): + self.input_size = 256 + self.target_face_scale = 1.0 + self.dl_framework = dl_framework + + # model + if self.dl_framework == "pytorch": + # conf + self.config = utility.get_config(args) + self.config.device_id = device_ids[0] + # set environment + utility.set_environment(self.config) + self.config.init_instance() + if self.config.logger is not None: + self.config.logger.info("Loaded configure file %s: %s" % (args.config_name, self.config.id)) + self.config.logger.info("\n" + "\n".join(["%s: %s" % item for item in self.config.__dict__.items()])) + + net = utility.get_net(self.config) + if device_ids == [-1]: + checkpoint = torch.load(model_path, map_location="cpu") + else: + checkpoint = torch.load(model_path) + net.load_state_dict(checkpoint["net"]) + net = net.to(self.config.device_id) + net.eval() + self.alignment = net + else: + assert False + + self.getCropMatrix = GetCropMatrix(image_size=self.input_size, target_face_scale=self.target_face_scale, + align_corners=True) + self.transformPerspective = TransformPerspective(image_size=self.input_size) + self.transformPoints2D = TransformPoints2D() + + def norm_points(self, points, align_corners=False): + if align_corners: + # [0, SIZE-1] -> [-1, +1] + return points / torch.tensor([self.input_size - 1, self.input_size - 1]).to(points).view(1, 1, 2) * 2 - 1 + else: + # [-0.5, SIZE-0.5] -> [-1, +1] + return (points * 2 + 1) / torch.tensor([self.input_size, self.input_size]).to(points).view(1, 1, 2) - 1 + + def denorm_points(self, points, align_corners=False): + if align_corners: + # [-1, +1] -> [0, SIZE-1] + return (points + 1) / 2 * torch.tensor([self.input_size - 1, self.input_size - 1]).to(points).view(1, 1, 2) + else: + # [-1, +1] -> [-0.5, SIZE-0.5] + return ((points + 1) * torch.tensor([self.input_size, self.input_size]).to(points).view(1, 1, 2) - 1) / 2 + + def preprocess(self, image, scale, center_w, center_h): + matrix = self.getCropMatrix.process(scale, center_w, center_h) + input_tensor = self.transformPerspective.process(image, matrix) + input_tensor = input_tensor[np.newaxis, :] + + input_tensor = torch.from_numpy(input_tensor) + input_tensor = input_tensor.float().permute(0, 3, 1, 2) + input_tensor = input_tensor / 255.0 * 2.0 - 1.0 + input_tensor = input_tensor.to(self.config.device_id) + return input_tensor, matrix + + def postprocess(self, srcPoints, coeff): + # dstPoints = self.transformPoints2D.process(srcPoints, coeff) + # matrix^(-1) * src = dst + # src = matrix * dst + dstPoints = np.zeros(srcPoints.shape, dtype=np.float32) + for i in range(srcPoints.shape[0]): + dstPoints[i][0] = coeff[0][0] * srcPoints[i][0] + coeff[0][1] * srcPoints[i][1] + coeff[0][2] + dstPoints[i][1] = coeff[1][0] * srcPoints[i][0] + coeff[1][1] * srcPoints[i][1] + coeff[1][2] + return dstPoints + + def analyze(self, image, scale, center_w, center_h): + input_tensor, matrix = self.preprocess(image, scale, center_w, center_h) + + if self.dl_framework == "pytorch": + with torch.no_grad(): + output = self.alignment(input_tensor) + landmarks = output[-1][0] + else: + assert False + + landmarks = self.denorm_points(landmarks) + landmarks = landmarks.data.cpu().numpy()[0] + landmarks = self.postprocess(landmarks, np.linalg.inv(matrix)) + + return landmarks + + +def draw_pts(img, pts, mode="pts", shift=4, color=(0, 255, 0), radius=1, thickness=1, save_path=None, dif=0, + scale=0.3, concat=False, ): + img_draw = copy.deepcopy(img) + for cnt, p in enumerate(pts): + if mode == "index": + cv2.putText(img_draw, str(cnt), (int(float(p[0] + dif)), int(float(p[1] + dif))), cv2.FONT_HERSHEY_SIMPLEX, + scale, color, thickness) + elif mode == 'pts': + if len(img_draw.shape) > 2: + # 此处来回切换是因为opencv的bug + img_draw = cv2.cvtColor(img_draw, cv2.COLOR_BGR2RGB) + img_draw = cv2.cvtColor(img_draw, cv2.COLOR_RGB2BGR) + cv2.circle(img_draw, (int(p[0] * (1 << shift)), int(p[1] * (1 << shift))), radius << shift, color, -1, + cv2.LINE_AA, shift=shift) + else: + raise NotImplementedError + if concat: + img_draw = np.concatenate((img, img_draw), axis=1) + if save_path is not None: + cv2.imwrite(save_path, img_draw) + return img_draw + + +def process(input_image): + image_draw = copy.deepcopy(input_image) + dets = detector(input_image, 1) + + num_faces = len(dets) + if num_faces == 0: + print("Sorry, there were no faces found in '{}'".format(face_file_path)) + exit() + + results = [] + for detection in dets: + face = sp(input_image, detection) + shape = [] + for i in range(68): + x = face.part(i).x + y = face.part(i).y + shape.append((x, y)) + shape = np.array(shape) + # image_draw = draw_pts(image_draw, shape) + x1, x2 = shape[:, 0].min(), shape[:, 0].max() + y1, y2 = shape[:, 1].min(), shape[:, 1].max() + scale = min(x2 - x1, y2 - y1) / 200 * 1.05 + center_w = (x2 + x1) / 2 + center_h = (y2 + y1) / 2 + + scale, center_w, center_h = float(scale), float(center_w), float(center_h) + landmarks_pv = alignment.analyze(input_image, scale, center_w, center_h) + results.append(landmarks_pv) + image_draw = draw_pts(image_draw, landmarks_pv) + return image_draw, results + + +if __name__ == '__main__': + # face detector + # could be downloaded in this repo: https://github.com/italojs/facial-landmarks-recognition/tree/master + predictor_path = '/path/to/shape_predictor_68_face_landmarks.dat' + detector = dlib.get_frontal_face_detector() + sp = dlib.shape_predictor(predictor_path) + + # facial landmark detector + args = argparse.Namespace() + args.config_name = 'alignment' + # could be downloaded here: https://drive.google.com/file/d/1aOx0wYEZUfBndYy_8IYszLPG_D2fhxrT/view + model_path = '/path/to/WFLW_STARLoss_NME_4_02_FR_2_32_AUC_0_605.pkl' + device_ids = '0' + device_ids = list(map(int, device_ids.split(","))) + alignment = Alignment(args, model_path, dl_framework="pytorch", device_ids=device_ids) + + # image: input image + # image_draw: draw the detected facial landmarks on image + # results: a list of detected facial landmarks + face_file_path = '/path/to/face/image/bald_guys.jpg' + image = cv2.imread(face_file_path) + image_draw, results = process(image) + + # visualize + img = cv2.cvtColor(image_draw, cv2.COLOR_BGR2RGB) + plt.imshow(img) + plt.show() + + # demo + # interface = gr.Interface(fn=process, inputs="image", outputs="image") + # interface.launch(share=True) diff --git a/models/STAR/evaluate.py b/models/STAR/evaluate.py new file mode 100644 index 0000000000000000000000000000000000000000..4c9b759381559cf32f9e461dc2e3c26b92d8dda0 --- /dev/null +++ b/models/STAR/evaluate.py @@ -0,0 +1,258 @@ +import os +import cv2 +import math +import argparse +import numpy as np +from tqdm import tqdm + +import torch + +# private package +from lib import utility + + + +class GetCropMatrix(): + """ + from_shape -> transform_matrix + """ + + def __init__(self, image_size, target_face_scale, align_corners=False): + self.image_size = image_size + self.target_face_scale = target_face_scale + self.align_corners = align_corners + + def _compose_rotate_and_scale(self, angle, scale, shift_xy, from_center, to_center): + cosv = math.cos(angle) + sinv = math.sin(angle) + + fx, fy = from_center + tx, ty = to_center + + acos = scale * cosv + asin = scale * sinv + + a0 = acos + a1 = -asin + a2 = tx - acos * fx + asin * fy + shift_xy[0] + + b0 = asin + b1 = acos + b2 = ty - asin * fx - acos * fy + shift_xy[1] + + rot_scale_m = np.array([ + [a0, a1, a2], + [b0, b1, b2], + [0.0, 0.0, 1.0] + ], np.float32) + return rot_scale_m + + def process(self, scale, center_w, center_h): + if self.align_corners: + to_w, to_h = self.image_size - 1, self.image_size - 1 + else: + to_w, to_h = self.image_size, self.image_size + + rot_mu = 0 + scale_mu = self.image_size / (scale * self.target_face_scale * 200.0) + shift_xy_mu = (0, 0) + matrix = self._compose_rotate_and_scale( + rot_mu, scale_mu, shift_xy_mu, + from_center=[center_w, center_h], + to_center=[to_w / 2.0, to_h / 2.0]) + return matrix + + +class TransformPerspective(): + """ + image, matrix3x3 -> transformed_image + """ + + def __init__(self, image_size): + self.image_size = image_size + + def process(self, image, matrix): + return cv2.warpPerspective( + image, matrix, dsize=(self.image_size, self.image_size), + flags=cv2.INTER_LINEAR, borderValue=0) + + +class TransformPoints2D(): + """ + points (nx2), matrix (3x3) -> points (nx2) + """ + + def process(self, srcPoints, matrix): + # nx3 + desPoints = np.concatenate([srcPoints, np.ones_like(srcPoints[:, [0]])], axis=1) + desPoints = desPoints @ np.transpose(matrix) # nx3 + desPoints = desPoints[:, :2] / desPoints[:, [2, 2]] + return desPoints.astype(srcPoints.dtype) + + +class Alignment: + def __init__(self, args, model_path, dl_framework, device_ids): + self.input_size = 256 + self.target_face_scale = 1.0 + self.dl_framework = dl_framework + + # model + if self.dl_framework == "pytorch": + # conf + self.config = utility.get_config(args) + self.config.device_id = device_ids[0] + # set environment + utility.set_environment(self.config) + self.config.init_instance() + if self.config.logger is not None: + self.config.logger.info("Loaded configure file %s: %s" % (args.config_name, self.config.id)) + self.config.logger.info("\n" + "\n".join(["%s: %s" % item for item in self.config.__dict__.items()])) + + net = utility.get_net(self.config) + if device_ids == [-1]: + checkpoint = torch.load(model_path, map_location="cpu") + else: + checkpoint = torch.load(model_path) + net.load_state_dict(checkpoint["net"]) + net = net.to(self.config.device_id) + net.eval() + self.alignment = net + else: + assert False + + self.getCropMatrix = GetCropMatrix(image_size=self.input_size, target_face_scale=self.target_face_scale, + align_corners=True) + self.transformPerspective = TransformPerspective(image_size=self.input_size) + self.transformPoints2D = TransformPoints2D() + + def norm_points(self, points, align_corners=False): + if align_corners: + # [0, SIZE-1] -> [-1, +1] + return points / torch.tensor([self.input_size - 1, self.input_size - 1]).to(points).view(1, 1, 2) * 2 - 1 + else: + # [-0.5, SIZE-0.5] -> [-1, +1] + return (points * 2 + 1) / torch.tensor([self.input_size, self.input_size]).to(points).view(1, 1, 2) - 1 + + def denorm_points(self, points, align_corners=False): + if align_corners: + # [-1, +1] -> [0, SIZE-1] + return (points + 1) / 2 * torch.tensor([self.input_size - 1, self.input_size - 1]).to(points).view(1, 1, 2) + else: + # [-1, +1] -> [-0.5, SIZE-0.5] + return ((points + 1) * torch.tensor([self.input_size, self.input_size]).to(points).view(1, 1, 2) - 1) / 2 + + def preprocess(self, image, scale, center_w, center_h): + matrix = self.getCropMatrix.process(scale, center_w, center_h) + input_tensor = self.transformPerspective.process(image, matrix) + input_tensor = input_tensor[np.newaxis, :] + + input_tensor = torch.from_numpy(input_tensor) + input_tensor = input_tensor.float().permute(0, 3, 1, 2) + input_tensor = input_tensor / 255.0 * 2.0 - 1.0 + input_tensor = input_tensor.to(self.config.device_id) + return input_tensor, matrix + + def postprocess(self, srcPoints, coeff): + # dstPoints = self.transformPoints2D.process(srcPoints, coeff) + # matrix^(-1) * src = dst + # src = matrix * dst + dstPoints = np.zeros(srcPoints.shape, dtype=np.float32) + for i in range(srcPoints.shape[0]): + dstPoints[i][0] = coeff[0][0] * srcPoints[i][0] + coeff[0][1] * srcPoints[i][1] + coeff[0][2] + dstPoints[i][1] = coeff[1][0] * srcPoints[i][0] + coeff[1][1] * srcPoints[i][1] + coeff[1][2] + return dstPoints + + def analyze(self, image, scale, center_w, center_h): + input_tensor, matrix = self.preprocess(image, scale, center_w, center_h) + + if self.dl_framework == "pytorch": + with torch.no_grad(): + output = self.alignment(input_tensor) + landmarks = output[-1][0] + else: + assert False + + landmarks = self.denorm_points(landmarks) + landmarks = landmarks.data.cpu().numpy()[0] + landmarks = self.postprocess(landmarks, np.linalg.inv(matrix)) + + return landmarks + + +def L2(p1, p2): + return np.linalg.norm(p1 - p2) + + +def NME(landmarks_gt, landmarks_pv): + pts_num = landmarks_gt.shape[0] + if pts_num == 29: + left_index = 16 + right_index = 17 + elif pts_num == 68: + left_index = 36 + right_index = 45 + elif pts_num == 98: + left_index = 60 + right_index = 72 + + nme = 0 + eye_span = L2(landmarks_gt[left_index], landmarks_gt[right_index]) + for i in range(pts_num): + error = L2(landmarks_pv[i], landmarks_gt[i]) + nme += error / eye_span + nme /= pts_num + return nme + + +def evaluate(args, model_path, metadata_path, device_ids, mode): + alignment = Alignment(args, model_path, dl_framework="pytorch", device_ids=device_ids) + config = alignment.config + nme_sum = 0 + with open(metadata_path, 'r') as f: + lines = f.readlines() + for k, line in enumerate(tqdm(lines)): + item = line.strip().split("\t") + image_name, landmarks_5pts, landmarks_gt, scale, center_w, center_h = item[:6] + # image & keypoints alignment + image_name = image_name.replace('\\', '/') + image_name = image_name.replace('//msr-facestore/Workspace/MSRA_EP_Allergan/users/yanghuan/training_data/wflw/rawImages/', '') + image_name = image_name.replace('./rawImages/', '') + image_path = os.path.join(config.image_dir, image_name) + landmarks_gt = np.array(list(map(float, landmarks_gt.split(","))), dtype=np.float32).reshape(-1, 2) + scale, center_w, center_h = float(scale), float(center_w), float(center_h) + + image = cv2.imread(image_path) + landmarks_pv = alignment.analyze(image, scale, center_w, center_h) + + # NME + if mode == "nme": + nme = NME(landmarks_gt, landmarks_pv) + nme_sum += nme + # print("Current NME(%d): %f" % (k + 1, (nme_sum / (k + 1)))) + else: + pass + + if mode == "nme": + print("Final NME: %f" % (nme_sum / (k + 1))) + else: + pass + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Evaluation script") + parser.add_argument("--config_name", type=str, default="alignment", help="set configure file name") + parser.add_argument("--model_path", type=str, default="./train.pkl", help="the path of model") + parser.add_argument("--data_definition", type=str, default='WFLW', help="COFW/300W/WFLW") + parser.add_argument("--metadata_path", type=str, default="", help="the path of metadata") + parser.add_argument("--image_dir", type=str, default="", help="the path of image") + parser.add_argument("--device_ids", type=str, default="0", help="set device ids, -1 means use cpu device, >= 0 means use gpu device") + parser.add_argument("--mode", type=str, default="nme", help="set the evaluate mode: nme") + args = parser.parse_args() + + device_ids = list(map(int, args.device_ids.split(","))) + evaluate( + args, + model_path=args.model_path, + metadata_path=args.metadata_path, + device_ids=device_ids, + mode=args.mode) diff --git a/models/STAR/lib/.DS_Store b/models/STAR/lib/.DS_Store new file mode 100644 index 0000000000000000000000000000000000000000..0a46003cbfa5d371dab2d128e02e050380ec0cec Binary files /dev/null and b/models/STAR/lib/.DS_Store differ diff --git a/models/STAR/lib/__init__.py b/models/STAR/lib/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..ff08a78fc3315c65bc061eb1a42a3e8cc90b9b8d --- /dev/null +++ b/models/STAR/lib/__init__.py @@ -0,0 +1,9 @@ +from .dataset import get_encoder, get_decoder +from .dataset import AlignmentDataset, Augmentation +from .backbone import StackedHGNetV1 +from .metric import NME, Accuracy +from .utils import time_print, time_string, time_for_file, time_string_short +from .utils import convert_secs2time, convert_size2str + +from .utility import get_dataloader, get_config, get_net, get_criterions +from .utility import get_optimizer, get_scheduler diff --git a/models/STAR/lib/backbone/__init__.py b/models/STAR/lib/backbone/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..b967103849100c921d4ecceb7a25142fa5ded6f6 --- /dev/null +++ b/models/STAR/lib/backbone/__init__.py @@ -0,0 +1,5 @@ +from .stackedHGNetV1 import StackedHGNetV1 + +__all__ = [ + "StackedHGNetV1", +] \ No newline at end of file diff --git a/models/STAR/lib/backbone/core/coord_conv.py b/models/STAR/lib/backbone/core/coord_conv.py new file mode 100644 index 0000000000000000000000000000000000000000..0eb8e2d6ec574b77576970b6e6aef20c62887228 --- /dev/null +++ b/models/STAR/lib/backbone/core/coord_conv.py @@ -0,0 +1,157 @@ +import torch +import torch.nn as nn + + +class AddCoordsTh(nn.Module): + def __init__(self, x_dim, y_dim, with_r=False, with_boundary=False): + super(AddCoordsTh, self).__init__() + self.x_dim = x_dim + self.y_dim = y_dim + self.with_r = with_r + self.with_boundary = with_boundary + + def forward(self, input_tensor, heatmap=None): + """ + input_tensor: (batch, c, x_dim, y_dim) + """ + batch_size_tensor = input_tensor.shape[0] + + xx_ones = torch.ones([1, self.y_dim], dtype=torch.int32).to(input_tensor) + xx_ones = xx_ones.unsqueeze(-1) + + xx_range = torch.arange(self.x_dim, dtype=torch.int32).unsqueeze(0).to(input_tensor) + xx_range = xx_range.unsqueeze(1) + + xx_channel = torch.matmul(xx_ones.float(), xx_range.float()) + xx_channel = xx_channel.unsqueeze(-1) + + yy_ones = torch.ones([1, self.x_dim], dtype=torch.int32).to(input_tensor) + yy_ones = yy_ones.unsqueeze(1) + + yy_range = torch.arange(self.y_dim, dtype=torch.int32).unsqueeze(0).to(input_tensor) + yy_range = yy_range.unsqueeze(-1) + + yy_channel = torch.matmul(yy_range.float(), yy_ones.float()) + yy_channel = yy_channel.unsqueeze(-1) + + xx_channel = xx_channel.permute(0, 3, 2, 1) + yy_channel = yy_channel.permute(0, 3, 2, 1) + + xx_channel = xx_channel / (self.x_dim - 1) + yy_channel = yy_channel / (self.y_dim - 1) + + xx_channel = xx_channel * 2 - 1 + yy_channel = yy_channel * 2 - 1 + + xx_channel = xx_channel.repeat(batch_size_tensor, 1, 1, 1) + yy_channel = yy_channel.repeat(batch_size_tensor, 1, 1, 1) + + if self.with_boundary and type(heatmap) != type(None): + boundary_channel = torch.clamp(heatmap[:, -1:, :, :], + 0.0, 1.0) + + zero_tensor = torch.zeros_like(xx_channel).to(xx_channel) + xx_boundary_channel = torch.where(boundary_channel>0.05, + xx_channel, zero_tensor) + yy_boundary_channel = torch.where(boundary_channel>0.05, + yy_channel, zero_tensor) + ret = torch.cat([input_tensor, xx_channel, yy_channel], dim=1) + + + if self.with_r: + rr = torch.sqrt(torch.pow(xx_channel, 2) + torch.pow(yy_channel, 2)) + rr = rr / torch.max(rr) + ret = torch.cat([ret, rr], dim=1) + + if self.with_boundary and type(heatmap) != type(None): + ret = torch.cat([ret, xx_boundary_channel, + yy_boundary_channel], dim=1) + return ret + + +class CoordConvTh(nn.Module): + """CoordConv layer as in the paper.""" + def __init__(self, x_dim, y_dim, with_r, with_boundary, + in_channels, out_channels, first_one=False, relu=False, bn=False, *args, **kwargs): + super(CoordConvTh, self).__init__() + self.addcoords = AddCoordsTh(x_dim=x_dim, y_dim=y_dim, with_r=with_r, + with_boundary=with_boundary) + in_channels += 2 + if with_r: + in_channels += 1 + if with_boundary and not first_one: + in_channels += 2 + self.conv = nn.Conv2d(in_channels=in_channels, out_channels=out_channels, *args, **kwargs) + self.relu = nn.ReLU() if relu else None + self.bn = nn.BatchNorm2d(out_channels) if bn else None + + self.with_boundary = with_boundary + self.first_one = first_one + + + def forward(self, input_tensor, heatmap=None): + assert (self.with_boundary and not self.first_one) == (heatmap is not None) + ret = self.addcoords(input_tensor, heatmap) + ret = self.conv(ret) + if self.bn is not None: + ret = self.bn(ret) + if self.relu is not None: + ret = self.relu(ret) + + return ret + + +''' +An alternative implementation for PyTorch with auto-infering the x-y dimensions. +''' +class AddCoords(nn.Module): + + def __init__(self, with_r=False): + super().__init__() + self.with_r = with_r + + def forward(self, input_tensor): + """ + Args: + input_tensor: shape(batch, channel, x_dim, y_dim) + """ + batch_size, _, x_dim, y_dim = input_tensor.size() + + xx_channel = torch.arange(x_dim).repeat(1, y_dim, 1).to(input_tensor) + yy_channel = torch.arange(y_dim).repeat(1, x_dim, 1).transpose(1, 2).to(input_tensor) + + xx_channel = xx_channel / (x_dim - 1) + yy_channel = yy_channel / (y_dim - 1) + + xx_channel = xx_channel * 2 - 1 + yy_channel = yy_channel * 2 - 1 + + xx_channel = xx_channel.repeat(batch_size, 1, 1, 1).transpose(2, 3) + yy_channel = yy_channel.repeat(batch_size, 1, 1, 1).transpose(2, 3) + + ret = torch.cat([ + input_tensor, + xx_channel.type_as(input_tensor), + yy_channel.type_as(input_tensor)], dim=1) + + if self.with_r: + rr = torch.sqrt(torch.pow(xx_channel - 0.5, 2) + torch.pow(yy_channel - 0.5, 2)) + ret = torch.cat([ret, rr], dim=1) + + return ret + + +class CoordConv(nn.Module): + + def __init__(self, in_channels, out_channels, with_r=False, **kwargs): + super().__init__() + self.addcoords = AddCoords(with_r=with_r) + in_channels += 2 + if with_r: + in_channels += 1 + self.conv = nn.Conv2d(in_channels, out_channels, **kwargs) + + def forward(self, x): + ret = self.addcoords(x) + ret = self.conv(ret) + return ret diff --git a/models/STAR/lib/backbone/stackedHGNetV1.py b/models/STAR/lib/backbone/stackedHGNetV1.py new file mode 100644 index 0000000000000000000000000000000000000000..471f8faadc534e8805dc1a072dfcf59d185f3117 --- /dev/null +++ b/models/STAR/lib/backbone/stackedHGNetV1.py @@ -0,0 +1,307 @@ +import numpy as np + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from .core.coord_conv import CoordConvTh +from models.STAR.lib.dataset import get_decoder + + + +class Activation(nn.Module): + def __init__(self, kind: str = 'relu', channel=None): + super().__init__() + self.kind = kind + + if '+' in kind: + norm_str, act_str = kind.split('+') + else: + norm_str, act_str = 'none', kind + + self.norm_fn = { + 'in': F.instance_norm, + 'bn': nn.BatchNorm2d(channel), + 'bn_noaffine': nn.BatchNorm2d(channel, affine=False, track_running_stats=True), + 'none': None + }[norm_str] + + self.act_fn = { + 'relu': F.relu, + 'softplus': nn.Softplus(), + 'exp': torch.exp, + 'sigmoid': torch.sigmoid, + 'tanh': torch.tanh, + 'none': None + }[act_str] + + self.channel = channel + + def forward(self, x): + if self.norm_fn is not None: + x = self.norm_fn(x) + if self.act_fn is not None: + x = self.act_fn(x) + return x + + def extra_repr(self): + return f'kind={self.kind}, channel={self.channel}' + + +class ConvBlock(nn.Module): + def __init__(self, inp_dim, out_dim, kernel_size=3, stride=1, bn=False, relu=True, groups=1): + super(ConvBlock, self).__init__() + self.inp_dim = inp_dim + self.conv = nn.Conv2d(inp_dim, out_dim, kernel_size, + stride, padding=(kernel_size - 1) // 2, groups=groups, bias=True) + self.relu = None + self.bn = None + if relu: + self.relu = nn.ReLU() + if bn: + self.bn = nn.BatchNorm2d(out_dim) + + def forward(self, x): + x = self.conv(x) + if self.bn is not None: + x = self.bn(x) + if self.relu is not None: + x = self.relu(x) + return x + + +class ResBlock(nn.Module): + def __init__(self, inp_dim, out_dim, mid_dim=None): + super(ResBlock, self).__init__() + if mid_dim is None: + mid_dim = out_dim // 2 + self.relu = nn.ReLU() + self.bn1 = nn.BatchNorm2d(inp_dim) + self.conv1 = ConvBlock(inp_dim, mid_dim, 1, relu=False) + self.bn2 = nn.BatchNorm2d(mid_dim) + self.conv2 = ConvBlock(mid_dim, mid_dim, 3, relu=False) + self.bn3 = nn.BatchNorm2d(mid_dim) + self.conv3 = ConvBlock(mid_dim, out_dim, 1, relu=False) + self.skip_layer = ConvBlock(inp_dim, out_dim, 1, relu=False) + if inp_dim == out_dim: + self.need_skip = False + else: + self.need_skip = True + + def forward(self, x): + if self.need_skip: + residual = self.skip_layer(x) + else: + residual = x + out = self.bn1(x) + out = self.relu(out) + out = self.conv1(out) + out = self.bn2(out) + out = self.relu(out) + out = self.conv2(out) + out = self.bn3(out) + out = self.relu(out) + out = self.conv3(out) + out += residual + return out + + +class Hourglass(nn.Module): + def __init__(self, n, f, increase=0, up_mode='nearest', + add_coord=False, first_one=False, x_dim=64, y_dim=64): + super(Hourglass, self).__init__() + nf = f + increase + + Block = ResBlock + + if add_coord: + self.coordconv = CoordConvTh(x_dim=x_dim, y_dim=y_dim, + with_r=True, with_boundary=True, + relu=False, bn=False, + in_channels=f, out_channels=f, + first_one=first_one, + kernel_size=1, + stride=1, padding=0) + else: + self.coordconv = None + self.up1 = Block(f, f) + + # Lower branch + self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2) + + self.low1 = Block(f, nf) + self.n = n + # Recursive hourglass + if self.n > 1: + self.low2 = Hourglass(n=n - 1, f=nf, increase=increase, up_mode=up_mode, add_coord=False) + else: + self.low2 = Block(nf, nf) + self.low3 = Block(nf, f) + self.up2 = nn.Upsample(scale_factor=2, mode=up_mode) + + def forward(self, x, heatmap=None): + if self.coordconv is not None: + x = self.coordconv(x, heatmap) + up1 = self.up1(x) + pool1 = self.pool1(x) + low1 = self.low1(pool1) + low2 = self.low2(low1) + low3 = self.low3(low2) + up2 = self.up2(low3) + return up1 + up2 + + +class E2HTransform(nn.Module): + def __init__(self, edge_info, num_points, num_edges): + super().__init__() + + e2h_matrix = np.zeros([num_points, num_edges]) + for edge_id, isclosed_indices in enumerate(edge_info): + is_closed, indices = isclosed_indices + for point_id in indices: + e2h_matrix[point_id, edge_id] = 1 + e2h_matrix = torch.from_numpy(e2h_matrix).float() + + # pn x en x 1 x 1. + self.register_buffer('weight', e2h_matrix.view( + e2h_matrix.size(0), e2h_matrix.size(1), 1, 1)) + + # some keypoints are not coverred by any edges, + # in these cases, we must add a constant bias to their heatmap weights. + bias = ((e2h_matrix @ torch.ones(e2h_matrix.size(1)).to( + e2h_matrix)) < 0.5).to(e2h_matrix) + # pn x 1. + self.register_buffer('bias', bias) + + def forward(self, edgemaps): + # input: batch_size x en x hw x hh. + # output: batch_size x pn x hw x hh. + return F.conv2d(edgemaps, weight=self.weight, bias=self.bias) + + +class StackedHGNetV1(nn.Module): + def __init__(self, config, classes_num, edge_info, + nstack=4, nlevels=4, in_channel=256, increase=0, + add_coord=True, decoder_type='default'): + super(StackedHGNetV1, self).__init__() + + self.cfg = config + self.coder_type = decoder_type + self.decoder = get_decoder(decoder_type=decoder_type) + self.nstack = nstack + self.add_coord = add_coord + + self.num_heats = classes_num[0] + + if self.add_coord: + convBlock = CoordConvTh(x_dim=self.cfg.width, y_dim=self.cfg.height, + with_r=True, with_boundary=False, + relu=True, bn=True, + in_channels=3, out_channels=64, + kernel_size=7, + stride=2, padding=3) + else: + convBlock = ConvBlock(3, 64, 7, 2, bn=True, relu=True) + + pool = nn.MaxPool2d(kernel_size=2, stride=2) + + Block = ResBlock + + self.pre = nn.Sequential( + convBlock, + Block(64, 128), + pool, + Block(128, 128), + Block(128, in_channel) + ) + + self.hgs = nn.ModuleList( + [Hourglass(n=nlevels, f=in_channel, increase=increase, add_coord=self.add_coord, first_one=(_ == 0), + x_dim=int(self.cfg.width / self.nstack), y_dim=int(self.cfg.height / self.nstack)) + for _ in range(nstack)]) + + self.features = nn.ModuleList([ + nn.Sequential( + Block(in_channel, in_channel), + ConvBlock(in_channel, in_channel, 1, bn=True, relu=True) + ) for _ in range(nstack)]) + + self.out_heatmaps = nn.ModuleList( + [ConvBlock(in_channel, self.num_heats, 1, relu=False, bn=False) + for _ in range(nstack)]) + + if self.cfg.use_AAM: + self.num_edges = classes_num[1] + self.num_points = classes_num[2] + + self.e2h_transform = E2HTransform(edge_info, self.num_points, self.num_edges) + self.out_edgemaps = nn.ModuleList( + [ConvBlock(in_channel, self.num_edges, 1, relu=False, bn=False) + for _ in range(nstack)]) + self.out_pointmaps = nn.ModuleList( + [ConvBlock(in_channel, self.num_points, 1, relu=False, bn=False) + for _ in range(nstack)]) + self.merge_edgemaps = nn.ModuleList( + [ConvBlock(self.num_edges, in_channel, 1, relu=False, bn=False) + for _ in range(nstack - 1)]) + self.merge_pointmaps = nn.ModuleList( + [ConvBlock(self.num_points, in_channel, 1, relu=False, bn=False) + for _ in range(nstack - 1)]) + self.edgemap_act = Activation("sigmoid", self.num_edges) + self.pointmap_act = Activation("sigmoid", self.num_points) + + self.merge_features = nn.ModuleList( + [ConvBlock(in_channel, in_channel, 1, relu=False, bn=False) + for _ in range(nstack - 1)]) + self.merge_heatmaps = nn.ModuleList( + [ConvBlock(self.num_heats, in_channel, 1, relu=False, bn=False) + for _ in range(nstack - 1)]) + + self.nstack = nstack + + self.heatmap_act = Activation("in+relu", self.num_heats) + + self.inference = False + + def set_inference(self, inference): + self.inference = inference + + def forward(self, x): + x = self.pre(x) + + y, fusionmaps = [], [] + heatmaps = None + for i in range(self.nstack): + hg = self.hgs[i](x, heatmap=heatmaps) + feature = self.features[i](hg) + + heatmaps0 = self.out_heatmaps[i](feature) + heatmaps = self.heatmap_act(heatmaps0) + + if self.cfg.use_AAM: + pointmaps0 = self.out_pointmaps[i](feature) + pointmaps = self.pointmap_act(pointmaps0) + edgemaps0 = self.out_edgemaps[i](feature) + edgemaps = self.edgemap_act(edgemaps0) + mask = self.e2h_transform(edgemaps) * pointmaps + fusion_heatmaps = mask * heatmaps + else: + fusion_heatmaps = heatmaps + + landmarks = self.decoder.get_coords_from_heatmap(fusion_heatmaps) + + if i < self.nstack - 1: + x = x + self.merge_features[i](feature) + \ + self.merge_heatmaps[i](heatmaps) + if self.cfg.use_AAM: + x += self.merge_pointmaps[i](pointmaps) + x += self.merge_edgemaps[i](edgemaps) + + y.append(landmarks) + if self.cfg.use_AAM: + y.append(pointmaps) + y.append(edgemaps) + + fusionmaps.append(fusion_heatmaps) + + return y, fusionmaps, landmarks \ No newline at end of file diff --git a/models/STAR/lib/dataset/__init__.py b/models/STAR/lib/dataset/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3380c4b601721c79c2076538d8e675badd05b5a2 --- /dev/null +++ b/models/STAR/lib/dataset/__init__.py @@ -0,0 +1,11 @@ +from .encoder import get_encoder +from .decoder import get_decoder +from .augmentation import Augmentation +from .alignmentDataset import AlignmentDataset + +__all__ = [ + "Augmentation", + "AlignmentDataset", + "get_encoder", + "get_decoder" +] diff --git a/models/STAR/lib/dataset/alignmentDataset.py b/models/STAR/lib/dataset/alignmentDataset.py new file mode 100644 index 0000000000000000000000000000000000000000..8ba888feae477174a98685e0d07eeadeadfaf6d7 --- /dev/null +++ b/models/STAR/lib/dataset/alignmentDataset.py @@ -0,0 +1,316 @@ +import os +import sys +import cv2 +import math +import copy +import hashlib +import imageio +import numpy as np +import pandas as pd +from scipy import interpolate +from PIL import Image, ImageEnhance, ImageFile + +import torch +import torch.nn.functional as F +from torch.utils.data import Dataset + +ImageFile.LOAD_TRUNCATED_IMAGES = True + +sys.path.append("./") +from models.STAR.lib.dataset.augmentation import Augmentation +from models.STAR.lib.dataset.encoder import get_encoder + + +class AlignmentDataset(Dataset): + + def __init__(self, tsv_flie, image_dir="", transform=None, + width=256, height=256, channels=3, + means=(127.5, 127.5, 127.5), scale=1 / 127.5, + classes_num=None, crop_op=True, aug_prob=0.0, edge_info=None, flip_mapping=None, is_train=True, + encoder_type='default', + ): + super(AlignmentDataset, self).__init__() + self.use_AAM = True + self.encoder_type = encoder_type + self.encoder = get_encoder(height, width, encoder_type=encoder_type) + self.items = pd.read_csv(tsv_flie, sep="\t") + self.image_dir = image_dir + self.landmark_num = classes_num[0] + self.transform = transform + + self.image_width = width + self.image_height = height + self.channels = channels + assert self.image_width == self.image_height + + self.means = means + self.scale = scale + + self.aug_prob = aug_prob + self.edge_info = edge_info + self.is_train = is_train + std_lmk_5pts = np.array([ + 196.0, 226.0, + 316.0, 226.0, + 256.0, 286.0, + 220.0, 360.4, + 292.0, 360.4], np.float32) / 256.0 - 1.0 + std_lmk_5pts = np.reshape(std_lmk_5pts, (5, 2)) # [-1 1] + target_face_scale = 1.0 if crop_op else 1.25 + + self.augmentation = Augmentation( + is_train=self.is_train, + aug_prob=self.aug_prob, + image_size=self.image_width, + crop_op=crop_op, + std_lmk_5pts=std_lmk_5pts, + target_face_scale=target_face_scale, + flip_rate=0.5, + flip_mapping=flip_mapping, + random_shift_sigma=0.05, + random_rot_sigma=math.pi / 180 * 18, + random_scale_sigma=0.1, + random_gray_rate=0.2, + random_occ_rate=0.4, + random_blur_rate=0.3, + random_gamma_rate=0.2, + random_nose_fusion_rate=0.2) + + def _circle(self, img, pt, sigma=1.0, label_type='Gaussian'): + # Check that any part of the gaussian is in-bounds + tmp_size = sigma * 3 + ul = [int(pt[0] - tmp_size), int(pt[1] - tmp_size)] + br = [int(pt[0] + tmp_size + 1), int(pt[1] + tmp_size + 1)] + if (ul[0] > img.shape[1] - 1 or ul[1] > img.shape[0] - 1 or + br[0] - 1 < 0 or br[1] - 1 < 0): + # If not, just return the image as is + return img + + # Generate gaussian + size = 2 * tmp_size + 1 + x = np.arange(0, size, 1, np.float32) + y = x[:, np.newaxis] + x0 = y0 = size // 2 + # The gaussian is not normalized, we want the center value to equal 1 + if label_type == 'Gaussian': + g = np.exp(- ((x - x0) ** 2 + (y - y0) ** 2) / (2 * sigma ** 2)) + else: + g = sigma / (((x - x0) ** 2 + (y - y0) ** 2 + sigma ** 2) ** 1.5) + + # Usable gaussian range + g_x = max(0, -ul[0]), min(br[0], img.shape[1]) - ul[0] + g_y = max(0, -ul[1]), min(br[1], img.shape[0]) - ul[1] + # Image range + img_x = max(0, ul[0]), min(br[0], img.shape[1]) + img_y = max(0, ul[1]), min(br[1], img.shape[0]) + + img[img_y[0]:img_y[1], img_x[0]:img_x[1]] = 255 * g[g_y[0]:g_y[1], g_x[0]:g_x[1]] + return img + + def _polylines(self, img, lmks, is_closed, color=255, thickness=1, draw_mode=cv2.LINE_AA, + interpolate_mode=cv2.INTER_AREA, scale=4): + h, w = img.shape + img_scale = cv2.resize(img, (w * scale, h * scale), interpolation=interpolate_mode) + lmks_scale = (lmks * scale + 0.5).astype(np.int32) + cv2.polylines(img_scale, [lmks_scale], is_closed, color, thickness * scale, draw_mode) + img = cv2.resize(img_scale, (w, h), interpolation=interpolate_mode) + return img + + def _generate_edgemap(self, points, scale=0.25, thickness=1): + h, w = self.image_height, self.image_width + edgemaps = [] + for is_closed, indices in self.edge_info: + edgemap = np.zeros([h, w], dtype=np.float32) + # align_corners: False. + part = copy.deepcopy(points[np.array(indices)]) + + part = self._fit_curve(part, is_closed) + part[:, 0] = np.clip(part[:, 0], 0, w - 1) + part[:, 1] = np.clip(part[:, 1], 0, h - 1) + edgemap = self._polylines(edgemap, part, is_closed, 255, thickness) + + edgemaps.append(edgemap) + edgemaps = np.stack(edgemaps, axis=0) / 255.0 + edgemaps = torch.from_numpy(edgemaps).float().unsqueeze(0) + edgemaps = F.interpolate(edgemaps, size=(int(w * scale), int(h * scale)), mode='bilinear', + align_corners=False).squeeze() + return edgemaps + + def _fit_curve(self, lmks, is_closed=False, density=5): + try: + x = lmks[:, 0].copy() + y = lmks[:, 1].copy() + if is_closed: + x = np.append(x, x[0]) + y = np.append(y, y[0]) + tck, u = interpolate.splprep([x, y], s=0, per=is_closed, k=3) + # bins = (x.shape[0] - 1) * density + 1 + # lmk_x, lmk_y = interpolate.splev(np.linspace(0, 1, bins), f) + intervals = np.array([]) + for i in range(len(u) - 1): + intervals = np.concatenate((intervals, np.linspace(u[i], u[i + 1], density, endpoint=False))) + if not is_closed: + intervals = np.concatenate((intervals, [u[-1]])) + lmk_x, lmk_y = interpolate.splev(intervals, tck, der=0) + # der_x, der_y = interpolate.splev(intervals, tck, der=1) + curve_lmks = np.stack([lmk_x, lmk_y], axis=-1) + # curve_ders = np.stack([der_x, der_y], axis=-1) + # origin_indices = np.arange(0, curve_lmks.shape[0], density) + + return curve_lmks + except: + return lmks + + def _image_id(self, image_path): + if not os.path.exists(image_path): + image_path = os.path.join(self.image_dir, image_path) + return hashlib.md5(open(image_path, "rb").read()).hexdigest() + + def _load_image(self, image_path): + if not os.path.exists(image_path): + image_path = os.path.join(self.image_dir, image_path) + + try: + # img = cv2.imdecode(np.fromfile(image_path, dtype=np.uint8), cv2.IMREAD_COLOR)#HWC, BGR, [0-255] + img = cv2.imread(image_path, cv2.IMREAD_COLOR) # HWC, BGR, [0-255] + assert img is not None and len(img.shape) == 3 and img.shape[2] == 3 + except: + try: + img = imageio.imread(image_path) # HWC, RGB, [0-255] + img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR) # HWC, BGR, [0-255] + assert img is not None and len(img.shape) == 3 and img.shape[2] == 3 + except: + try: + gifImg = imageio.mimread(image_path) # BHWC, RGB, [0-255] + img = gifImg[0] # HWC, RGB, [0-255] + img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR) # HWC, BGR, [0-255] + assert img is not None and len(img.shape) == 3 and img.shape[2] == 3 + except: + img = None + return img + + def _compose_rotate_and_scale(self, angle, scale, shift_xy, from_center, to_center): + cosv = math.cos(angle) + sinv = math.sin(angle) + + fx, fy = from_center + tx, ty = to_center + + acos = scale * cosv + asin = scale * sinv + + a0 = acos + a1 = -asin + a2 = tx - acos * fx + asin * fy + shift_xy[0] + + b0 = asin + b1 = acos + b2 = ty - asin * fx - acos * fy + shift_xy[1] + + rot_scale_m = np.array([ + [a0, a1, a2], + [b0, b1, b2], + [0.0, 0.0, 1.0] + ], np.float32) + return rot_scale_m + + def _transformPoints2D(self, points, matrix): + """ + points (nx2), matrix (3x3) -> points (nx2) + """ + dtype = points.dtype + + # nx3 + points = np.concatenate([points, np.ones_like(points[:, [0]])], axis=1) + points = points @ np.transpose(matrix) # nx3 + points = points[:, :2] / points[:, [2, 2]] + return points.astype(dtype) + + def _transformPerspective(self, image, matrix, target_shape): + """ + image, matrix3x3 -> transformed_image + """ + return cv2.warpPerspective( + image, matrix, + dsize=(target_shape[1], target_shape[0]), + flags=cv2.INTER_LINEAR, borderValue=0) + + def _norm_points(self, points, h, w, align_corners=False): + if align_corners: + # [0, SIZE-1] -> [-1, +1] + des_points = points / torch.tensor([w - 1, h - 1]).to(points).view(1, 2) * 2 - 1 + else: + # [-0.5, SIZE-0.5] -> [-1, +1] + des_points = (points * 2 + 1) / torch.tensor([w, h]).to(points).view(1, 2) - 1 + des_points = torch.clamp(des_points, -1, 1) + return des_points + + def _denorm_points(self, points, h, w, align_corners=False): + if align_corners: + # [-1, +1] -> [0, SIZE-1] + des_points = (points + 1) / 2 * torch.tensor([w - 1, h - 1]).to(points).view(1, 1, 2) + else: + # [-1, +1] -> [-0.5, SIZE-0.5] + des_points = ((points + 1) * torch.tensor([w, h]).to(points).view(1, 1, 2) - 1) / 2 + return des_points + + def __len__(self): + return len(self.items) + + def __getitem__(self, index): + sample = dict() + + image_path = self.items.iloc[index, 0] + landmarks_5pts = self.items.iloc[index, 1] + landmarks_5pts = np.array(list(map(float, landmarks_5pts.split(","))), dtype=np.float32).reshape(5, 2) + landmarks_target = self.items.iloc[index, 2] + landmarks_target = np.array(list(map(float, landmarks_target.split(","))), dtype=np.float32).reshape( + self.landmark_num, 2) + scale = float(self.items.iloc[index, 3]) + center_w, center_h = float(self.items.iloc[index, 4]), float(self.items.iloc[index, 5]) + if len(self.items.iloc[index]) > 6: + tags = np.array(list(map(lambda x: int(float(x)), self.items.iloc[index, 6].split(",")))) + else: + tags = np.array([]) + + # image & keypoints alignment + image_path = image_path.replace('\\', '/') + # wflw testset + image_path = image_path.replace( + '//msr-facestore/Workspace/MSRA_EP_Allergan/users/yanghuan/training_data/wflw/rawImages/', '') + # trainset + image_path = image_path.replace('./rawImages/', '') + image_path = os.path.join(self.image_dir, image_path) + + # image path + sample["image_path"] = image_path + + img = self._load_image(image_path) # HWC, BGR, [0, 255] + assert img is not None + + # augmentation + # landmarks_target = [-0.5, edge-0.5] + img, landmarks_target, matrix = \ + self.augmentation.process(img, landmarks_target, landmarks_5pts, scale, center_w, center_h) + + landmarks = self._norm_points(torch.from_numpy(landmarks_target), self.image_height, self.image_width) + + sample["label"] = [landmarks, ] + + if self.use_AAM: + pointmap = self.encoder.generate_heatmap(landmarks_target) + edgemap = self._generate_edgemap(landmarks_target) + sample["label"] += [pointmap, edgemap] + + sample['matrix'] = matrix + + # image normalization + img = img.transpose(2, 0, 1).astype(np.float32) # CHW, BGR, [0, 255] + img[0, :, :] = (img[0, :, :] - self.means[0]) * self.scale + img[1, :, :] = (img[1, :, :] - self.means[1]) * self.scale + img[2, :, :] = (img[2, :, :] - self.means[2]) * self.scale + sample["data"] = torch.from_numpy(img) # CHW, BGR, [-1, 1] + + sample["tags"] = tags + + return sample diff --git a/models/STAR/lib/dataset/augmentation.py b/models/STAR/lib/dataset/augmentation.py new file mode 100644 index 0000000000000000000000000000000000000000..0694d316b8d20ccbb6fda6f4b9750c851f83995a --- /dev/null +++ b/models/STAR/lib/dataset/augmentation.py @@ -0,0 +1,355 @@ +import os +import cv2 +import math +import random +import numpy as np +from skimage import transform + + +class Augmentation: + def __init__(self, + is_train=True, + aug_prob=1.0, + image_size=256, + crop_op=True, + std_lmk_5pts=None, + target_face_scale=1.0, + flip_rate=0.5, + flip_mapping=None, + random_shift_sigma=0.05, + random_rot_sigma=math.pi/180*18, + random_scale_sigma=0.1, + random_gray_rate=0.2, + random_occ_rate=0.4, + random_blur_rate=0.3, + random_gamma_rate=0.2, + random_nose_fusion_rate=0.2): + self.is_train = is_train + self.aug_prob = aug_prob + self.crop_op = crop_op + self._flip = Flip(flip_mapping, flip_rate) + if self.crop_op: + self._cropMatrix = GetCropMatrix( + image_size=image_size, + target_face_scale=target_face_scale, + align_corners=True) + else: + self._alignMatrix = GetAlignMatrix( + image_size=image_size, + target_face_scale=target_face_scale, + std_lmk_5pts=std_lmk_5pts) + self._randomGeometryMatrix = GetRandomGeometryMatrix( + target_shape=(image_size, image_size), + from_shape=(image_size, image_size), + shift_sigma=random_shift_sigma, + rot_sigma=random_rot_sigma, + scale_sigma=random_scale_sigma, + align_corners=True) + self._transform = Transform(image_size=image_size) + self._randomTexture = RandomTexture( + random_gray_rate=random_gray_rate, + random_occ_rate=random_occ_rate, + random_blur_rate=random_blur_rate, + random_gamma_rate=random_gamma_rate, + random_nose_fusion_rate=random_nose_fusion_rate) + + def process(self, img, lmk, lmk_5pts=None, scale=1.0, center_w=0, center_h=0, is_train=True): + if self.is_train and random.random() < self.aug_prob: + img, lmk, lmk_5pts, center_w, center_h = self._flip.process(img, lmk, lmk_5pts, center_w, center_h) + matrix_geoaug = self._randomGeometryMatrix.process() + if self.crop_op: + matrix_pre = self._cropMatrix.process(scale, center_w, center_h) + else: + matrix_pre = self._alignMatrix.process(lmk_5pts) + matrix = matrix_geoaug @ matrix_pre + aug_img, aug_lmk = self._transform.process(img, lmk, matrix) + aug_img = self._randomTexture.process(aug_img) + else: + if self.crop_op: + matrix = self._cropMatrix.process(scale, center_w, center_h) + else: + matrix = self._alignMatrix.process(lmk_5pts) + aug_img, aug_lmk = self._transform.process(img, lmk, matrix) + return aug_img, aug_lmk, matrix + + +class GetCropMatrix: + def __init__(self, image_size, target_face_scale, align_corners=False): + self.image_size = image_size + self.target_face_scale = target_face_scale + self.align_corners = align_corners + + def _compose_rotate_and_scale(self, angle, scale, shift_xy, from_center, to_center): + cosv = math.cos(angle) + sinv = math.sin(angle) + + fx, fy = from_center + tx, ty = to_center + + acos = scale * cosv + asin = scale * sinv + + a0 = acos + a1 = -asin + a2 = tx - acos * fx + asin * fy + shift_xy[0] + + b0 = asin + b1 = acos + b2 = ty - asin * fx - acos * fy + shift_xy[1] + + rot_scale_m = np.array([ + [a0, a1, a2], + [b0, b1, b2], + [0.0, 0.0, 1.0] + ], np.float32) + return rot_scale_m + + def process(self, scale, center_w, center_h): + if self.align_corners: + to_w, to_h = self.image_size-1, self.image_size-1 + else: + to_w, to_h = self.image_size, self.image_size + + rot_mu = 0 + scale_mu = self.image_size / (scale * self.target_face_scale * 200.0) + shift_xy_mu = (0, 0) + matrix = self._compose_rotate_and_scale( + rot_mu, scale_mu, shift_xy_mu, + from_center=[center_w, center_h], + to_center=[to_w/2.0, to_h/2.0]) + return matrix + + +class GetAlignMatrix: + def __init__(self, image_size, target_face_scale, std_lmk_5pts): + """ + points in std_lmk_5pts range from -1 to 1. + """ + self.std_lmk_5pts = (std_lmk_5pts * target_face_scale + 1) * \ + np.array([image_size, image_size], np.float32) / 2.0 + + def process(self, lmk_5pts): + assert lmk_5pts.shape[-2:] == (5, 2) + tform = transform.SimilarityTransform() + tform.estimate(lmk_5pts, self.std_lmk_5pts) + return tform.params + + +class GetRandomGeometryMatrix: + def __init__(self, target_shape, from_shape, + shift_sigma=0.1, rot_sigma=18*math.pi/180, scale_sigma=0.1, + shift_mu=0.0, rot_mu=0.0, scale_mu=1.0, + shift_normal=True, rot_normal=True, scale_normal=True, + align_corners=False): + self.target_shape = target_shape + self.from_shape = from_shape + self.shift_config = (shift_mu, shift_sigma, shift_normal) + self.rot_config = (rot_mu, rot_sigma, rot_normal) + self.scale_config = (scale_mu, scale_sigma, scale_normal) + self.align_corners = align_corners + + def _compose_rotate_and_scale(self, angle, scale, shift_xy, from_center, to_center): + cosv = math.cos(angle) + sinv = math.sin(angle) + + fx, fy = from_center + tx, ty = to_center + + acos = scale * cosv + asin = scale * sinv + + a0 = acos + a1 = -asin + a2 = tx - acos * fx + asin * fy + shift_xy[0] + + b0 = asin + b1 = acos + b2 = ty - asin * fx - acos * fy + shift_xy[1] + + rot_scale_m = np.array([ + [a0, a1, a2], + [b0, b1, b2], + [0.0, 0.0, 1.0] + ], np.float32) + return rot_scale_m + + def _random(self, mu_sigma_normal, size=None): + mu, sigma, is_normal = mu_sigma_normal + if is_normal: + return np.random.normal(mu, sigma, size=size) + else: + return np.random.uniform(low=mu-sigma, high=mu+sigma, size=size) + + def process(self): + if self.align_corners: + from_w, from_h = self.from_shape[1]-1, self.from_shape[0]-1 + to_w, to_h = self.target_shape[1]-1, self.target_shape[0]-1 + else: + from_w, from_h = self.from_shape[1], self.from_shape[0] + to_w, to_h = self.target_shape[1], self.target_shape[0] + + if self.shift_config[:2] != (0.0, 0.0) or \ + self.rot_config[:2] != (0.0, 0.0) or \ + self.scale_config[:2] != (1.0, 0.0): + shift_xy = self._random(self.shift_config, size=[2]) * \ + min(to_h, to_w) + rot_angle = self._random(self.rot_config) + scale = self._random(self.scale_config) + matrix_geoaug = self._compose_rotate_and_scale( + rot_angle, scale, shift_xy, + from_center=[from_w/2.0, from_h/2.0], + to_center=[to_w/2.0, to_h/2.0]) + + return matrix_geoaug + + +class Transform: + def __init__(self, image_size): + self.image_size = image_size + + def _transformPoints2D(self, points, matrix): + """ + points (nx2), matrix (3x3) -> points (nx2) + """ + dtype = points.dtype + + # nx3 + points = np.concatenate([points, np.ones_like(points[:, [0]])], axis=1) + points = points @ np.transpose(matrix) + points = points[:, :2] / points[:, [2, 2]] + return points.astype(dtype) + + def _transformPerspective(self, image, matrix): + """ + image, matrix3x3 -> transformed_image + """ + return cv2.warpPerspective( + image, matrix, + dsize=(self.image_size, self.image_size), + flags=cv2.INTER_LINEAR, borderValue=0) + + def process(self, image, landmarks, matrix): + t_landmarks = self._transformPoints2D(landmarks, matrix) + t_image = self._transformPerspective(image, matrix) + return t_image, t_landmarks + + +class RandomTexture: + def __init__(self, random_gray_rate=0, random_occ_rate=0, random_blur_rate=0, random_gamma_rate=0, random_nose_fusion_rate=0): + self.random_gray_rate = random_gray_rate + self.random_occ_rate = random_occ_rate + self.random_blur_rate = random_blur_rate + self.random_gamma_rate = random_gamma_rate + self.random_nose_fusion_rate = random_nose_fusion_rate + self.texture_augs = ( + (self.add_occ, self.random_occ_rate), + (self.add_blur, self.random_blur_rate), + (self.add_gamma, self.random_gamma_rate), + (self.add_nose_fusion, self.random_nose_fusion_rate) + ) + + def add_gray(self, image): + assert image.ndim == 3 and image.shape[-1] == 3 + image = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY) + image = np.tile(np.expand_dims(image, -1), [1, 1, 3]) + return image + + def add_occ(self, image): + h, w, c = image.shape + rh = 0.2 + 0.6 * random.random() # [0.2, 0.8] + rw = rh - 0.2 + 0.4 * random.random() + cx = int((h - 1) * random.random()) + cy = int((w - 1) * random.random()) + dh = int(h / 2 * rh) + dw = int(w / 2 * rw) + x0 = max(0, cx - dw // 2) + y0 = max(0, cy - dh // 2) + x1 = min(w - 1, cx + dw // 2) + y1 = min(h - 1, cy + dh // 2) + image[y0:y1+1, x0:x1+1] = 0 + return image + + def add_blur(self, image): + blur_kratio = 0.05 * random.random() + blur_ksize = int((image.shape[0] + image.shape[1]) / 2 * blur_kratio) + if blur_ksize > 1: + image = cv2.blur(image, (blur_ksize, blur_ksize)) + return image + + def add_gamma(self, image): + if random.random() < 0.5: + gamma = 0.25 + 0.75 * random.random() + else: + gamma = 1.0 + 3.0 * random.random() + image = (((image / 255.0) ** gamma) * 255).astype("uint8") + return image + + def add_nose_fusion(self, image): + h, w, c = image.shape + nose = np.array(bytearray(os.urandom(h * w * c)), dtype=image.dtype).reshape(h, w, c) + alpha = 0.5 * random.random() + image = (1 - alpha) * image + alpha * nose + return image.astype(np.uint8) + + def process(self, image): + image = image.copy() + if random.random() < self.random_occ_rate: + image = self.add_occ(image) + if random.random() < self.random_blur_rate: + image = self.add_blur(image) + if random.random() < self.random_gamma_rate: + image = self.add_gamma(image) + if random.random() < self.random_nose_fusion_rate: + image = self.add_nose_fusion(image) + """ + orders = list(range(len(self.texture_augs))) + random.shuffle(orders) + for order in orders: + if random.random() < self.texture_augs[order][1]: + image = self.texture_augs[order][0](image) + """ + + if random.random() < self.random_gray_rate: + image = self.add_gray(image) + + return image + + +class Flip: + def __init__(self, flip_mapping, random_rate): + self.flip_mapping = flip_mapping + self.random_rate = random_rate + + def process(self, image, landmarks, landmarks_5pts, center_w, center_h): + if random.random() >= self.random_rate or self.flip_mapping is None: + return image, landmarks, landmarks_5pts, center_w, center_h + + # COFW + if landmarks.shape[0] == 29: + flip_offset = 0 + # 300W, WFLW + elif landmarks.shape[0] in (68, 98): + flip_offset = -1 + else: + flip_offset = -1 + + h, w, _ = image.shape + #image_flip = cv2.flip(image, 1) + image_flip = np.fliplr(image).copy() + landmarks_flip = landmarks.copy() + for i, j in self.flip_mapping: + landmarks_flip[i] = landmarks[j] + landmarks_flip[j] = landmarks[i] + landmarks_flip[:, 0] = w + flip_offset - landmarks_flip[:, 0] + if landmarks_5pts is not None: + flip_mapping = ([0, 1], [3, 4]) + landmarks_5pts_flip = landmarks_5pts.copy() + for i, j in flip_mapping: + landmarks_5pts_flip[i] = landmarks_5pts[j] + landmarks_5pts_flip[j] = landmarks_5pts[i] + landmarks_5pts_flip[:, 0] = w + flip_offset - landmarks_5pts_flip[:, 0] + else: + landmarks_5pts_flip = None + + center_w = w + flip_offset - center_w + return image_flip, landmarks_flip, landmarks_5pts_flip, center_w, center_h diff --git a/models/STAR/lib/dataset/decoder/__init__.py b/models/STAR/lib/dataset/decoder/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..2315040c50fbccacc9cff3e0b8e92c1691a4600b --- /dev/null +++ b/models/STAR/lib/dataset/decoder/__init__.py @@ -0,0 +1,8 @@ +from .decoder_default import decoder_default + +def get_decoder(decoder_type='default'): + if decoder_type == 'default': + decoder = decoder_default() + else: + raise NotImplementedError + return decoder \ No newline at end of file diff --git a/models/STAR/lib/dataset/decoder/decoder_default.py b/models/STAR/lib/dataset/decoder/decoder_default.py new file mode 100644 index 0000000000000000000000000000000000000000..19b981e3392e17b065a12a96861c64e600d790bd --- /dev/null +++ b/models/STAR/lib/dataset/decoder/decoder_default.py @@ -0,0 +1,38 @@ +import torch + + +class decoder_default: + def __init__(self, weight=1, use_weight_map=False): + self.weight = weight + self.use_weight_map = use_weight_map + + def _make_grid(self, h, w): + yy, xx = torch.meshgrid( + torch.arange(h).float() / (h - 1) * 2 - 1, + torch.arange(w).float() / (w - 1) * 2 - 1) + return yy, xx + + def get_coords_from_heatmap(self, heatmap): + """ + inputs: + - heatmap: batch x npoints x h x w + + outputs: + - coords: batch x npoints x 2 (x,y), [-1, +1] + - radius_sq: batch x npoints + """ + batch, npoints, h, w = heatmap.shape + if self.use_weight_map: + heatmap = heatmap * self.weight + + yy, xx = self._make_grid(h, w) + yy = yy.view(1, 1, h, w).to(heatmap) + xx = xx.view(1, 1, h, w).to(heatmap) + + heatmap_sum = torch.clamp(heatmap.sum([2, 3]), min=1e-6) + + yy_coord = (yy * heatmap).sum([2, 3]) / heatmap_sum # batch x npoints + xx_coord = (xx * heatmap).sum([2, 3]) / heatmap_sum # batch x npoints + coords = torch.stack([xx_coord, yy_coord], dim=-1) + + return coords diff --git a/models/STAR/lib/dataset/encoder/__init__.py b/models/STAR/lib/dataset/encoder/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..b80fe999ee83928d6abde95c51ef03c6bd052098 --- /dev/null +++ b/models/STAR/lib/dataset/encoder/__init__.py @@ -0,0 +1,8 @@ +from .encoder_default import encoder_default + +def get_encoder(image_height, image_width, scale=0.25, sigma=1.5, encoder_type='default'): + if encoder_type == 'default': + encoder = encoder_default(image_height, image_width, scale, sigma) + else: + raise NotImplementedError + return encoder diff --git a/models/STAR/lib/dataset/encoder/encoder_default.py b/models/STAR/lib/dataset/encoder/encoder_default.py new file mode 100644 index 0000000000000000000000000000000000000000..6662a94bdc5788ebb00cc2eb0c4103ba5a10d653 --- /dev/null +++ b/models/STAR/lib/dataset/encoder/encoder_default.py @@ -0,0 +1,63 @@ +import copy +import numpy as np + +import torch +import torch.nn.functional as F + + +class encoder_default: + def __init__(self, image_height, image_width, scale=0.25, sigma=1.5): + self.image_height = image_height + self.image_width = image_width + self.scale = scale + self.sigma = sigma + + def generate_heatmap(self, points): + # points = (num_pts, 2) + h, w = self.image_height, self.image_width + pointmaps = [] + for i in range(len(points)): + pointmap = np.zeros([h, w], dtype=np.float32) + # align_corners: False. + point = copy.deepcopy(points[i]) + point[0] = max(0, min(w - 1, point[0])) + point[1] = max(0, min(h - 1, point[1])) + pointmap = self._circle(pointmap, point, sigma=self.sigma) + + pointmaps.append(pointmap) + pointmaps = np.stack(pointmaps, axis=0) / 255.0 + pointmaps = torch.from_numpy(pointmaps).float().unsqueeze(0) + pointmaps = F.interpolate(pointmaps, size=(int(w * self.scale), int(h * self.scale)), mode='bilinear', + align_corners=False).squeeze() + return pointmaps + + def _circle(self, img, pt, sigma=1.0, label_type='Gaussian'): + # Check that any part of the gaussian is in-bounds + tmp_size = sigma * 3 + ul = [int(pt[0] - tmp_size), int(pt[1] - tmp_size)] + br = [int(pt[0] + tmp_size + 1), int(pt[1] + tmp_size + 1)] + if (ul[0] > img.shape[1] - 1 or ul[1] > img.shape[0] - 1 or + br[0] - 1 < 0 or br[1] - 1 < 0): + # If not, just return the image as is + return img + + # Generate gaussian + size = 2 * tmp_size + 1 + x = np.arange(0, size, 1, np.float32) + y = x[:, np.newaxis] + x0 = y0 = size // 2 + # The gaussian is not normalized, we want the center value to equal 1 + if label_type == 'Gaussian': + g = np.exp(- ((x - x0) ** 2 + (y - y0) ** 2) / (2 * sigma ** 2)) + else: + g = sigma / (((x - x0) ** 2 + (y - y0) ** 2 + sigma ** 2) ** 1.5) + + # Usable gaussian range + g_x = max(0, -ul[0]), min(br[0], img.shape[1]) - ul[0] + g_y = max(0, -ul[1]), min(br[1], img.shape[0]) - ul[1] + # Image range + img_x = max(0, ul[0]), min(br[0], img.shape[1]) + img_y = max(0, ul[1]), min(br[1], img.shape[0]) + + img[img_y[0]:img_y[1], img_x[0]:img_x[1]] = 255 * g[g_y[0]:g_y[1], g_x[0]:g_x[1]] + return img diff --git a/models/STAR/lib/loss/__init__.py b/models/STAR/lib/loss/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f71a33bddba049b10a2a94ee5e51342f50878637 --- /dev/null +++ b/models/STAR/lib/loss/__init__.py @@ -0,0 +1,14 @@ +from .awingLoss import AWingLoss +from .smoothL1Loss import SmoothL1Loss +from .wingLoss import WingLoss +from .starLoss import STARLoss +from .starLoss_v2 import STARLoss_v2 + +__all__ = [ + "AWingLoss", + "SmoothL1Loss", + "WingLoss", + "STARLoss", + + "STARLoss_v2", +] diff --git a/models/STAR/lib/loss/awingLoss.py b/models/STAR/lib/loss/awingLoss.py new file mode 100644 index 0000000000000000000000000000000000000000..a5bfc579bd52702506e0dd0c4a6f2932d0698586 --- /dev/null +++ b/models/STAR/lib/loss/awingLoss.py @@ -0,0 +1,39 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F + + +class AWingLoss(nn.Module): + def __init__(self, omega=14, theta=0.5, epsilon=1, alpha=2.1, use_weight_map=True): + super(AWingLoss, self).__init__() + self.omega = omega + self.theta = theta + self.epsilon = epsilon + self.alpha = alpha + self.use_weight_map = use_weight_map + + def __repr__(self): + return "AWingLoss()" + + def generate_weight_map(self, heatmap, k_size=3, w=10): + dilate = F.max_pool2d(heatmap, kernel_size=k_size, stride=1, padding=1) + weight_map = torch.where(dilate < 0.2, torch.zeros_like(heatmap), torch.ones_like(heatmap)) + return w * weight_map + 1 + + def forward(self, output, groundtruth): + """ + input: b x n x h x w + output: b x n x h x w => 1 + """ + delta = (output - groundtruth).abs() + A = self.omega * (1 / (1 + torch.pow(self.theta / self.epsilon, self.alpha - groundtruth))) * (self.alpha - groundtruth) * \ + (torch.pow(self.theta / self.epsilon, self.alpha - groundtruth - 1)) * (1 / self.epsilon) + C = self.theta * A - self.omega * \ + torch.log(1 + torch.pow(self.theta / self.epsilon, self.alpha - groundtruth)) + loss = torch.where(delta < self.theta, + self.omega * torch.log(1 + torch.pow(delta / self.epsilon, self.alpha - groundtruth)), + (A * delta - C)) + if self.use_weight_map: + weight = self.generate_weight_map(groundtruth) + loss = loss * weight + return loss.mean() diff --git a/models/STAR/lib/loss/smoothL1Loss.py b/models/STAR/lib/loss/smoothL1Loss.py new file mode 100644 index 0000000000000000000000000000000000000000..e81104da2f7a3134103bea6a1463a924f4f6ce20 --- /dev/null +++ b/models/STAR/lib/loss/smoothL1Loss.py @@ -0,0 +1,36 @@ +import torch +import torch.nn as nn + + +class SmoothL1Loss(nn.Module): + def __init__(self, scale=0.01): + super(SmoothL1Loss, self).__init__() + self.scale = scale + self.EPSILON = 1e-10 + + def __repr__(self): + return "SmoothL1Loss()" + + def forward(self, output: torch.Tensor, groundtruth: torch.Tensor, reduction='mean'): + """ + input: b x n x 2 + output: b x n x 1 => 1 + """ + if output.dim() == 4: + shape = output.shape + groundtruth = groundtruth.reshape(shape[0], shape[1], 1, shape[3]) + + delta_2 = (output - groundtruth).pow(2).sum(dim=-1, keepdim=False) + delta = delta_2.clamp(min=1e-6).sqrt() + # delta = torch.sqrt(delta_2 + self.EPSILON) + loss = torch.where( \ + delta_2 < self.scale * self.scale, \ + 0.5 / self.scale * delta_2, \ + delta - 0.5 * self.scale) + + if reduction == 'mean': + loss = loss.mean() + elif reduction == 'sum': + loss = loss.sum() + + return loss diff --git a/models/STAR/lib/loss/starLoss.py b/models/STAR/lib/loss/starLoss.py new file mode 100644 index 0000000000000000000000000000000000000000..bfd43782daba6f01d7d5ad9647c7028de111dd13 --- /dev/null +++ b/models/STAR/lib/loss/starLoss.py @@ -0,0 +1,140 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.autograd import Variable + +from .smoothL1Loss import SmoothL1Loss +from .wingLoss import WingLoss + + +def get_channel_sum(input): + temp = torch.sum(input, dim=3) + output = torch.sum(temp, dim=2) + return output + + +def expand_two_dimensions_at_end(input, dim1, dim2): + input = input.unsqueeze(-1).unsqueeze(-1) + input = input.expand(-1, -1, dim1, dim2) + return input + + +class STARLoss(nn.Module): + def __init__(self, w=1, dist='smoothl1', num_dim_image=2, EPSILON=1e-5): + super(STARLoss, self).__init__() + self.w = w + self.num_dim_image = num_dim_image + self.EPSILON = EPSILON + self.dist = dist + if self.dist == 'smoothl1': + self.dist_func = SmoothL1Loss() + elif self.dist == 'l1': + self.dist_func = F.l1_loss + elif self.dist == 'l2': + self.dist_func = F.mse_loss + elif self.dist == 'wing': + self.dist_func = WingLoss() + else: + raise NotImplementedError + + def __repr__(self): + return "STARLoss()" + + def _make_grid(self, h, w): + yy, xx = torch.meshgrid( + torch.arange(h).float() / (h - 1) * 2 - 1, + torch.arange(w).float() / (w - 1) * 2 - 1) + return yy, xx + + def weighted_mean(self, heatmap): + batch, npoints, h, w = heatmap.shape + + yy, xx = self._make_grid(h, w) + yy = yy.view(1, 1, h, w).to(heatmap) + xx = xx.view(1, 1, h, w).to(heatmap) + + yy_coord = (yy * heatmap).sum([2, 3]) # batch x npoints + xx_coord = (xx * heatmap).sum([2, 3]) # batch x npoints + coords = torch.stack([xx_coord, yy_coord], dim=-1) + return coords + + def unbiased_weighted_covariance(self, htp, means, num_dim_image=2, EPSILON=1e-5): + batch_size, num_points, height, width = htp.shape + + yv, xv = self._make_grid(height, width) + xv = Variable(xv) + yv = Variable(yv) + + if htp.is_cuda: + xv = xv.cuda() + yv = yv.cuda() + + xmean = means[:, :, 0] + xv_minus_mean = xv.expand(batch_size, num_points, -1, -1) - expand_two_dimensions_at_end(xmean, height, + width) # [batch_size, 68, 64, 64] + ymean = means[:, :, 1] + yv_minus_mean = yv.expand(batch_size, num_points, -1, -1) - expand_two_dimensions_at_end(ymean, height, + width) # [batch_size, 68, 64, 64] + wt_xv_minus_mean = xv_minus_mean + wt_yv_minus_mean = yv_minus_mean + + wt_xv_minus_mean = wt_xv_minus_mean.view(batch_size * num_points, height * width) # [batch_size*68, 4096] + wt_xv_minus_mean = wt_xv_minus_mean.view(batch_size * num_points, 1, height * width) # [batch_size*68, 1, 4096] + wt_yv_minus_mean = wt_yv_minus_mean.view(batch_size * num_points, height * width) # [batch_size*68, 4096] + wt_yv_minus_mean = wt_yv_minus_mean.view(batch_size * num_points, 1, height * width) # [batch_size*68, 1, 4096] + vec_concat = torch.cat((wt_xv_minus_mean, wt_yv_minus_mean), 1) # [batch_size*68, 2, 4096] + + htp_vec = htp.view(batch_size * num_points, 1, height * width) + htp_vec = htp_vec.expand(-1, 2, -1) + + covariance = torch.bmm(htp_vec * vec_concat, vec_concat.transpose(1, 2)) # [batch_size*68, 2, 2] + covariance = covariance.view(batch_size, num_points, num_dim_image, num_dim_image) # [batch_size, 68, 2, 2] + + V_1 = htp.sum([2, 3]) + EPSILON # [batch_size, 68] + V_2 = torch.pow(htp, 2).sum([2, 3]) + EPSILON # [batch_size, 68] + + denominator = V_1 - (V_2 / V_1) + covariance = covariance / expand_two_dimensions_at_end(denominator, num_dim_image, num_dim_image) + + return covariance + + def ambiguity_guided_decompose(self, pts, eigenvalues, eigenvectors): + batch_size, npoints = pts.shape[:2] + rotate = torch.matmul(pts.view(batch_size, npoints, 1, 2), eigenvectors.transpose(-1, -2)) + scale = rotate.view(batch_size, npoints, 2) / torch.sqrt(eigenvalues + self.EPSILON) + return scale + + def eigenvalue_restriction(self, evalues, batch, npoints): + eigen_loss = torch.abs(evalues.view(batch * npoints, 2)).sum(-1) + return eigen_loss.mean() + + def forward(self, heatmap, groundtruth): + """ + heatmap: b x n x 64 x 64 + groundtruth: b x n x 2 + output: b x n x 1 => 1 + """ + # normalize + bs, npoints, h, w = heatmap.shape + heatmap_sum = torch.clamp(heatmap.sum([2, 3]), min=1e-6) + heatmap = heatmap / heatmap_sum.view(bs, npoints, 1, 1) + + means = self.weighted_mean(heatmap) # [bs, 68, 2] + covars = self.unbiased_weighted_covariance(heatmap, means) # covars [bs, 68, 2, 2] + + # TODO: GPU-based eigen-decomposition + # https://github.com/pytorch/pytorch/issues/60537 + _covars = covars.view(bs * npoints, 2, 2).cpu() + evalues, evectors = _covars.symeig(eigenvectors=True) # evalues [bs * 68, 2], evectors [bs * 68, 2, 2] + evalues = evalues.view(bs, npoints, 2).to(heatmap) + evectors = evectors.view(bs, npoints, 2, 2).to(heatmap) + + # STAR Loss + # Ambiguity-guided Decomposition + error = self.ambiguity_guided_decompose(groundtruth - means, evalues, evectors) + loss_trans = self.dist_func(torch.zeros_like(error).to(error), error) + # Eigenvalue Restriction + loss_eigen = self.eigenvalue_restriction(evalues, bs, npoints) + star_loss = loss_trans + self.w * loss_eigen + + return star_loss diff --git a/models/STAR/lib/loss/starLoss_v2.py b/models/STAR/lib/loss/starLoss_v2.py new file mode 100644 index 0000000000000000000000000000000000000000..c182ff81b24eb0f46130485ef16baff57632810d --- /dev/null +++ b/models/STAR/lib/loss/starLoss_v2.py @@ -0,0 +1,150 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.autograd import Variable + +from .smoothL1Loss import SmoothL1Loss +from .wingLoss import WingLoss + + +def get_channel_sum(input): + temp = torch.sum(input, dim=3) + output = torch.sum(temp, dim=2) + return output + + +def expand_two_dimensions_at_end(input, dim1, dim2): + input = input.unsqueeze(-1).unsqueeze(-1) + input = input.expand(-1, -1, dim1, dim2) + return input + + +class STARLoss_v2(nn.Module): + def __init__(self, w=1, dist='smoothl1', num_dim_image=2, EPSILON=1e-5): + super(STARLoss_v2, self).__init__() + self.w = w + self.num_dim_image = num_dim_image + self.EPSILON = EPSILON + self.dist = dist + if self.dist == 'smoothl1': + self.dist_func = SmoothL1Loss() + elif self.dist == 'l1': + self.dist_func = F.l1_loss + elif self.dist == 'l2': + self.dist_func = F.mse_loss + elif self.dist == 'wing': + self.dist_func = WingLoss() + else: + raise NotImplementedError + + def __repr__(self): + return "STARLoss()" + + def _make_grid(self, h, w): + yy, xx = torch.meshgrid( + torch.arange(h).float() / (h - 1) * 2 - 1, + torch.arange(w).float() / (w - 1) * 2 - 1) + return yy, xx + + def weighted_mean(self, heatmap): + batch, npoints, h, w = heatmap.shape + + yy, xx = self._make_grid(h, w) + yy = yy.view(1, 1, h, w).to(heatmap) + xx = xx.view(1, 1, h, w).to(heatmap) + + yy_coord = (yy * heatmap).sum([2, 3]) # batch x npoints + xx_coord = (xx * heatmap).sum([2, 3]) # batch x npoints + coords = torch.stack([xx_coord, yy_coord], dim=-1) + return coords + + def unbiased_weighted_covariance(self, htp, means, num_dim_image=2, EPSILON=1e-5): + batch_size, num_points, height, width = htp.shape + + yv, xv = self._make_grid(height, width) + xv = Variable(xv) + yv = Variable(yv) + + if htp.is_cuda: + xv = xv.cuda() + yv = yv.cuda() + + xmean = means[:, :, 0] + xv_minus_mean = xv.expand(batch_size, num_points, -1, -1) - expand_two_dimensions_at_end(xmean, height, + width) # [batch_size, 68, 64, 64] + ymean = means[:, :, 1] + yv_minus_mean = yv.expand(batch_size, num_points, -1, -1) - expand_two_dimensions_at_end(ymean, height, + width) # [batch_size, 68, 64, 64] + wt_xv_minus_mean = xv_minus_mean + wt_yv_minus_mean = yv_minus_mean + + wt_xv_minus_mean = wt_xv_minus_mean.view(batch_size * num_points, height * width) # [batch_size*68, 4096] + wt_xv_minus_mean = wt_xv_minus_mean.view(batch_size * num_points, 1, height * width) # [batch_size*68, 1, 4096] + wt_yv_minus_mean = wt_yv_minus_mean.view(batch_size * num_points, height * width) # [batch_size*68, 4096] + wt_yv_minus_mean = wt_yv_minus_mean.view(batch_size * num_points, 1, height * width) # [batch_size*68, 1, 4096] + vec_concat = torch.cat((wt_xv_minus_mean, wt_yv_minus_mean), 1) # [batch_size*68, 2, 4096] + + htp_vec = htp.view(batch_size * num_points, 1, height * width) + htp_vec = htp_vec.expand(-1, 2, -1) + + covariance = torch.bmm(htp_vec * vec_concat, vec_concat.transpose(1, 2)) # [batch_size*68, 2, 2] + covariance = covariance.view(batch_size, num_points, num_dim_image, num_dim_image) # [batch_size, 68, 2, 2] + + V_1 = htp.sum([2, 3]) + EPSILON # [batch_size, 68] + V_2 = torch.pow(htp, 2).sum([2, 3]) + EPSILON # [batch_size, 68] + + denominator = V_1 - (V_2 / V_1) + covariance = covariance / expand_two_dimensions_at_end(denominator, num_dim_image, num_dim_image) + + return covariance + + def ambiguity_guided_decompose(self, error, evalues, evectors): + bs, npoints = error.shape[:2] + normal_vector = evectors[:, :, 0] + tangent_vector = evectors[:, :, 1] + normal_error = torch.matmul(normal_vector.unsqueeze(-2), error.unsqueeze(-1)) + tangent_error = torch.matmul(tangent_vector.unsqueeze(-2), error.unsqueeze(-1)) + normal_error = normal_error.squeeze(dim=-1) + tangent_error = tangent_error.squeeze(dim=-1) + normal_dist = self.dist_func(normal_error, torch.zeros_like(normal_error).to(normal_error), reduction='none') + tangent_dist = self.dist_func(tangent_error, torch.zeros_like(tangent_error).to(tangent_error), reduction='none') + normal_dist = normal_dist.reshape(bs, npoints, 1) + tangent_dist = tangent_dist.reshape(bs, npoints, 1) + dist = torch.cat((normal_dist, tangent_dist), dim=-1) + scale_dist = dist / torch.sqrt(evalues + self.EPSILON) + scale_dist = scale_dist.sum(-1) + return scale_dist + + def eigenvalue_restriction(self, evalues, batch, npoints): + eigen_loss = torch.abs(evalues.view(batch, npoints, 2)).sum(-1) + return eigen_loss + + def forward(self, heatmap, groundtruth): + """ + heatmap: b x n x 64 x 64 + groundtruth: b x n x 2 + output: b x n x 1 => 1 + """ + # normalize + bs, npoints, h, w = heatmap.shape + heatmap_sum = torch.clamp(heatmap.sum([2, 3]), min=1e-6) + heatmap = heatmap / heatmap_sum.view(bs, npoints, 1, 1) + + means = self.weighted_mean(heatmap) # [bs, 68, 2] + covars = self.unbiased_weighted_covariance(heatmap, means) # covars [bs, 68, 2, 2] + + # TODO: GPU-based eigen-decomposition + # https://github.com/pytorch/pytorch/issues/60537 + _covars = covars.view(bs * npoints, 2, 2).cpu() + evalues, evectors = _covars.symeig(eigenvectors=True) # evalues [bs * 68, 2], evectors [bs * 68, 2, 2] + evalues = evalues.view(bs, npoints, 2).to(heatmap) + evectors = evectors.view(bs, npoints, 2, 2).to(heatmap) + + # STAR Loss + # Ambiguity-guided Decomposition + loss_trans = self.ambiguity_guided_decompose(groundtruth - means, evalues, evectors) + # Eigenvalue Restriction + loss_eigen = self.eigenvalue_restriction(evalues, bs, npoints) + star_loss = loss_trans + self.w * loss_eigen + + return star_loss.mean() diff --git a/models/STAR/lib/loss/wingLoss.py b/models/STAR/lib/loss/wingLoss.py new file mode 100644 index 0000000000000000000000000000000000000000..578f71cbf9947ce585b2aef6c060a78f1f90993f --- /dev/null +++ b/models/STAR/lib/loss/wingLoss.py @@ -0,0 +1,27 @@ +# -*- coding: utf-8 -*- + +import math +import torch +from torch import nn + + +# torch.log and math.log is e based +class WingLoss(nn.Module): + def __init__(self, omega=0.01, epsilon=2): + super(WingLoss, self).__init__() + self.omega = omega + self.epsilon = epsilon + + def forward(self, pred, target): + y = target + y_hat = pred + delta_2 = (y - y_hat).pow(2).sum(dim=-1, keepdim=False) + # delta = delta_2.sqrt() + delta = delta_2.clamp(min=1e-6).sqrt() + C = self.omega - self.omega * math.log(1 + self.omega / self.epsilon) + loss = torch.where( + delta < self.omega, + self.omega * torch.log(1 + delta / self.epsilon), + delta - C + ) + return loss.mean() diff --git a/models/STAR/lib/metric/__init__.py b/models/STAR/lib/metric/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e843d42131dbd60a09f8951ddb652f9ce358b3f9 --- /dev/null +++ b/models/STAR/lib/metric/__init__.py @@ -0,0 +1,11 @@ +from .nme import NME +from .accuracy import Accuracy +from .fr_and_auc import FR_AUC +from .params import count_parameters_in_MB + +__all__ = [ + "NME", + "Accuracy", + "FR_AUC", + 'count_parameters_in_MB', +] diff --git a/models/STAR/lib/metric/accuracy.py b/models/STAR/lib/metric/accuracy.py new file mode 100644 index 0000000000000000000000000000000000000000..d007da207bdd203ad4dbff24cfa5d5e0733c7b76 --- /dev/null +++ b/models/STAR/lib/metric/accuracy.py @@ -0,0 +1,21 @@ +import torch +import torch.nn.functional as F + +class Accuracy: + def __init__(self): + pass + + def __repr__(self): + return "Accuracy()" + + def test(self, label_pd, label_gt, ignore_label=-1): + correct_cnt = 0 + total_cnt = 0 + with torch.no_grad(): + label_pd = F.softmax(label_pd, dim=1) + label_pd = torch.max(label_pd, 1)[1] + label_gt = label_gt.long() + c = (label_pd == label_gt) + correct_cnt = torch.sum(c).item() + total_cnt = c.size(0) - torch.sum(label_gt==ignore_label).item() + return correct_cnt, total_cnt diff --git a/models/STAR/lib/metric/fr_and_auc.py b/models/STAR/lib/metric/fr_and_auc.py new file mode 100644 index 0000000000000000000000000000000000000000..b4ceec47500a04eb1e42aa02ba9b67aeaf44e46b --- /dev/null +++ b/models/STAR/lib/metric/fr_and_auc.py @@ -0,0 +1,25 @@ +import numpy as np +from scipy.integrate import simps + + +class FR_AUC: + def __init__(self, data_definition): + self.data_definition = data_definition + if data_definition == '300W': + self.thresh = 0.05 + else: + self.thresh = 0.1 + + def __repr__(self): + return "FR_AUC()" + + def test(self, nmes, thres=None, step=0.0001): + if thres is None: + thres = self.thresh + + num_data = len(nmes) + xs = np.arange(0, thres + step, step) + ys = np.array([np.count_nonzero(nmes <= x) for x in xs]) / float(num_data) + fr = 1.0 - ys[-1] + auc = simps(ys, x=xs) / thres + return [round(fr, 4), round(auc, 6)] diff --git a/models/STAR/lib/metric/nme.py b/models/STAR/lib/metric/nme.py new file mode 100644 index 0000000000000000000000000000000000000000..2da6b07dc48bd1863bbc6d8365fe0e821f6f8783 --- /dev/null +++ b/models/STAR/lib/metric/nme.py @@ -0,0 +1,39 @@ +import torch +import numpy as np + +class NME: + def __init__(self, nme_left_index, nme_right_index): + self.nme_left_index = nme_left_index + self.nme_right_index = nme_right_index + + def __repr__(self): + return "NME()" + + def get_norm_distance(self, landmarks): + assert isinstance(self.nme_right_index, list), 'the nme_right_index is not list.' + assert isinstance(self.nme_left_index, list), 'the nme_left, index is not list.' + right_pupil = landmarks[self.nme_right_index, :].mean(0) + left_pupil = landmarks[self.nme_left_index, :].mean(0) + norm_distance = np.linalg.norm(right_pupil - left_pupil) + return norm_distance + + def test(self, label_pd, label_gt): + nme_list = [] + label_pd = label_pd.data.cpu().numpy() + label_gt = label_gt.data.cpu().numpy() + + for i in range(label_gt.shape[0]): + landmarks_gt = label_gt[i] + landmarks_pv = label_pd[i] + if isinstance(self.nme_right_index, list): + norm_distance = self.get_norm_distance(landmarks_gt) + elif isinstance(self.nme_right_index, int): + norm_distance = np.linalg.norm(landmarks_gt[self.nme_left_index] - landmarks_gt[self.nme_right_index]) + else: + raise NotImplementedError + landmarks_delta = landmarks_pv - landmarks_gt + nme = (np.linalg.norm(landmarks_delta, axis=1) / norm_distance).mean() + nme_list.append(nme) + # sum_nme += nme + # total_cnt += 1 + return nme_list diff --git a/models/STAR/lib/metric/params.py b/models/STAR/lib/metric/params.py new file mode 100644 index 0000000000000000000000000000000000000000..7b5552085951f85279bf386242a235063f9f77ea --- /dev/null +++ b/models/STAR/lib/metric/params.py @@ -0,0 +1,7 @@ +import torch.nn as nn + +def count_parameters_in_MB(model): + if isinstance(model, nn.Module): + return sum(v.numel() for v in model.parameters()) / 1e6 + else: + return sum(v.numel() for v in model) / 1e6 \ No newline at end of file diff --git a/models/STAR/lib/utility.py b/models/STAR/lib/utility.py new file mode 100644 index 0000000000000000000000000000000000000000..46fb27e32c0fffd25708117220462b8ffa6e9765 --- /dev/null +++ b/models/STAR/lib/utility.py @@ -0,0 +1,364 @@ +import json +import os.path as osp +import time +import torch +import numpy as np +from tqdm import tqdm + +import torchvision.transforms as transforms +from torch.utils.data import DataLoader, DistributedSampler +import torch.optim as optim +import torch.optim.lr_scheduler as lr_scheduler +import torch.nn.functional as F + +# private package +from models.STAR.conf import * +from models.STAR.lib.dataset import AlignmentDataset +from models.STAR.lib.backbone import StackedHGNetV1 +from models.STAR.lib.loss import * +from models.STAR.lib.metric import NME, FR_AUC +from models.STAR.lib.utils import convert_secs2time +from models.STAR.lib.utils import AverageMeter + + +def get_config(args): + config = None + config_name = args.config_name + if config_name == "alignment": + config = Alignment(args) + else: + assert NotImplementedError + + return config + + +def get_dataset(config, tsv_file, image_dir, loader_type, is_train): + dataset = None + if loader_type == "alignment": + dataset = AlignmentDataset( + tsv_file, + image_dir, + transforms.Compose([transforms.ToTensor()]), + config.width, + config.height, + config.channels, + config.means, + config.scale, + config.classes_num, + config.crop_op, + config.aug_prob, + config.edge_info, + config.flip_mapping, + is_train, + encoder_type=config.encoder_type + ) + else: + assert False + return dataset + + +def get_dataloader(config, data_type, world_rank=0, world_size=1): + loader = None + if data_type == "train": + dataset = get_dataset( + config, + config.train_tsv_file, + config.train_pic_dir, + config.loader_type, + is_train=True) + if world_size > 1: + sampler = DistributedSampler(dataset, rank=world_rank, num_replicas=world_size, shuffle=True) + loader = DataLoader(dataset, sampler=sampler, batch_size=config.batch_size // world_size, + num_workers=config.train_num_workers, pin_memory=True, drop_last=True) + else: + loader = DataLoader(dataset, batch_size=config.batch_size, shuffle=True, + num_workers=config.train_num_workers) + elif data_type == "val": + dataset = get_dataset( + config, + config.val_tsv_file, + config.val_pic_dir, + config.loader_type, + is_train=False) + loader = DataLoader(dataset, shuffle=False, batch_size=config.val_batch_size, + num_workers=config.val_num_workers) + elif data_type == "test": + dataset = get_dataset( + config, + config.test_tsv_file, + config.test_pic_dir, + config.loader_type, + is_train=False) + loader = DataLoader(dataset, shuffle=False, batch_size=config.test_batch_size, + num_workers=config.test_num_workers) + else: + assert False + return loader + + +def get_optimizer(config, net): + params = net.parameters() + + optimizer = None + if config.optimizer == "sgd": + optimizer = optim.SGD( + params, + lr=config.learn_rate, + momentum=config.momentum, + weight_decay=config.weight_decay, + nesterov=config.nesterov) + elif config.optimizer == "adam": + optimizer = optim.Adam( + params, + lr=config.learn_rate) + elif config.optimizer == "rmsprop": + optimizer = optim.RMSprop( + params, + lr=config.learn_rate, + momentum=config.momentum, + alpha=config.alpha, + eps=config.epsilon, + weight_decay=config.weight_decay + ) + else: + assert False + return optimizer + + +def get_scheduler(config, optimizer): + if config.scheduler == "MultiStepLR": + scheduler = lr_scheduler.MultiStepLR(optimizer, milestones=config.milestones, gamma=config.gamma) + else: + assert False + return scheduler + + +def get_net(config): + net = None + if config.net == "stackedHGnet_v1": + net = StackedHGNetV1(config=config, + classes_num=config.classes_num, + edge_info=config.edge_info, + nstack=config.nstack, + add_coord=config.add_coord, + decoder_type=config.decoder_type) + else: + assert False + return net + + +def get_criterions(config): + criterions = list() + for k in range(config.label_num): + if config.criterions[k] == "AWingLoss": + criterion = AWingLoss() + elif config.criterions[k] == "smoothl1": + criterion = SmoothL1Loss() + elif config.criterions[k] == "l1": + criterion = F.l1_loss + elif config.criterions[k] == 'l2': + criterion = F.mse_loss + elif config.criterions[k] == "STARLoss": + criterion = STARLoss(dist=config.star_dist, w=config.star_w) + elif config.criterions[k] == "STARLoss_v2": + criterion = STARLoss_v2(dist=config.star_dist, w=config.star_w) + else: + assert False + criterions.append(criterion) + return criterions + + +def set_environment(config): + if config.device_id >= 0: + assert torch.cuda.is_available() and torch.cuda.device_count() > config.device_id + torch.cuda.empty_cache() + config.device = torch.device("cuda", config.device_id) + config.use_gpu = True + else: + config.device = torch.device("cpu") + config.use_gpu = False + + torch.set_default_dtype(torch.float32) + torch.set_default_tensor_type(torch.FloatTensor) + torch.set_flush_denormal(True) # ignore extremely small value + torch.backends.cudnn.benchmark = True # This flag allows you to enable the inbuilt cudnn auto-tuner to find the best algorithm to use for your hardware. + torch.autograd.set_detect_anomaly(True) + + +def forward(config, test_loader, net): + # ave_metrics = [[0, 0] for i in range(config.label_num)] + list_nmes = [[] for i in range(config.label_num)] + metric_nme = NME(nme_left_index=config.nme_left_index, nme_right_index=config.nme_right_index) + metric_fr_auc = FR_AUC(data_definition=config.data_definition) + + output_pd = None + + net = net.float().to(config.device) + net.eval() + dataset_size = len(test_loader.dataset) + batch_size = test_loader.batch_size + if config.logger is not None: + config.logger.info("Forward process, Dataset size: %d, Batch size: %d" % (dataset_size, batch_size)) + for i, sample in enumerate(tqdm(test_loader)): + input = sample["data"].float().to(config.device, non_blocking=True) + labels = list() + if isinstance(sample["label"], list): + for label in sample["label"]: + label = label.float().to(config.device, non_blocking=True) + labels.append(label) + else: + label = sample["label"].float().to(config.device, non_blocking=True) + for k in range(label.shape[1]): + labels.append(label[:, k]) + labels = config.nstack * labels + + with torch.no_grad(): + output, heatmap, landmarks = net(input) + + # metrics + for k in range(config.label_num): + if config.metrics[k] is not None: + list_nmes[k] += metric_nme.test(output[k], labels[k]) + + metrics = [[np.mean(nmes), ] + metric_fr_auc.test(nmes) for nmes in list_nmes] + + return output_pd, metrics + + +def compute_loss(config, criterions, output, labels, heatmap=None, landmarks=None): + batch_weight = 1.0 + sum_loss = 0 + losses = list() + for k in range(config.label_num): + if config.criterions[k] in ['smoothl1', 'l1', 'l2', 'WingLoss', 'AWingLoss']: + loss = criterions[k](output[k], labels[k]) + elif config.criterions[k] in ["STARLoss", "STARLoss_v2"]: + _k = int(k / 3) if config.use_AAM else k + loss = criterions[k](heatmap[_k], labels[k]) + else: + assert NotImplementedError + loss = batch_weight * loss + sum_loss += config.loss_weights[k] * loss + loss = float(loss.data.cpu().item()) + losses.append(loss) + return losses, sum_loss + + +def forward_backward(config, train_loader, net_module, net, net_ema, criterions, optimizer, epoch): + train_model_time = AverageMeter() + ave_losses = [0] * config.label_num + + net_module = net_module.float().to(config.device) + net_module.train(True) + dataset_size = len(train_loader.dataset) + batch_size = config.batch_size # train_loader.batch_size + batch_num = max(dataset_size / max(batch_size, 1), 1) + if config.logger is not None: + config.logger.info(config.note) + config.logger.info("Forward Backward process, Dataset size: %d, Batch size: %d" % (dataset_size, batch_size)) + + iter_num = len(train_loader) + epoch_start_time = time.time() + if net_module != net: + train_loader.sampler.set_epoch(epoch) + for iter, sample in enumerate(train_loader): + iter_start_time = time.time() + # input + input = sample["data"].float().to(config.device, non_blocking=True) + # labels + labels = list() + if isinstance(sample["label"], list): + for label in sample["label"]: + label = label.float().to(config.device, non_blocking=True) + labels.append(label) + else: + label = sample["label"].float().to(config.device, non_blocking=True) + for k in range(label.shape[1]): + labels.append(label[:, k]) + labels = config.nstack * labels + # forward + output, heatmaps, landmarks = net_module(input) + + # loss + losses, sum_loss = compute_loss(config, criterions, output, labels, heatmaps, landmarks) + ave_losses = list(map(sum, zip(ave_losses, losses))) + + # backward + optimizer.zero_grad() + with torch.autograd.detect_anomaly(): + sum_loss.backward() + # torch.nn.utils.clip_grad_norm_(net_module.parameters(), 128.0) + optimizer.step() + + if net_ema is not None: + accumulate_net(net_ema, net, 0.5 ** (config.batch_size / 10000.0)) + # accumulate_net(net_ema, net, 0.5 ** (8 / 10000.0)) + + # output + train_model_time.update(time.time() - iter_start_time) + last_time = convert_secs2time(train_model_time.avg * (iter_num - iter - 1), True) + if iter % config.display_iteration == 0 or iter + 1 == len(train_loader): + if config.logger is not None: + losses_str = ' Average Loss: {:.6f}'.format(sum(losses) / len(losses)) + for k, loss in enumerate(losses): + losses_str += ', L{}: {:.3f}'.format(k, loss) + config.logger.info( + ' -->>[{:03d}/{:03d}][{:03d}/{:03d}]'.format(epoch, config.max_epoch, iter, iter_num) \ + + last_time + losses_str) + + epoch_end_time = time.time() + epoch_total_time = epoch_end_time - epoch_start_time + epoch_load_data_time = epoch_total_time - train_model_time.sum + if config.logger is not None: + config.logger.info("Train/Epoch: %d/%d, Average total time cost per iteration in this epoch: %.6f" % ( + epoch, config.max_epoch, epoch_total_time / iter_num)) + config.logger.info("Train/Epoch: %d/%d, Average loading data time cost per iteration in this epoch: %.6f" % ( + epoch, config.max_epoch, epoch_load_data_time / iter_num)) + config.logger.info("Train/Epoch: %d/%d, Average training model time cost per iteration in this epoch: %.6f" % ( + epoch, config.max_epoch, train_model_time.avg)) + + ave_losses = [loss / iter_num for loss in ave_losses] + if config.logger is not None: + config.logger.info("Train/Epoch: %d/%d, Average Loss in this epoch: %.6f" % ( + epoch, config.max_epoch, sum(ave_losses) / len(ave_losses))) + for k, ave_loss in enumerate(ave_losses): + if config.logger is not None: + config.logger.info("Train/Loss%03d in this epoch: %.6f" % (k, ave_loss)) + + +def accumulate_net(model1, model2, decay): + """ + operation: model1 = model1 * decay + model2 * (1 - decay) + """ + par1 = dict(model1.named_parameters()) + par2 = dict(model2.named_parameters()) + for k in par1.keys(): + par1[k].data.mul_(decay).add_( + other=par2[k].data.to(par1[k].data.device), + alpha=1 - decay) + + par1 = dict(model1.named_buffers()) + par2 = dict(model2.named_buffers()) + for k in par1.keys(): + if par1[k].data.is_floating_point(): + par1[k].data.mul_(decay).add_( + other=par2[k].data.to(par1[k].data.device), + alpha=1 - decay) + else: + par1[k].data = par2[k].data.to(par1[k].data.device) + + +def save_model(config, epoch, net, net_ema, optimizer, scheduler, pytorch_model_path): + # save pytorch model + state = { + "net": net.state_dict(), + "optimizer": optimizer.state_dict(), + "scheduler": scheduler.state_dict(), + "epoch": epoch + } + if config.ema: + state["net_ema"] = net_ema.state_dict() + + torch.save(state, pytorch_model_path) + if config.logger is not None: + config.logger.info("Epoch: %d/%d, model saved in this epoch" % (epoch, config.max_epoch)) diff --git a/models/STAR/lib/utils/__init__.py b/models/STAR/lib/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..8cf0cbd5b0bf523d300e6f35167eadf85a1e1a46 --- /dev/null +++ b/models/STAR/lib/utils/__init__.py @@ -0,0 +1,16 @@ +from .meter import AverageMeter +from .time_utils import time_print, time_string, time_string_short, time_for_file +from .time_utils import convert_secs2time, convert_size2str +from .vis_utils import plot_points + +__all__ = [ + "AverageMeter", + "time_print", + "time_string", + "time_string_short", + "time_for_file", + "convert_size2str", + "convert_secs2time", + + "plot_points", +] diff --git a/models/STAR/lib/utils/dist_utils.py b/models/STAR/lib/utils/dist_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..ed54cab62100d985509ec1070257a790fdffac9c --- /dev/null +++ b/models/STAR/lib/utils/dist_utils.py @@ -0,0 +1,183 @@ +import torch +from torch.autograd import Variable +import matplotlib.pyplot as plt +import seaborn as sns + + +def get_channel_sum(input): + """ + Generates the sum of each channel of the input + input = batch_size x 68 x 64 x 64 + output = batch_size x 68 + """ + temp = torch.sum(input, dim=3) + output = torch.sum(temp, dim=2) + + return output + + +def expand_two_dimensions_at_end(input, dim1, dim2): + """ + Adds two more dimensions to the end of the input + input = batch_size x 68 + output= batch_size x 68 x dim1 x dim2 + """ + input = input.unsqueeze(-1).unsqueeze(-1) + input = input.expand(-1, -1, dim1, dim2) + + return input + + +class Distribution(object): + def __init__(self, heatmaps, num_dim_dist=2, EPSILON=1e-5, is_normalize=True): + self.heatmaps = heatmaps + self.num_dim_dist = num_dim_dist + self.EPSILON = EPSILON + self.is_normalize = is_normalize + batch, npoints, h, w = heatmaps.shape + # normalize + heatmap_sum = torch.clamp(heatmaps.sum([2, 3]), min=1e-6) + self.heatmaps = heatmaps / heatmap_sum.view(batch, npoints, 1, 1) + + # means [batch_size x 68 x 2] + self.mean = self.get_spatial_mean(self.heatmaps) + # covars [batch_size x 68 x 2 x 2] + self.covars = self.get_covariance_matrix(self.heatmaps, self.mean) + + _covars = self.covars.view(batch * npoints, 2, 2).cpu() + evalues, evectors = _covars.symeig(eigenvectors=True) + # eigenvalues [batch_size x 68 x 2] + self.evalues = evalues.view(batch, npoints, 2).to(heatmaps) + # eignvectors [batch_size x 68 x 2 x 2] + self.evectors = evectors.view(batch, npoints, 2, 2).to(heatmaps) + + def __repr__(self): + return "Distribution()" + + def plot(self, heatmap, mean, evalues, evectors): + # heatmap is not normalized + plt.figure(0) + if heatmap.is_cuda: + heatmap, mean = heatmap.cpu(), mean.cpu() + evalues, evectors = evalues.cpu(), evectors.cpu() + sns.heatmap(heatmap, cmap="RdBu_r") + for evalue, evector in zip(evalues, evectors): + plt.arrow(mean[0], mean[1], evalue * evector[0], evalue * evector[1], + width=0.2, shape="full") + plt.show() + + def easy_plot(self, index): + # index = (num of batch_size, num of num_points) + num_bs, num_p = index + heatmap = self.heatmaps[num_bs, num_p] + mean = self.mean[num_bs, num_p] + evalues = self.evalues[num_bs, num_p] + evectors = self.evectors[num_bs, num_p] + self.plot(heatmap, mean, evalues, evectors) + + def project_and_scale(self, pts, eigenvalues, eigenvectors): + batch_size, npoints, _ = pts.shape + proj_pts = torch.matmul(pts.view(batch_size, npoints, 1, 2), eigenvectors) + scale_proj_pts = proj_pts.view(batch_size, npoints, 2) / torch.sqrt(eigenvalues) + return scale_proj_pts + + def _make_grid(self, h, w): + if self.is_normalize: + yy, xx = torch.meshgrid( + torch.arange(h).float() / (h - 1) * 2 - 1, + torch.arange(w).float() / (w - 1) * 2 - 1) + else: + yy, xx = torch.meshgrid( + torch.arange(h).float(), + torch.arange(w).float() + ) + + return yy, xx + + def get_spatial_mean(self, heatmap): + batch, npoints, h, w = heatmap.shape + + yy, xx = self._make_grid(h, w) + yy = yy.view(1, 1, h, w).to(heatmap) + xx = xx.view(1, 1, h, w).to(heatmap) + + yy_coord = (yy * heatmap).sum([2, 3]) # batch x npoints + xx_coord = (xx * heatmap).sum([2, 3]) # batch x npoints + coords = torch.stack([xx_coord, yy_coord], dim=-1) + return coords + + def get_covariance_matrix(self, htp, means): + """ + Covariance calculation from the normalized heatmaps + Reference https://en.wikipedia.org/wiki/Weighted_arithmetic_mean#Weighted_sample_covariance + The unbiased estimate is given by + Unbiased covariance = + ___ + \ + /__ w_i (x_i - \mu_i)^T (x_i - \mu_i) + + ___________________________________________ + + V_1 - (V_2/V_1) + + ___ ___ + \ \ + where V_1 = /__ w_i and V_2 = /__ w_i^2 + + + Input: + htp = batch_size x 68 x 64 x 64 + means = batch_size x 68 x 2 + + Output: + covariance = batch_size x 68 x 2 x 2 + """ + batch_size = htp.shape[0] + num_points = htp.shape[1] + height = htp.shape[2] + width = htp.shape[3] + + yv, xv = self._make_grid(height, width) + xv = Variable(xv) + yv = Variable(yv) + + if htp.is_cuda: + xv = xv.cuda() + yv = yv.cuda() + + xmean = means[:, :, 0] + xv_minus_mean = xv.expand(batch_size, num_points, -1, -1) - expand_two_dimensions_at_end(xmean, height, + width) # batch_size x 68 x 64 x 64 + ymean = means[:, :, 1] + yv_minus_mean = yv.expand(batch_size, num_points, -1, -1) - expand_two_dimensions_at_end(ymean, height, + width) # batch_size x 68 x 64 x 64 + + # These are the unweighted versions + wt_xv_minus_mean = xv_minus_mean + wt_yv_minus_mean = yv_minus_mean + + wt_xv_minus_mean = wt_xv_minus_mean.view(batch_size * num_points, height * width) # batch_size*68 x 4096 + wt_xv_minus_mean = wt_xv_minus_mean.view(batch_size * num_points, 1, + height * width) # batch_size*68 x 1 x 4096 + wt_yv_minus_mean = wt_yv_minus_mean.view(batch_size * num_points, height * width) # batch_size*68 x 4096 + wt_yv_minus_mean = wt_yv_minus_mean.view(batch_size * num_points, 1, + height * width) # batch_size*68 x 1 x 4096 + vec_concat = torch.cat((wt_xv_minus_mean, wt_yv_minus_mean), 1) # batch_size*68 x 2 x 4096 + + htp_vec = htp.view(batch_size * num_points, 1, height * width) + htp_vec = htp_vec.expand(-1, 2, -1) + + # Torch batch matrix multiplication + # https://pytorch.org/docs/stable/torch.html#torch.bmm + # Also use the heatmap as the weights at one place now + covariance = torch.bmm(htp_vec * vec_concat, vec_concat.transpose(1, 2)) # batch_size*68 x 2 x 2 + covariance = covariance.view(batch_size, num_points, self.num_dim_dist, + self.num_dim_dist) # batch_size x 68 x 2 x 2 + + V_1 = get_channel_sum(htp) + self.EPSILON # batch_size x 68 + V_2 = get_channel_sum(torch.pow(htp, 2)) # batch_size x 68 + denominator = V_1 - (V_2 / V_1) + + covariance = covariance / expand_two_dimensions_at_end(denominator, self.num_dim_dist, self.num_dim_dist) + + return (covariance) diff --git a/models/STAR/lib/utils/meter.py b/models/STAR/lib/utils/meter.py new file mode 100644 index 0000000000000000000000000000000000000000..4ba5f27e79f85b02fd7dffad665cdf826731d462 --- /dev/null +++ b/models/STAR/lib/utils/meter.py @@ -0,0 +1,20 @@ +class AverageMeter(object): + """Computes and stores the average and current value""" + + def __init__(self): + self.reset() + + def reset(self): + self.val = 0.0 + self.avg = 0.0 + self.sum = 0.0 + self.count = 0.0 + + def update(self, val, n=1): + self.val = val + self.sum += val + self.count += n + self.avg = self.sum / self.count + + def __repr__(self): + return ('{name}(val={val}, avg={avg}, count={count})'.format(name=self.__class__.__name__, **self.__dict__)) \ No newline at end of file diff --git a/models/STAR/lib/utils/time_utils.py b/models/STAR/lib/utils/time_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..d177aafea369fbd5e006a3a6fdbe01fe51207e6f --- /dev/null +++ b/models/STAR/lib/utils/time_utils.py @@ -0,0 +1,49 @@ +# Copyright (c) Facebook, Inc. and its affiliates. +# All rights reserved. +# +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. +# +import time, sys +import numpy as np + + +def time_for_file(): + ISOTIMEFORMAT = '%d-%h-at-%H-%M-%S' + return '{}'.format(time.strftime(ISOTIMEFORMAT, time.gmtime(time.time()))) + + +def time_string(): + ISOTIMEFORMAT = '%Y-%m-%d %X' + string = '[{}]'.format(time.strftime(ISOTIMEFORMAT, time.gmtime(time.time()))) + return string + + +def time_string_short(): + ISOTIMEFORMAT = '%Y%m%d' + string = '{}'.format(time.strftime(ISOTIMEFORMAT, time.gmtime(time.time()))) + return string + + +def time_print(string, is_print=True): + if (is_print): + print('{} : {}'.format(time_string(), string)) + + +def convert_size2str(torch_size): + dims = len(torch_size) + string = '[' + for idim in range(dims): + string = string + ' {}'.format(torch_size[idim]) + return string + ']' + + +def convert_secs2time(epoch_time, return_str=False): + need_hour = int(epoch_time / 3600) + need_mins = int((epoch_time - 3600 * need_hour) / 60) + need_secs = int(epoch_time - 3600 * need_hour - 60 * need_mins) + if return_str: + str = '[Time Left: {:02d}:{:02d}:{:02d}]'.format(need_hour, need_mins, need_secs) + return str + else: + return need_hour, need_mins, need_secs diff --git a/models/STAR/lib/utils/vis_utils.py b/models/STAR/lib/utils/vis_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..99b5ed1471d9eb75328bf4ae4396d4e3b2544b4e --- /dev/null +++ b/models/STAR/lib/utils/vis_utils.py @@ -0,0 +1,31 @@ +import cv2 +import numpy as np +import numbers + + +def plot_points(vis, points, radius=1, color=(255, 255, 0), shift=4, indexes=0, is_index=False): + if isinstance(points, list): + num_point = len(points) + elif isinstance(points, np.numarray): + num_point = points.shape[0] + else: + raise NotImplementedError + if isinstance(radius, numbers.Number): + radius = np.zeros((num_point)) + radius + + if isinstance(indexes, numbers.Number): + indexes = [indexes + i for i in range(num_point)] + elif isinstance(indexes, list): + pass + else: + raise NotImplementedError + + factor = (1 << shift) + for (index, p, s) in zip(indexes, points, radius): + cv2.circle(vis, (int(p[0] * factor + 0.5), int(p[1] * factor + 0.5)), + int(s * factor), color, 1, cv2.LINE_AA, shift=shift) + if is_index: + vis = cv2.putText(vis, str(index), (int(p[0]), int(p[1])), cv2.FONT_HERSHEY_SIMPLEX, 0.2, + (255, 255, 255), 1) + + return vis diff --git a/models/STAR/main.py b/models/STAR/main.py new file mode 100644 index 0000000000000000000000000000000000000000..e5147d7a77e98a0a15720a54557f40524caf4c27 --- /dev/null +++ b/models/STAR/main.py @@ -0,0 +1,67 @@ +import argparse +from trainer import train +from tester import test + + +def add_data_options(parser): + group = parser.add_argument_group("dataset") + group.add_argument("--image_dir", type=str, default=None, help="the directory of image") + group.add_argument("--annot_dir", type=str, default=None, help="the directory of annot") + + +def add_base_options(parser): + group = parser.add_argument_group("base") + group.add_argument("--mode", type=str, default="train", help="train or test") + group.add_argument("--config_name", type=str, default="alignment", help="set configure file name") + group.add_argument('--device_ids', type=str, default="0,1,2,3", + help="set device ids, -1 means use cpu device, >= 0 means use gpu device") + group.add_argument('--data_definition', type=str, default='WFLW', help="COFW, 300W, WFLW") + group.add_argument('--learn_rate', type=float, default=0.001, help='learning rate') + group.add_argument("--batch_size", type=int, default=128, help="the batch size in train process") + group.add_argument('--width', type=int, default=256, help='the width of input image') + group.add_argument('--height', type=int, default=256, help='the height of input image') + + +def add_train_options(parser): + group = parser.add_argument_group('train') + group.add_argument("--train_num_workers", type=int, default=None, help="the num of workers in train process") + group.add_argument('--loss_func', type=str, default='STARLoss_v2', help="loss function") + group.add_argument("--val_batch_size", type=int, default=None, help="the batch size in val process") + group.add_argument("--val_num_workers", type=int, default=None, help="the num of workers in val process") + + +def add_eval_options(parser): + group = parser.add_argument_group("eval") + group.add_argument("--pretrained_weight", type=str, default=None, + help="set pretrained model file name, if ignored then train the network without pretrain model") + group.add_argument('--norm_type', type=str, default='default', help='default, ocular, pupil') + group.add_argument('--test_file', type=str, default="test.tsv", help='for wflw, test.tsv/test_xx_metadata.tsv') + + +def add_starloss_options(parser): + group = parser.add_argument_group('starloss') + group.add_argument('--star_w', type=float, default=1, help="regular loss ratio") + group.add_argument('--star_dist', type=str, default='smoothl1', help='STARLoss distance function') + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Entry Function") + add_base_options(parser) + add_data_options(parser) + add_train_options(parser) + add_eval_options(parser) + add_starloss_options(parser) + + args = parser.parse_args() + print(args) + + print( + "mode is %s, config_name is %s, pretrained_weight is %s, image_dir is %s, annot_dir is %s, device_ids is %s" % ( + args.mode, args.config_name, args.pretrained_weight, args.image_dir, args.annot_dir, args.device_ids)) + args.device_ids = list(map(int, args.device_ids.split(","))) + if args.mode == "train": + train(args) + elif args.mode == "test": + test(args) + else: + print("unknown running mode") diff --git a/models/STAR/requirements.txt b/models/STAR/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..2e61114224b460979bc877eff7a0d96571831199 --- /dev/null +++ b/models/STAR/requirements.txt @@ -0,0 +1,19 @@ +tqdm +torch==1.6.0 +torchvision==0.7.0 +python-gflags==3.1.2 +pandas==0.24.2 +pillow==6.0.0 +numpy==1.16.4 +opencv-python==4.1.0.25 +imageio==2.5.0 +imgaug==0.2.9 +lmdb==0.98 +lxml==4.5.0 +tensorboard==2.4.1 +protobuf==3.20 +tensorboardX==1.8 +# pyarrow==0.17.1 +# wandb==0.10.25 +# https://pytorch.org/get-started/previous-versions/ +# pip install torch==1.6.0+cpu torchvision==0.7.0+cpu -f https://download.pytorch.org/whl/torch_stable.html diff --git a/models/STAR/tester.py b/models/STAR/tester.py new file mode 100644 index 0000000000000000000000000000000000000000..2b79b2c2b7da39b28ed4572fe3861e84903b661b --- /dev/null +++ b/models/STAR/tester.py @@ -0,0 +1,49 @@ +import os +import torch +from lib import utility + + +def test(args): + # conf + config = utility.get_config(args) + config.device_id = args.device_ids[0] + + # set environment + utility.set_environment(config) + config.init_instance() + if config.logger is not None: + config.logger.info("Loaded configure file %s: %s" % (args.config_name, config.id)) + config.logger.info("\n" + "\n".join(["%s: %s" % item for item in config.__dict__.items()])) + + # model + net = utility.get_net(config) + model_path = os.path.join(config.model_dir, + "train.pkl") if args.pretrained_weight is None else args.pretrained_weight + if args.device_ids == [-1]: + checkpoint = torch.load(model_path, map_location="cpu") + else: + checkpoint = torch.load(model_path) + + net.load_state_dict(checkpoint["net"]) + + if config.logger is not None: + config.logger.info("Loaded network") + # config.logger.info('Net flops: {} G, params: {} MB'.format(flops/1e9, params/1e6)) + + # data - test + test_loader = utility.get_dataloader(config, "test") + + if config.logger is not None: + config.logger.info("Loaded data from {:}".format(config.test_tsv_file)) + + # inference + result, metrics = utility.forward(config, test_loader, net) + if config.logger is not None: + config.logger.info("Finished inference") + + # output + for k, metric in enumerate(metrics): + if config.logger is not None and len(metric) != 0: + config.logger.info( + "Tested {} dataset, the Size is {}, Metric: [NME {:.6f}, FR {:.6f}, AUC {:.6f}]".format( + config.type, len(test_loader.dataset), metric[0], metric[1], metric[2])) diff --git a/models/STAR/tools/analysis_motivation.py b/models/STAR/tools/analysis_motivation.py new file mode 100644 index 0000000000000000000000000000000000000000..bbcbdd3a30510c07c3571675e0452eba70de442f --- /dev/null +++ b/models/STAR/tools/analysis_motivation.py @@ -0,0 +1,220 @@ +import glob +import json +import os.path as osp +import numpy as np +from tqdm import tqdm +import matplotlib.pyplot as plt +import seaborn as sns +from pandas import DataFrame +import pandas as pd + + +def L2(p1, p2): + return np.linalg.norm(p1 - p2) + + +def NME(landmarks_gt, landmarks_pv): + pts_num = landmarks_gt.shape[0] + if pts_num == 29: + left_index = 16 + right_index = 17 + elif pts_num == 68: + left_index = 36 + right_index = 45 + elif pts_num == 98: + left_index = 60 + right_index = 72 + + nme = 0 + eye_span = L2(landmarks_gt[left_index], landmarks_gt[right_index]) + nmeList = [] + for i in range(pts_num): + error = L2(landmarks_pv[i], landmarks_gt[i]) + _nme = error / eye_span + nmeList.append(_nme) + nme += _nme + nme /= pts_num + return nme, nmeList + + +def NME_analysis(listA): + for jsonA in listA: + pred = np.array(jsonA['pred']) + gt = np.array(jsonA['gt']) + nme, nmeList = NME(gt, pred) + jsonA['nme'] = nme + jsonA['nmeList'] = nmeList + return listA + + +def nme_analysis(listA): + bdy_nmeList = [] + scene_nmeList = [] + for jsonA in tqdm(listA): + nme = jsonA['nmeList'] + nme = np.array(nme) + bdy_nme = np.mean(nme[:33]) + scene_nme = np.mean(nme[33:]) + # scene_nme = np.mean(nme[[33, 35, 40, 38, + # 60, 62, 96, 66, 64, + # 50, 44, 48, 46, + # 68, 70, 97, 74, 72, + # 54, 55, 57, 59, + # 76, 82, 79, 90, 94, 85, 16]]) + bdy_nmeList.append(bdy_nme) + scene_nmeList.append(scene_nme) + print('bdy nme: {:.4f}'.format(np.mean(bdy_nmeList))) + print('scene_nmeList: {:.4f}'.format(np.mean(scene_nmeList))) + + +def Energy_analysis(listA, easyThresh=0.02, easyNum=10, hardThresh=0.07, hardNum=10): + easyDict = {'energy': [], 'nme': []} + hardDict = {'energy': [], 'nme': []} + + _easyNum, _hardNum = 0, 0 + + def cal_energy(evalues): + evalues = np.array(evalues) + # _energy = _energy.max(1) + eccentricity = evalues.max(1) / evalues.min(1) + # _energy = _energy.sum() / 2 + _energy = np.mean(eccentricity) + return _energy + + for jsonA in tqdm(listA): + nme = jsonA['nme'] + evalues = jsonA['evalues'] + + if _easyNum == easyNum and _hardNum == hardNum: + break + + if nme < easyThresh and _easyNum < easyNum: + energy = cal_energy(evalues) + easyDict['energy'].append(energy) + easyDict['nme'].append(nme) + _easyNum += 1 + elif nme > hardThresh and _hardNum < hardNum: + energy = cal_energy(evalues) + hardDict['energy'].append(energy) + hardDict['nme'].append(nme) + _hardNum += 1 + + print('easyThresh: < {}; hardThresh > {}'.format(easyThresh, hardThresh)) + print(' |nme |energy |num |') + print('easy samples: |{:.4f} |{:.4f} |{} |'.format(np.mean(easyDict['nme']), + np.mean(easyDict['energy']), + len(easyDict['energy']))) + print('hard samples: |{:.4f} |{:.4f} |{} |'.format(np.mean(hardDict['nme']), + np.mean(hardDict['energy']), + len(hardDict['energy']))) + + return easyDict, hardDict + + +def Eccentricity_analysis(listA): + eyecornerList = [] + boundaryList = [] + for jsonA in listA: + evalues = np.array(jsonA['evalues']) + eccentricity = evalues.max(1) / evalues.min(1) + + eyecorner = np.mean(eccentricity[[60, 64, 68, 72]]) + boundary = np.mean(eccentricity[0:33]) + eyecornerList.append(eyecorner) + boundaryList.append(boundary) + + print('eyecorner: {:.4f}'.format(np.mean(eyecornerList))) + print('boundary: {:.4f}'.format(np.mean(boundaryList))) + return eyecornerList, boundaryList + + +def plot_bar(dataList): + x = list(range(98)) + assert len(x) == len(dataList) + _x = 'Landmark Index' + # _y = 'elliptical eccentricity (λ1/λ2)' + _y = 'PCA Analyze (λ1/λ2)' + data = { + _x: x, + _y: dataList + } + df = DataFrame(data) + plt.figure(figsize=(10, 4)) + sns.barplot(x=_x, y=_y, data=df) + plt.show() + + +def Eccentricity_analysis2(listA, is_vis=False): + landmarksList = [[] for i in range(98)] + for jsonA in listA: + evalues = np.array(jsonA['evalues']) + eccentricity = evalues.max(1) / evalues.min(1) + for i, e in enumerate(eccentricity): + landmarksList[i].append(e) + print('Mean value: {:.4f}'.format(np.mean(np.array(landmarksList)))) + landmarksList = [np.mean(l) for l in landmarksList] + if is_vis: + plot_bar(landmarksList) + return landmarksList + + +def std_analysis2(): + save_dir = '/apdcephfs/share_1134483/charlinzhou/experiment/cvpr-23/wflw_results' + # l2_npy = glob.glob(osp.join(save_dir, '*DSNT*.npy')) + l2_npy = glob.glob(osp.join(save_dir, '*MHNLoss_v2_l2*.npy')) + + def npy2std(npyList): + datas = [np.load(npy)[np.newaxis, :] for npy in npyList] + datas = np.concatenate(datas, axis=0) + # denormalization + datas = (datas + 1) * 256 / 2 + mean = datas.mean(axis=0)[np.newaxis, :] + dist = np.linalg.norm(datas - mean, axis=-1) + std = np.std(dist, 0) + print('min: {}, max:{}, mean:{}'.format(std.min(), std.max(), std.mean())) + return std + + std1 = npy2std(l2_npy) + std1 = std1.mean(0) + # plot_bar(std1) + bdy_std = np.mean(std1[:33]) + cofw_std = np.mean(std1[[33, 35, 40, 38, + 60, 62, 96, 66, 64, + 50, 44, 48, 46, + 68, 70, 97, 74, 72, + 54, 55, 57, 59, + 76, 82, 79, 90, 94, 85, 16]]) + print('bdy_std: {:.4f}, cofw_std: {:.4f}'.format(bdy_std, cofw_std)) + print('the ratio of Boundary std and ALL std: {:.4f} / {:.4f}'.format(np.sum(std1[:33]), np.sum(std1))) + + +if __name__ == '__main__': + # 4.29模型 + json_path = '/apdcephfs/share_1134483/charlinzhou/ckpts/STAR/WFLW/WFLW_256x256_adam_ep500_lr0.001_bs128_STARLoss_smoothl1_1_b0183746-161a-4b76-9cb9-8a2059090233/results.json' + # 无初始化 + # json_path = '/apdcephfs/share_1134483/charlinzhou/ckpts/STAR/WFLW/WFLW_256x256_adam_ep500_lr0.001_bs128_STARLoss_smoothl1_1_9cff3656-8ca8-4c3d-a95d-da76f9f76ea5/results.json' + # 4.02模型 + # json_path = '/apdcephfs/share_1134483/charlinzhou/ckpts/STAR/WFLW/WFLW_256x256_adam_ep500_lr0.001_bs128_STARLoss_smoothl1_1_AAM_2d2bb70e-6fdb-459c-baf7-18c89e7a165f/results.json' + listA = json.load(open(json_path, 'r')) + print('Load Done!') + listA = NME_analysis(listA) + print('NME analysis Done!') + # Exp1: 分析简单样本和困难样本的能量差异 + easyDict, hardDict = Energy_analysis(listA, easyNum=2500, hardNum=2500, easyThresh=0.03, hardThresh=0.08) + + # Exp2.1: 分析眼角点和轮廓点的斜率差异 + # eyecornerList, boundaryList = Eccentricity_analysis(listA) + + # Exp2.2: 可视化所有点的斜率分布 + # landmarksList = Eccentricity_analysis2(listA, is_vis=True) + + # Exp2.3: 可视化所有点的方差分布 + # std_analysis2() + + # Exp3: 五官和轮廓NME分析 + # nme_analysis(listA) + # print(easyDict) + # print(hardDict) + + # nmeList = [jsonA['nme'] for jsonA in listA] + # print(len(nmeList)) diff --git a/models/STAR/tools/infinite_loop.py b/models/STAR/tools/infinite_loop.py new file mode 100644 index 0000000000000000000000000000000000000000..510011e4ddcbef5c6b67dce38edce3e2236b694d --- /dev/null +++ b/models/STAR/tools/infinite_loop.py @@ -0,0 +1,4 @@ +import time + +while True: + time.sleep(1) diff --git a/models/STAR/tools/infinite_loop_gpu.py b/models/STAR/tools/infinite_loop_gpu.py new file mode 100644 index 0000000000000000000000000000000000000000..6bfc2a50fede08ce1412f03aee5db9132ca285ff --- /dev/null +++ b/models/STAR/tools/infinite_loop_gpu.py @@ -0,0 +1,21 @@ +# -*- coding: utf-8 -*- + +import os +import time +import torch +import argparse + +parser = argparse.ArgumentParser(description='inf') +parser.add_argument('--gpu', default='1', type=str, help='index of gpu to use') +args = parser.parse_args() + +os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu + +n = 1000 + +x = torch.zeros(4, n, n).cuda() +rest_time = 0.0000000000001 +while True: + y = x * x + time.sleep(rest_time) + y1 = x * x diff --git a/models/STAR/tools/split_wflw.py b/models/STAR/tools/split_wflw.py new file mode 100644 index 0000000000000000000000000000000000000000..0337f4236e43c48dd8ccd0f3041f274448d712f1 --- /dev/null +++ b/models/STAR/tools/split_wflw.py @@ -0,0 +1,38 @@ +import csv +import os.path as osp +import numpy as np +import pandas as pd +from tqdm import tqdm + +tsv_file = '/apdcephfs/share_1134483/charlinzhou/datas/ADNet/WFLW/test.tsv' +save_folder = '/apdcephfs/share_1134483/charlinzhou/datas/ADNet/_WFLW/' + +save_tags = ['largepose', 'expression', 'illumination', 'makeup', 'occlusion', 'blur'] +save_tags = ['test_{}_metadata.tsv'.format(t) for t in save_tags] +save_files = [osp.join(save_folder, t) for t in save_tags] +save_files = [open(f, 'w', newline='') for f in save_files] + +landmark_num = 98 +items = pd.read_csv(tsv_file, sep="\t") + +items_num = len(items) +for index in tqdm(range(items_num)): + image_path = items.iloc[index, 0] + landmarks_5pts = items.iloc[index, 1] + # landmarks_5pts = np.array(list(map(float, landmarks_5pts.split(","))), dtype=np.float32).reshape(5, 2) + landmarks_target = items.iloc[index, 2] + # landmarks_target = np.array(list(map(float, landmarks_target.split(","))), dtype=np.float32).reshape(landmark_num, 2) + scale = items.iloc[index, 3] + center_w, center_h = items.iloc[index, 4], items.iloc[index, 5] + if len(items.iloc[index]) > 6: + tags = np.array(list(map(lambda x: int(float(x)), items.iloc[index, 6].split(",")))) + else: + tags = np.array([]) + assert len(tags) == 6, '{} v.s. 6'.format(len(tags)) + for k, tag in enumerate(tags): + if tag == 1: + save_file = save_files[k] + tsv_w = csv.writer(save_file, delimiter='\t') + tsv_w.writerow([image_path, landmarks_5pts, landmarks_target, scale, center_w, center_h]) + +print('Done!') diff --git a/models/STAR/tools/testtime_pca.py b/models/STAR/tools/testtime_pca.py new file mode 100644 index 0000000000000000000000000000000000000000..c231a96719b986354b9e60d783eacc0467a413f2 --- /dev/null +++ b/models/STAR/tools/testtime_pca.py @@ -0,0 +1,107 @@ +import torch +import torch.nn as nn +from torch.autograd import Variable + + +def get_channel_sum(input): + temp = torch.sum(input, dim=3) + output = torch.sum(temp, dim=2) + return output + + +def expand_two_dimensions_at_end(input, dim1, dim2): + input = input.unsqueeze(-1).unsqueeze(-1) + input = input.expand(-1, -1, dim1, dim2) + return input + + +class TestTimePCA(nn.Module): + def __init__(self): + super(TestTimePCA, self).__init__() + + def _make_grid(self, h, w): + yy, xx = torch.meshgrid( + torch.arange(h).float() / (h - 1) * 2 - 1, + torch.arange(w).float() / (w - 1) * 2 - 1) + return yy, xx + + def weighted_mean(self, heatmap): + batch, npoints, h, w = heatmap.shape + + yy, xx = self._make_grid(h, w) + yy = yy.view(1, 1, h, w).to(heatmap) + xx = xx.view(1, 1, h, w).to(heatmap) + + yy_coord = (yy * heatmap).sum([2, 3]) # batch x npoints + xx_coord = (xx * heatmap).sum([2, 3]) # batch x npoints + coords = torch.stack([xx_coord, yy_coord], dim=-1) + return coords + + def unbiased_weighted_covariance(self, htp, means, num_dim_image=2, EPSILON=1e-5): + batch_size, num_points, height, width = htp.shape + + yv, xv = self._make_grid(height, width) + xv = Variable(xv) + yv = Variable(yv) + + if htp.is_cuda: + xv = xv.cuda() + yv = yv.cuda() + + xmean = means[:, :, 0] + xv_minus_mean = xv.expand(batch_size, num_points, -1, -1) - expand_two_dimensions_at_end(xmean, height, + width) # [batch_size, 68, 64, 64] + ymean = means[:, :, 1] + yv_minus_mean = yv.expand(batch_size, num_points, -1, -1) - expand_two_dimensions_at_end(ymean, height, + width) # [batch_size, 68, 64, 64] + wt_xv_minus_mean = xv_minus_mean + wt_yv_minus_mean = yv_minus_mean + + wt_xv_minus_mean = wt_xv_minus_mean.view(batch_size * num_points, height * width) # [batch_size*68, 4096] + wt_xv_minus_mean = wt_xv_minus_mean.view(batch_size * num_points, 1, height * width) # [batch_size*68, 1, 4096] + wt_yv_minus_mean = wt_yv_minus_mean.view(batch_size * num_points, height * width) # [batch_size*68, 4096] + wt_yv_minus_mean = wt_yv_minus_mean.view(batch_size * num_points, 1, height * width) # [batch_size*68, 1, 4096] + vec_concat = torch.cat((wt_xv_minus_mean, wt_yv_minus_mean), 1) # [batch_size*68, 2, 4096] + + htp_vec = htp.view(batch_size * num_points, 1, height * width) + htp_vec = htp_vec.expand(-1, 2, -1) + + covariance = torch.bmm(htp_vec * vec_concat, vec_concat.transpose(1, 2)) # [batch_size*68, 2, 2] + covariance = covariance.view(batch_size, num_points, num_dim_image, num_dim_image) # [batch_size, 68, 2, 2] + + V_1 = htp.sum([2, 3]) + EPSILON # [batch_size, 68] + V_2 = torch.pow(htp, 2).sum([2, 3]) + EPSILON # [batch_size, 68] + + denominator = V_1 - (V_2 / V_1) + covariance = covariance / expand_two_dimensions_at_end(denominator, num_dim_image, num_dim_image) + + return covariance + + def forward(self, heatmap, groudtruth): + + batch, npoints, h, w = heatmap.shape + + heatmap_sum = torch.clamp(heatmap.sum([2, 3]), min=1e-6) + heatmap = heatmap / heatmap_sum.view(batch, npoints, 1, 1) + + # means [batch_size, 68, 2] + means = self.weighted_mean(heatmap) + + # covars [batch_size, 68, 2, 2] + covars = self.unbiased_weighted_covariance(heatmap, means) + + # eigenvalues [batch_size * 68, 2] , eigenvectors [batch_size * 68, 2, 2] + covars = covars.view(batch * npoints, 2, 2).cpu() + evalues, evectors = covars.symeig(eigenvectors=True) + evalues = evalues.view(batch, npoints, 2) + evectors = evectors.view(batch, npoints, 2, 2) + means = means.cpu() + + results = [dict() for _ in range(batch)] + for i in range(batch): + results[i]['pred'] = means[i].numpy().tolist() + results[i]['gt'] = groudtruth[i].cpu().numpy().tolist() + results[i]['evalues'] = evalues[i].numpy().tolist() + results[i]['evectors'] = evectors[i].numpy().tolist() + + return results diff --git a/models/STAR/trainer.py b/models/STAR/trainer.py new file mode 100644 index 0000000000000000000000000000000000000000..ff7bce92c866521643715ca22236d0b94b5f860e --- /dev/null +++ b/models/STAR/trainer.py @@ -0,0 +1,209 @@ +import os +import sys +import time +import argparse +import traceback +import torch +import torch.nn as nn +from lib import utility +from lib.utils import AverageMeter, convert_secs2time + +os.environ["MKL_THREADING_LAYER"] = "GNU" + + +# os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,3" + +def train(args): + device_ids = args.device_ids + nprocs = len(device_ids) + if nprocs > 1: + torch.multiprocessing.spawn( + train_worker, args=(nprocs, 1, args), nprocs=nprocs, + join=True) + elif nprocs == 1: + train_worker(device_ids[0], nprocs, 1, args) + else: + assert False + + +def train_worker(world_rank, world_size, nodes_size, args): + # initialize config. + config = utility.get_config(args) + config.device_id = world_rank if nodes_size == 1 else world_rank % torch.cuda.device_count() + # set environment + utility.set_environment(config) + # initialize instances, such as writer, logger and wandb. + if world_rank == 0: + config.init_instance() + + if config.logger is not None: + config.logger.info("\n" + "\n".join(["%s: %s" % item for item in config.__dict__.items()])) + config.logger.info("Loaded configure file %s: %s" % (config.type, config.id)) + + # worker communication + if world_size > 1: + torch.distributed.init_process_group( + backend="nccl", init_method="tcp://localhost:23456" if nodes_size == 1 else "env://", + rank=world_rank, world_size=world_size) + torch.cuda.set_device(config.device) + + # model + net = utility.get_net(config) + if world_size > 1: + net = torch.nn.SyncBatchNorm.convert_sync_batchnorm(net) + net = net.float().to(config.device) + net.train(True) + if config.ema and world_rank == 0: + net_ema = utility.get_net(config) + if world_size > 1: + net_ema = torch.nn.SyncBatchNorm.convert_sync_batchnorm(net_ema) + net_ema = net_ema.float().to(config.device) + net_ema.eval() + utility.accumulate_net(net_ema, net, 0) + else: + net_ema = None + + # multi-GPU training + if world_size > 1: + net_module = nn.parallel.DistributedDataParallel(net, device_ids=[config.device_id], + output_device=config.device_id, find_unused_parameters=True) + else: + net_module = net + + criterions = utility.get_criterions(config) + optimizer = utility.get_optimizer(config, net_module) + scheduler = utility.get_scheduler(config, optimizer) + + # load pretrain model + if args.pretrained_weight is not None: + if not os.path.exists(args.pretrained_weight): + pretrained_weight = os.path.join(config.work_dir, args.pretrained_weight) + else: + pretrained_weight = args.pretrained_weight + + try: + checkpoint = torch.load(pretrained_weight) + net.load_state_dict(checkpoint["net"], strict=False) + if net_ema is not None: + net_ema.load_state_dict(checkpoint["net_ema"], strict=False) + if config.logger is not None: + config.logger.warn("Successed to load pretrain model %s." % pretrained_weight) + start_epoch = checkpoint["epoch"] + optimizer.load_state_dict(checkpoint["optimizer"]) + scheduler.load_state_dict(checkpoint["scheduler"]) + except: + start_epoch = 0 + if config.logger is not None: + config.logger.warn("Failed to load pretrain model %s." % pretrained_weight) + else: + start_epoch = 0 + + if config.logger is not None: + config.logger.info("Loaded network") + + # data - train, val + train_loader = utility.get_dataloader(config, "train", world_rank, world_size) + if world_rank == 0: + val_loader = utility.get_dataloader(config, "val") + if config.logger is not None: + config.logger.info("Loaded data") + + # forward & backward + if config.logger is not None: + config.logger.info("Optimizer type %s. Start training..." % (config.optimizer)) + if not os.path.exists(config.model_dir) and world_rank == 0: + os.makedirs(config.model_dir) + + # training + best_metric, best_net = None, None + epoch_time, eval_time = AverageMeter(), AverageMeter() + for i_epoch, epoch in enumerate(range(config.max_epoch + 1)): + try: + epoch_start_time = time.time() + if epoch >= start_epoch: + # forward and backward + if epoch != start_epoch: + utility.forward_backward(config, train_loader, net_module, net, net_ema, criterions, optimizer, + epoch) + + if world_size > 1: + torch.distributed.barrier() + + # validating + if epoch % config.val_epoch == 0 and epoch != 0 and world_rank == 0: + eval_start_time = time.time() + epoch_nets = {"net": net, "net_ema": net_ema} + for net_name, epoch_net in epoch_nets.items(): + if epoch_net is None: + continue + result, metrics = utility.forward(config, val_loader, epoch_net) + for k, metric in enumerate(metrics): + if config.logger is not None and len(metric) != 0: + config.logger.info( + "Val_{}/Metric{:3d} in this epoch: [NME {:.6f}, FR {:.6f}, AUC {:.6f}]".format( + net_name, k, metric[0], metric[1], metric[2])) + + # update best model. + cur_metric = metrics[config.key_metric_index][0] + if best_metric is None or best_metric > cur_metric: + best_metric = cur_metric + best_net = epoch_net + current_pytorch_model_path = os.path.join(config.model_dir, "best_model.pkl") + # current_onnx_model_path = os.path.join(config.model_dir, "train.onnx") + utility.save_model( + config, + epoch, + best_net, + net_ema, + optimizer, + scheduler, + current_pytorch_model_path) + if best_metric is not None: + config.logger.info( + "Val/Best_Metric%03d in this epoch: %.6f" % (config.key_metric_index, best_metric)) + eval_time.update(time.time() - eval_start_time) + + # saving model + if epoch == config.max_epoch and world_rank == 0: + current_pytorch_model_path = os.path.join(config.model_dir, "last_model.pkl") + # current_onnx_model_path = os.path.join(config.model_dir, "model_epoch_%s.onnx" % epoch) + utility.save_model( + config, + epoch, + net, + net_ema, + optimizer, + scheduler, + current_pytorch_model_path) + + if world_size > 1: + torch.distributed.barrier() + + # adjusting learning rate + if epoch > 0: + scheduler.step() + epoch_time.update(time.time() - epoch_start_time) + last_time = convert_secs2time(epoch_time.avg * (config.max_epoch - i_epoch), True) + if config.logger is not None: + config.logger.info( + "Train/Epoch: %d/%d, Learning rate decays to %s, " % ( + epoch, config.max_epoch, str(scheduler.get_last_lr())) \ + + last_time + 'eval_time: {:4.2f}, '.format(eval_time.avg) + '\n\n') + + except: + traceback.print_exc() + config.logger.error("Exception happened in training steps") + + if config.logger is not None: + config.logger.info("Training finished") + + try: + if config.logger is not None and best_metric is not None: + new_folder_name = config.folder + '-fin-{:.4f}'.format(best_metric) + new_work_dir = os.path.join(config.ckpt_dir, config.data_definition, new_folder_name) + os.system('mv {} {}'.format(config.work_dir, new_work_dir)) + except: + traceback.print_exc() + + if world_size > 1: + torch.distributed.destroy_process_group() diff --git a/models/__init__.py b/models/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/models/encoder4editing/.DS_Store b/models/encoder4editing/.DS_Store new file mode 100644 index 0000000000000000000000000000000000000000..eaef0e7fe24cc28841e2843076297d8c7fdb582e Binary files /dev/null and b/models/encoder4editing/.DS_Store differ diff --git a/models/encoder4editing/LICENSE b/models/encoder4editing/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..88ba9d421ea8acea9b4e3937535e72c282b3d4e6 --- /dev/null +++ b/models/encoder4editing/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2021 omertov + +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. diff --git a/models/encoder4editing/README.md b/models/encoder4editing/README.md new file mode 100644 index 0000000000000000000000000000000000000000..639e7347a730fea1e1fbf595cbbaa0e03d79ac6b --- /dev/null +++ b/models/encoder4editing/README.md @@ -0,0 +1,143 @@ +# Designing an Encoder for StyleGAN Image Manipulation (SIGGRAPH 2021) + + + [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](http://colab.research.google.com/github/omertov/encoder4editing/blob/main/notebooks/inference_playground.ipynb) + +> Recently, there has been a surge of diverse methods for performing image editing by employing pre-trained unconditional generators. Applying these methods on real images, however, remains a challenge, as it necessarily requires the inversion of the images into their latent space. To successfully invert a real image, one needs to find a latent code that reconstructs the input image accurately, and more importantly, allows for its meaningful manipulation. In this paper, we carefully study the latent space of StyleGAN, the state-of-the-art unconditional generator. We identify and analyze the existence of a distortion-editability tradeoff and a distortion-perception tradeoff within the StyleGAN latent space. We then suggest two principles for designing encoders in a manner that allows one to control the proximity of the inversions to regions that StyleGAN was originally trained on. We present an encoder based on our two principles that is specifically designed for facilitating editing on real images by balancing these tradeoffs. By evaluating its performance qualitatively and quantitatively on numerous challenging domains, including cars and horses, we show that our inversion method, followed by common editing techniques, achieves superior real-image editing quality, with only a small reconstruction accuracy drop. + +

+ +

+ +## Description +Official Implementation of "Designing an Encoder for StyleGAN Image Manipulation" paper for both training and evaluation. +The e4e encoder is specifically designed to complement existing image manipulation techniques performed over StyleGAN's latent space. + +## Recent Updates +`2021.08.17`: Add single style code encoder (use `--encoder_type SingleStyleCodeEncoder`).
+`2021.03.25`: Add pose editing direction. + +## Getting Started +### Prerequisites +- Linux or macOS +- NVIDIA GPU + CUDA CuDNN (CPU may be possible with some modifications, but is not inherently supported) +- Python 3 + +### Installation +- Clone the repository: +``` +git clone https://github.com/omertov/encoder4editing.git +cd encoder4editing +``` +- Dependencies: +We recommend running this repository using [Anaconda](https://docs.anaconda.com/anaconda/install/). +All dependencies for defining the environment are provided in `environment/e4e_env.yaml`. + +### Inference Notebook +We provide a Jupyter notebook found in `notebooks/inference_playground.ipynb` that allows one to encode and perform several editings on real images using StyleGAN. + +### Pretrained Models +Please download the pre-trained models from the following links. Each e4e model contains the entire pSp framework architecture, including the encoder and decoder weights. +| Path | Description +| :--- | :---------- +|[FFHQ Inversion](https://drive.google.com/file/d/1cUv_reLE6k3604or78EranS7XzuVMWeO/view?usp=sharing) | FFHQ e4e encoder. +|[Cars Inversion](https://drive.google.com/file/d/17faPqBce2m1AQeLCLHUVXaDfxMRU2QcV/view?usp=sharing) | Cars e4e encoder. +|[Horse Inversion](https://drive.google.com/file/d/1TkLLnuX86B_BMo2ocYD0kX9kWh53rUVX/view?usp=sharing) | Horse e4e encoder. +|[Church Inversion](https://drive.google.com/file/d/1-L0ZdnQLwtdy6-A_Ccgq5uNJGTqE7qBa/view?usp=sharing) | Church e4e encoder. + +If you wish to use one of the pretrained models for training or inference, you may do so using the flag `--checkpoint_path`. + +In addition, we provide various auxiliary models needed for training your own e4e model from scratch. +| Path | Description +| :--- | :---------- +|[FFHQ StyleGAN](https://drive.google.com/file/d/1EM87UquaoQmk17Q8d5kYIAHqu0dkYqdT/view?usp=sharing) | StyleGAN model pretrained on FFHQ taken from [rosinality](https://github.com/rosinality/stylegan2-pytorch) with 1024x1024 output resolution. +|[IR-SE50 Model](https://drive.google.com/file/d/1KW7bjndL3QG3sxBbZxreGHigcCCpsDgn/view?usp=sharing) | Pretrained IR-SE50 model taken from [TreB1eN](https://github.com/TreB1eN/InsightFace_Pytorch) for use in our ID loss during training. +|[MOCOv2 Model](https://drive.google.com/file/d/18rLcNGdteX5LwT7sv_F7HWr12HpVEzVe/view?usp=sharing) | Pretrained ResNet-50 model trained using MOCOv2 for use in our simmilarity loss for domains other then human faces during training. + +By default, we assume that all auxiliary models are downloaded and saved to the directory `pretrained_models`. However, you may use your own paths by changing the necessary values in `configs/path_configs.py`. + +## Training +To train the e4e encoder, make sure the paths to the required models, as well as training and testing data is configured in `configs/path_configs.py` and `configs/data_configs.py`. +#### **Training the e4e Encoder** +``` +python scripts/train.py \ +--dataset_type cars_encode \ +--exp_dir new/experiment/directory \ +--start_from_latent_avg \ +--use_w_pool \ +--w_discriminator_lambda 0.1 \ +--progressive_start 20000 \ +--id_lambda 0.5 \ +--val_interval 10000 \ +--max_steps 200000 \ +--stylegan_size 512 \ +--stylegan_weights path/to/pretrained/stylegan.pt \ +--workers 8 \ +--batch_size 8 \ +--test_batch_size 4 \ +--test_workers 4 +``` + +#### Training on your own dataset +In order to train the e4e encoder on a custom dataset, perform the following adjustments: +1. Insert the paths to your train and test data into the `dataset_paths` variable defined in `configs/paths_config.py`: +``` +dataset_paths = { + 'my_train_data': '/path/to/train/images/directory', + 'my_test_data': '/path/to/test/images/directory' +} +``` +2. Configure a new dataset under the DATASETS variable defined in `configs/data_configs.py`: +``` +DATASETS = { + 'my_data_encode': { + 'transforms': transforms_config.EncodeTransforms, + 'train_source_root': dataset_paths['my_train_data'], + 'train_target_root': dataset_paths['my_train_data'], + 'test_source_root': dataset_paths['my_test_data'], + 'test_target_root': dataset_paths['my_test_data'] + } +} +``` +Refer to `configs/transforms_config.py` for the transformations applied to the train and test images during training. + +3. Finally, run a training session with `--dataset_type my_data_encode`. + +## Inference +Having trained your model, you can use `scripts/inference.py` to apply the model on a set of images. +For example, +``` +python scripts/inference.py \ +--images_dir=/path/to/images/directory \ +--save_dir=/path/to/saving/directory \ +path/to/checkpoint.pt +``` + +## Latent Editing Consistency (LEC) +As described in the paper, we suggest a new metric, Latent Editing Consistency (LEC), for evaluating the encoder's +performance. +We provide an example for calculating the metric over the FFHQ StyleGAN using the aging editing direction in +`metrics/LEC.py`. + +To run the example: +``` +cd metrics +python LEC.py \ +--images_dir=/path/to/images/directory \ +path/to/checkpoint.pt +``` + +## Acknowledgments +This code borrows heavily from [pixel2style2pixel](https://github.com/eladrich/pixel2style2pixel) + +## Citation +If you use this code for your research, please cite our paper Designing an Encoder for StyleGAN Image Manipulation: + +``` +@article{tov2021designing, + title={Designing an Encoder for StyleGAN Image Manipulation}, + author={Tov, Omer and Alaluf, Yuval and Nitzan, Yotam and Patashnik, Or and Cohen-Or, Daniel}, + journal={arXiv preprint arXiv:2102.02766}, + year={2021} +} +``` diff --git a/models/encoder4editing/configs/__init__.py b/models/encoder4editing/configs/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/models/encoder4editing/configs/data_configs.py b/models/encoder4editing/configs/data_configs.py new file mode 100644 index 0000000000000000000000000000000000000000..deccb0b1c266ad4b6abaef53d67ec1ed0ddbd462 --- /dev/null +++ b/models/encoder4editing/configs/data_configs.py @@ -0,0 +1,41 @@ +from configs import transforms_config +from configs.paths_config import dataset_paths + + +DATASETS = { + 'ffhq_encode': { + 'transforms': transforms_config.EncodeTransforms, + 'train_source_root': dataset_paths['ffhq'], + 'train_target_root': dataset_paths['ffhq'], + 'test_source_root': dataset_paths['celeba_test'], + 'test_target_root': dataset_paths['celeba_test'], + }, + 'cars_encode': { + 'transforms': transforms_config.CarsEncodeTransforms, + 'train_source_root': dataset_paths['cars_train'], + 'train_target_root': dataset_paths['cars_train'], + 'test_source_root': dataset_paths['cars_test'], + 'test_target_root': dataset_paths['cars_test'], + }, + 'horse_encode': { + 'transforms': transforms_config.EncodeTransforms, + 'train_source_root': dataset_paths['horse_train'], + 'train_target_root': dataset_paths['horse_train'], + 'test_source_root': dataset_paths['horse_test'], + 'test_target_root': dataset_paths['horse_test'], + }, + 'church_encode': { + 'transforms': transforms_config.EncodeTransforms, + 'train_source_root': dataset_paths['church_train'], + 'train_target_root': dataset_paths['church_train'], + 'test_source_root': dataset_paths['church_test'], + 'test_target_root': dataset_paths['church_test'], + }, + 'cats_encode': { + 'transforms': transforms_config.EncodeTransforms, + 'train_source_root': dataset_paths['cats_train'], + 'train_target_root': dataset_paths['cats_train'], + 'test_source_root': dataset_paths['cats_test'], + 'test_target_root': dataset_paths['cats_test'], + } +} diff --git a/models/encoder4editing/configs/paths_config.py b/models/encoder4editing/configs/paths_config.py new file mode 100644 index 0000000000000000000000000000000000000000..8bc45c4a7ca33b1d074bd5582d57dc83caf77027 --- /dev/null +++ b/models/encoder4editing/configs/paths_config.py @@ -0,0 +1,28 @@ +dataset_paths = { + # Face Datasets (In the paper: FFHQ - scripts, CelebAHQ - test) + 'ffhq': '', + 'celeba_test': '', + + # Cars Dataset (In the paper: Stanford cars) + 'cars_train': '', + 'cars_test': '', + + # Horse Dataset (In the paper: LSUN Horse) + 'horse_train': '', + 'horse_test': '', + + # Church Dataset (In the paper: LSUN Church) + 'church_train': '', + 'church_test': '', + + # Cats Dataset (In the paper: LSUN Cat) + 'cats_train': '', + 'cats_test': '' +} + +model_paths = { + 'stylegan_ffhq': 'pretrained_models/stylegan2-ffhq-config-f.pt', + 'ir_se50': 'pretrained_models/model_ir_se50.pth', + 'shape_predictor': 'pretrained_models/shape_predictor_68_face_landmarks.dat', + 'moco': 'pretrained_models/moco_v2_800ep_pretrain.pth' +} diff --git a/models/encoder4editing/configs/transforms_config.py b/models/encoder4editing/configs/transforms_config.py new file mode 100644 index 0000000000000000000000000000000000000000..ac12b5d5ba0571f21715e0f6b24b9c1ebe84bf72 --- /dev/null +++ b/models/encoder4editing/configs/transforms_config.py @@ -0,0 +1,62 @@ +from abc import abstractmethod +import torchvision.transforms as transforms + + +class TransformsConfig(object): + + def __init__(self, opts): + self.opts = opts + + @abstractmethod + def get_transforms(self): + pass + + +class EncodeTransforms(TransformsConfig): + + def __init__(self, opts): + super(EncodeTransforms, self).__init__(opts) + + def get_transforms(self): + transforms_dict = { + 'transform_gt_train': transforms.Compose([ + transforms.Resize((256, 256)), + transforms.RandomHorizontalFlip(0.5), + transforms.ToTensor(), + transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])]), + 'transform_source': None, + 'transform_test': transforms.Compose([ + transforms.Resize((256, 256)), + transforms.ToTensor(), + transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])]), + 'transform_inference': transforms.Compose([ + transforms.Resize((256, 256)), + transforms.ToTensor(), + transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])]) + } + return transforms_dict + + +class CarsEncodeTransforms(TransformsConfig): + + def __init__(self, opts): + super(CarsEncodeTransforms, self).__init__(opts) + + def get_transforms(self): + transforms_dict = { + 'transform_gt_train': transforms.Compose([ + transforms.Resize((192, 256)), + transforms.RandomHorizontalFlip(0.5), + transforms.ToTensor(), + transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])]), + 'transform_source': None, + 'transform_test': transforms.Compose([ + transforms.Resize((192, 256)), + transforms.ToTensor(), + transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])]), + 'transform_inference': transforms.Compose([ + transforms.Resize((192, 256)), + transforms.ToTensor(), + transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])]) + } + return transforms_dict diff --git a/models/encoder4editing/criteria/__init__.py b/models/encoder4editing/criteria/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/models/encoder4editing/criteria/id_loss.py b/models/encoder4editing/criteria/id_loss.py new file mode 100644 index 0000000000000000000000000000000000000000..bab806172eff18c0630536ae96817508c3197b8b --- /dev/null +++ b/models/encoder4editing/criteria/id_loss.py @@ -0,0 +1,47 @@ +import torch +from torch import nn +from configs.paths_config import model_paths +from models.encoders.model_irse import Backbone + + +class IDLoss(nn.Module): + def __init__(self): + super(IDLoss, self).__init__() + print('Loading ResNet ArcFace') + self.facenet = Backbone(input_size=112, num_layers=50, drop_ratio=0.6, mode='ir_se') + self.facenet.load_state_dict(torch.load(model_paths['ir_se50'])) + self.face_pool = torch.nn.AdaptiveAvgPool2d((112, 112)) + self.facenet.eval() + for module in [self.facenet, self.face_pool]: + for param in module.parameters(): + param.requires_grad = False + + def extract_feats(self, x): + x = x[:, :, 35:223, 32:220] # Crop interesting region + x = self.face_pool(x) + x_feats = self.facenet(x) + return x_feats + + def forward(self, y_hat, y, x): + n_samples = x.shape[0] + x_feats = self.extract_feats(x) + y_feats = self.extract_feats(y) # Otherwise use the feature from there + y_hat_feats = self.extract_feats(y_hat) + y_feats = y_feats.detach() + loss = 0 + sim_improvement = 0 + id_logs = [] + count = 0 + for i in range(n_samples): + diff_target = y_hat_feats[i].dot(y_feats[i]) + diff_input = y_hat_feats[i].dot(x_feats[i]) + diff_views = y_feats[i].dot(x_feats[i]) + id_logs.append({'diff_target': float(diff_target), + 'diff_input': float(diff_input), + 'diff_views': float(diff_views)}) + loss += 1 - diff_target + id_diff = float(diff_target) - float(diff_views) + sim_improvement += id_diff + count += 1 + + return loss / count, sim_improvement / count, id_logs diff --git a/models/encoder4editing/criteria/lpips/__init__.py b/models/encoder4editing/criteria/lpips/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/models/encoder4editing/criteria/lpips/lpips.py b/models/encoder4editing/criteria/lpips/lpips.py new file mode 100644 index 0000000000000000000000000000000000000000..1add6acc84c1c04cfcb536cf31ec5acdf24b716b --- /dev/null +++ b/models/encoder4editing/criteria/lpips/lpips.py @@ -0,0 +1,35 @@ +import torch +import torch.nn as nn + +from criteria.lpips.networks import get_network, LinLayers +from criteria.lpips.utils import get_state_dict + + +class LPIPS(nn.Module): + r"""Creates a criterion that measures + Learned Perceptual Image Patch Similarity (LPIPS). + Arguments: + net_type (str): the network type to compare the features: + 'alex' | 'squeeze' | 'vgg'. Default: 'alex'. + version (str): the version of LPIPS. Default: 0.1. + """ + def __init__(self, net_type: str = 'alex', version: str = '0.1'): + + assert version in ['0.1'], 'v0.1 is only supported now' + + super(LPIPS, self).__init__() + + # pretrained network + self.net = get_network(net_type).to("cuda") + + # linear layers + self.lin = LinLayers(self.net.n_channels_list).to("cuda") + self.lin.load_state_dict(get_state_dict(net_type, version)) + + def forward(self, x: torch.Tensor, y: torch.Tensor): + feat_x, feat_y = self.net(x), self.net(y) + + diff = [(fx - fy) ** 2 for fx, fy in zip(feat_x, feat_y)] + res = [l(d).mean((2, 3), True) for d, l in zip(diff, self.lin)] + + return torch.sum(torch.cat(res, 0)) / x.shape[0] diff --git a/models/encoder4editing/criteria/lpips/networks.py b/models/encoder4editing/criteria/lpips/networks.py new file mode 100644 index 0000000000000000000000000000000000000000..3a0d13ad2d560278f16586da68d3a5eadb26e746 --- /dev/null +++ b/models/encoder4editing/criteria/lpips/networks.py @@ -0,0 +1,96 @@ +from typing import Sequence + +from itertools import chain + +import torch +import torch.nn as nn +from torchvision import models + +from criteria.lpips.utils import normalize_activation + + +def get_network(net_type: str): + if net_type == 'alex': + return AlexNet() + elif net_type == 'squeeze': + return SqueezeNet() + elif net_type == 'vgg': + return VGG16() + else: + raise NotImplementedError('choose net_type from [alex, squeeze, vgg].') + + +class LinLayers(nn.ModuleList): + def __init__(self, n_channels_list: Sequence[int]): + super(LinLayers, self).__init__([ + nn.Sequential( + nn.Identity(), + nn.Conv2d(nc, 1, 1, 1, 0, bias=False) + ) for nc in n_channels_list + ]) + + for param in self.parameters(): + param.requires_grad = False + + +class BaseNet(nn.Module): + def __init__(self): + super(BaseNet, self).__init__() + + # register buffer + self.register_buffer( + 'mean', torch.Tensor([-.030, -.088, -.188])[None, :, None, None]) + self.register_buffer( + 'std', torch.Tensor([.458, .448, .450])[None, :, None, None]) + + def set_requires_grad(self, state: bool): + for param in chain(self.parameters(), self.buffers()): + param.requires_grad = state + + def z_score(self, x: torch.Tensor): + return (x - self.mean) / self.std + + def forward(self, x: torch.Tensor): + x = self.z_score(x) + + output = [] + for i, (_, layer) in enumerate(self.layers._modules.items(), 1): + x = layer(x) + if i in self.target_layers: + output.append(normalize_activation(x)) + if len(output) == len(self.target_layers): + break + return output + + +class SqueezeNet(BaseNet): + def __init__(self): + super(SqueezeNet, self).__init__() + + self.layers = models.squeezenet1_1(True).features + self.target_layers = [2, 5, 8, 10, 11, 12, 13] + self.n_channels_list = [64, 128, 256, 384, 384, 512, 512] + + self.set_requires_grad(False) + + +class AlexNet(BaseNet): + def __init__(self): + super(AlexNet, self).__init__() + + self.layers = models.alexnet(True).features + self.target_layers = [2, 5, 8, 10, 12] + self.n_channels_list = [64, 192, 384, 256, 256] + + self.set_requires_grad(False) + + +class VGG16(BaseNet): + def __init__(self): + super(VGG16, self).__init__() + + self.layers = models.vgg16(True).features + self.target_layers = [4, 9, 16, 23, 30] + self.n_channels_list = [64, 128, 256, 512, 512] + + self.set_requires_grad(False) \ No newline at end of file diff --git a/models/encoder4editing/criteria/lpips/utils.py b/models/encoder4editing/criteria/lpips/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..3d15a0983775810ef6239c561c67939b2b9ee3b5 --- /dev/null +++ b/models/encoder4editing/criteria/lpips/utils.py @@ -0,0 +1,30 @@ +from collections import OrderedDict + +import torch + + +def normalize_activation(x, eps=1e-10): + norm_factor = torch.sqrt(torch.sum(x ** 2, dim=1, keepdim=True)) + return x / (norm_factor + eps) + + +def get_state_dict(net_type: str = 'alex', version: str = '0.1'): + # build url + url = 'https://raw.githubusercontent.com/richzhang/PerceptualSimilarity/' \ + + f'master/lpips/weights/v{version}/{net_type}.pth' + + # download + old_state_dict = torch.hub.load_state_dict_from_url( + url, progress=True, + map_location=None if torch.cuda.is_available() else torch.device('cpu') + ) + + # rename keys + new_state_dict = OrderedDict() + for key, val in old_state_dict.items(): + new_key = key + new_key = new_key.replace('lin', '') + new_key = new_key.replace('model.', '') + new_state_dict[new_key] = val + + return new_state_dict diff --git a/models/encoder4editing/criteria/moco_loss.py b/models/encoder4editing/criteria/moco_loss.py new file mode 100644 index 0000000000000000000000000000000000000000..8fb13fbd426202cff9014c876c85b0d5c4ec6a9d --- /dev/null +++ b/models/encoder4editing/criteria/moco_loss.py @@ -0,0 +1,71 @@ +import torch +from torch import nn +import torch.nn.functional as F + +from configs.paths_config import model_paths + + +class MocoLoss(nn.Module): + + def __init__(self, opts): + super(MocoLoss, self).__init__() + print("Loading MOCO model from path: {}".format(model_paths["moco"])) + self.model = self.__load_model() + self.model.eval() + for param in self.model.parameters(): + param.requires_grad = False + + @staticmethod + def __load_model(): + import torchvision.models as models + model = models.__dict__["resnet50"]() + # freeze all layers but the last fc + for name, param in model.named_parameters(): + if name not in ['fc.weight', 'fc.bias']: + param.requires_grad = False + checkpoint = torch.load(model_paths['moco'], map_location="cpu") + state_dict = checkpoint['state_dict'] + # rename moco pre-trained keys + for k in list(state_dict.keys()): + # retain only encoder_q up to before the embedding layer + if k.startswith('module.encoder_q') and not k.startswith('module.encoder_q.fc'): + # remove prefix + state_dict[k[len("module.encoder_q."):]] = state_dict[k] + # delete renamed or unused k + del state_dict[k] + msg = model.load_state_dict(state_dict, strict=False) + assert set(msg.missing_keys) == {"fc.weight", "fc.bias"} + # remove output layer + model = nn.Sequential(*list(model.children())[:-1]).cuda() + return model + + def extract_feats(self, x): + x = F.interpolate(x, size=224) + x_feats = self.model(x) + x_feats = nn.functional.normalize(x_feats, dim=1) + x_feats = x_feats.squeeze() + return x_feats + + def forward(self, y_hat, y, x): + n_samples = x.shape[0] + x_feats = self.extract_feats(x) + y_feats = self.extract_feats(y) + y_hat_feats = self.extract_feats(y_hat) + y_feats = y_feats.detach() + loss = 0 + sim_improvement = 0 + sim_logs = [] + count = 0 + for i in range(n_samples): + diff_target = y_hat_feats[i].dot(y_feats[i]) + diff_input = y_hat_feats[i].dot(x_feats[i]) + diff_views = y_feats[i].dot(x_feats[i]) + sim_logs.append({'diff_target': float(diff_target), + 'diff_input': float(diff_input), + 'diff_views': float(diff_views)}) + loss += 1 - diff_target + sim_diff = float(diff_target) - float(diff_views) + sim_improvement += sim_diff + count += 1 + + return loss / count, sim_improvement / count, sim_logs diff --git a/models/encoder4editing/criteria/w_norm.py b/models/encoder4editing/criteria/w_norm.py new file mode 100644 index 0000000000000000000000000000000000000000..a45ab6f67d8a3f7051be4b7236fa2f38446fd2c1 --- /dev/null +++ b/models/encoder4editing/criteria/w_norm.py @@ -0,0 +1,14 @@ +import torch +from torch import nn + + +class WNormLoss(nn.Module): + + def __init__(self, start_from_latent_avg=True): + super(WNormLoss, self).__init__() + self.start_from_latent_avg = start_from_latent_avg + + def forward(self, latent, latent_avg=None): + if self.start_from_latent_avg: + latent = latent - latent_avg + return torch.sum(latent.norm(2, dim=(1, 2))) / latent.shape[0] diff --git a/models/encoder4editing/datasets/__init__.py b/models/encoder4editing/datasets/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/models/encoder4editing/datasets/gt_res_dataset.py b/models/encoder4editing/datasets/gt_res_dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..c0beacfee5335aa10aa7e8b7cabe206d7f9a56f7 --- /dev/null +++ b/models/encoder4editing/datasets/gt_res_dataset.py @@ -0,0 +1,32 @@ +#!/usr/bin/python +# encoding: utf-8 +import os +from torch.utils.data import Dataset +from PIL import Image +import torch + +class GTResDataset(Dataset): + + def __init__(self, root_path, gt_dir=None, transform=None, transform_train=None): + self.pairs = [] + for f in os.listdir(root_path): + image_path = os.path.join(root_path, f) + gt_path = os.path.join(gt_dir, f) + if f.endswith(".jpg") or f.endswith(".png"): + self.pairs.append([image_path, gt_path.replace('.png', '.jpg'), None]) + self.transform = transform + self.transform_train = transform_train + + def __len__(self): + return len(self.pairs) + + def __getitem__(self, index): + from_path, to_path, _ = self.pairs[index] + from_im = Image.open(from_path).convert('RGB') + to_im = Image.open(to_path).convert('RGB') + + if self.transform: + to_im = self.transform(to_im) + from_im = self.transform(from_im) + + return from_im, to_im diff --git a/models/encoder4editing/datasets/images_dataset.py b/models/encoder4editing/datasets/images_dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..8387117ece2643b356e7c792c1da8ec6d05fa772 --- /dev/null +++ b/models/encoder4editing/datasets/images_dataset.py @@ -0,0 +1,33 @@ +from torch.utils.data import Dataset +from PIL import Image +from ..utils import data_utils + + +class ImagesDataset(Dataset): + + def __init__(self, source_root, target_root, opts, target_transform=None, source_transform=None): + self.source_paths = sorted(data_utils.make_dataset(source_root)) + self.target_paths = sorted(data_utils.make_dataset(target_root)) + self.source_transform = source_transform + self.target_transform = target_transform + self.opts = opts + + def __len__(self): + return len(self.source_paths) + + def __getitem__(self, index): + from_path = self.source_paths[index] + from_im = Image.open(from_path) + from_im = from_im.convert('RGB') + + to_path = self.target_paths[index] + to_im = Image.open(to_path).convert('RGB') + if self.target_transform: + to_im = self.target_transform(to_im) + + if self.source_transform: + from_im = self.source_transform(from_im) + else: + from_im = to_im + + return from_im, to_im diff --git a/models/encoder4editing/datasets/inference_dataset.py b/models/encoder4editing/datasets/inference_dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..cfccc5adbdf20e5b90739ec4d4772cc5f48c0357 --- /dev/null +++ b/models/encoder4editing/datasets/inference_dataset.py @@ -0,0 +1,25 @@ +from torch.utils.data import Dataset +from PIL import Image +from ..utils import data_utils + + +class InferenceDataset(Dataset): + + def __init__(self, root, opts, transform=None, preprocess=None): + self.paths = sorted(data_utils.make_dataset(root)) + self.transform = transform + self.preprocess = preprocess + self.opts = opts + + def __len__(self): + return len(self.paths) + + def __getitem__(self, index): + from_path = self.paths[index] + if self.preprocess is not None: + from_im = self.preprocess(from_path) + else: + from_im = Image.open(from_path).convert('RGB') + if self.transform: + from_im = self.transform(from_im) + return from_im diff --git a/models/encoder4editing/editings/ganspace.py b/models/encoder4editing/editings/ganspace.py new file mode 100644 index 0000000000000000000000000000000000000000..0c286a421280c542e9776a75e64bb65409da8fc7 --- /dev/null +++ b/models/encoder4editing/editings/ganspace.py @@ -0,0 +1,22 @@ +import torch + + +def edit(latents, pca, edit_directions): + edit_latents = [] + for latent in latents: + for pca_idx, start, end, strength in edit_directions: + delta = get_delta(pca, latent, pca_idx, strength) + delta_padded = torch.zeros(latent.shape).to('cuda') + delta_padded[start:end] += delta.repeat(end - start, 1) + edit_latents.append(latent + delta_padded) + return torch.stack(edit_latents) + + +def get_delta(pca, latent, idx, strength): + # pca: ganspace checkpoint. latent: (16, 512) w+ + w_centered = latent - pca['mean'].to('cuda') + lat_comp = pca['comp'].to('cuda') + lat_std = pca['std'].to('cuda') + w_coord = torch.sum(w_centered[0].reshape(-1)*lat_comp[idx].reshape(-1)) / lat_std[idx] + delta = (strength - w_coord)*lat_comp[idx]*lat_std[idx] + return delta diff --git a/models/encoder4editing/editings/latent_editor.py b/models/encoder4editing/editings/latent_editor.py new file mode 100644 index 0000000000000000000000000000000000000000..4bebca2f5c86f71b58fa1f30d24bfcb0da06d88f --- /dev/null +++ b/models/encoder4editing/editings/latent_editor.py @@ -0,0 +1,45 @@ +import torch +import sys +sys.path.append(".") +sys.path.append("..") +from editings import ganspace, sefa +from utils.common import tensor2im + + +class LatentEditor(object): + def __init__(self, stylegan_generator, is_cars=False): + self.generator = stylegan_generator + self.is_cars = is_cars # Since the cars StyleGAN output is 384x512, there is a need to crop the 512x512 output. + + def apply_ganspace(self, latent, ganspace_pca, edit_directions): + edit_latents = ganspace.edit(latent, ganspace_pca, edit_directions) + return self._latents_to_image(edit_latents) + + def apply_interfacegan(self, latent, direction, factor=1, factor_range=None): + edit_latents = [] + if factor_range is not None: # Apply a range of editing factors. for example, (-5, 5) + for f in range(*factor_range): + edit_latent = latent + f * direction + edit_latents.append(edit_latent) + edit_latents = torch.cat(edit_latents) + else: + edit_latents = latent + factor * direction + return self._latents_to_image(edit_latents) + + def apply_sefa(self, latent, indices=[2, 3, 4, 5], **kwargs): + edit_latents = sefa.edit(self.generator, latent, indices, **kwargs) + return self._latents_to_image(edit_latents) + + # Currently, in order to apply StyleFlow editings, one should run inference, + # save the latent codes and load them form the official StyleFlow repository. + # def apply_styleflow(self): + # pass + + def _latents_to_image(self, latents): + with torch.no_grad(): + images, _ = self.generator([latents], randomize_noise=False, input_is_latent=True) + if self.is_cars: + images = images[:, :, 64:448, :] # 512x512 -> 384x512 + horizontal_concat_image = torch.cat(list(images), 2) + final_image = tensor2im(horizontal_concat_image) + return final_image diff --git a/models/encoder4editing/editings/sefa.py b/models/encoder4editing/editings/sefa.py new file mode 100644 index 0000000000000000000000000000000000000000..db7083ce463b765a7cf452807883a3b85fb63fa5 --- /dev/null +++ b/models/encoder4editing/editings/sefa.py @@ -0,0 +1,46 @@ +import torch +import numpy as np +from tqdm import tqdm + + +def edit(generator, latents, indices, semantics=1, start_distance=-15.0, end_distance=15.0, num_samples=1, step=11): + + layers, boundaries, values = factorize_weight(generator, indices) + codes = latents.detach().cpu().numpy() # (1,18,512) + + # Generate visualization pages. + distances = np.linspace(start_distance, end_distance, step) + num_sam = num_samples + num_sem = semantics + + edited_latents = [] + for sem_id in tqdm(range(num_sem), desc='Semantic ', leave=False): + boundary = boundaries[sem_id:sem_id + 1] + for sam_id in tqdm(range(num_sam), desc='Sample ', leave=False): + code = codes[sam_id:sam_id + 1] + for col_id, d in enumerate(distances, start=1): + temp_code = code.copy() + temp_code[:, layers, :] += boundary * d + edited_latents.append(torch.from_numpy(temp_code).float().cuda()) + return torch.cat(edited_latents) + + +def factorize_weight(g_ema, layers='all'): + + weights = [] + if layers == 'all' or 0 in layers: + weight = g_ema.conv1.conv.modulation.weight.T + weights.append(weight.cpu().detach().numpy()) + + if layers == 'all': + layers = list(range(g_ema.num_layers - 1)) + else: + layers = [l - 1 for l in layers if l != 0] + + for idx in layers: + weight = g_ema.convs[idx].conv.modulation.weight.T + weights.append(weight.cpu().detach().numpy()) + weight = np.concatenate(weights, axis=1).astype(np.float32) + weight = weight / np.linalg.norm(weight, axis=0, keepdims=True) + eigen_values, eigen_vectors = np.linalg.eig(weight.dot(weight.T)) + return layers, eigen_vectors.T, eigen_values diff --git a/models/encoder4editing/environment/e4e_env.yaml b/models/encoder4editing/environment/e4e_env.yaml new file mode 100644 index 0000000000000000000000000000000000000000..4f537615ebb47afd74b5a9856fb9cbea2e0c4bf4 --- /dev/null +++ b/models/encoder4editing/environment/e4e_env.yaml @@ -0,0 +1,73 @@ +name: e4e_env +channels: + - conda-forge + - defaults +dependencies: + - _libgcc_mutex=0.1=main + - ca-certificates=2020.4.5.1=hecc5488_0 + - certifi=2020.4.5.1=py36h9f0ad1d_0 + - libedit=3.1.20181209=hc058e9b_0 + - libffi=3.2.1=hd88cf55_4 + - libgcc-ng=9.1.0=hdf63c60_0 + - libstdcxx-ng=9.1.0=hdf63c60_0 + - ncurses=6.2=he6710b0_1 + - ninja=1.10.0=hc9558a2_0 + - openssl=1.1.1g=h516909a_0 + - pip=20.0.2=py36_3 + - python=3.6.7=h0371630_0 + - python_abi=3.6=1_cp36m + - readline=7.0=h7b6447c_5 + - setuptools=46.4.0=py36_0 + - sqlite=3.31.1=h62c20be_1 + - tk=8.6.8=hbc83047_0 + - wheel=0.34.2=py36_0 + - xz=5.2.5=h7b6447c_0 + - zlib=1.2.11=h7b6447c_3 + - pip: + - absl-py==0.9.0 + - cachetools==4.1.0 + - chardet==3.0.4 + - cycler==0.10.0 + - decorator==4.4.2 + - future==0.18.2 + - google-auth==1.15.0 + - google-auth-oauthlib==0.4.1 + - grpcio==1.29.0 + - idna==2.9 + - imageio==2.8.0 + - importlib-metadata==1.6.0 + - kiwisolver==1.2.0 + - markdown==3.2.2 + - matplotlib==3.2.1 + - mxnet==1.6.0 + - networkx==2.4 + - numpy==1.18.4 + - oauthlib==3.1.0 + - opencv-python==4.2.0.34 + - pillow==7.1.2 + - protobuf==3.12.1 + - pyasn1==0.4.8 + - pyasn1-modules==0.2.8 + - pyparsing==2.4.7 + - python-dateutil==2.8.1 + - pytorch-lightning==0.7.1 + - pywavelets==1.1.1 + - requests==2.23.0 + - requests-oauthlib==1.3.0 + - rsa==4.0 + - scikit-image==0.17.2 + - scipy==1.4.1 + - six==1.15.0 + - tensorboard==2.2.1 + - tensorboard-plugin-wit==1.6.0.post3 + - tensorboardx==1.9 + - tifffile==2020.5.25 + - torch==1.6.0 + - torchvision==0.7.1 + - tqdm==4.46.0 + - urllib3==1.25.9 + - werkzeug==1.0.1 + - zipp==3.1.0 + - pyaml +prefix: ~/anaconda3/envs/e4e_env + diff --git a/models/encoder4editing/metrics/LEC.py b/models/encoder4editing/metrics/LEC.py new file mode 100644 index 0000000000000000000000000000000000000000..3eef2d2f00a4d757a56b6e845a8fde16aab306ab --- /dev/null +++ b/models/encoder4editing/metrics/LEC.py @@ -0,0 +1,134 @@ +import sys +import argparse +import torch +import numpy as np +from torch.utils.data import DataLoader + +sys.path.append(".") +sys.path.append("..") + +from configs import data_configs +from datasets.images_dataset import ImagesDataset +from utils.model_utils import setup_model + + +class LEC: + def __init__(self, net, is_cars=False): + """ + Latent Editing Consistency metric as proposed in the main paper. + :param net: e4e model loaded over the pSp framework. + :param is_cars: An indication as to whether or not to crop the middle of the StyleGAN's output images. + """ + self.net = net + self.is_cars = is_cars + + def _encode(self, images): + """ + Encodes the given images into StyleGAN's latent space. + :param images: Tensor of shape NxCxHxW representing the images to be encoded. + :return: Tensor of shape NxKx512 representing the latent space embeddings of the given image (in W(K, *) space). + """ + codes = self.net.encoder(images) + assert codes.ndim == 3, f"Invalid latent codes shape, should be NxKx512 but is {codes.shape}" + # normalize with respect to the center of an average face + if self.net.opts.start_from_latent_avg: + codes = codes + self.net.latent_avg.repeat(codes.shape[0], 1, 1) + return codes + + def _generate(self, codes): + """ + Generate the StyleGAN2 images of the given codes + :param codes: Tensor of shape NxKx512 representing the StyleGAN's latent codes (in W(K, *) space). + :return: Tensor of shape NxCxHxW representing the generated images. + """ + images, _ = self.net.decoder([codes], input_is_latent=True, randomize_noise=False, return_latents=True) + images = self.net.face_pool(images) + if self.is_cars: + images = images[:, :, 32:224, :] + return images + + @staticmethod + def _filter_outliers(arr): + arr = np.array(arr) + + lo = np.percentile(arr, 1, interpolation="lower") + hi = np.percentile(arr, 99, interpolation="higher") + return np.extract( + np.logical_and(lo <= arr, arr <= hi), arr + ) + + def calculate_metric(self, data_loader, edit_function, inverse_edit_function): + """ + Calculate the LEC metric score. + :param data_loader: An iterable that returns a tuple of (images, _), similar to the training data loader. + :param edit_function: A function that receives latent codes and performs a semantically meaningful edit in the + latent space. + :param inverse_edit_function: A function that receives latent codes and performs the inverse edit of the + `edit_function` parameter. + :return: The LEC metric score. + """ + distances = [] + with torch.no_grad(): + for batch in data_loader: + x, _ = batch + inputs = x.to(device).float() + + codes = self._encode(inputs) + edited_codes = edit_function(codes) + edited_image = self._generate(edited_codes) + edited_image_inversion_codes = self._encode(edited_image) + inverse_edit_codes = inverse_edit_function(edited_image_inversion_codes) + + dist = (codes - inverse_edit_codes).norm(2, dim=(1, 2)).mean() + distances.append(dist.to("cpu").numpy()) + + distances = self._filter_outliers(distances) + return distances.mean() + + +if __name__ == "__main__": + device = "cuda" + + parser = argparse.ArgumentParser(description="LEC metric calculator") + + parser.add_argument("--batch", type=int, default=8, help="batch size for the models") + parser.add_argument("--images_dir", type=str, default=None, + help="Path to the images directory on which we calculate the LEC score") + parser.add_argument("ckpt", metavar="CHECKPOINT", help="path to the model checkpoints") + + args = parser.parse_args() + print(args) + + net, opts = setup_model(args.ckpt, device) + dataset_args = data_configs.DATASETS[opts.dataset_type] + transforms_dict = dataset_args['transforms'](opts).get_transforms() + + images_directory = dataset_args['test_source_root'] if args.images_dir is None else args.images_dir + test_dataset = ImagesDataset(source_root=images_directory, + target_root=images_directory, + source_transform=transforms_dict['transform_source'], + target_transform=transforms_dict['transform_test'], + opts=opts) + + data_loader = DataLoader(test_dataset, + batch_size=args.batch, + shuffle=False, + num_workers=2, + drop_last=True) + + print(f'dataset length: {len(test_dataset)}') + + # In the following example, we are using an InterfaceGAN based editing to calculate the LEC metric. + # Change the provided example according to your domain and needs. + direction = torch.load('../editings/interfacegan_directions/age.pt').to(device) + + def edit_func_example(codes): + return codes + 3 * direction + + + def inverse_edit_func_example(codes): + return codes - 3 * direction + + lec = LEC(net, is_cars='car' in opts.dataset_type) + result = lec.calculate_metric(data_loader, edit_func_example, inverse_edit_func_example) + print(f"LEC: {result}") diff --git a/models/encoder4editing/models/.DS_Store b/models/encoder4editing/models/.DS_Store new file mode 100644 index 0000000000000000000000000000000000000000..0eb37726d705b0ac8e70751bdbc647781f065da7 Binary files /dev/null and b/models/encoder4editing/models/.DS_Store differ diff --git a/models/encoder4editing/models/__init__.py b/models/encoder4editing/models/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/models/encoder4editing/models/discriminator.py b/models/encoder4editing/models/discriminator.py new file mode 100644 index 0000000000000000000000000000000000000000..16bf3722c7f2e35cdc9bd177a33ed0975e67200d --- /dev/null +++ b/models/encoder4editing/models/discriminator.py @@ -0,0 +1,20 @@ +from torch import nn + + +class LatentCodesDiscriminator(nn.Module): + def __init__(self, style_dim, n_mlp): + super().__init__() + + self.style_dim = style_dim + + layers = [] + for i in range(n_mlp-1): + layers.append( + nn.Linear(style_dim, style_dim) + ) + layers.append(nn.LeakyReLU(0.2)) + layers.append(nn.Linear(512, 1)) + self.mlp = nn.Sequential(*layers) + + def forward(self, w): + return self.mlp(w) diff --git a/models/encoder4editing/models/encoders/__init__.py b/models/encoder4editing/models/encoders/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/models/encoder4editing/models/encoders/helpers.py b/models/encoder4editing/models/encoders/helpers.py new file mode 100644 index 0000000000000000000000000000000000000000..c4a58b34ea5ca6912fe53c63dede0a8696f5c024 --- /dev/null +++ b/models/encoder4editing/models/encoders/helpers.py @@ -0,0 +1,140 @@ +from collections import namedtuple +import torch +import torch.nn.functional as F +from torch.nn import Conv2d, BatchNorm2d, PReLU, ReLU, Sigmoid, MaxPool2d, AdaptiveAvgPool2d, Sequential, Module + +""" +ArcFace implementation from [TreB1eN](https://github.com/TreB1eN/InsightFace_Pytorch) +""" + + +class Flatten(Module): + def forward(self, input): + return input.view(input.size(0), -1) + + +def l2_norm(input, axis=1): + norm = torch.norm(input, 2, axis, True) + output = torch.div(input, norm) + return output + + +class Bottleneck(namedtuple('Block', ['in_channel', 'depth', 'stride'])): + """ A named tuple describing a ResNet block. """ + + +def get_block(in_channel, depth, num_units, stride=2): + return [Bottleneck(in_channel, depth, stride)] + [Bottleneck(depth, depth, 1) for i in range(num_units - 1)] + + +def get_blocks(num_layers): + if num_layers == 50: + blocks = [ + get_block(in_channel=64, depth=64, num_units=3), + get_block(in_channel=64, depth=128, num_units=4), + get_block(in_channel=128, depth=256, num_units=14), + get_block(in_channel=256, depth=512, num_units=3) + ] + elif num_layers == 100: + blocks = [ + get_block(in_channel=64, depth=64, num_units=3), + get_block(in_channel=64, depth=128, num_units=13), + get_block(in_channel=128, depth=256, num_units=30), + get_block(in_channel=256, depth=512, num_units=3) + ] + elif num_layers == 152: + blocks = [ + get_block(in_channel=64, depth=64, num_units=3), + get_block(in_channel=64, depth=128, num_units=8), + get_block(in_channel=128, depth=256, num_units=36), + get_block(in_channel=256, depth=512, num_units=3) + ] + else: + raise ValueError("Invalid number of layers: {}. Must be one of [50, 100, 152]".format(num_layers)) + return blocks + + +class SEModule(Module): + def __init__(self, channels, reduction): + super(SEModule, self).__init__() + self.avg_pool = AdaptiveAvgPool2d(1) + self.fc1 = Conv2d(channels, channels // reduction, kernel_size=1, padding=0, bias=False) + self.relu = ReLU(inplace=True) + self.fc2 = Conv2d(channels // reduction, channels, kernel_size=1, padding=0, bias=False) + self.sigmoid = Sigmoid() + + def forward(self, x): + module_input = x + x = self.avg_pool(x) + x = self.fc1(x) + x = self.relu(x) + x = self.fc2(x) + x = self.sigmoid(x) + return module_input * x + + +class bottleneck_IR(Module): + def __init__(self, in_channel, depth, stride): + super(bottleneck_IR, self).__init__() + if in_channel == depth: + self.shortcut_layer = MaxPool2d(1, stride) + else: + self.shortcut_layer = Sequential( + Conv2d(in_channel, depth, (1, 1), stride, bias=False), + BatchNorm2d(depth) + ) + self.res_layer = Sequential( + BatchNorm2d(in_channel), + Conv2d(in_channel, depth, (3, 3), (1, 1), 1, bias=False), PReLU(depth), + Conv2d(depth, depth, (3, 3), stride, 1, bias=False), BatchNorm2d(depth) + ) + + def forward(self, x): + shortcut = self.shortcut_layer(x) + res = self.res_layer(x) + return res + shortcut + + +class bottleneck_IR_SE(Module): + def __init__(self, in_channel, depth, stride): + super(bottleneck_IR_SE, self).__init__() + if in_channel == depth: + self.shortcut_layer = MaxPool2d(1, stride) + else: + self.shortcut_layer = Sequential( + Conv2d(in_channel, depth, (1, 1), stride, bias=False), + BatchNorm2d(depth) + ) + self.res_layer = Sequential( + BatchNorm2d(in_channel), + Conv2d(in_channel, depth, (3, 3), (1, 1), 1, bias=False), + PReLU(depth), + Conv2d(depth, depth, (3, 3), stride, 1, bias=False), + BatchNorm2d(depth), + SEModule(depth, 16) + ) + + def forward(self, x): + shortcut = self.shortcut_layer(x) + res = self.res_layer(x) + return res + shortcut + + +def _upsample_add(x, y): + """Upsample and add two feature maps. + Args: + x: (Variable) top feature map to be upsampled. + y: (Variable) lateral feature map. + Returns: + (Variable) added feature map. + Note in PyTorch, when input size is odd, the upsampled feature map + with `F.upsample(..., scale_factor=2, mode='nearest')` + maybe not equal to the lateral feature map size. + e.g. + original input size: [N,_,15,15] -> + conv2d feature map size: [N,_,8,8] -> + upsampled feature map size: [N,_,16,16] + So we choose bilinear upsample which supports arbitrary output sizes. + """ + _, _, H, W = y.size() + return F.interpolate(x, size=(H, W), mode='bilinear', align_corners=True) + y diff --git a/models/encoder4editing/models/encoders/model_irse.py b/models/encoder4editing/models/encoders/model_irse.py new file mode 100644 index 0000000000000000000000000000000000000000..ea6c6091c1e71279ff0bc7e013b0cea287cb01b3 --- /dev/null +++ b/models/encoder4editing/models/encoders/model_irse.py @@ -0,0 +1,84 @@ +from torch.nn import Linear, Conv2d, BatchNorm1d, BatchNorm2d, PReLU, Dropout, Sequential, Module +from models.encoders.helpers import get_blocks, Flatten, bottleneck_IR, bottleneck_IR_SE, l2_norm + +""" +Modified Backbone implementation from [TreB1eN](https://github.com/TreB1eN/InsightFace_Pytorch) +""" + + +class Backbone(Module): + def __init__(self, input_size, num_layers, mode='ir', drop_ratio=0.4, affine=True): + super(Backbone, self).__init__() + assert input_size in [112, 224], "input_size should be 112 or 224" + assert num_layers in [50, 100, 152], "num_layers should be 50, 100 or 152" + assert mode in ['ir', 'ir_se'], "mode should be ir or ir_se" + blocks = get_blocks(num_layers) + if mode == 'ir': + unit_module = bottleneck_IR + elif mode == 'ir_se': + unit_module = bottleneck_IR_SE + self.input_layer = Sequential(Conv2d(3, 64, (3, 3), 1, 1, bias=False), + BatchNorm2d(64), + PReLU(64)) + if input_size == 112: + self.output_layer = Sequential(BatchNorm2d(512), + Dropout(drop_ratio), + Flatten(), + Linear(512 * 7 * 7, 512), + BatchNorm1d(512, affine=affine)) + else: + self.output_layer = Sequential(BatchNorm2d(512), + Dropout(drop_ratio), + Flatten(), + Linear(512 * 14 * 14, 512), + BatchNorm1d(512, affine=affine)) + + modules = [] + for block in blocks: + for bottleneck in block: + modules.append(unit_module(bottleneck.in_channel, + bottleneck.depth, + bottleneck.stride)) + self.body = Sequential(*modules) + + def forward(self, x): + x = self.input_layer(x) + x = self.body(x) + x = self.output_layer(x) + return l2_norm(x) + + +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 diff --git a/models/encoder4editing/models/encoders/psp_encoders.py b/models/encoder4editing/models/encoders/psp_encoders.py new file mode 100644 index 0000000000000000000000000000000000000000..183561be36853c6d35ff0ff338e28b45b0468c09 --- /dev/null +++ b/models/encoder4editing/models/encoders/psp_encoders.py @@ -0,0 +1,235 @@ +from enum import Enum +import math +import numpy as np +import torch +from torch import nn +from torch.nn import Conv2d, BatchNorm2d, PReLU, Sequential, Module + +from models.encoder4editing.models.encoders.helpers import get_blocks, bottleneck_IR, bottleneck_IR_SE, _upsample_add +from models.encoder4editing.models.stylegan2.model import EqualLinear + + +class ProgressiveStage(Enum): + WTraining = 0 + Delta1Training = 1 + Delta2Training = 2 + Delta3Training = 3 + Delta4Training = 4 + Delta5Training = 5 + Delta6Training = 6 + Delta7Training = 7 + Delta8Training = 8 + Delta9Training = 9 + Delta10Training = 10 + Delta11Training = 11 + Delta12Training = 12 + Delta13Training = 13 + Delta14Training = 14 + Delta15Training = 15 + Delta16Training = 16 + Delta17Training = 17 + Inference = 18 + + +class GradualStyleBlock(Module): + def __init__(self, in_c, out_c, spatial): + super(GradualStyleBlock, self).__init__() + self.out_c = out_c + self.spatial = spatial + num_pools = int(np.log2(spatial)) + modules = [] + modules += [Conv2d(in_c, out_c, kernel_size=3, stride=2, padding=1), + nn.LeakyReLU()] + for i in range(num_pools - 1): + modules += [ + Conv2d(out_c, out_c, kernel_size=3, stride=2, padding=1), + nn.LeakyReLU() + ] + self.convs = nn.Sequential(*modules) + self.linear = EqualLinear(out_c, out_c, lr_mul=1) + + def forward(self, x): + x = self.convs(x) + x = x.view(-1, self.out_c) + x = self.linear(x) + return x + + +class GradualStyleEncoder(Module): + def __init__(self, num_layers, mode='ir', opts=None): + super(GradualStyleEncoder, self).__init__() + assert num_layers in [50, 100, 152], 'num_layers should be 50,100, or 152' + assert mode in ['ir', 'ir_se'], 'mode should be ir or ir_se' + blocks = get_blocks(num_layers) + if mode == 'ir': + unit_module = bottleneck_IR + elif mode == 'ir_se': + unit_module = bottleneck_IR_SE + self.input_layer = Sequential(Conv2d(3, 64, (3, 3), 1, 1, bias=False), + BatchNorm2d(64), + PReLU(64)) + modules = [] + for block in blocks: + for bottleneck in block: + modules.append(unit_module(bottleneck.in_channel, + bottleneck.depth, + bottleneck.stride)) + self.body = Sequential(*modules) + + self.styles = nn.ModuleList() + log_size = int(math.log(opts.stylegan_size, 2)) + self.style_count = 2 * log_size - 2 + self.coarse_ind = 3 + self.middle_ind = 7 + for i in range(self.style_count): + if i < self.coarse_ind: + style = GradualStyleBlock(512, 512, 16) + elif i < self.middle_ind: + style = GradualStyleBlock(512, 512, 32) + else: + style = GradualStyleBlock(512, 512, 64) + self.styles.append(style) + self.latlayer1 = nn.Conv2d(256, 512, kernel_size=1, stride=1, padding=0) + self.latlayer2 = nn.Conv2d(128, 512, kernel_size=1, stride=1, padding=0) + + def forward(self, x): + x = self.input_layer(x) + + latents = [] + modulelist = list(self.body._modules.values()) + for i, l in enumerate(modulelist): + x = l(x) + if i == 6: + c1 = x + elif i == 20: + c2 = x + elif i == 23: + c3 = x + + for j in range(self.coarse_ind): + latents.append(self.styles[j](c3)) + + p2 = _upsample_add(c3, self.latlayer1(c2)) + for j in range(self.coarse_ind, self.middle_ind): + latents.append(self.styles[j](p2)) + + p1 = _upsample_add(p2, self.latlayer2(c1)) + for j in range(self.middle_ind, self.style_count): + latents.append(self.styles[j](p1)) + + out = torch.stack(latents, dim=1) + return out + + +class Encoder4Editing(Module): + def __init__(self, num_layers, mode='ir', opts=None): + super(Encoder4Editing, self).__init__() + assert num_layers in [50, 100, 152], 'num_layers should be 50,100, or 152' + assert mode in ['ir', 'ir_se'], 'mode should be ir or ir_se' + blocks = get_blocks(num_layers) + if mode == 'ir': + unit_module = bottleneck_IR + elif mode == 'ir_se': + unit_module = bottleneck_IR_SE + self.input_layer = Sequential(Conv2d(3, 64, (3, 3), 1, 1, bias=False), + BatchNorm2d(64), + PReLU(64)) + modules = [] + for block in blocks: + for bottleneck in block: + modules.append(unit_module(bottleneck.in_channel, + bottleneck.depth, + bottleneck.stride)) + self.body = Sequential(*modules) + + self.styles = nn.ModuleList() + log_size = int(math.log(opts.stylegan_size, 2)) + self.style_count = 2 * log_size - 2 + self.coarse_ind = 3 + self.middle_ind = 7 + + for i in range(self.style_count): + if i < self.coarse_ind: + style = GradualStyleBlock(512, 512, 16) + elif i < self.middle_ind: + style = GradualStyleBlock(512, 512, 32) + else: + style = GradualStyleBlock(512, 512, 64) + self.styles.append(style) + + self.latlayer1 = nn.Conv2d(256, 512, kernel_size=1, stride=1, padding=0) + self.latlayer2 = nn.Conv2d(128, 512, kernel_size=1, stride=1, padding=0) + + self.progressive_stage = ProgressiveStage.Inference + + def get_deltas_starting_dimensions(self): + ''' Get a list of the initial dimension of every delta from which it is applied ''' + return list(range(self.style_count)) # Each dimension has a delta applied to it + + def set_progressive_stage(self, new_stage: ProgressiveStage): + self.progressive_stage = new_stage + print('Changed progressive stage to: ', new_stage) + + def forward(self, x): + x = self.input_layer(x) + + modulelist = list(self.body._modules.values()) + for i, l in enumerate(modulelist): + x = l(x) + if i == 6: + c1 = x + elif i == 20: + c2 = x + elif i == 23: + c3 = x + + # Infer main W and duplicate it + w0 = self.styles[0](c3) + w = w0.repeat(self.style_count, 1, 1).permute(1, 0, 2) + stage = self.progressive_stage.value + features = c3 + for i in range(1, min(stage + 1, self.style_count)): # Infer additional deltas + if i == self.coarse_ind: + p2 = _upsample_add(c3, self.latlayer1(c2)) # FPN's middle features + features = p2 + elif i == self.middle_ind: + p1 = _upsample_add(p2, self.latlayer2(c1)) # FPN's fine features + features = p1 + delta_i = self.styles[i](features) + w[:, i] += delta_i + return w + + +class BackboneEncoderUsingLastLayerIntoW(Module): + def __init__(self, num_layers, mode='ir', opts=None): + super(BackboneEncoderUsingLastLayerIntoW, self).__init__() + print('Using BackboneEncoderUsingLastLayerIntoW') + 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)) + self.output_pool = torch.nn.AdaptiveAvgPool2d((1, 1)) + self.linear = EqualLinear(512, 512, lr_mul=1) + modules = [] + for block in blocks: + for bottleneck in block: + modules.append(unit_module(bottleneck.in_channel, + bottleneck.depth, + bottleneck.stride)) + self.body = Sequential(*modules) + log_size = int(math.log(opts.stylegan_size, 2)) + self.style_count = 2 * log_size - 2 + + def forward(self, x): + x = self.input_layer(x) + x = self.body(x) + x = self.output_pool(x) + x = x.view(-1, 512) + x = self.linear(x) + return x.repeat(self.style_count, 1, 1).permute(1, 0, 2) diff --git a/models/encoder4editing/models/latent_codes_pool.py b/models/encoder4editing/models/latent_codes_pool.py new file mode 100644 index 0000000000000000000000000000000000000000..0281d4b5e80f8eb26e824fa35b4f908dcb6634e6 --- /dev/null +++ b/models/encoder4editing/models/latent_codes_pool.py @@ -0,0 +1,55 @@ +import random +import torch + + +class LatentCodesPool: + """This class implements latent codes buffer that stores previously generated w latent codes. + This buffer enables us to update discriminators using a history of generated w's + rather than the ones produced by the latest encoder. + """ + + def __init__(self, pool_size): + """Initialize the ImagePool class + Parameters: + pool_size (int) -- the size of image buffer, if pool_size=0, no buffer will be created + """ + self.pool_size = pool_size + if self.pool_size > 0: # create an empty pool + self.num_ws = 0 + self.ws = [] + + def query(self, ws): + """Return w's from the pool. + Parameters: + ws: the latest generated w's from the generator + Returns w's from the buffer. + By 50/100, the buffer will return input w's. + By 50/100, the buffer will return w's previously stored in the buffer, + and insert the current w's to the buffer. + """ + if self.pool_size == 0: # if the buffer size is 0, do nothing + return ws + return_ws = [] + for w in ws: # ws.shape: (batch, 512) or (batch, n_latent, 512) + # w = torch.unsqueeze(image.data, 0) + if w.ndim == 2: + i = random.randint(0, len(w) - 1) # apply a random latent index as a candidate + w = w[i] + self.handle_w(w, return_ws) + return_ws = torch.stack(return_ws, 0) # collect all the images and return + return return_ws + + def handle_w(self, w, return_ws): + if self.num_ws < self.pool_size: # if the buffer is not full; keep inserting current codes to the buffer + self.num_ws = self.num_ws + 1 + self.ws.append(w) + return_ws.append(w) + else: + p = random.uniform(0, 1) + if p > 0.5: # by 50% chance, the buffer will return a previously stored latent code, and insert the current code into the buffer + random_id = random.randint(0, self.pool_size - 1) # randint is inclusive + tmp = self.ws[random_id].clone() + self.ws[random_id] = w + return_ws.append(tmp) + else: # by another 50% chance, the buffer will return the current image + return_ws.append(w) diff --git a/models/encoder4editing/models/psp.py b/models/encoder4editing/models/psp.py new file mode 100644 index 0000000000000000000000000000000000000000..40b45c72022312fbf85076c2fb81f1d758729798 --- /dev/null +++ b/models/encoder4editing/models/psp.py @@ -0,0 +1,104 @@ +import matplotlib + +matplotlib.use('Agg') +import torch +from torch import nn +from models.encoder4editing.models.encoders import psp_encoders +from models.encoder4editing.models.stylegan2.model import Generator +from models.encoder4editing.configs.paths_config import model_paths + + +def get_keys(d, name): + if 'state_dict' in d: + d = d['state_dict'] + d_filt = {k[len(name) + 1:]: v for k, v in d.items() if k[:len(name)] == name} + return d_filt + + +class pSp(nn.Module): + + def __init__(self, opts): + super(pSp, self).__init__() + self.opts = opts + # Define architecture + self.encoder = self.set_encoder() + self.decoder = Generator(opts.stylegan_size, 512, 8, channel_multiplier=2) + self.face_pool = torch.nn.AdaptiveAvgPool2d((256, 256)) + # Load weights if needed + self.load_weights() + + def set_encoder(self): + if self.opts.encoder_type == 'GradualStyleEncoder': + encoder = psp_encoders.GradualStyleEncoder(50, 'ir_se', self.opts) + elif self.opts.encoder_type == 'Encoder4Editing': + encoder = psp_encoders.Encoder4Editing(50, 'ir_se', self.opts) + elif self.opts.encoder_type == 'SingleStyleCodeEncoder': + encoder = psp_encoders.BackboneEncoderUsingLastLayerIntoW(50, 'ir_se', self.opts) + else: + raise Exception('{} is not a valid encoders'.format(self.opts.encoder_type)) + return encoder + + def load_weights(self): + if self.opts.checkpoint_path is not None: + print('Loading e4e over the pSp framework from checkpoint: {}'.format(self.opts.checkpoint_path)) + ckpt = torch.load(self.opts.checkpoint_path, map_location='cpu') + self.encoder.load_state_dict(get_keys(ckpt, 'encoder'), strict=True) + self.decoder.load_state_dict(get_keys(ckpt, 'decoder'), strict=True) + self.__load_latent_avg(ckpt) + else: + print('Loading encoders weights from irse50!') + encoder_ckpt = torch.load(model_paths['ir_se50']) + self.encoder.load_state_dict(encoder_ckpt, strict=False) + print('Loading decoder weights from pretrained!') + ckpt = torch.load(self.opts.stylegan_weights) + self.decoder.load_state_dict(ckpt['g_ema'], strict=False) + self.__load_latent_avg(ckpt, repeat=self.encoder.style_count) + + def forward(self, x, resize=True, latent_mask=None, input_code=False, randomize_noise=True, + inject_latent=None, return_latents=False, alpha=None): + if input_code: + codes = x + else: + codes = self.encoder(x) + # normalize with respect to the center of an average face + if self.opts.start_from_latent_avg: + if codes.ndim == 2: + codes = codes + self.latent_avg.repeat(codes.shape[0], 1, 1)[:, 0, :] + else: + codes = codes + self.latent_avg.repeat(codes.shape[0], 1, 1) + + if latent_mask is not None: + for i in latent_mask: + if inject_latent is not None: + if alpha is not None: + codes[:, i] = alpha * inject_latent[:, i] + (1 - alpha) * codes[:, i] + else: + codes[:, i] = inject_latent[:, i] + else: + codes[:, i] = 0 + + input_is_latent = not input_code + images, result_latent = self.decoder([codes], + input_is_latent=input_is_latent, + randomize_noise=randomize_noise, + return_latents=return_latents) + + if resize: + images = self.face_pool(images) + + if return_latents: + return images, result_latent + else: + return images + + def __load_latent_avg(self, ckpt, repeat=None): + if 'latent_avg' in ckpt: + self.latent_avg = ckpt['latent_avg'].to(self.opts.device) + elif self.opts.start_from_latent_avg: + # Compute mean code based on a large number of latents (10,000 here) + with torch.no_grad(): + self.latent_avg = self.decoder.mean_latent(10000).to(self.opts.device) + else: + self.latent_avg = None + if repeat is not None and self.latent_avg is not None: + self.latent_avg = self.latent_avg.repeat(repeat, 1) diff --git a/models/encoder4editing/models/stylegan2/__init__.py b/models/encoder4editing/models/stylegan2/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/models/encoder4editing/models/stylegan2/model.py b/models/encoder4editing/models/stylegan2/model.py new file mode 100644 index 0000000000000000000000000000000000000000..54870486c6ef5a0d34e8e63b94ba5e3ac6e68944 --- /dev/null +++ b/models/encoder4editing/models/stylegan2/model.py @@ -0,0 +1,673 @@ +import math +import random +import torch +from torch import nn +from torch.nn import functional as F + +from models.stylegan2.op import FusedLeakyReLU, fused_leaky_relu, upfirdn2d + + +class PixelNorm(nn.Module): + def __init__(self): + super().__init__() + + def forward(self, input): + return input * torch.rsqrt(torch.mean(input ** 2, dim=1, keepdim=True) + 1e-8) + + +def make_kernel(k): + k = torch.tensor(k, dtype=torch.float32) + + if k.ndim == 1: + k = k[None, :] * k[:, None] + + k /= k.sum() + + return k + + +class Upsample(nn.Module): + def __init__(self, kernel, factor=2): + super().__init__() + + self.factor = factor + kernel = make_kernel(kernel) * (factor ** 2) + self.register_buffer('kernel', kernel) + + p = kernel.shape[0] - factor + + pad0 = (p + 1) // 2 + factor - 1 + pad1 = p // 2 + + self.pad = (pad0, pad1) + + def forward(self, input): + out = upfirdn2d(input, self.kernel, up=self.factor, down=1, pad=self.pad) + + return out + + +class Downsample(nn.Module): + def __init__(self, kernel, factor=2): + super().__init__() + + self.factor = factor + kernel = make_kernel(kernel) + self.register_buffer('kernel', kernel) + + p = kernel.shape[0] - factor + + pad0 = (p + 1) // 2 + pad1 = p // 2 + + self.pad = (pad0, pad1) + + def forward(self, input): + out = upfirdn2d(input, self.kernel, up=1, down=self.factor, pad=self.pad) + + return out + + +class Blur(nn.Module): + def __init__(self, kernel, pad, upsample_factor=1): + super().__init__() + + kernel = make_kernel(kernel) + + if upsample_factor > 1: + kernel = kernel * (upsample_factor ** 2) + + self.register_buffer('kernel', kernel) + + self.pad = pad + + def forward(self, input): + out = upfirdn2d(input, self.kernel, pad=self.pad) + + return out + + +class EqualConv2d(nn.Module): + def __init__( + self, in_channel, out_channel, kernel_size, stride=1, padding=0, bias=True + ): + super().__init__() + + self.weight = nn.Parameter( + torch.randn(out_channel, in_channel, kernel_size, kernel_size) + ) + self.scale = 1 / math.sqrt(in_channel * kernel_size ** 2) + + self.stride = stride + self.padding = padding + + if bias: + self.bias = nn.Parameter(torch.zeros(out_channel)) + + else: + self.bias = None + + def forward(self, input): + out = F.conv2d( + input, + self.weight * self.scale, + bias=self.bias, + stride=self.stride, + padding=self.padding, + ) + + return out + + def __repr__(self): + return ( + f'{self.__class__.__name__}({self.weight.shape[1]}, {self.weight.shape[0]},' + f' {self.weight.shape[2]}, stride={self.stride}, padding={self.padding})' + ) + + +class EqualLinear(nn.Module): + def __init__( + self, in_dim, out_dim, bias=True, bias_init=0, lr_mul=1, activation=None + ): + super().__init__() + + self.weight = nn.Parameter(torch.randn(out_dim, in_dim).div_(lr_mul)) + + if bias: + self.bias = nn.Parameter(torch.zeros(out_dim).fill_(bias_init)) + + else: + self.bias = None + + self.activation = activation + + self.scale = (1 / math.sqrt(in_dim)) * lr_mul + self.lr_mul = lr_mul + + def forward(self, input): + if self.activation: + out = F.linear(input, self.weight * self.scale) + out = fused_leaky_relu(out, self.bias * self.lr_mul) + + else: + out = F.linear( + input, self.weight * self.scale, bias=self.bias * self.lr_mul + ) + + return out + + def __repr__(self): + return ( + f'{self.__class__.__name__}({self.weight.shape[1]}, {self.weight.shape[0]})' + ) + + +class ScaledLeakyReLU(nn.Module): + def __init__(self, negative_slope=0.2): + super().__init__() + + self.negative_slope = negative_slope + + def forward(self, input): + out = F.leaky_relu(input, negative_slope=self.negative_slope) + + return out * math.sqrt(2) + + +class ModulatedConv2d(nn.Module): + def __init__( + self, + in_channel, + out_channel, + kernel_size, + style_dim, + demodulate=True, + upsample=False, + downsample=False, + blur_kernel=[1, 3, 3, 1], + ): + super().__init__() + + self.eps = 1e-8 + self.kernel_size = kernel_size + self.in_channel = in_channel + self.out_channel = out_channel + self.upsample = upsample + self.downsample = downsample + + if upsample: + factor = 2 + p = (len(blur_kernel) - factor) - (kernel_size - 1) + pad0 = (p + 1) // 2 + factor - 1 + pad1 = p // 2 + 1 + + self.blur = Blur(blur_kernel, pad=(pad0, pad1), upsample_factor=factor) + + if downsample: + factor = 2 + p = (len(blur_kernel) - factor) + (kernel_size - 1) + pad0 = (p + 1) // 2 + pad1 = p // 2 + + self.blur = Blur(blur_kernel, pad=(pad0, pad1)) + + fan_in = in_channel * kernel_size ** 2 + self.scale = 1 / math.sqrt(fan_in) + self.padding = kernel_size // 2 + + self.weight = nn.Parameter( + torch.randn(1, out_channel, in_channel, kernel_size, kernel_size) + ) + + self.modulation = EqualLinear(style_dim, in_channel, bias_init=1) + + self.demodulate = demodulate + + def __repr__(self): + return ( + f'{self.__class__.__name__}({self.in_channel}, {self.out_channel}, {self.kernel_size}, ' + f'upsample={self.upsample}, downsample={self.downsample})' + ) + + def forward(self, input, style): + batch, in_channel, height, width = input.shape + + style = self.modulation(style).view(batch, 1, in_channel, 1, 1) + weight = self.scale * self.weight * style + + if self.demodulate: + demod = torch.rsqrt(weight.pow(2).sum([2, 3, 4]) + 1e-8) + weight = weight * demod.view(batch, self.out_channel, 1, 1, 1) + + weight = weight.view( + batch * self.out_channel, in_channel, self.kernel_size, self.kernel_size + ) + + if self.upsample: + input = input.view(1, batch * in_channel, height, width) + weight = weight.view( + batch, self.out_channel, in_channel, self.kernel_size, self.kernel_size + ) + weight = weight.transpose(1, 2).reshape( + batch * in_channel, self.out_channel, self.kernel_size, self.kernel_size + ) + out = F.conv_transpose2d(input, weight, padding=0, stride=2, groups=batch) + _, _, height, width = out.shape + out = out.view(batch, self.out_channel, height, width) + out = self.blur(out) + + elif self.downsample: + input = self.blur(input) + _, _, height, width = input.shape + input = input.view(1, batch * in_channel, height, width) + out = F.conv2d(input, weight, padding=0, stride=2, groups=batch) + _, _, height, width = out.shape + out = out.view(batch, self.out_channel, height, width) + + else: + input = input.view(1, batch * in_channel, height, width) + out = F.conv2d(input, weight, padding=self.padding, groups=batch) + _, _, height, width = out.shape + out = out.view(batch, self.out_channel, height, width) + + return out + + +class NoiseInjection(nn.Module): + def __init__(self): + super().__init__() + + self.weight = nn.Parameter(torch.zeros(1)) + + def forward(self, image, noise=None): + if noise is None: + batch, _, height, width = image.shape + noise = image.new_empty(batch, 1, height, width).normal_() + + return image + self.weight * noise + + +class ConstantInput(nn.Module): + def __init__(self, channel, size=4): + super().__init__() + + self.input = nn.Parameter(torch.randn(1, channel, size, size)) + + def forward(self, input): + batch = input.shape[0] + out = self.input.repeat(batch, 1, 1, 1) + + return out + + +class StyledConv(nn.Module): + def __init__( + self, + in_channel, + out_channel, + kernel_size, + style_dim, + upsample=False, + blur_kernel=[1, 3, 3, 1], + demodulate=True, + ): + super().__init__() + + self.conv = ModulatedConv2d( + in_channel, + out_channel, + kernel_size, + style_dim, + upsample=upsample, + blur_kernel=blur_kernel, + demodulate=demodulate, + ) + + self.noise = NoiseInjection() + # self.bias = nn.Parameter(torch.zeros(1, out_channel, 1, 1)) + # self.activate = ScaledLeakyReLU(0.2) + self.activate = FusedLeakyReLU(out_channel) + + def forward(self, input, style, noise=None): + out = self.conv(input, style) + out = self.noise(out, noise=noise) + # out = out + self.bias + out = self.activate(out) + + return out + + +class ToRGB(nn.Module): + def __init__(self, in_channel, style_dim, upsample=True, blur_kernel=[1, 3, 3, 1]): + super().__init__() + + if upsample: + self.upsample = Upsample(blur_kernel) + + self.conv = ModulatedConv2d(in_channel, 3, 1, style_dim, demodulate=False) + self.bias = nn.Parameter(torch.zeros(1, 3, 1, 1)) + + def forward(self, input, style, skip=None): + out = self.conv(input, style) + out = out + self.bias + + if skip is not None: + skip = self.upsample(skip) + + out = out + skip + + return out + + +class Generator(nn.Module): + def __init__( + self, + size, + style_dim, + n_mlp, + channel_multiplier=2, + blur_kernel=[1, 3, 3, 1], + lr_mlp=0.01, + ): + super().__init__() + + self.size = size + + self.style_dim = style_dim + + layers = [PixelNorm()] + + for i in range(n_mlp): + layers.append( + EqualLinear( + style_dim, style_dim, lr_mul=lr_mlp, activation='fused_lrelu' + ) + ) + + self.style = nn.Sequential(*layers) + + self.channels = { + 4: 512, + 8: 512, + 16: 512, + 32: 512, + 64: 256 * channel_multiplier, + 128: 128 * channel_multiplier, + 256: 64 * channel_multiplier, + 512: 32 * channel_multiplier, + 1024: 16 * channel_multiplier, + } + + self.input = ConstantInput(self.channels[4]) + self.conv1 = StyledConv( + self.channels[4], self.channels[4], 3, style_dim, blur_kernel=blur_kernel + ) + self.to_rgb1 = ToRGB(self.channels[4], style_dim, upsample=False) + + self.log_size = int(math.log(size, 2)) + self.num_layers = (self.log_size - 2) * 2 + 1 + + self.convs = nn.ModuleList() + self.upsamples = nn.ModuleList() + self.to_rgbs = nn.ModuleList() + self.noises = nn.Module() + + in_channel = self.channels[4] + + for layer_idx in range(self.num_layers): + res = (layer_idx + 5) // 2 + shape = [1, 1, 2 ** res, 2 ** res] + self.noises.register_buffer(f'noise_{layer_idx}', torch.randn(*shape)) + + for i in range(3, self.log_size + 1): + out_channel = self.channels[2 ** i] + + self.convs.append( + StyledConv( + in_channel, + out_channel, + 3, + style_dim, + upsample=True, + blur_kernel=blur_kernel, + ) + ) + + self.convs.append( + StyledConv( + out_channel, out_channel, 3, style_dim, blur_kernel=blur_kernel + ) + ) + + self.to_rgbs.append(ToRGB(out_channel, style_dim)) + + in_channel = out_channel + + self.n_latent = self.log_size * 2 - 2 + + def make_noise(self): + device = self.input.input.device + + noises = [torch.randn(1, 1, 2 ** 2, 2 ** 2, device=device)] + + for i in range(3, self.log_size + 1): + for _ in range(2): + noises.append(torch.randn(1, 1, 2 ** i, 2 ** i, device=device)) + + return noises + + def mean_latent(self, n_latent): + latent_in = torch.randn( + n_latent, self.style_dim, device=self.input.input.device + ) + latent = self.style(latent_in).mean(0, keepdim=True) + + return latent + + def get_latent(self, input): + return self.style(input) + + def forward( + self, + styles, + return_latents=False, + return_features=False, + inject_index=None, + truncation=1, + truncation_latent=None, + input_is_latent=False, + noise=None, + randomize_noise=True, + ): + if not input_is_latent: + styles = [self.style(s) for s in styles] + + if noise is None: + if randomize_noise: + noise = [None] * self.num_layers + else: + noise = [ + getattr(self.noises, f'noise_{i}') for i in range(self.num_layers) + ] + + if truncation < 1: + style_t = [] + + for style in styles: + style_t.append( + truncation_latent + truncation * (style - truncation_latent) + ) + + styles = style_t + + if len(styles) < 2: + inject_index = self.n_latent + + if styles[0].ndim < 3: + latent = styles[0].unsqueeze(1).repeat(1, inject_index, 1) + else: + latent = styles[0] + + else: + if inject_index is None: + inject_index = random.randint(1, self.n_latent - 1) + + latent = styles[0].unsqueeze(1).repeat(1, inject_index, 1) + latent2 = styles[1].unsqueeze(1).repeat(1, self.n_latent - inject_index, 1) + + latent = torch.cat([latent, latent2], 1) + + out = self.input(latent) + out = self.conv1(out, latent[:, 0], noise=noise[0]) + + skip = self.to_rgb1(out, latent[:, 1]) + + i = 1 + for conv1, conv2, noise1, noise2, to_rgb in zip( + self.convs[::2], self.convs[1::2], noise[1::2], noise[2::2], self.to_rgbs + ): + out = conv1(out, latent[:, i], noise=noise1) + out = conv2(out, latent[:, i + 1], noise=noise2) + skip = to_rgb(out, latent[:, i + 2], skip) + + i += 2 + + image = skip + + if return_latents: + return image, latent + elif return_features: + return image, out + else: + return image, None + + +class ConvLayer(nn.Sequential): + def __init__( + self, + in_channel, + out_channel, + kernel_size, + downsample=False, + blur_kernel=[1, 3, 3, 1], + bias=True, + activate=True, + ): + layers = [] + + if downsample: + factor = 2 + p = (len(blur_kernel) - factor) + (kernel_size - 1) + pad0 = (p + 1) // 2 + pad1 = p // 2 + + layers.append(Blur(blur_kernel, pad=(pad0, pad1))) + + stride = 2 + self.padding = 0 + + else: + stride = 1 + self.padding = kernel_size // 2 + + layers.append( + EqualConv2d( + in_channel, + out_channel, + kernel_size, + padding=self.padding, + stride=stride, + bias=bias and not activate, + ) + ) + + if activate: + if bias: + layers.append(FusedLeakyReLU(out_channel)) + + else: + layers.append(ScaledLeakyReLU(0.2)) + + super().__init__(*layers) + + +class ResBlock(nn.Module): + def __init__(self, in_channel, out_channel, blur_kernel=[1, 3, 3, 1]): + super().__init__() + + self.conv1 = ConvLayer(in_channel, in_channel, 3) + self.conv2 = ConvLayer(in_channel, out_channel, 3, downsample=True) + + self.skip = ConvLayer( + in_channel, out_channel, 1, downsample=True, activate=False, bias=False + ) + + def forward(self, input): + out = self.conv1(input) + out = self.conv2(out) + + skip = self.skip(input) + out = (out + skip) / math.sqrt(2) + + return out + + +class Discriminator(nn.Module): + def __init__(self, size, channel_multiplier=2, blur_kernel=[1, 3, 3, 1]): + super().__init__() + + channels = { + 4: 512, + 8: 512, + 16: 512, + 32: 512, + 64: 256 * channel_multiplier, + 128: 128 * channel_multiplier, + 256: 64 * channel_multiplier, + 512: 32 * channel_multiplier, + 1024: 16 * channel_multiplier, + } + + convs = [ConvLayer(3, channels[size], 1)] + + log_size = int(math.log(size, 2)) + + in_channel = channels[size] + + for i in range(log_size, 2, -1): + out_channel = channels[2 ** (i - 1)] + + convs.append(ResBlock(in_channel, out_channel, blur_kernel)) + + in_channel = out_channel + + self.convs = nn.Sequential(*convs) + + self.stddev_group = 4 + self.stddev_feat = 1 + + self.final_conv = ConvLayer(in_channel + 1, channels[4], 3) + self.final_linear = nn.Sequential( + EqualLinear(channels[4] * 4 * 4, channels[4], activation='fused_lrelu'), + EqualLinear(channels[4], 1), + ) + + def forward(self, input): + out = self.convs(input) + + batch, channel, height, width = out.shape + group = min(batch, self.stddev_group) + stddev = out.view( + group, -1, self.stddev_feat, channel // self.stddev_feat, height, width + ) + stddev = torch.sqrt(stddev.var(0, unbiased=False) + 1e-8) + stddev = stddev.mean([2, 3, 4], keepdims=True).squeeze(2) + stddev = stddev.repeat(group, 1, height, width) + out = torch.cat([out, stddev], 1) + + out = self.final_conv(out) + + out = out.view(batch, -1) + out = self.final_linear(out) + + return out diff --git a/models/encoder4editing/models/stylegan2/op/__init__.py b/models/encoder4editing/models/stylegan2/op/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d0918d92285955855be89f00096b888ee5597ce3 --- /dev/null +++ b/models/encoder4editing/models/stylegan2/op/__init__.py @@ -0,0 +1,2 @@ +from .fused_act import FusedLeakyReLU, fused_leaky_relu +from .upfirdn2d import upfirdn2d diff --git a/models/encoder4editing/models/stylegan2/op/fused_act.py b/models/encoder4editing/models/stylegan2/op/fused_act.py new file mode 100644 index 0000000000000000000000000000000000000000..973a84fffde53668d31397da5fb993bbc95f7be0 --- /dev/null +++ b/models/encoder4editing/models/stylegan2/op/fused_act.py @@ -0,0 +1,85 @@ +import os + +import torch +from torch import nn +from torch.autograd import Function +from torch.utils.cpp_extension import load + +module_path = os.path.dirname(__file__) +fused = load( + 'fused', + sources=[ + os.path.join(module_path, 'fused_bias_act.cpp'), + os.path.join(module_path, 'fused_bias_act_kernel.cu'), + ], +) + + +class FusedLeakyReLUFunctionBackward(Function): + @staticmethod + def forward(ctx, grad_output, out, negative_slope, scale): + ctx.save_for_backward(out) + ctx.negative_slope = negative_slope + ctx.scale = scale + + empty = grad_output.new_empty(0) + + grad_input = fused.fused_bias_act( + grad_output, empty, out, 3, 1, negative_slope, scale + ) + + dim = [0] + + if grad_input.ndim > 2: + dim += list(range(2, grad_input.ndim)) + + grad_bias = grad_input.sum(dim).detach() + + return grad_input, grad_bias + + @staticmethod + def backward(ctx, gradgrad_input, gradgrad_bias): + out, = ctx.saved_tensors + gradgrad_out = fused.fused_bias_act( + gradgrad_input, gradgrad_bias, out, 3, 1, ctx.negative_slope, ctx.scale + ) + + return gradgrad_out, None, None, None + + +class FusedLeakyReLUFunction(Function): + @staticmethod + def forward(ctx, input, bias, negative_slope, scale): + empty = input.new_empty(0) + out = fused.fused_bias_act(input, bias, empty, 3, 0, negative_slope, scale) + ctx.save_for_backward(out) + ctx.negative_slope = negative_slope + ctx.scale = scale + + return out + + @staticmethod + def backward(ctx, grad_output): + out, = ctx.saved_tensors + + grad_input, grad_bias = FusedLeakyReLUFunctionBackward.apply( + grad_output, out, ctx.negative_slope, ctx.scale + ) + + return grad_input, grad_bias, None, None + + +class FusedLeakyReLU(nn.Module): + def __init__(self, channel, negative_slope=0.2, scale=2 ** 0.5): + super().__init__() + + self.bias = nn.Parameter(torch.zeros(channel)) + self.negative_slope = negative_slope + self.scale = scale + + def forward(self, input): + return fused_leaky_relu(input, self.bias, self.negative_slope, self.scale) + + +def fused_leaky_relu(input, bias, negative_slope=0.2, scale=2 ** 0.5): + return FusedLeakyReLUFunction.apply(input, bias, negative_slope, scale) diff --git a/models/encoder4editing/models/stylegan2/op/fused_bias_act.cpp b/models/encoder4editing/models/stylegan2/op/fused_bias_act.cpp new file mode 100644 index 0000000000000000000000000000000000000000..02be898f970bcc8ea297867fcaa4e71b24b3d949 --- /dev/null +++ b/models/encoder4editing/models/stylegan2/op/fused_bias_act.cpp @@ -0,0 +1,21 @@ +#include + + +torch::Tensor fused_bias_act_op(const torch::Tensor& input, const torch::Tensor& bias, const torch::Tensor& refer, + int act, int grad, float alpha, float scale); + +#define CHECK_CUDA(x) TORCH_CHECK(x.type().is_cuda(), #x " must be a CUDA tensor") +#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous") +#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x) + +torch::Tensor fused_bias_act(const torch::Tensor& input, const torch::Tensor& bias, const torch::Tensor& refer, + int act, int grad, float alpha, float scale) { + CHECK_CUDA(input); + CHECK_CUDA(bias); + + return fused_bias_act_op(input, bias, refer, act, grad, alpha, scale); +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("fused_bias_act", &fused_bias_act, "fused bias act (CUDA)"); +} \ No newline at end of file diff --git a/models/encoder4editing/models/stylegan2/op/fused_bias_act_kernel.cu b/models/encoder4editing/models/stylegan2/op/fused_bias_act_kernel.cu new file mode 100644 index 0000000000000000000000000000000000000000..c9fa56fea7ede7072dc8925cfb0148f136eb85b8 --- /dev/null +++ b/models/encoder4editing/models/stylegan2/op/fused_bias_act_kernel.cu @@ -0,0 +1,99 @@ +// Copyright (c) 2019, NVIDIA Corporation. All rights reserved. +// +// This work is made available under the Nvidia Source Code License-NC. +// To view a copy of this license, visit +// https://nvlabs.github.io/stylegan2/license.html + +#include + +#include +#include +#include +#include + +#include +#include + + +template +static __global__ void fused_bias_act_kernel(scalar_t* out, const scalar_t* p_x, const scalar_t* p_b, const scalar_t* p_ref, + int act, int grad, scalar_t alpha, scalar_t scale, int loop_x, int size_x, int step_b, int size_b, int use_bias, int use_ref) { + int xi = blockIdx.x * loop_x * blockDim.x + threadIdx.x; + + scalar_t zero = 0.0; + + for (int loop_idx = 0; loop_idx < loop_x && xi < size_x; loop_idx++, xi += blockDim.x) { + scalar_t x = p_x[xi]; + + if (use_bias) { + x += p_b[(xi / step_b) % size_b]; + } + + scalar_t ref = use_ref ? p_ref[xi] : zero; + + scalar_t y; + + switch (act * 10 + grad) { + default: + case 10: y = x; break; + case 11: y = x; break; + case 12: y = 0.0; break; + + case 30: y = (x > 0.0) ? x : x * alpha; break; + case 31: y = (ref > 0.0) ? x : x * alpha; break; + case 32: y = 0.0; break; + } + + out[xi] = y * scale; + } +} + + +torch::Tensor fused_bias_act_op(const torch::Tensor& input, const torch::Tensor& bias, const torch::Tensor& refer, + int act, int grad, float alpha, float scale) { + int curDevice = -1; + cudaGetDevice(&curDevice); + cudaStream_t stream = at::cuda::getCurrentCUDAStream(curDevice); + + auto x = input.contiguous(); + auto b = bias.contiguous(); + auto ref = refer.contiguous(); + + int use_bias = b.numel() ? 1 : 0; + int use_ref = ref.numel() ? 1 : 0; + + int size_x = x.numel(); + int size_b = b.numel(); + int step_b = 1; + + for (int i = 1 + 1; i < x.dim(); i++) { + step_b *= x.size(i); + } + + int loop_x = 4; + int block_size = 4 * 32; + int grid_size = (size_x - 1) / (loop_x * block_size) + 1; + + auto y = torch::empty_like(x); + + AT_DISPATCH_FLOATING_TYPES_AND_HALF(x.scalar_type(), "fused_bias_act_kernel", [&] { + fused_bias_act_kernel<<>>( + y.data_ptr(), + x.data_ptr(), + b.data_ptr(), + ref.data_ptr(), + act, + grad, + alpha, + scale, + loop_x, + size_x, + step_b, + size_b, + use_bias, + use_ref + ); + }); + + return y; +} \ No newline at end of file diff --git a/models/encoder4editing/models/stylegan2/op/upfirdn2d.cpp b/models/encoder4editing/models/stylegan2/op/upfirdn2d.cpp new file mode 100644 index 0000000000000000000000000000000000000000..d2e633dc896433c205e18bc3e455539192ff968e --- /dev/null +++ b/models/encoder4editing/models/stylegan2/op/upfirdn2d.cpp @@ -0,0 +1,23 @@ +#include + + +torch::Tensor upfirdn2d_op(const torch::Tensor& input, const torch::Tensor& kernel, + int up_x, int up_y, int down_x, int down_y, + int pad_x0, int pad_x1, int pad_y0, int pad_y1); + +#define CHECK_CUDA(x) TORCH_CHECK(x.type().is_cuda(), #x " must be a CUDA tensor") +#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous") +#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x) + +torch::Tensor upfirdn2d(const torch::Tensor& input, const torch::Tensor& kernel, + int up_x, int up_y, int down_x, int down_y, + int pad_x0, int pad_x1, int pad_y0, int pad_y1) { + CHECK_CUDA(input); + CHECK_CUDA(kernel); + + return upfirdn2d_op(input, kernel, up_x, up_y, down_x, down_y, pad_x0, pad_x1, pad_y0, pad_y1); +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("upfirdn2d", &upfirdn2d, "upfirdn2d (CUDA)"); +} \ No newline at end of file diff --git a/models/encoder4editing/models/stylegan2/op/upfirdn2d.py b/models/encoder4editing/models/stylegan2/op/upfirdn2d.py new file mode 100644 index 0000000000000000000000000000000000000000..7bc5a1e331c2bbb1893ac748cfd0f144ff0651b4 --- /dev/null +++ b/models/encoder4editing/models/stylegan2/op/upfirdn2d.py @@ -0,0 +1,184 @@ +import os + +import torch +from torch.autograd import Function +from torch.utils.cpp_extension import load + +module_path = os.path.dirname(__file__) +upfirdn2d_op = load( + 'upfirdn2d', + sources=[ + os.path.join(module_path, 'upfirdn2d.cpp'), + os.path.join(module_path, 'upfirdn2d_kernel.cu'), + ], +) + + +class UpFirDn2dBackward(Function): + @staticmethod + def forward( + ctx, grad_output, kernel, grad_kernel, up, down, pad, g_pad, in_size, out_size + ): + up_x, up_y = up + down_x, down_y = down + g_pad_x0, g_pad_x1, g_pad_y0, g_pad_y1 = g_pad + + grad_output = grad_output.reshape(-1, out_size[0], out_size[1], 1) + + grad_input = upfirdn2d_op.upfirdn2d( + grad_output, + grad_kernel, + down_x, + down_y, + up_x, + up_y, + g_pad_x0, + g_pad_x1, + g_pad_y0, + g_pad_y1, + ) + grad_input = grad_input.view(in_size[0], in_size[1], in_size[2], in_size[3]) + + ctx.save_for_backward(kernel) + + pad_x0, pad_x1, pad_y0, pad_y1 = pad + + ctx.up_x = up_x + ctx.up_y = up_y + ctx.down_x = down_x + ctx.down_y = down_y + ctx.pad_x0 = pad_x0 + ctx.pad_x1 = pad_x1 + ctx.pad_y0 = pad_y0 + ctx.pad_y1 = pad_y1 + ctx.in_size = in_size + ctx.out_size = out_size + + return grad_input + + @staticmethod + def backward(ctx, gradgrad_input): + kernel, = ctx.saved_tensors + + gradgrad_input = gradgrad_input.reshape(-1, ctx.in_size[2], ctx.in_size[3], 1) + + gradgrad_out = upfirdn2d_op.upfirdn2d( + gradgrad_input, + kernel, + ctx.up_x, + ctx.up_y, + ctx.down_x, + ctx.down_y, + ctx.pad_x0, + ctx.pad_x1, + ctx.pad_y0, + ctx.pad_y1, + ) + # gradgrad_out = gradgrad_out.view(ctx.in_size[0], ctx.out_size[0], ctx.out_size[1], ctx.in_size[3]) + gradgrad_out = gradgrad_out.view( + ctx.in_size[0], ctx.in_size[1], ctx.out_size[0], ctx.out_size[1] + ) + + return gradgrad_out, None, None, None, None, None, None, None, None + + +class UpFirDn2d(Function): + @staticmethod + def forward(ctx, input, kernel, up, down, pad): + up_x, up_y = up + down_x, down_y = down + pad_x0, pad_x1, pad_y0, pad_y1 = pad + + kernel_h, kernel_w = kernel.shape + batch, channel, in_h, in_w = input.shape + ctx.in_size = input.shape + + input = input.reshape(-1, in_h, in_w, 1) + + ctx.save_for_backward(kernel, torch.flip(kernel, [0, 1])) + + out_h = (in_h * up_y + pad_y0 + pad_y1 - kernel_h) // down_y + 1 + out_w = (in_w * up_x + pad_x0 + pad_x1 - kernel_w) // down_x + 1 + ctx.out_size = (out_h, out_w) + + ctx.up = (up_x, up_y) + ctx.down = (down_x, down_y) + ctx.pad = (pad_x0, pad_x1, pad_y0, pad_y1) + + g_pad_x0 = kernel_w - pad_x0 - 1 + g_pad_y0 = kernel_h - pad_y0 - 1 + g_pad_x1 = in_w * up_x - out_w * down_x + pad_x0 - up_x + 1 + g_pad_y1 = in_h * up_y - out_h * down_y + pad_y0 - up_y + 1 + + ctx.g_pad = (g_pad_x0, g_pad_x1, g_pad_y0, g_pad_y1) + + out = upfirdn2d_op.upfirdn2d( + input, kernel, up_x, up_y, down_x, down_y, pad_x0, pad_x1, pad_y0, pad_y1 + ) + # out = out.view(major, out_h, out_w, minor) + out = out.view(-1, channel, out_h, out_w) + + return out + + @staticmethod + def backward(ctx, grad_output): + kernel, grad_kernel = ctx.saved_tensors + + grad_input = UpFirDn2dBackward.apply( + grad_output, + kernel, + grad_kernel, + ctx.up, + ctx.down, + ctx.pad, + ctx.g_pad, + ctx.in_size, + ctx.out_size, + ) + + return grad_input, None, None, None, None + + +def upfirdn2d(input, kernel, up=1, down=1, pad=(0, 0)): + out = UpFirDn2d.apply( + input, kernel, (up, up), (down, down), (pad[0], pad[1], pad[0], pad[1]) + ) + + return out + + +def upfirdn2d_native( + input, kernel, up_x, up_y, down_x, down_y, pad_x0, pad_x1, pad_y0, pad_y1 +): + _, in_h, in_w, minor = input.shape + kernel_h, kernel_w = kernel.shape + + out = input.view(-1, in_h, 1, in_w, 1, minor) + out = F.pad(out, [0, 0, 0, up_x - 1, 0, 0, 0, up_y - 1]) + out = out.view(-1, in_h * up_y, in_w * up_x, minor) + + out = F.pad( + out, [0, 0, max(pad_x0, 0), max(pad_x1, 0), max(pad_y0, 0), max(pad_y1, 0)] + ) + out = out[ + :, + max(-pad_y0, 0): out.shape[1] - max(-pad_y1, 0), + max(-pad_x0, 0): out.shape[2] - max(-pad_x1, 0), + :, + ] + + out = out.permute(0, 3, 1, 2) + out = out.reshape( + [-1, 1, in_h * up_y + pad_y0 + pad_y1, in_w * up_x + pad_x0 + pad_x1] + ) + w = torch.flip(kernel, [0, 1]).view(1, 1, kernel_h, kernel_w) + out = F.conv2d(out, w) + out = out.reshape( + -1, + minor, + in_h * up_y + pad_y0 + pad_y1 - kernel_h + 1, + in_w * up_x + pad_x0 + pad_x1 - kernel_w + 1, + ) + out = out.permute(0, 2, 3, 1) + + return out[:, ::down_y, ::down_x, :] diff --git a/models/encoder4editing/models/stylegan2/op/upfirdn2d_kernel.cu b/models/encoder4editing/models/stylegan2/op/upfirdn2d_kernel.cu new file mode 100644 index 0000000000000000000000000000000000000000..2a710aa6adc3d43ac93136a1814e3c39970e1c7e --- /dev/null +++ b/models/encoder4editing/models/stylegan2/op/upfirdn2d_kernel.cu @@ -0,0 +1,272 @@ +// Copyright (c) 2019, NVIDIA Corporation. All rights reserved. +// +// This work is made available under the Nvidia Source Code License-NC. +// To view a copy of this license, visit +// https://nvlabs.github.io/stylegan2/license.html + +#include + +#include +#include +#include +#include + +#include +#include + + +static __host__ __device__ __forceinline__ int floor_div(int a, int b) { + int c = a / b; + + if (c * b > a) { + c--; + } + + return c; +} + + +struct UpFirDn2DKernelParams { + int up_x; + int up_y; + int down_x; + int down_y; + int pad_x0; + int pad_x1; + int pad_y0; + int pad_y1; + + int major_dim; + int in_h; + int in_w; + int minor_dim; + int kernel_h; + int kernel_w; + int out_h; + int out_w; + int loop_major; + int loop_x; +}; + + +template +__global__ void upfirdn2d_kernel(scalar_t* out, const scalar_t* input, const scalar_t* kernel, const UpFirDn2DKernelParams p) { + const int tile_in_h = ((tile_out_h - 1) * down_y + kernel_h - 1) / up_y + 1; + const int tile_in_w = ((tile_out_w - 1) * down_x + kernel_w - 1) / up_x + 1; + + __shared__ volatile float sk[kernel_h][kernel_w]; + __shared__ volatile float sx[tile_in_h][tile_in_w]; + + int minor_idx = blockIdx.x; + int tile_out_y = minor_idx / p.minor_dim; + minor_idx -= tile_out_y * p.minor_dim; + tile_out_y *= tile_out_h; + int tile_out_x_base = blockIdx.y * p.loop_x * tile_out_w; + int major_idx_base = blockIdx.z * p.loop_major; + + if (tile_out_x_base >= p.out_w | tile_out_y >= p.out_h | major_idx_base >= p.major_dim) { + return; + } + + for (int tap_idx = threadIdx.x; tap_idx < kernel_h * kernel_w; tap_idx += blockDim.x) { + int ky = tap_idx / kernel_w; + int kx = tap_idx - ky * kernel_w; + scalar_t v = 0.0; + + if (kx < p.kernel_w & ky < p.kernel_h) { + v = kernel[(p.kernel_h - 1 - ky) * p.kernel_w + (p.kernel_w - 1 - kx)]; + } + + sk[ky][kx] = v; + } + + for (int loop_major = 0, major_idx = major_idx_base; loop_major < p.loop_major & major_idx < p.major_dim; loop_major++, major_idx++) { + for (int loop_x = 0, tile_out_x = tile_out_x_base; loop_x < p.loop_x & tile_out_x < p.out_w; loop_x++, tile_out_x += tile_out_w) { + int tile_mid_x = tile_out_x * down_x + up_x - 1 - p.pad_x0; + int tile_mid_y = tile_out_y * down_y + up_y - 1 - p.pad_y0; + int tile_in_x = floor_div(tile_mid_x, up_x); + int tile_in_y = floor_div(tile_mid_y, up_y); + + __syncthreads(); + + for (int in_idx = threadIdx.x; in_idx < tile_in_h * tile_in_w; in_idx += blockDim.x) { + int rel_in_y = in_idx / tile_in_w; + int rel_in_x = in_idx - rel_in_y * tile_in_w; + int in_x = rel_in_x + tile_in_x; + int in_y = rel_in_y + tile_in_y; + + scalar_t v = 0.0; + + if (in_x >= 0 & in_y >= 0 & in_x < p.in_w & in_y < p.in_h) { + v = input[((major_idx * p.in_h + in_y) * p.in_w + in_x) * p.minor_dim + minor_idx]; + } + + sx[rel_in_y][rel_in_x] = v; + } + + __syncthreads(); + for (int out_idx = threadIdx.x; out_idx < tile_out_h * tile_out_w; out_idx += blockDim.x) { + int rel_out_y = out_idx / tile_out_w; + int rel_out_x = out_idx - rel_out_y * tile_out_w; + int out_x = rel_out_x + tile_out_x; + int out_y = rel_out_y + tile_out_y; + + int mid_x = tile_mid_x + rel_out_x * down_x; + int mid_y = tile_mid_y + rel_out_y * down_y; + int in_x = floor_div(mid_x, up_x); + int in_y = floor_div(mid_y, up_y); + int rel_in_x = in_x - tile_in_x; + int rel_in_y = in_y - tile_in_y; + int kernel_x = (in_x + 1) * up_x - mid_x - 1; + int kernel_y = (in_y + 1) * up_y - mid_y - 1; + + scalar_t v = 0.0; + + #pragma unroll + for (int y = 0; y < kernel_h / up_y; y++) + #pragma unroll + for (int x = 0; x < kernel_w / up_x; x++) + v += sx[rel_in_y + y][rel_in_x + x] * sk[kernel_y + y * up_y][kernel_x + x * up_x]; + + if (out_x < p.out_w & out_y < p.out_h) { + out[((major_idx * p.out_h + out_y) * p.out_w + out_x) * p.minor_dim + minor_idx] = v; + } + } + } + } +} + + +torch::Tensor upfirdn2d_op(const torch::Tensor& input, const torch::Tensor& kernel, + int up_x, int up_y, int down_x, int down_y, + int pad_x0, int pad_x1, int pad_y0, int pad_y1) { + int curDevice = -1; + cudaGetDevice(&curDevice); + cudaStream_t stream = at::cuda::getCurrentCUDAStream(curDevice); + + UpFirDn2DKernelParams p; + + auto x = input.contiguous(); + auto k = kernel.contiguous(); + + p.major_dim = x.size(0); + p.in_h = x.size(1); + p.in_w = x.size(2); + p.minor_dim = x.size(3); + p.kernel_h = k.size(0); + p.kernel_w = k.size(1); + p.up_x = up_x; + p.up_y = up_y; + p.down_x = down_x; + p.down_y = down_y; + p.pad_x0 = pad_x0; + p.pad_x1 = pad_x1; + p.pad_y0 = pad_y0; + p.pad_y1 = pad_y1; + + p.out_h = (p.in_h * p.up_y + p.pad_y0 + p.pad_y1 - p.kernel_h + p.down_y) / p.down_y; + p.out_w = (p.in_w * p.up_x + p.pad_x0 + p.pad_x1 - p.kernel_w + p.down_x) / p.down_x; + + auto out = at::empty({p.major_dim, p.out_h, p.out_w, p.minor_dim}, x.options()); + + int mode = -1; + + int tile_out_h; + int tile_out_w; + + if (p.up_x == 1 && p.up_y == 1 && p.down_x == 1 && p.down_y == 1 && p.kernel_h <= 4 && p.kernel_w <= 4) { + mode = 1; + tile_out_h = 16; + tile_out_w = 64; + } + + if (p.up_x == 1 && p.up_y == 1 && p.down_x == 1 && p.down_y == 1 && p.kernel_h <= 3 && p.kernel_w <= 3) { + mode = 2; + tile_out_h = 16; + tile_out_w = 64; + } + + if (p.up_x == 2 && p.up_y == 2 && p.down_x == 1 && p.down_y == 1 && p.kernel_h <= 4 && p.kernel_w <= 4) { + mode = 3; + tile_out_h = 16; + tile_out_w = 64; + } + + if (p.up_x == 2 && p.up_y == 2 && p.down_x == 1 && p.down_y == 1 && p.kernel_h <= 2 && p.kernel_w <= 2) { + mode = 4; + tile_out_h = 16; + tile_out_w = 64; + } + + if (p.up_x == 1 && p.up_y == 1 && p.down_x == 2 && p.down_y == 2 && p.kernel_h <= 4 && p.kernel_w <= 4) { + mode = 5; + tile_out_h = 8; + tile_out_w = 32; + } + + if (p.up_x == 1 && p.up_y == 1 && p.down_x == 2 && p.down_y == 2 && p.kernel_h <= 2 && p.kernel_w <= 2) { + mode = 6; + tile_out_h = 8; + tile_out_w = 32; + } + + dim3 block_size; + dim3 grid_size; + + if (tile_out_h > 0 && tile_out_w) { + p.loop_major = (p.major_dim - 1) / 16384 + 1; + p.loop_x = 1; + block_size = dim3(32 * 8, 1, 1); + grid_size = dim3(((p.out_h - 1) / tile_out_h + 1) * p.minor_dim, + (p.out_w - 1) / (p.loop_x * tile_out_w) + 1, + (p.major_dim - 1) / p.loop_major + 1); + } + + AT_DISPATCH_FLOATING_TYPES_AND_HALF(x.scalar_type(), "upfirdn2d_cuda", [&] { + switch (mode) { + case 1: + upfirdn2d_kernel<<>>( + out.data_ptr(), x.data_ptr(), k.data_ptr(), p + ); + + break; + + case 2: + upfirdn2d_kernel<<>>( + out.data_ptr(), x.data_ptr(), k.data_ptr(), p + ); + + break; + + case 3: + upfirdn2d_kernel<<>>( + out.data_ptr(), x.data_ptr(), k.data_ptr(), p + ); + + break; + + case 4: + upfirdn2d_kernel<<>>( + out.data_ptr(), x.data_ptr(), k.data_ptr(), p + ); + + break; + + case 5: + upfirdn2d_kernel<<>>( + out.data_ptr(), x.data_ptr(), k.data_ptr(), p + ); + + break; + + case 6: + upfirdn2d_kernel<<>>( + out.data_ptr(), x.data_ptr(), k.data_ptr(), p + ); + + break; + } + }); + + return out; +} \ No newline at end of file diff --git a/models/encoder4editing/options/__init__.py b/models/encoder4editing/options/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/models/encoder4editing/options/train_options.py b/models/encoder4editing/options/train_options.py new file mode 100644 index 0000000000000000000000000000000000000000..33c015c76949d20aaf656498fb472bc0d33038b5 --- /dev/null +++ b/models/encoder4editing/options/train_options.py @@ -0,0 +1,84 @@ +from argparse import ArgumentParser +from configs.paths_config import model_paths + + +class TrainOptions: + + def __init__(self): + self.parser = ArgumentParser() + self.initialize() + + def initialize(self): + self.parser.add_argument('--exp_dir', type=str, help='Path to experiment output directory') + self.parser.add_argument('--dataset_type', default='ffhq_encode', type=str, + help='Type of dataset/experiment to run') + self.parser.add_argument('--encoder_type', default='Encoder4Editing', type=str, help='Which encoder to use') + + self.parser.add_argument('--batch_size', default=4, type=int, help='Batch size for training') + self.parser.add_argument('--test_batch_size', default=2, type=int, help='Batch size for testing and inference') + self.parser.add_argument('--workers', default=4, type=int, help='Number of scripts dataloader workers') + self.parser.add_argument('--test_workers', default=2, type=int, + help='Number of test/inference dataloader workers') + + self.parser.add_argument('--learning_rate', default=0.0001, type=float, help='Optimizer learning rate') + self.parser.add_argument('--optim_name', default='ranger', type=str, help='Which optimizer to use') + self.parser.add_argument('--train_decoder', default=False, type=bool, help='Whether to scripts the decoder model') + self.parser.add_argument('--start_from_latent_avg', action='store_true', + help='Whether to add average latent vector to generate codes from encoder.') + self.parser.add_argument('--lpips_type', default='alex', type=str, help='LPIPS backbone') + + self.parser.add_argument('--lpips_lambda', default=0.8, type=float, help='LPIPS loss multiplier factor') + self.parser.add_argument('--id_lambda', default=0.1, type=float, help='ID loss multiplier factor') + self.parser.add_argument('--l2_lambda', default=1.0, type=float, help='L2 loss multiplier factor') + + self.parser.add_argument('--stylegan_weights', default=model_paths['stylegan_ffhq'], type=str, + help='Path to StyleGAN model weights') + self.parser.add_argument('--stylegan_size', default=1024, type=int, + help='size of pretrained StyleGAN Generator') + self.parser.add_argument('--checkpoint_path', default=None, type=str, help='Path to pSp model checkpoint') + + self.parser.add_argument('--max_steps', default=500000, type=int, help='Maximum number of training steps') + self.parser.add_argument('--image_interval', default=100, type=int, + help='Interval for logging scripts images during training') + self.parser.add_argument('--board_interval', default=50, type=int, + help='Interval for logging metrics to tensorboard') + self.parser.add_argument('--val_interval', default=1000, type=int, help='Validation interval') + self.parser.add_argument('--save_interval', default=None, type=int, help='Model checkpoint interval') + + # Discriminator flags + self.parser.add_argument('--w_discriminator_lambda', default=0, type=float, help='Dw loss multiplier') + self.parser.add_argument('--w_discriminator_lr', default=2e-5, type=float, help='Dw learning rate') + self.parser.add_argument("--r1", type=float, default=10, help="weight of the r1 regularization") + self.parser.add_argument("--d_reg_every", type=int, default=16, + help="interval for applying r1 regularization") + self.parser.add_argument('--use_w_pool', action='store_true', + help='Whether to store a latnet codes pool for the discriminator\'s training') + self.parser.add_argument("--w_pool_size", type=int, default=50, + help="W\'s pool size, depends on --use_w_pool") + + # e4e specific + self.parser.add_argument('--delta_norm', type=int, default=2, help="norm type of the deltas") + self.parser.add_argument('--delta_norm_lambda', type=float, default=2e-4, help="lambda for delta norm loss") + + # Progressive training + self.parser.add_argument('--progressive_steps', nargs='+', type=int, default=None, + help="The training steps of training new deltas. steps[i] starts the delta_i training") + self.parser.add_argument('--progressive_start', type=int, default=None, + help="The training step to start training the deltas, overrides progressive_steps") + self.parser.add_argument('--progressive_step_every', type=int, default=2_000, + help="Amount of training steps for each progressive step") + + # Save additional training info to enable future training continuation from produced checkpoints + self.parser.add_argument('--save_training_data', action='store_true', + help='Save intermediate training data to resume training from the checkpoint') + self.parser.add_argument('--sub_exp_dir', default=None, type=str, help='Name of sub experiment directory') + self.parser.add_argument('--keep_optimizer', action='store_true', + help='Whether to continue from the checkpoint\'s optimizer') + self.parser.add_argument('--resume_training_from_ckpt', default=None, type=str, + help='Path to training checkpoint, works when --save_training_data was set to True') + self.parser.add_argument('--update_param_list', nargs='+', type=str, default=None, + help="Name of training parameters to update the loaded training checkpoint") + + def parse(self): + opts = self.parser.parse_args() + return opts diff --git a/models/encoder4editing/scripts/calc_losses_on_images.py b/models/encoder4editing/scripts/calc_losses_on_images.py new file mode 100644 index 0000000000000000000000000000000000000000..32b6bcee854da7ae357daf82bd986f30db9fb72c --- /dev/null +++ b/models/encoder4editing/scripts/calc_losses_on_images.py @@ -0,0 +1,87 @@ +from argparse import ArgumentParser +import os +import json +import sys +from tqdm import tqdm +import numpy as np +import torch +from torch.utils.data import DataLoader +import torchvision.transforms as transforms + +sys.path.append(".") +sys.path.append("..") + +from criteria.lpips.lpips import LPIPS +from datasets.gt_res_dataset import GTResDataset + + +def parse_args(): + parser = ArgumentParser(add_help=False) + parser.add_argument('--mode', type=str, default='lpips', choices=['lpips', 'l2']) + parser.add_argument('--data_path', type=str, default='results') + parser.add_argument('--gt_path', type=str, default='gt_images') + parser.add_argument('--workers', type=int, default=4) + parser.add_argument('--batch_size', type=int, default=4) + parser.add_argument('--is_cars', action='store_true') + args = parser.parse_args() + return args + + +def run(args): + resize_dims = (256, 256) + if args.is_cars: + resize_dims = (192, 256) + transform = transforms.Compose([transforms.Resize(resize_dims), + transforms.ToTensor(), + transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])]) + + print('Loading dataset') + dataset = GTResDataset(root_path=args.data_path, + gt_dir=args.gt_path, + transform=transform) + + dataloader = DataLoader(dataset, + batch_size=args.batch_size, + shuffle=False, + num_workers=int(args.workers), + drop_last=True) + + if args.mode == 'lpips': + loss_func = LPIPS(net_type='alex') + elif args.mode == 'l2': + loss_func = torch.nn.MSELoss() + else: + raise Exception('Not a valid mode!') + loss_func.cuda() + + global_i = 0 + scores_dict = {} + all_scores = [] + for result_batch, gt_batch in tqdm(dataloader): + for i in range(args.batch_size): + loss = float(loss_func(result_batch[i:i + 1].cuda(), gt_batch[i:i + 1].cuda())) + all_scores.append(loss) + im_path = dataset.pairs[global_i][0] + scores_dict[os.path.basename(im_path)] = loss + global_i += 1 + + all_scores = list(scores_dict.values()) + mean = np.mean(all_scores) + std = np.std(all_scores) + result_str = 'Average loss is {:.2f}+-{:.2f}'.format(mean, std) + print('Finished with ', args.data_path) + print(result_str) + + out_path = os.path.join(os.path.dirname(args.data_path), 'inference_metrics') + if not os.path.exists(out_path): + os.makedirs(out_path) + + with open(os.path.join(out_path, 'stat_{}.txt'.format(args.mode)), 'w') as f: + f.write(result_str) + with open(os.path.join(out_path, 'scores_{}.json'.format(args.mode)), 'w') as f: + json.dump(scores_dict, f) + + +if __name__ == '__main__': + args = parse_args() + run(args) diff --git a/models/encoder4editing/scripts/inference.py b/models/encoder4editing/scripts/inference.py new file mode 100644 index 0000000000000000000000000000000000000000..4d2ad37dcf3eda5ec0c9b6a77da0cc223c6ff452 --- /dev/null +++ b/models/encoder4editing/scripts/inference.py @@ -0,0 +1,134 @@ +import argparse + +import torch +import numpy as np +import sys +import os +import dlib + +sys.path.append(".") +sys.path.append("..") + +from configs import data_configs, paths_config +from datasets.inference_dataset import InferenceDataset +from torch.utils.data import DataLoader +from utils.model_utils import setup_model +from utils.common import tensor2im +from utils.alignment import align_face +from PIL import Image + + +def main(args): + net, opts = setup_model(args.ckpt, device) + is_cars = 'cars_' in opts.dataset_type + generator = net.decoder + generator.eval() + args, data_loader = setup_data_loader(args, opts) + + # Check if latents exist + latents_file_path = os.path.join(args.save_dir, 'latents.pt') + if os.path.exists(latents_file_path): + latent_codes = torch.load(latents_file_path).to(device) + else: + latent_codes = get_all_latents(net, data_loader, args.n_sample, is_cars=is_cars) + torch.save(latent_codes, latents_file_path) + + if not args.latents_only: + generate_inversions(args, generator, latent_codes, is_cars=is_cars) + + +def setup_data_loader(args, opts): + dataset_args = data_configs.DATASETS[opts.dataset_type] + transforms_dict = dataset_args['transforms'](opts).get_transforms() + images_path = args.images_dir if args.images_dir is not None else dataset_args['test_source_root'] + print(f"images path: {images_path}") + align_function = None + if args.align: + align_function = run_alignment + test_dataset = InferenceDataset(root=images_path, + transform=transforms_dict['transform_test'], + preprocess=align_function, + opts=opts) + + data_loader = DataLoader(test_dataset, + batch_size=args.batch, + shuffle=False, + num_workers=2, + drop_last=True) + + print(f'dataset length: {len(test_dataset)}') + + if args.n_sample is None: + args.n_sample = len(test_dataset) + return args, data_loader + + +def get_latents(net, x, is_cars=False): + codes = net.encoder(x) + if net.opts.start_from_latent_avg: + if codes.ndim == 2: + codes = codes + net.latent_avg.repeat(codes.shape[0], 1, 1)[:, 0, :] + else: + codes = codes + net.latent_avg.repeat(codes.shape[0], 1, 1) + if codes.shape[1] == 18 and is_cars: + codes = codes[:, :16, :] + return codes + + +def get_all_latents(net, data_loader, n_images=None, is_cars=False): + all_latents = [] + i = 0 + with torch.no_grad(): + for batch in data_loader: + if n_images is not None and i > n_images: + break + x = batch + inputs = x.to(device).float() + latents = get_latents(net, inputs, is_cars) + all_latents.append(latents) + i += len(latents) + return torch.cat(all_latents) + + +def save_image(img, save_dir, idx): + result = tensor2im(img) + im_save_path = os.path.join(save_dir, f"{idx:05d}.jpg") + Image.fromarray(np.array(result)).save(im_save_path) + + +@torch.no_grad() +def generate_inversions(args, g, latent_codes, is_cars): + print('Saving inversion images') + inversions_directory_path = os.path.join(args.save_dir, 'inversions') + os.makedirs(inversions_directory_path, exist_ok=True) + for i in range(min(args.n_sample, len(latent_codes))): + imgs, _ = g([latent_codes[i].unsqueeze(0)], input_is_latent=True, randomize_noise=False, return_latents=True) + if is_cars: + imgs = imgs[:, :, 64:448, :] + save_image(imgs[0], inversions_directory_path, i + 1) + + +def run_alignment(image_path): + predictor = dlib.shape_predictor(paths_config.model_paths['shape_predictor']) + aligned_image = align_face(filepath=image_path, predictor=predictor) + print("Aligned image has shape: {}".format(aligned_image.size)) + return aligned_image + + +if __name__ == "__main__": + device = "cuda" + + parser = argparse.ArgumentParser(description="Inference") + parser.add_argument("--images_dir", type=str, default=None, + help="The directory of the images to be inverted") + parser.add_argument("--save_dir", type=str, default=None, + help="The directory to save the latent codes and inversion images. (default: images_dir") + parser.add_argument("--batch", type=int, default=1, help="batch size for the generator") + parser.add_argument("--n_sample", type=int, default=None, help="number of the samples to infer.") + parser.add_argument("--latents_only", action="store_true", help="infer only the latent codes of the directory") + parser.add_argument("--align", action="store_true", help="align face images before inference") + parser.add_argument("ckpt", metavar="CHECKPOINT", help="path to generator checkpoint") + + args = parser.parse_args() + print(args) + main(args) diff --git a/models/encoder4editing/scripts/train.py b/models/encoder4editing/scripts/train.py new file mode 100644 index 0000000000000000000000000000000000000000..d885cfde49a0b21140e663e475918698d5e51ee3 --- /dev/null +++ b/models/encoder4editing/scripts/train.py @@ -0,0 +1,88 @@ +""" +This file runs the main training/val loop +""" +import os +import json +import math +import sys +import pprint +import torch +from argparse import Namespace + +sys.path.append(".") +sys.path.append("..") + +from options.train_options import TrainOptions +from training.coach import Coach + + +def main(): + opts = TrainOptions().parse() + previous_train_ckpt = None + if opts.resume_training_from_ckpt: + opts, previous_train_ckpt = load_train_checkpoint(opts) + else: + setup_progressive_steps(opts) + create_initial_experiment_dir(opts) + + coach = Coach(opts, previous_train_ckpt) + coach.train() + + +def load_train_checkpoint(opts): + train_ckpt_path = opts.resume_training_from_ckpt + previous_train_ckpt = torch.load(opts.resume_training_from_ckpt, map_location='cpu') + new_opts_dict = vars(opts) + opts = previous_train_ckpt['opts'] + opts['resume_training_from_ckpt'] = train_ckpt_path + update_new_configs(opts, new_opts_dict) + pprint.pprint(opts) + opts = Namespace(**opts) + if opts.sub_exp_dir is not None: + sub_exp_dir = opts.sub_exp_dir + opts.exp_dir = os.path.join(opts.exp_dir, sub_exp_dir) + create_initial_experiment_dir(opts) + return opts, previous_train_ckpt + + +def setup_progressive_steps(opts): + log_size = int(math.log(opts.stylegan_size, 2)) + num_style_layers = 2*log_size - 2 + num_deltas = num_style_layers - 1 + if opts.progressive_start is not None: # If progressive delta training + opts.progressive_steps = [0] + next_progressive_step = opts.progressive_start + for i in range(num_deltas): + opts.progressive_steps.append(next_progressive_step) + next_progressive_step += opts.progressive_step_every + + assert opts.progressive_steps is None or is_valid_progressive_steps(opts, num_style_layers), \ + "Invalid progressive training input" + + +def is_valid_progressive_steps(opts, num_style_layers): + return len(opts.progressive_steps) == num_style_layers and opts.progressive_steps[0] == 0 + + +def create_initial_experiment_dir(opts): + if os.path.exists(opts.exp_dir): + raise Exception('Oops... {} already exists'.format(opts.exp_dir)) + os.makedirs(opts.exp_dir) + + opts_dict = vars(opts) + pprint.pprint(opts_dict) + with open(os.path.join(opts.exp_dir, 'opt.json'), 'w') as f: + json.dump(opts_dict, f, indent=4, sort_keys=True) + + +def update_new_configs(ckpt_opts, new_opts): + for k, v in new_opts.items(): + if k not in ckpt_opts: + ckpt_opts[k] = v + if new_opts['update_param_list']: + for param in new_opts['update_param_list']: + ckpt_opts[param] = new_opts[param] + + +if __name__ == '__main__': + main() diff --git a/models/encoder4editing/training/__init__.py b/models/encoder4editing/training/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/models/encoder4editing/training/coach.py b/models/encoder4editing/training/coach.py new file mode 100644 index 0000000000000000000000000000000000000000..d0270123bd001f36f90a1f8dbf1a84207093f3cf --- /dev/null +++ b/models/encoder4editing/training/coach.py @@ -0,0 +1,439 @@ +import os +import random +import matplotlib +import matplotlib.pyplot as plt + +matplotlib.use('Agg') + +import torch +from torch import nn, autograd +from torch.utils.data import DataLoader +from torch.utils.tensorboard import SummaryWriter +import torch.nn.functional as F + +from utils import common, train_utils +from criteria import id_loss, moco_loss +from configs import data_configs +from datasets.images_dataset import ImagesDataset +from criteria.lpips.lpips import LPIPS +from models.psp import pSp +from models.latent_codes_pool import LatentCodesPool +from models.discriminator import LatentCodesDiscriminator +from models.encoders.psp_encoders import ProgressiveStage +from training.ranger import Ranger + +random.seed(0) +torch.manual_seed(0) + + +class Coach: + def __init__(self, opts, prev_train_checkpoint=None): + self.opts = opts + + self.global_step = 0 + + self.device = 'cuda:0' + self.opts.device = self.device + # Initialize network + self.net = pSp(self.opts).to(self.device) + + # Initialize loss + if self.opts.lpips_lambda > 0: + self.lpips_loss = LPIPS(net_type=self.opts.lpips_type).to(self.device).eval() + if self.opts.id_lambda > 0: + if 'ffhq' in self.opts.dataset_type or 'celeb' in self.opts.dataset_type: + self.id_loss = id_loss.IDLoss().to(self.device).eval() + else: + self.id_loss = moco_loss.MocoLoss(opts).to(self.device).eval() + self.mse_loss = nn.MSELoss().to(self.device).eval() + + # Initialize optimizer + self.optimizer = self.configure_optimizers() + + # Initialize discriminator + if self.opts.w_discriminator_lambda > 0: + self.discriminator = LatentCodesDiscriminator(512, 4).to(self.device) + self.discriminator_optimizer = torch.optim.Adam(list(self.discriminator.parameters()), + lr=opts.w_discriminator_lr) + self.real_w_pool = LatentCodesPool(self.opts.w_pool_size) + self.fake_w_pool = LatentCodesPool(self.opts.w_pool_size) + + # Initialize dataset + self.train_dataset, self.test_dataset = self.configure_datasets() + self.train_dataloader = DataLoader(self.train_dataset, + batch_size=self.opts.batch_size, + shuffle=True, + num_workers=int(self.opts.workers), + drop_last=True) + self.test_dataloader = DataLoader(self.test_dataset, + batch_size=self.opts.test_batch_size, + shuffle=False, + num_workers=int(self.opts.test_workers), + drop_last=True) + + # Initialize logger + log_dir = os.path.join(opts.exp_dir, 'logs') + os.makedirs(log_dir, exist_ok=True) + self.logger = SummaryWriter(log_dir=log_dir) + + # Initialize checkpoint dir + self.checkpoint_dir = os.path.join(opts.exp_dir, 'checkpoints') + os.makedirs(self.checkpoint_dir, exist_ok=True) + self.best_val_loss = None + if self.opts.save_interval is None: + self.opts.save_interval = self.opts.max_steps + + if prev_train_checkpoint is not None: + self.load_from_train_checkpoint(prev_train_checkpoint) + prev_train_checkpoint = None + + def load_from_train_checkpoint(self, ckpt): + print('Loading previous training data...') + self.global_step = ckpt['global_step'] + 1 + self.best_val_loss = ckpt['best_val_loss'] + self.net.load_state_dict(ckpt['state_dict']) + + if self.opts.keep_optimizer: + self.optimizer.load_state_dict(ckpt['optimizer']) + if self.opts.w_discriminator_lambda > 0: + self.discriminator.load_state_dict(ckpt['discriminator_state_dict']) + self.discriminator_optimizer.load_state_dict(ckpt['discriminator_optimizer_state_dict']) + if self.opts.progressive_steps: + self.check_for_progressive_training_update(is_resume_from_ckpt=True) + print(f'Resuming training from step {self.global_step}') + + def train(self): + self.net.train() + if self.opts.progressive_steps: + self.check_for_progressive_training_update() + while self.global_step < self.opts.max_steps: + for batch_idx, batch in enumerate(self.train_dataloader): + loss_dict = {} + if self.is_training_discriminator(): + loss_dict = self.train_discriminator(batch) + x, y, y_hat, latent = self.forward(batch) + loss, encoder_loss_dict, id_logs = self.calc_loss(x, y, y_hat, latent) + loss_dict = {**loss_dict, **encoder_loss_dict} + self.optimizer.zero_grad() + loss.backward() + self.optimizer.step() + + # Logging related + if self.global_step % self.opts.image_interval == 0 or ( + self.global_step < 1000 and self.global_step % 25 == 0): + self.parse_and_log_images(id_logs, x, y, y_hat, title='images/scripts/faces') + if self.global_step % self.opts.board_interval == 0: + self.print_metrics(loss_dict, prefix='scripts') + self.log_metrics(loss_dict, prefix='scripts') + + # Validation related + val_loss_dict = None + if self.global_step % self.opts.val_interval == 0 or self.global_step == self.opts.max_steps: + val_loss_dict = self.validate() + if val_loss_dict and (self.best_val_loss is None or val_loss_dict['loss'] < self.best_val_loss): + self.best_val_loss = val_loss_dict['loss'] + self.checkpoint_me(val_loss_dict, is_best=True) + + if self.global_step % self.opts.save_interval == 0 or self.global_step == self.opts.max_steps: + if val_loss_dict is not None: + self.checkpoint_me(val_loss_dict, is_best=False) + else: + self.checkpoint_me(loss_dict, is_best=False) + + if self.global_step == self.opts.max_steps: + print('OMG, finished training!') + break + + self.global_step += 1 + if self.opts.progressive_steps: + self.check_for_progressive_training_update() + + def check_for_progressive_training_update(self, is_resume_from_ckpt=False): + for i in range(len(self.opts.progressive_steps)): + if is_resume_from_ckpt and self.global_step >= self.opts.progressive_steps[i]: # Case checkpoint + self.net.encoder.set_progressive_stage(ProgressiveStage(i)) + if self.global_step == self.opts.progressive_steps[i]: # Case training reached progressive step + self.net.encoder.set_progressive_stage(ProgressiveStage(i)) + + def validate(self): + self.net.eval() + agg_loss_dict = [] + for batch_idx, batch in enumerate(self.test_dataloader): + cur_loss_dict = {} + if self.is_training_discriminator(): + cur_loss_dict = self.validate_discriminator(batch) + with torch.no_grad(): + x, y, y_hat, latent = self.forward(batch) + loss, cur_encoder_loss_dict, id_logs = self.calc_loss(x, y, y_hat, latent) + cur_loss_dict = {**cur_loss_dict, **cur_encoder_loss_dict} + agg_loss_dict.append(cur_loss_dict) + + # Logging related + self.parse_and_log_images(id_logs, x, y, y_hat, + title='images/test/faces', + subscript='{:04d}'.format(batch_idx)) + + # For first step just do sanity test on small amount of data + if self.global_step == 0 and batch_idx >= 4: + self.net.train() + return None # Do not log, inaccurate in first batch + + loss_dict = train_utils.aggregate_loss_dict(agg_loss_dict) + self.log_metrics(loss_dict, prefix='test') + self.print_metrics(loss_dict, prefix='test') + + self.net.train() + return loss_dict + + def checkpoint_me(self, loss_dict, is_best): + save_name = 'best_model.pt' if is_best else 'iteration_{}.pt'.format(self.global_step) + save_dict = self.__get_save_dict() + checkpoint_path = os.path.join(self.checkpoint_dir, save_name) + torch.save(save_dict, checkpoint_path) + with open(os.path.join(self.checkpoint_dir, 'timestamp.txt'), 'a') as f: + if is_best: + f.write( + '**Best**: Step - {}, Loss - {:.3f} \n{}\n'.format(self.global_step, self.best_val_loss, loss_dict)) + else: + f.write('Step - {}, \n{}\n'.format(self.global_step, loss_dict)) + + def configure_optimizers(self): + params = list(self.net.encoder.parameters()) + if self.opts.train_decoder: + params += list(self.net.decoder.parameters()) + else: + self.requires_grad(self.net.decoder, False) + if self.opts.optim_name == 'adam': + optimizer = torch.optim.Adam(params, lr=self.opts.learning_rate) + else: + optimizer = Ranger(params, lr=self.opts.learning_rate) + return optimizer + + def configure_datasets(self): + if self.opts.dataset_type not in data_configs.DATASETS.keys(): + Exception('{} is not a valid dataset_type'.format(self.opts.dataset_type)) + print('Loading dataset for {}'.format(self.opts.dataset_type)) + dataset_args = data_configs.DATASETS[self.opts.dataset_type] + transforms_dict = dataset_args['transforms'](self.opts).get_transforms() + train_dataset = ImagesDataset(source_root=dataset_args['train_source_root'], + target_root=dataset_args['train_target_root'], + source_transform=transforms_dict['transform_source'], + target_transform=transforms_dict['transform_gt_train'], + opts=self.opts) + test_dataset = ImagesDataset(source_root=dataset_args['test_source_root'], + target_root=dataset_args['test_target_root'], + source_transform=transforms_dict['transform_source'], + target_transform=transforms_dict['transform_test'], + opts=self.opts) + print("Number of training samples: {}".format(len(train_dataset))) + print("Number of test samples: {}".format(len(test_dataset))) + return train_dataset, test_dataset + + def calc_loss(self, x, y, y_hat, latent): + loss_dict = {} + loss = 0.0 + id_logs = None + if self.is_training_discriminator(): # Adversarial loss + loss_disc = 0. + dims_to_discriminate = self.get_dims_to_discriminate() if self.is_progressive_training() else \ + list(range(self.net.decoder.n_latent)) + + for i in dims_to_discriminate: + w = latent[:, i, :] + fake_pred = self.discriminator(w) + loss_disc += F.softplus(-fake_pred).mean() + loss_disc /= len(dims_to_discriminate) + loss_dict['encoder_discriminator_loss'] = float(loss_disc) + loss += self.opts.w_discriminator_lambda * loss_disc + + if self.opts.progressive_steps and self.net.encoder.progressive_stage.value != 18: # delta regularization loss + total_delta_loss = 0 + deltas_latent_dims = self.net.encoder.get_deltas_starting_dimensions() + + first_w = latent[:, 0, :] + for i in range(1, self.net.encoder.progressive_stage.value + 1): + curr_dim = deltas_latent_dims[i] + delta = latent[:, curr_dim, :] - first_w + delta_loss = torch.norm(delta, self.opts.delta_norm, dim=1).mean() + loss_dict[f"delta{i}_loss"] = float(delta_loss) + total_delta_loss += delta_loss + loss_dict['total_delta_loss'] = float(total_delta_loss) + loss += self.opts.delta_norm_lambda * total_delta_loss + + if self.opts.id_lambda > 0: # Similarity loss + loss_id, sim_improvement, id_logs = self.id_loss(y_hat, y, x) + loss_dict['loss_id'] = float(loss_id) + loss_dict['id_improve'] = float(sim_improvement) + loss += loss_id * self.opts.id_lambda + if self.opts.l2_lambda > 0: + loss_l2 = F.mse_loss(y_hat, y) + loss_dict['loss_l2'] = float(loss_l2) + loss += loss_l2 * self.opts.l2_lambda + if self.opts.lpips_lambda > 0: + loss_lpips = self.lpips_loss(y_hat, y) + loss_dict['loss_lpips'] = float(loss_lpips) + loss += loss_lpips * self.opts.lpips_lambda + loss_dict['loss'] = float(loss) + return loss, loss_dict, id_logs + + def forward(self, batch): + x, y = batch + x, y = x.to(self.device).float(), y.to(self.device).float() + y_hat, latent = self.net.forward(x, return_latents=True) + if self.opts.dataset_type == "cars_encode": + y_hat = y_hat[:, :, 32:224, :] + return x, y, y_hat, latent + + def log_metrics(self, metrics_dict, prefix): + for key, value in metrics_dict.items(): + self.logger.add_scalar('{}/{}'.format(prefix, key), value, self.global_step) + + def print_metrics(self, metrics_dict, prefix): + print('Metrics for {}, step {}'.format(prefix, self.global_step)) + for key, value in metrics_dict.items(): + print('\t{} = '.format(key), value) + + def parse_and_log_images(self, id_logs, x, y, y_hat, title, subscript=None, display_count=2): + im_data = [] + for i in range(display_count): + cur_im_data = { + 'input_face': common.log_input_image(x[i], self.opts), + 'target_face': common.tensor2im(y[i]), + 'output_face': common.tensor2im(y_hat[i]), + } + if id_logs is not None: + for key in id_logs[i]: + cur_im_data[key] = id_logs[i][key] + im_data.append(cur_im_data) + self.log_images(title, im_data=im_data, subscript=subscript) + + def log_images(self, name, im_data, subscript=None, log_latest=False): + fig = common.vis_faces(im_data) + step = self.global_step + if log_latest: + step = 0 + if subscript: + path = os.path.join(self.logger.log_dir, name, '{}_{:04d}.jpg'.format(subscript, step)) + else: + path = os.path.join(self.logger.log_dir, name, '{:04d}.jpg'.format(step)) + os.makedirs(os.path.dirname(path), exist_ok=True) + fig.savefig(path) + plt.close(fig) + + def __get_save_dict(self): + save_dict = { + 'state_dict': self.net.state_dict(), + 'opts': vars(self.opts) + } + # save the latent avg in state_dict for inference if truncation of w was used during training + if self.opts.start_from_latent_avg: + save_dict['latent_avg'] = self.net.latent_avg + + if self.opts.save_training_data: # Save necessary information to enable training continuation from checkpoint + save_dict['global_step'] = self.global_step + save_dict['optimizer'] = self.optimizer.state_dict() + save_dict['best_val_loss'] = self.best_val_loss + if self.opts.w_discriminator_lambda > 0: + save_dict['discriminator_state_dict'] = self.discriminator.state_dict() + save_dict['discriminator_optimizer_state_dict'] = self.discriminator_optimizer.state_dict() + return save_dict + + def get_dims_to_discriminate(self): + deltas_starting_dimensions = self.net.encoder.get_deltas_starting_dimensions() + return deltas_starting_dimensions[:self.net.encoder.progressive_stage.value + 1] + + def is_progressive_training(self): + return self.opts.progressive_steps is not None + +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ Discriminator ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # + + def is_training_discriminator(self): + return self.opts.w_discriminator_lambda > 0 + + @staticmethod + def discriminator_loss(real_pred, fake_pred, loss_dict): + real_loss = F.softplus(-real_pred).mean() + fake_loss = F.softplus(fake_pred).mean() + + loss_dict['d_real_loss'] = float(real_loss) + loss_dict['d_fake_loss'] = float(fake_loss) + + return real_loss + fake_loss + + @staticmethod + def discriminator_r1_loss(real_pred, real_w): + grad_real, = autograd.grad( + outputs=real_pred.sum(), inputs=real_w, create_graph=True + ) + grad_penalty = grad_real.pow(2).reshape(grad_real.shape[0], -1).sum(1).mean() + + return grad_penalty + + @staticmethod + def requires_grad(model, flag=True): + for p in model.parameters(): + p.requires_grad = flag + + def train_discriminator(self, batch): + loss_dict = {} + x, _ = batch + x = x.to(self.device).float() + self.requires_grad(self.discriminator, True) + + with torch.no_grad(): + real_w, fake_w = self.sample_real_and_fake_latents(x) + real_pred = self.discriminator(real_w) + fake_pred = self.discriminator(fake_w) + loss = self.discriminator_loss(real_pred, fake_pred, loss_dict) + loss_dict['discriminator_loss'] = float(loss) + + self.discriminator_optimizer.zero_grad() + loss.backward() + self.discriminator_optimizer.step() + + # r1 regularization + d_regularize = self.global_step % self.opts.d_reg_every == 0 + if d_regularize: + real_w = real_w.detach() + real_w.requires_grad = True + real_pred = self.discriminator(real_w) + r1_loss = self.discriminator_r1_loss(real_pred, real_w) + + self.discriminator.zero_grad() + r1_final_loss = self.opts.r1 / 2 * r1_loss * self.opts.d_reg_every + 0 * real_pred[0] + r1_final_loss.backward() + self.discriminator_optimizer.step() + loss_dict['discriminator_r1_loss'] = float(r1_final_loss) + + # Reset to previous state + self.requires_grad(self.discriminator, False) + + return loss_dict + + def validate_discriminator(self, test_batch): + with torch.no_grad(): + loss_dict = {} + x, _ = test_batch + x = x.to(self.device).float() + real_w, fake_w = self.sample_real_and_fake_latents(x) + real_pred = self.discriminator(real_w) + fake_pred = self.discriminator(fake_w) + loss = self.discriminator_loss(real_pred, fake_pred, loss_dict) + loss_dict['discriminator_loss'] = float(loss) + return loss_dict + + def sample_real_and_fake_latents(self, x): + sample_z = torch.randn(self.opts.batch_size, 512, device=self.device) + real_w = self.net.decoder.get_latent(sample_z) + fake_w = self.net.encoder(x) + if self.opts.start_from_latent_avg: + fake_w = fake_w + self.net.latent_avg.repeat(fake_w.shape[0], 1, 1) + if self.is_progressive_training(): # When progressive training, feed only unique w's + dims_to_discriminate = self.get_dims_to_discriminate() + fake_w = fake_w[:, dims_to_discriminate, :] + if self.opts.use_w_pool: + real_w = self.real_w_pool.query(real_w) + fake_w = self.fake_w_pool.query(fake_w) + if fake_w.ndim == 3: + fake_w = fake_w[:, 0, :] + return real_w, fake_w diff --git a/models/encoder4editing/training/ranger.py b/models/encoder4editing/training/ranger.py new file mode 100644 index 0000000000000000000000000000000000000000..3d63264dda6df0ee40cac143440f0b5f8977a9ad --- /dev/null +++ b/models/encoder4editing/training/ranger.py @@ -0,0 +1,164 @@ +# Ranger deep learning optimizer - RAdam + Lookahead + Gradient Centralization, combined into one optimizer. + +# https://github.com/lessw2020/Ranger-Deep-Learning-Optimizer +# and/or +# https://github.com/lessw2020/Best-Deep-Learning-Optimizers + +# Ranger has now been used to capture 12 records on the FastAI leaderboard. + +# This version = 20.4.11 + +# Credits: +# Gradient Centralization --> https://arxiv.org/abs/2004.01461v2 (a new optimization technique for DNNs), github: https://github.com/Yonghongwei/Gradient-Centralization +# RAdam --> https://github.com/LiyuanLucasLiu/RAdam +# Lookahead --> rewritten by lessw2020, but big thanks to Github @LonePatient and @RWightman for ideas from their code. +# Lookahead paper --> MZhang,G Hinton https://arxiv.org/abs/1907.08610 + +# summary of changes: +# 4/11/20 - add gradient centralization option. Set new testing benchmark for accuracy with it, toggle with use_gc flag at init. +# full code integration with all updates at param level instead of group, moves slow weights into state dict (from generic weights), +# supports group learning rates (thanks @SHolderbach), fixes sporadic load from saved model issues. +# changes 8/31/19 - fix references to *self*.N_sma_threshold; +# changed eps to 1e-5 as better default than 1e-8. + +import math +import torch +from torch.optim.optimizer import Optimizer + + +class Ranger(Optimizer): + + def __init__(self, params, lr=1e-3, # lr + alpha=0.5, k=6, N_sma_threshhold=5, # Ranger options + betas=(.95, 0.999), eps=1e-5, weight_decay=0, # Adam options + use_gc=True, gc_conv_only=False + # Gradient centralization on or off, applied to conv layers only or conv + fc layers + ): + + # parameter checks + if not 0.0 <= alpha <= 1.0: + raise ValueError(f'Invalid slow update rate: {alpha}') + if not 1 <= k: + raise ValueError(f'Invalid lookahead steps: {k}') + if not lr > 0: + raise ValueError(f'Invalid Learning Rate: {lr}') + if not eps > 0: + raise ValueError(f'Invalid eps: {eps}') + + # parameter comments: + # beta1 (momentum) of .95 seems to work better than .90... + # N_sma_threshold of 5 seems better in testing than 4. + # In both cases, worth testing on your dataset (.90 vs .95, 4 vs 5) to make sure which works best for you. + + # prep defaults and init torch.optim base + defaults = dict(lr=lr, alpha=alpha, k=k, step_counter=0, betas=betas, N_sma_threshhold=N_sma_threshhold, + eps=eps, weight_decay=weight_decay) + super().__init__(params, defaults) + + # adjustable threshold + self.N_sma_threshhold = N_sma_threshhold + + # look ahead params + + self.alpha = alpha + self.k = k + + # radam buffer for state + self.radam_buffer = [[None, None, None] for ind in range(10)] + + # gc on or off + self.use_gc = use_gc + + # level of gradient centralization + self.gc_gradient_threshold = 3 if gc_conv_only else 1 + + def __setstate__(self, state): + super(Ranger, self).__setstate__(state) + + def step(self, closure=None): + loss = None + + # Evaluate averages and grad, update param tensors + for group in self.param_groups: + + for p in group['params']: + if p.grad is None: + continue + grad = p.grad.data.float() + + if grad.is_sparse: + raise RuntimeError('Ranger optimizer does not support sparse gradients') + + p_data_fp32 = p.data.float() + + state = self.state[p] # get state dict for this param + + if len(state) == 0: # if first time to run...init dictionary with our desired entries + # if self.first_run_check==0: + # self.first_run_check=1 + # print("Initializing slow buffer...should not see this at load from saved model!") + state['step'] = 0 + state['exp_avg'] = torch.zeros_like(p_data_fp32) + state['exp_avg_sq'] = torch.zeros_like(p_data_fp32) + + # look ahead weight storage now in state dict + state['slow_buffer'] = torch.empty_like(p.data) + state['slow_buffer'].copy_(p.data) + + else: + state['exp_avg'] = state['exp_avg'].type_as(p_data_fp32) + state['exp_avg_sq'] = state['exp_avg_sq'].type_as(p_data_fp32) + + # begin computations + exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq'] + beta1, beta2 = group['betas'] + + # GC operation for Conv layers and FC layers + if grad.dim() > self.gc_gradient_threshold: + grad.add_(-grad.mean(dim=tuple(range(1, grad.dim())), keepdim=True)) + + state['step'] += 1 + + # compute variance mov avg + exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad) + # compute mean moving avg + exp_avg.mul_(beta1).add_(1 - beta1, grad) + + buffered = self.radam_buffer[int(state['step'] % 10)] + + if state['step'] == buffered[0]: + N_sma, step_size = buffered[1], buffered[2] + else: + buffered[0] = state['step'] + beta2_t = beta2 ** state['step'] + N_sma_max = 2 / (1 - beta2) - 1 + N_sma = N_sma_max - 2 * state['step'] * beta2_t / (1 - beta2_t) + buffered[1] = N_sma + if N_sma > self.N_sma_threshhold: + step_size = math.sqrt( + (1 - beta2_t) * (N_sma - 4) / (N_sma_max - 4) * (N_sma - 2) / N_sma * N_sma_max / ( + N_sma_max - 2)) / (1 - beta1 ** state['step']) + else: + step_size = 1.0 / (1 - beta1 ** state['step']) + buffered[2] = step_size + + if group['weight_decay'] != 0: + p_data_fp32.add_(-group['weight_decay'] * group['lr'], p_data_fp32) + + # apply lr + if N_sma > self.N_sma_threshhold: + denom = exp_avg_sq.sqrt().add_(group['eps']) + p_data_fp32.addcdiv_(-step_size * group['lr'], exp_avg, denom) + else: + p_data_fp32.add_(-step_size * group['lr'], exp_avg) + + p.data.copy_(p_data_fp32) + + # integrated look ahead... + # we do it at the param level instead of group level + if state['step'] % group['k'] == 0: + slow_p = state['slow_buffer'] # get access to slow param tensor + slow_p.add_(self.alpha, p.data - slow_p) # (fast weights - slow weights) * alpha + p.data.copy_(slow_p) # copy interpolated weights to RAdam param tensor + + return loss \ No newline at end of file diff --git a/models/encoder4editing/utils/__init__.py b/models/encoder4editing/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/models/encoder4editing/utils/alignment.py b/models/encoder4editing/utils/alignment.py new file mode 100644 index 0000000000000000000000000000000000000000..a02798f0f7c9fdcc319f7884a491b9e6580cc8aa --- /dev/null +++ b/models/encoder4editing/utils/alignment.py @@ -0,0 +1,115 @@ +import numpy as np +import PIL +import PIL.Image +import scipy +import scipy.ndimage +import dlib + + +def get_landmark(filepath, predictor): + """get landmark with dlib + :return: np.array shape=(68, 2) + """ + detector = dlib.get_frontal_face_detector() + + img = dlib.load_rgb_image(filepath) + dets = detector(img, 1) + + for k, d in enumerate(dets): + shape = predictor(img, d) + + t = list(shape.parts()) + a = [] + for tt in t: + a.append([tt.x, tt.y]) + lm = np.array(a) + return lm + + +def align_face(filepath, predictor): + """ + :param filepath: str + :return: PIL Image + """ + + lm = get_landmark(filepath, predictor) + + lm_chin = lm[0: 17] # left-right + lm_eyebrow_left = lm[17: 22] # left-right + lm_eyebrow_right = lm[22: 27] # left-right + lm_nose = lm[27: 31] # top-down + lm_nostrils = lm[31: 36] # top-down + lm_eye_left = lm[36: 42] # left-clockwise + lm_eye_right = lm[42: 48] # left-clockwise + lm_mouth_outer = lm[48: 60] # left-clockwise + lm_mouth_inner = lm[60: 68] # left-clockwise + + # Calculate auxiliary vectors. + eye_left = np.mean(lm_eye_left, axis=0) + eye_right = np.mean(lm_eye_right, axis=0) + eye_avg = (eye_left + eye_right) * 0.5 + eye_to_eye = eye_right - eye_left + mouth_left = lm_mouth_outer[0] + mouth_right = lm_mouth_outer[6] + mouth_avg = (mouth_left + mouth_right) * 0.5 + eye_to_mouth = mouth_avg - eye_avg + + # Choose oriented crop rectangle. + x = eye_to_eye - np.flipud(eye_to_mouth) * [-1, 1] + x /= np.hypot(*x) + x *= max(np.hypot(*eye_to_eye) * 2.0, np.hypot(*eye_to_mouth) * 1.8) + y = np.flipud(x) * [-1, 1] + c = eye_avg + eye_to_mouth * 0.1 + quad = np.stack([c - x - y, c - x + y, c + x + y, c + x - y]) + qsize = np.hypot(*x) * 2 + + # read image + img = PIL.Image.open(filepath) + + output_size = 256 + transform_size = 256 + enable_padding = True + + # Shrink. + shrink = int(np.floor(qsize / output_size * 0.5)) + if shrink > 1: + rsize = (int(np.rint(float(img.size[0]) / shrink)), int(np.rint(float(img.size[1]) / shrink))) + img = img.resize(rsize, PIL.Image.ANTIALIAS) + quad /= shrink + qsize /= shrink + + # Crop. + border = max(int(np.rint(qsize * 0.1)), 3) + crop = (int(np.floor(min(quad[:, 0]))), int(np.floor(min(quad[:, 1]))), int(np.ceil(max(quad[:, 0]))), + int(np.ceil(max(quad[:, 1])))) + crop = (max(crop[0] - border, 0), max(crop[1] - border, 0), min(crop[2] + border, img.size[0]), + min(crop[3] + border, img.size[1])) + if crop[2] - crop[0] < img.size[0] or crop[3] - crop[1] < img.size[1]: + img = img.crop(crop) + quad -= crop[0:2] + + # Pad. + pad = (int(np.floor(min(quad[:, 0]))), int(np.floor(min(quad[:, 1]))), int(np.ceil(max(quad[:, 0]))), + int(np.ceil(max(quad[:, 1])))) + pad = (max(-pad[0] + border, 0), max(-pad[1] + border, 0), max(pad[2] - img.size[0] + border, 0), + max(pad[3] - img.size[1] + border, 0)) + if enable_padding and max(pad) > border - 4: + pad = np.maximum(pad, int(np.rint(qsize * 0.3))) + img = np.pad(np.float32(img), ((pad[1], pad[3]), (pad[0], pad[2]), (0, 0)), 'reflect') + h, w, _ = img.shape + y, x, _ = np.ogrid[:h, :w, :1] + mask = np.maximum(1.0 - np.minimum(np.float32(x) / pad[0], np.float32(w - 1 - x) / pad[2]), + 1.0 - np.minimum(np.float32(y) / pad[1], np.float32(h - 1 - y) / pad[3])) + blur = qsize * 0.02 + img += (scipy.ndimage.gaussian_filter(img, [blur, blur, 0]) - img) * np.clip(mask * 3.0 + 1.0, 0.0, 1.0) + img += (np.median(img, axis=(0, 1)) - img) * np.clip(mask, 0.0, 1.0) + img = PIL.Image.fromarray(np.uint8(np.clip(np.rint(img), 0, 255)), 'RGB') + quad += pad[:2] + + # Transform. + img = img.transform((transform_size, transform_size), PIL.Image.QUAD, (quad + 0.5).flatten(), PIL.Image.BILINEAR) + if output_size < transform_size: + img = img.resize((output_size, output_size), PIL.Image.ANTIALIAS) + + # Return aligned image. + return img diff --git a/models/encoder4editing/utils/common.py b/models/encoder4editing/utils/common.py new file mode 100644 index 0000000000000000000000000000000000000000..b19e18ddcb78b06678fa18e4a76da44fc511b789 --- /dev/null +++ b/models/encoder4editing/utils/common.py @@ -0,0 +1,55 @@ +from PIL import Image +import matplotlib.pyplot as plt + + +# Log images +def log_input_image(x, opts): + return tensor2im(x) + + +def tensor2im(var): + # var shape: (3, H, W) + var = var.cpu().detach().transpose(0, 2).transpose(0, 1).numpy() + var = ((var + 1) / 2) + var[var < 0] = 0 + var[var > 1] = 1 + var = var * 255 + return Image.fromarray(var.astype('uint8')) + + +def vis_faces(log_hooks): + display_count = len(log_hooks) + fig = plt.figure(figsize=(8, 4 * display_count)) + gs = fig.add_gridspec(display_count, 3) + for i in range(display_count): + hooks_dict = log_hooks[i] + fig.add_subplot(gs[i, 0]) + if 'diff_input' in hooks_dict: + vis_faces_with_id(hooks_dict, fig, gs, i) + else: + vis_faces_no_id(hooks_dict, fig, gs, i) + plt.tight_layout() + return fig + + +def vis_faces_with_id(hooks_dict, fig, gs, i): + plt.imshow(hooks_dict['input_face']) + plt.title('Input\nOut Sim={:.2f}'.format(float(hooks_dict['diff_input']))) + fig.add_subplot(gs[i, 1]) + plt.imshow(hooks_dict['target_face']) + plt.title('Target\nIn={:.2f}, Out={:.2f}'.format(float(hooks_dict['diff_views']), + float(hooks_dict['diff_target']))) + fig.add_subplot(gs[i, 2]) + plt.imshow(hooks_dict['output_face']) + plt.title('Output\n Target Sim={:.2f}'.format(float(hooks_dict['diff_target']))) + + +def vis_faces_no_id(hooks_dict, fig, gs, i): + plt.imshow(hooks_dict['input_face'], cmap="gray") + plt.title('Input') + fig.add_subplot(gs[i, 1]) + plt.imshow(hooks_dict['target_face']) + plt.title('Target') + fig.add_subplot(gs[i, 2]) + plt.imshow(hooks_dict['output_face']) + plt.title('Output') diff --git a/models/encoder4editing/utils/data_utils.py b/models/encoder4editing/utils/data_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..f1ba79f4a2d5cc2b97dce76d87bf6e7cdebbc257 --- /dev/null +++ b/models/encoder4editing/utils/data_utils.py @@ -0,0 +1,25 @@ +""" +Code adopted from pix2pixHD: +https://github.com/NVIDIA/pix2pixHD/blob/master/data/image_folder.py +""" +import os + +IMG_EXTENSIONS = [ + '.jpg', '.JPG', '.jpeg', '.JPEG', + '.png', '.PNG', '.ppm', '.PPM', '.bmp', '.BMP', '.tiff' +] + + +def is_image_file(filename): + return any(filename.endswith(extension) for extension in IMG_EXTENSIONS) + + +def make_dataset(dir): + images = [] + assert os.path.isdir(dir), '%s is not a valid directory' % dir + for root, _, fnames in sorted(os.walk(dir)): + for fname in fnames: + if is_image_file(fname): + path = os.path.join(root, fname) + images.append(path) + return images diff --git a/models/encoder4editing/utils/model_utils.py b/models/encoder4editing/utils/model_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..ee3acd4faa80d0e44e9d41af39540c94ea95ef53 --- /dev/null +++ b/models/encoder4editing/utils/model_utils.py @@ -0,0 +1,45 @@ +import torch +import argparse +from models.encoder4editing.models.psp import pSp +from models.encoder4editing.models.encoders.psp_encoders import Encoder4Editing + + +def get_latents(net, x): + codes = net.encoder(x) + if net.opts.start_from_latent_avg: + if codes.ndim == 2: + codes = codes + net.latent_avg.repeat(codes.shape[0], 1, 1)[:, 0, :] + else: + codes = codes + net.latent_avg.repeat(codes.shape[0], 1, 1) + return codes + + +def setup_model(checkpoint_path, device='cuda'): + ckpt = torch.load(checkpoint_path, map_location='cpu') + opts = ckpt['opts'] + + opts['checkpoint_path'] = checkpoint_path + opts['device'] = device + opts = argparse.Namespace(**opts) + + net = pSp(opts) + net.eval() + net = net.to(device) + return net, opts + + +def load_e4e_standalone(checkpoint_path, device='cuda'): + ckpt = torch.load(checkpoint_path, map_location='cpu') + opts = argparse.Namespace(**ckpt['opts']) + e4e = Encoder4Editing(50, 'ir_se', opts) + e4e_dict = {k.replace('encoder.', ''): v for k, v in ckpt['state_dict'].items() if k.startswith('encoder.')} + e4e.load_state_dict(e4e_dict) + e4e.eval() + e4e = e4e.to(device) + latent_avg = ckpt['latent_avg'].to(device) + + def add_latent_avg(model, inputs, outputs): + return outputs + latent_avg.repeat(outputs.shape[0], 1, 1) + + e4e.register_forward_hook(add_latent_avg) + return e4e diff --git a/models/encoder4editing/utils/train_utils.py b/models/encoder4editing/utils/train_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..0c55177f7442010bc1fcc64de3d142585c22adc0 --- /dev/null +++ b/models/encoder4editing/utils/train_utils.py @@ -0,0 +1,13 @@ + +def aggregate_loss_dict(agg_loss_dict): + mean_vals = {} + for output in agg_loss_dict: + for key in output: + mean_vals[key] = mean_vals.setdefault(key, []) + [output[key]] + for key in mean_vals: + if len(mean_vals[key]) > 0: + mean_vals[key] = sum(mean_vals[key]) / len(mean_vals[key]) + else: + print('{} has no value'.format(key)) + mean_vals[key] = 0 + return mean_vals diff --git a/models/face_parsing/README.md b/models/face_parsing/README.md new file mode 100644 index 0000000000000000000000000000000000000000..849d55e2789c8852e01707d1ff755dc74e63a7f5 --- /dev/null +++ b/models/face_parsing/README.md @@ -0,0 +1,68 @@ +# face-parsing.PyTorch + +

+ + + +

+ +### Contents +- [Training](#training) +- [Demo](#Demo) +- [References](#references) + +## Training + +1. Prepare training data: + -- download [CelebAMask-HQ dataset](https://github.com/switchablenorms/CelebAMask-HQ) + + -- change file path in the `prepropess_data.py` and run +```Shell +python prepropess_data.py +``` + +2. Train the model using CelebAMask-HQ dataset: +Just run the train script: +``` + $ CUDA_VISIBLE_DEVICES=0,1 python -m torch.distributed.launch --nproc_per_node=2 train.py +``` + +If you do not wish to train the model, you can download [our pre-trained model](https://drive.google.com/open?id=154JgKpzCPW82qINcVieuPH3fZ2e0P812) and save it in `res/cp`. + + +## Demo +1. Evaluate the trained model using: +```Shell +# evaluate using GPU +python test.py +``` + +## Face makeup using parsing maps +[**face-makeup.PyTorch**](https://github.com/zllrunning/face-makeup.PyTorch) + + + + + + + + + + + + + + + + + + + + + + +
 HairLip
Original InputOriginal InputOriginal Input
ColorColorColor
+ + +## References +- [BiSeNet](https://github.com/CoinCheung/BiSeNet) \ No newline at end of file diff --git a/models/face_parsing/face_dataset.py b/models/face_parsing/face_dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..847715d3fd60faa1d8f435e02dc4d001b96b0d03 --- /dev/null +++ b/models/face_parsing/face_dataset.py @@ -0,0 +1,106 @@ +# #!/usr/bin/python +# # -*- encoding: utf-8 -*- +# +# import torch +# from torch.utils.data import Dataset +# import torchvision.transforms as transforms +# +# import os.path as osp +# import os +# from PIL import Image +# import numpy as np +# import json +# import cv2 +# +# from .transform import * +# +# +# +# class FaceMask(Dataset): +# def __init__(self, rootpth, cropsize=(640, 480), mode='scripts', *args, **kwargs): +# super(FaceMask, self).__init__(*args, **kwargs) +# assert mode in ('scripts', 'val', 'test') +# self.mode = mode +# self.ignore_lb = 255 +# self.rootpth = rootpth +# +# self.imgs = os.listdir(os.path.join(self.rootpth, 'CelebA-HQ-img')) +# +# # pre-processing +# self.to_tensor = transforms.Compose([ +# transforms.ToTensor(), +# transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)), +# ]) +# self.trans_train = Compose([ +# ColorJitter( +# brightness=0.5, +# contrast=0.5, +# saturation=0.5), +# HorizontalFlip(), +# RandomScale((0.75, 1.0, 1.25, 1.5, 1.75, 2.0)), +# RandomCrop(cropsize) +# ]) +# +# def __getitem__(self, idx): +# impth = self.imgs[idx] +# img = Image.open(osp.join(self.rootpth, 'CelebA-HQ-img', impth)) +# img = img.resize((512, 512), Image.BILINEAR) +# label = Image.open(osp.join(self.rootpth, 'Mask_less', impth[:-3]+'png')).convert('P') +# # print(np.unique(np.array(label))) +# if self.mode == 'scripts': +# im_lb = dict(im=img, lb=label) +# im_lb = self.trans_train(im_lb) +# img, label = im_lb['im'], im_lb['lb'] +# img = self.to_tensor(img) +# label = np.array(label).astype(np.int64)[np.newaxis, :] +# return img, label +# +# def __len__(self): +# return len(self.imgs) +# +# +# if __name__ == "__main__": +# face_data = '/home/zll/data/CelebAMask-HQ/CelebA-HQ-img' +# face_sep_mask = '/home/zll/data/CelebAMask-HQ/CelebAMask-HQ-mask-anno' +# mask_path = '/home/zll/data/CelebAMask-HQ/mask' +# counter = 0 +# total = 0 +# for i in range(15): +# # files = os.listdir(osp.join(face_sep_mask, str(i))) +# +# atts = ['skin', 'nose', 'eye_g', 'l_eye', 'r_eye', 'l_brow', 'r_brow', 'l_ear', 'r_ear', 'mouth', 'u_lip', +# 'l_lip', 'hair', 'hat', 'ear_r', 'neck_l', 'neck', 'cloth'] +# +# for j in range(i*2000, (i+1)*2000): +# +# mask = np.zeros((512, 512)) +# +# for l, att in enumerate(atts, 1): +# total += 1 +# file_name = ''.join([str(j).rjust(5, '0'), '_', att, '.png']) +# path = osp.join(face_sep_mask, str(i), file_name) +# +# if os.path.exists(path): +# counter += 1 +# sep_mask = np.array(Image.open(path).convert('P')) +# # print(np.unique(sep_mask)) +# +# mask[sep_mask == 225] = l +# cv2.imwrite('{}/{}.png'.format(mask_path, j), mask) +# print(j) +# +# print(counter, total) +# +# +# +# +# +# +# +# +# +# +# +# +# +# diff --git a/models/face_parsing/makeup.py b/models/face_parsing/makeup.py new file mode 100644 index 0000000000000000000000000000000000000000..b03f1414a7463d056ac1452743e2ad8b272c4796 --- /dev/null +++ b/models/face_parsing/makeup.py @@ -0,0 +1,130 @@ +import cv2 +import os +import numpy as np +from skimage.filters import gaussian + + +def sharpen(img): + img = img * 1.0 + gauss_out = gaussian(img, sigma=5, multichannel=True) + + alpha = 1.5 + img_out = (img - gauss_out) * alpha + img + + img_out = img_out / 255.0 + + mask_1 = img_out < 0 + mask_2 = img_out > 1 + + img_out = img_out * (1 - mask_1) + img_out = img_out * (1 - mask_2) + mask_2 + img_out = np.clip(img_out, 0, 1) + img_out = img_out * 255 + return np.array(img_out, dtype=np.uint8) + + +def hair(image, parsing, part=17, color=[230, 50, 20]): + b, g, r = color #[10, 50, 250] # [10, 250, 10] + tar_color = np.zeros_like(image) + tar_color[:, :, 0] = b + tar_color[:, :, 1] = g + tar_color[:, :, 2] = r + + image_hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV) + tar_hsv = cv2.cvtColor(tar_color, cv2.COLOR_BGR2HSV) + + if part == 12 or part == 13: + image_hsv[:, :, 0:2] = tar_hsv[:, :, 0:2] + else: + image_hsv[:, :, 0:1] = tar_hsv[:, :, 0:1] + + changed = cv2.cvtColor(image_hsv, cv2.COLOR_HSV2BGR) + + if part == 17: + changed = sharpen(changed) + + changed[parsing != part] = image[parsing != part] + # changed = cv2.resize(changed, (512, 512)) + return changed + +# +# def lip(image, parsing, part=17, color=[230, 50, 20]): +# b, g, r = color #[10, 50, 250] # [10, 250, 10] +# tar_color = np.zeros_like(image) +# tar_color[:, :, 0] = b +# tar_color[:, :, 1] = g +# tar_color[:, :, 2] = r +# +# image_lab = cv2.cvtColor(image, cv2.COLOR_BGR2Lab) +# il, ia, ib = cv2.split(image_lab) +# +# tar_lab = cv2.cvtColor(tar_color, cv2.COLOR_BGR2Lab) +# tl, ta, tb = cv2.split(tar_lab) +# +# image_lab[:, :, 0] = np.clip(il - np.mean(il) + tl, 0, 100) +# image_lab[:, :, 1] = np.clip(ia - np.mean(ia) + ta, -127, 128) +# image_lab[:, :, 2] = np.clip(ib - np.mean(ib) + tb, -127, 128) +# +# +# changed = cv2.cvtColor(image_lab, cv2.COLOR_Lab2BGR) +# +# if part == 17: +# changed = sharpen(changed) +# +# changed[parsing != part] = image[parsing != part] +# # changed = cv2.resize(changed, (512, 512)) +# return changed + + +if __name__ == '__main__': + # 1 face + # 10 nose + # 11 teeth + # 12 upper lip + # 13 lower lip + # 17 hair + num = 116 + table = { + 'hair': 17, + 'upper_lip': 12, + 'lower_lip': 13 + } + image_path = '/home/zll/data/CelebAMask-HQ/test-img/{}.jpg'.format(num) + parsing_path = 'res/test_res/{}.png'.format(num) + + image = cv2.imread(image_path) + ori = image.copy() + parsing = np.array(cv2.imread(parsing_path, 0)) + parsing = cv2.resize(parsing, image.shape[0:2], interpolation=cv2.INTER_NEAREST) + + parts = [table['hair'], table['upper_lip'], table['lower_lip']] + # colors = [[20, 20, 200], [100, 100, 230], [100, 100, 230]] + colors = [[100, 200, 100]] + for part, color in zip(parts, colors): + image = hair(image, parsing, part, color) + cv2.imwrite('res/makeup/116_ori.png', cv2.resize(ori, (512, 512))) + cv2.imwrite('res/makeup/116_2.png', cv2.resize(image, (512, 512))) + + cv2.imshow('image', cv2.resize(ori, (512, 512))) + cv2.imshow('color', cv2.resize(image, (512, 512))) + + # cv2.imshow('image', ori) + # cv2.imshow('color', image) + + cv2.waitKey(0) + cv2.destroyAllWindows() + + + + + + + + + + + + + + + diff --git a/models/face_parsing/model.py b/models/face_parsing/model.py new file mode 100644 index 0000000000000000000000000000000000000000..e8c3dc86c1e528767f28ad5461de41be40eaf0ac --- /dev/null +++ b/models/face_parsing/model.py @@ -0,0 +1,289 @@ +#!/usr/bin/python +# -*- encoding: utf-8 -*- + + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torchvision + +from .resnet import Resnet18 +# from modules.bn import InPlaceABNSync as BatchNorm2d + +import numpy as np + +seg_mean = torch.from_numpy(np.array([[0.485, 0.456, 0.406]])).float().cuda().reshape(1,3,1,1) +seg_std = torch.from_numpy(np.array([[0.229, 0.224, 0.225]])).float().cuda().reshape(1,3,1,1) +seg_criterion = nn.CrossEntropyLoss() + + +class ConvBNReLU(nn.Module): + def __init__(self, in_chan, out_chan, ks=3, stride=1, padding=1, *args, **kwargs): + super(ConvBNReLU, self).__init__() + self.conv = nn.Conv2d(in_chan, + out_chan, + kernel_size = ks, + stride = stride, + padding = padding, + bias = False) + self.bn = nn.BatchNorm2d(out_chan) + self.init_weight() + + def forward(self, x): + x = self.conv(x) + x = F.relu(self.bn(x)) + return x + + def init_weight(self): + for ly in self.children(): + if isinstance(ly, nn.Conv2d): + nn.init.kaiming_normal_(ly.weight, a=1) + if not ly.bias is None: nn.init.constant_(ly.bias, 0) + +class BiSeNetOutput(nn.Module): + def __init__(self, in_chan, mid_chan, n_classes, *args, **kwargs): + super(BiSeNetOutput, self).__init__() + self.conv = ConvBNReLU(in_chan, mid_chan, ks=3, stride=1, padding=1) + self.conv_out = nn.Conv2d(mid_chan, n_classes, kernel_size=1, bias=False) + self.init_weight() + + def forward(self, x): + x = self.conv(x) + x = self.conv_out(x) + return x + + def init_weight(self): + for ly in self.children(): + if isinstance(ly, nn.Conv2d): + nn.init.kaiming_normal_(ly.weight, a=1) + if not ly.bias is None: nn.init.constant_(ly.bias, 0) + + def get_params(self): + wd_params, nowd_params = [], [] + for name, module in self.named_modules(): + if isinstance(module, nn.Linear) or isinstance(module, nn.Conv2d): + wd_params.append(module.weight) + if not module.bias is None: + nowd_params.append(module.bias) + elif isinstance(module, nn.BatchNorm2d): + nowd_params += list(module.parameters()) + return wd_params, nowd_params + + +class AttentionRefinementModule(nn.Module): + def __init__(self, in_chan, out_chan, *args, **kwargs): + super(AttentionRefinementModule, self).__init__() + self.conv = ConvBNReLU(in_chan, out_chan, ks=3, stride=1, padding=1) + self.conv_atten = nn.Conv2d(out_chan, out_chan, kernel_size= 1, bias=False) + self.bn_atten = nn.BatchNorm2d(out_chan) + self.sigmoid_atten = nn.Sigmoid() + self.init_weight() + + def forward(self, x): + feat = self.conv(x) + atten = F.avg_pool2d(feat, feat.size()[2:]) + atten = self.conv_atten(atten) + atten = self.bn_atten(atten) + atten = self.sigmoid_atten(atten) + out = torch.mul(feat, atten) + return out + + def init_weight(self): + for ly in self.children(): + if isinstance(ly, nn.Conv2d): + nn.init.kaiming_normal_(ly.weight, a=1) + if not ly.bias is None: nn.init.constant_(ly.bias, 0) + + +class ContextPath(nn.Module): + def __init__(self, *args, **kwargs): + super(ContextPath, self).__init__() + self.resnet = Resnet18() + self.arm16 = AttentionRefinementModule(256, 128) + self.arm32 = AttentionRefinementModule(512, 128) + self.conv_head32 = ConvBNReLU(128, 128, ks=3, stride=1, padding=1) + self.conv_head16 = ConvBNReLU(128, 128, ks=3, stride=1, padding=1) + self.conv_avg = ConvBNReLU(512, 128, ks=1, stride=1, padding=0) + + self.init_weight() + + def forward(self, x): + H0, W0 = x.size()[2:] + feat8, feat16, feat32 = self.resnet(x) + H8, W8 = feat8.size()[2:] + H16, W16 = feat16.size()[2:] + H32, W32 = feat32.size()[2:] + + avg = F.avg_pool2d(feat32, feat32.size()[2:]) + avg = self.conv_avg(avg) + avg_up = F.interpolate(avg, (H32, W32), mode='nearest') + + feat32_arm = self.arm32(feat32) + feat32_sum = feat32_arm + avg_up + feat32_up = F.interpolate(feat32_sum, (H16, W16), mode='nearest') + feat32_up = self.conv_head32(feat32_up) + + feat16_arm = self.arm16(feat16) + feat16_sum = feat16_arm + feat32_up + feat16_up = F.interpolate(feat16_sum, (H8, W8), mode='nearest') + feat16_up = self.conv_head16(feat16_up) + + return feat8, feat16_up, feat32_up # x8, x8, x16 + + def init_weight(self): + for ly in self.children(): + if isinstance(ly, nn.Conv2d): + nn.init.kaiming_normal_(ly.weight, a=1) + if not ly.bias is None: nn.init.constant_(ly.bias, 0) + + def get_params(self): + wd_params, nowd_params = [], [] + for name, module in self.named_modules(): + if isinstance(module, (nn.Linear, nn.Conv2d)): + wd_params.append(module.weight) + if not module.bias is None: + nowd_params.append(module.bias) + elif isinstance(module, nn.BatchNorm2d): + nowd_params += list(module.parameters()) + return wd_params, nowd_params + + +### This is not used, since I replace this with the resnet feature with the same size +class SpatialPath(nn.Module): + def __init__(self, *args, **kwargs): + super(SpatialPath, self).__init__() + self.conv1 = ConvBNReLU(3, 64, ks=7, stride=2, padding=3) + self.conv2 = ConvBNReLU(64, 64, ks=3, stride=2, padding=1) + self.conv3 = ConvBNReLU(64, 64, ks=3, stride=2, padding=1) + self.conv_out = ConvBNReLU(64, 128, ks=1, stride=1, padding=0) + self.init_weight() + + def forward(self, x): + feat = self.conv1(x) + feat = self.conv2(feat) + feat = self.conv3(feat) + feat = self.conv_out(feat) + return feat + + def init_weight(self): + for ly in self.children(): + if isinstance(ly, nn.Conv2d): + nn.init.kaiming_normal_(ly.weight, a=1) + if not ly.bias is None: nn.init.constant_(ly.bias, 0) + + def get_params(self): + wd_params, nowd_params = [], [] + for name, module in self.named_modules(): + if isinstance(module, nn.Linear) or isinstance(module, nn.Conv2d): + wd_params.append(module.weight) + if not module.bias is None: + nowd_params.append(module.bias) + elif isinstance(module, nn.BatchNorm2d): + nowd_params += list(module.parameters()) + return wd_params, nowd_params + + +class FeatureFusionModule(nn.Module): + def __init__(self, in_chan, out_chan, *args, **kwargs): + super(FeatureFusionModule, self).__init__() + self.convblk = ConvBNReLU(in_chan, out_chan, ks=1, stride=1, padding=0) + self.conv1 = nn.Conv2d(out_chan, + out_chan//4, + kernel_size = 1, + stride = 1, + padding = 0, + bias = False) + self.conv2 = nn.Conv2d(out_chan//4, + out_chan, + kernel_size = 1, + stride = 1, + padding = 0, + bias = False) + self.relu = nn.ReLU(inplace=True) + self.sigmoid = nn.Sigmoid() + self.init_weight() + + def forward(self, fsp, fcp): + fcat = torch.cat([fsp, fcp], dim=1) + feat = self.convblk(fcat) + atten = F.avg_pool2d(feat, feat.size()[2:]) + atten = self.conv1(atten) + atten = self.relu(atten) + atten = self.conv2(atten) + atten = self.sigmoid(atten) + feat_atten = torch.mul(feat, atten) + feat_out = feat_atten + feat + return feat_out + + def init_weight(self): + for ly in self.children(): + if isinstance(ly, nn.Conv2d): + nn.init.kaiming_normal_(ly.weight, a=1) + if not ly.bias is None: nn.init.constant_(ly.bias, 0) + + def get_params(self): + wd_params, nowd_params = [], [] + for name, module in self.named_modules(): + if isinstance(module, nn.Linear) or isinstance(module, nn.Conv2d): + wd_params.append(module.weight) + if not module.bias is None: + nowd_params.append(module.bias) + elif isinstance(module, nn.BatchNorm2d): + nowd_params += list(module.parameters()) + return wd_params, nowd_params + + +class BiSeNet(nn.Module): + def __init__(self, n_classes, *args, **kwargs): + super(BiSeNet, self).__init__() + self.cp = ContextPath() + ## here self.sp is deleted + self.ffm = FeatureFusionModule(256, 256) + self.conv_out = BiSeNetOutput(256, 256, n_classes) + self.conv_out16 = BiSeNetOutput(128, 64, n_classes) + self.conv_out32 = BiSeNetOutput(128, 64, n_classes) + self.init_weight() + + def forward(self, x): + H, W = x.size()[2:] + feat_res8, feat_cp8, feat_cp16 = self.cp(x) # here return res3b1 feature + feat_sp = feat_res8 # use res3b1 feature to replace spatial path feature + feat_fuse = self.ffm(feat_sp, feat_cp8) + + feat_out = self.conv_out(feat_fuse) + feat_out16 = self.conv_out16(feat_cp8) + feat_out32 = self.conv_out32(feat_cp16) + + feat_out = F.interpolate(feat_out, (H, W), mode='bilinear', align_corners=True) + feat_out16 = F.interpolate(feat_out16, (H, W), mode='bilinear', align_corners=True) + feat_out32 = F.interpolate(feat_out32, (H, W), mode='bilinear', align_corners=True) + return feat_out, feat_out16, feat_out32 + + def init_weight(self): + for ly in self.children(): + if isinstance(ly, nn.Conv2d): + nn.init.kaiming_normal_(ly.weight, a=1) + if not ly.bias is None: nn.init.constant_(ly.bias, 0) + + def get_params(self): + wd_params, nowd_params, lr_mul_wd_params, lr_mul_nowd_params = [], [], [], [] + for name, child in self.named_children(): + child_wd_params, child_nowd_params = child.get_params() + if isinstance(child, FeatureFusionModule) or isinstance(child, BiSeNetOutput): + lr_mul_wd_params += child_wd_params + lr_mul_nowd_params += child_nowd_params + else: + wd_params += child_wd_params + nowd_params += child_nowd_params + return wd_params, nowd_params, lr_mul_wd_params, lr_mul_nowd_params + + +if __name__ == "__main__": + net = BiSeNet(16) + net.cuda() + net.eval() + in_ten = torch.randn(16, 3, 640, 480).cuda() + out, out16, out32 = net(in_ten) + print(out.shape) + + net.get_params() diff --git a/models/face_parsing/modules/__init__.py b/models/face_parsing/modules/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..8a098dee5911f3613d320d23db37bc401cf57fa4 --- /dev/null +++ b/models/face_parsing/modules/__init__.py @@ -0,0 +1,5 @@ +from .bn import ABN, InPlaceABN, InPlaceABNSync +from .functions import ACT_RELU, ACT_LEAKY_RELU, ACT_ELU, ACT_NONE +from .misc import GlobalAvgPool2d, SingleGPU +from .residual import IdentityResidualBlock +from .dense import DenseModule diff --git a/models/face_parsing/modules/bn.py b/models/face_parsing/modules/bn.py new file mode 100644 index 0000000000000000000000000000000000000000..cd3928bccfd3f70233414d837876b323217864c8 --- /dev/null +++ b/models/face_parsing/modules/bn.py @@ -0,0 +1,130 @@ +import torch +import torch.nn as nn +import torch.nn.functional as functional + +try: + from queue import Queue +except ImportError: + from Queue import Queue + +from .functions import * + + +class ABN(nn.Module): + """Activated Batch Normalization + + This gathers a `BatchNorm2d` and an activation function in a single module + """ + + def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=True, activation="leaky_relu", slope=0.01): + """Creates an Activated Batch Normalization module + + Parameters + ---------- + num_features : int + Number of feature channels in the input and output. + eps : float + Small constant to prevent numerical issues. + momentum : float + Momentum factor applied to compute running statistics as. + affine : bool + If `True` apply learned scale and shift transformation after normalization. + activation : str + Name of the activation functions, one of: `leaky_relu`, `elu` or `none`. + slope : float + Negative slope for the `leaky_relu` activation. + """ + super(ABN, self).__init__() + self.num_features = num_features + self.affine = affine + self.eps = eps + self.momentum = momentum + self.activation = activation + self.slope = slope + if self.affine: + self.weight = nn.Parameter(torch.ones(num_features)) + self.bias = nn.Parameter(torch.zeros(num_features)) + else: + self.register_parameter('weight', None) + self.register_parameter('bias', None) + self.register_buffer('running_mean', torch.zeros(num_features)) + self.register_buffer('running_var', torch.ones(num_features)) + self.reset_parameters() + + def reset_parameters(self): + nn.init.constant_(self.running_mean, 0) + nn.init.constant_(self.running_var, 1) + if self.affine: + nn.init.constant_(self.weight, 1) + nn.init.constant_(self.bias, 0) + + def forward(self, x): + x = functional.batch_norm(x, self.running_mean, self.running_var, self.weight, self.bias, + self.training, self.momentum, self.eps) + + if self.activation == ACT_RELU: + return functional.relu(x, inplace=True) + elif self.activation == ACT_LEAKY_RELU: + return functional.leaky_relu(x, negative_slope=self.slope, inplace=True) + elif self.activation == ACT_ELU: + return functional.elu(x, inplace=True) + else: + return x + + def __repr__(self): + rep = '{name}({num_features}, eps={eps}, momentum={momentum},' \ + ' affine={affine}, activation={activation}' + if self.activation == "leaky_relu": + rep += ', slope={slope})' + else: + rep += ')' + return rep.format(name=self.__class__.__name__, **self.__dict__) + + +class InPlaceABN(ABN): + """InPlace Activated Batch Normalization""" + + def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=True, activation="leaky_relu", slope=0.01): + """Creates an InPlace Activated Batch Normalization module + + Parameters + ---------- + num_features : int + Number of feature channels in the input and output. + eps : float + Small constant to prevent numerical issues. + momentum : float + Momentum factor applied to compute running statistics as. + affine : bool + If `True` apply learned scale and shift transformation after normalization. + activation : str + Name of the activation functions, one of: `leaky_relu`, `elu` or `none`. + slope : float + Negative slope for the `leaky_relu` activation. + """ + super(InPlaceABN, self).__init__(num_features, eps, momentum, affine, activation, slope) + + def forward(self, x): + return inplace_abn(x, self.weight, self.bias, self.running_mean, self.running_var, + self.training, self.momentum, self.eps, self.activation, self.slope) + + +class InPlaceABNSync(ABN): + """InPlace Activated Batch Normalization with cross-GPU synchronization + This assumes that it will be replicated across GPUs using the same mechanism as in `nn.DistributedDataParallel`. + """ + + def forward(self, x): + return inplace_abn_sync(x, self.weight, self.bias, self.running_mean, self.running_var, + self.training, self.momentum, self.eps, self.activation, self.slope) + + def __repr__(self): + rep = '{name}({num_features}, eps={eps}, momentum={momentum},' \ + ' affine={affine}, activation={activation}' + if self.activation == "leaky_relu": + rep += ', slope={slope})' + else: + rep += ')' + return rep.format(name=self.__class__.__name__, **self.__dict__) + + diff --git a/models/face_parsing/modules/deeplab.py b/models/face_parsing/modules/deeplab.py new file mode 100644 index 0000000000000000000000000000000000000000..fd25b78369b27ef02c183a0b17b9bf8354c5f7c3 --- /dev/null +++ b/models/face_parsing/modules/deeplab.py @@ -0,0 +1,84 @@ +import torch +import torch.nn as nn +import torch.nn.functional as functional + +from models._util import try_index +from .bn import ABN + + +class DeeplabV3(nn.Module): + def __init__(self, + in_channels, + out_channels, + hidden_channels=256, + dilations=(12, 24, 36), + norm_act=ABN, + pooling_size=None): + super(DeeplabV3, self).__init__() + self.pooling_size = pooling_size + + self.map_convs = nn.ModuleList([ + nn.Conv2d(in_channels, hidden_channels, 1, bias=False), + nn.Conv2d(in_channels, hidden_channels, 3, bias=False, dilation=dilations[0], padding=dilations[0]), + nn.Conv2d(in_channels, hidden_channels, 3, bias=False, dilation=dilations[1], padding=dilations[1]), + nn.Conv2d(in_channels, hidden_channels, 3, bias=False, dilation=dilations[2], padding=dilations[2]) + ]) + self.map_bn = norm_act(hidden_channels * 4) + + self.global_pooling_conv = nn.Conv2d(in_channels, hidden_channels, 1, bias=False) + self.global_pooling_bn = norm_act(hidden_channels) + + self.red_conv = nn.Conv2d(hidden_channels * 4, out_channels, 1, bias=False) + self.pool_red_conv = nn.Conv2d(hidden_channels, out_channels, 1, bias=False) + self.red_bn = norm_act(out_channels) + + self.reset_parameters(self.map_bn.activation, self.map_bn.slope) + + def reset_parameters(self, activation, slope): + gain = nn.init.calculate_gain(activation, slope) + for m in self.modules(): + if isinstance(m, nn.Conv2d): + nn.init.xavier_normal_(m.weight.data, gain) + if hasattr(m, "bias") and m.bias is not None: + nn.init.constant_(m.bias, 0) + elif isinstance(m, ABN): + if hasattr(m, "weight") and m.weight is not None: + nn.init.constant_(m.weight, 1) + if hasattr(m, "bias") and m.bias is not None: + nn.init.constant_(m.bias, 0) + + def forward(self, x): + # Map convolutions + out = torch.cat([m(x) for m in self.map_convs], dim=1) + out = self.map_bn(out) + out = self.red_conv(out) + + # Global pooling + pool = self._global_pooling(x) + pool = self.global_pooling_conv(pool) + pool = self.global_pooling_bn(pool) + pool = self.pool_red_conv(pool) + if self.training or self.pooling_size is None: + pool = pool.repeat(1, 1, x.size(2), x.size(3)) + + out += pool + out = self.red_bn(out) + return out + + def _global_pooling(self, x): + if self.training or self.pooling_size is None: + pool = x.view(x.size(0), x.size(1), -1).mean(dim=-1) + pool = pool.view(x.size(0), x.size(1), 1, 1) + else: + pooling_size = (min(try_index(self.pooling_size, 0), x.shape[2]), + min(try_index(self.pooling_size, 1), x.shape[3])) + padding = ( + (pooling_size[1] - 1) // 2, + (pooling_size[1] - 1) // 2 if pooling_size[1] % 2 == 1 else (pooling_size[1] - 1) // 2 + 1, + (pooling_size[0] - 1) // 2, + (pooling_size[0] - 1) // 2 if pooling_size[0] % 2 == 1 else (pooling_size[0] - 1) // 2 + 1 + ) + + pool = functional.avg_pool2d(x, pooling_size, stride=1) + pool = functional.pad(pool, pad=padding, mode="replicate") + return pool diff --git a/models/face_parsing/modules/dense.py b/models/face_parsing/modules/dense.py new file mode 100644 index 0000000000000000000000000000000000000000..9638d6e86d2ae838550fefa9002a984af52e6cc8 --- /dev/null +++ b/models/face_parsing/modules/dense.py @@ -0,0 +1,42 @@ +from collections import OrderedDict + +import torch +import torch.nn as nn + +from .bn import ABN + + +class DenseModule(nn.Module): + def __init__(self, in_channels, growth, layers, bottleneck_factor=4, norm_act=ABN, dilation=1): + super(DenseModule, self).__init__() + self.in_channels = in_channels + self.growth = growth + self.layers = layers + + self.convs1 = nn.ModuleList() + self.convs3 = nn.ModuleList() + for i in range(self.layers): + self.convs1.append(nn.Sequential(OrderedDict([ + ("bn", norm_act(in_channels)), + ("conv", nn.Conv2d(in_channels, self.growth * bottleneck_factor, 1, bias=False)) + ]))) + self.convs3.append(nn.Sequential(OrderedDict([ + ("bn", norm_act(self.growth * bottleneck_factor)), + ("conv", nn.Conv2d(self.growth * bottleneck_factor, self.growth, 3, padding=dilation, bias=False, + dilation=dilation)) + ]))) + in_channels += self.growth + + @property + def out_channels(self): + return self.in_channels + self.growth * self.layers + + def forward(self, x): + inputs = [x] + for i in range(self.layers): + x = torch.cat(inputs, dim=1) + x = self.convs1[i](x) + x = self.convs3[i](x) + inputs += [x] + + return torch.cat(inputs, dim=1) diff --git a/models/face_parsing/modules/functions.py b/models/face_parsing/modules/functions.py new file mode 100644 index 0000000000000000000000000000000000000000..093615ff4f383e95712c96b57286338ec3b28f3b --- /dev/null +++ b/models/face_parsing/modules/functions.py @@ -0,0 +1,234 @@ +from os import path +import torch +import torch.distributed as dist +import torch.autograd as autograd +import torch.cuda.comm as comm +from torch.autograd.function import once_differentiable +from torch.utils.cpp_extension import load + +_src_path = path.join(path.dirname(path.abspath(__file__)), "src") +_backend = load(name="inplace_abn", + extra_cflags=["-O3"], + sources=[path.join(_src_path, f) for f in [ + "inplace_abn.cpp", + "inplace_abn_cpu.cpp", + "inplace_abn_cuda.cu", + "inplace_abn_cuda_half.cu" + ]], + extra_cuda_cflags=["--expt-extended-lambda"]) + +# Activation names +ACT_RELU = "relu" +ACT_LEAKY_RELU = "leaky_relu" +ACT_ELU = "elu" +ACT_NONE = "none" + + +def _check(fn, *args, **kwargs): + success = fn(*args, **kwargs) + if not success: + raise RuntimeError("CUDA Error encountered in {}".format(fn)) + + +def _broadcast_shape(x): + out_size = [] + for i, s in enumerate(x.size()): + if i != 1: + out_size.append(1) + else: + out_size.append(s) + return out_size + + +def _reduce(x): + if len(x.size()) == 2: + return x.sum(dim=0) + else: + n, c = x.size()[0:2] + return x.contiguous().view((n, c, -1)).sum(2).sum(0) + + +def _count_samples(x): + count = 1 + for i, s in enumerate(x.size()): + if i != 1: + count *= s + return count + + +def _act_forward(ctx, x): + if ctx.activation == ACT_LEAKY_RELU: + _backend.leaky_relu_forward(x, ctx.slope) + elif ctx.activation == ACT_ELU: + _backend.elu_forward(x) + elif ctx.activation == ACT_NONE: + pass + + +def _act_backward(ctx, x, dx): + if ctx.activation == ACT_LEAKY_RELU: + _backend.leaky_relu_backward(x, dx, ctx.slope) + elif ctx.activation == ACT_ELU: + _backend.elu_backward(x, dx) + elif ctx.activation == ACT_NONE: + pass + + +class InPlaceABN(autograd.Function): + @staticmethod + def forward(ctx, x, weight, bias, running_mean, running_var, + training=True, momentum=0.1, eps=1e-05, activation=ACT_LEAKY_RELU, slope=0.01): + # Save context + ctx.training = training + ctx.momentum = momentum + ctx.eps = eps + ctx.activation = activation + ctx.slope = slope + ctx.affine = weight is not None and bias is not None + + # Prepare inputs + count = _count_samples(x) + x = x.contiguous() + weight = weight.contiguous() if ctx.affine else x.new_empty(0) + bias = bias.contiguous() if ctx.affine else x.new_empty(0) + + if ctx.training: + mean, var = _backend.mean_var(x) + + # Update running stats + running_mean.mul_((1 - ctx.momentum)).add_(ctx.momentum * mean) + running_var.mul_((1 - ctx.momentum)).add_(ctx.momentum * var * count / (count - 1)) + + # Mark in-place modified tensors + ctx.mark_dirty(x, running_mean, running_var) + else: + mean, var = running_mean.contiguous(), running_var.contiguous() + ctx.mark_dirty(x) + + # BN forward + activation + _backend.forward(x, mean, var, weight, bias, ctx.affine, ctx.eps) + _act_forward(ctx, x) + + # Output + ctx.var = var + ctx.save_for_backward(x, var, weight, bias) + return x + + @staticmethod + @once_differentiable + def backward(ctx, dz): + z, var, weight, bias = ctx.saved_tensors + dz = dz.contiguous() + + # Undo activation + _act_backward(ctx, z, dz) + + if ctx.training: + edz, eydz = _backend.edz_eydz(z, dz, weight, bias, ctx.affine, ctx.eps) + else: + # TODO: implement simplified CUDA backward for inference mode + edz = dz.new_zeros(dz.size(1)) + eydz = dz.new_zeros(dz.size(1)) + + dx = _backend.backward(z, dz, var, weight, bias, edz, eydz, ctx.affine, ctx.eps) + dweight = eydz * weight.sign() if ctx.affine else None + dbias = edz if ctx.affine else None + + return dx, dweight, dbias, None, None, None, None, None, None, None + +class InPlaceABNSync(autograd.Function): + @classmethod + def forward(cls, ctx, x, weight, bias, running_mean, running_var, + training=True, momentum=0.1, eps=1e-05, activation=ACT_LEAKY_RELU, slope=0.01, equal_batches=True): + # Save context + ctx.training = training + ctx.momentum = momentum + ctx.eps = eps + ctx.activation = activation + ctx.slope = slope + ctx.affine = weight is not None and bias is not None + + # Prepare inputs + ctx.world_size = dist.get_world_size() if dist.is_initialized() else 1 + + #count = _count_samples(x) + batch_size = x.new_tensor([x.shape[0]],dtype=torch.long) + + x = x.contiguous() + weight = weight.contiguous() if ctx.affine else x.new_empty(0) + bias = bias.contiguous() if ctx.affine else x.new_empty(0) + + if ctx.training: + mean, var = _backend.mean_var(x) + if ctx.world_size>1: + # get global batch size + if equal_batches: + batch_size *= ctx.world_size + else: + dist.all_reduce(batch_size, dist.ReduceOp.SUM) + + ctx.factor = x.shape[0]/float(batch_size.item()) + + mean_all = mean.clone() * ctx.factor + dist.all_reduce(mean_all, dist.ReduceOp.SUM) + + var_all = (var + (mean - mean_all) ** 2) * ctx.factor + dist.all_reduce(var_all, dist.ReduceOp.SUM) + + mean = mean_all + var = var_all + + # Update running stats + running_mean.mul_((1 - ctx.momentum)).add_(ctx.momentum * mean) + count = batch_size.item() * x.view(x.shape[0],x.shape[1],-1).shape[-1] + running_var.mul_((1 - ctx.momentum)).add_(ctx.momentum * var * (float(count) / (count - 1))) + + # Mark in-place modified tensors + ctx.mark_dirty(x, running_mean, running_var) + else: + mean, var = running_mean.contiguous(), running_var.contiguous() + ctx.mark_dirty(x) + + # BN forward + activation + _backend.forward(x, mean, var, weight, bias, ctx.affine, ctx.eps) + _act_forward(ctx, x) + + # Output + ctx.var = var + ctx.save_for_backward(x, var, weight, bias) + return x + + @staticmethod + @once_differentiable + def backward(ctx, dz): + z, var, weight, bias = ctx.saved_tensors + dz = dz.contiguous() + + # Undo activation + _act_backward(ctx, z, dz) + + if ctx.training: + edz, eydz = _backend.edz_eydz(z, dz, weight, bias, ctx.affine, ctx.eps) + edz_local = edz.clone() + eydz_local = eydz.clone() + + if ctx.world_size>1: + edz *= ctx.factor + dist.all_reduce(edz, dist.ReduceOp.SUM) + + eydz *= ctx.factor + dist.all_reduce(eydz, dist.ReduceOp.SUM) + else: + edz_local = edz = dz.new_zeros(dz.size(1)) + eydz_local = eydz = dz.new_zeros(dz.size(1)) + + dx = _backend.backward(z, dz, var, weight, bias, edz, eydz, ctx.affine, ctx.eps) + dweight = eydz_local * weight.sign() if ctx.affine else None + dbias = edz_local if ctx.affine else None + + return dx, dweight, dbias, None, None, None, None, None, None, None + +inplace_abn = InPlaceABN.apply +inplace_abn_sync = InPlaceABNSync.apply + +__all__ = ["inplace_abn", "inplace_abn_sync", "ACT_RELU", "ACT_LEAKY_RELU", "ACT_ELU", "ACT_NONE"] diff --git a/models/face_parsing/modules/misc.py b/models/face_parsing/modules/misc.py new file mode 100644 index 0000000000000000000000000000000000000000..3c50b69b38c950801baacba8b3684ffd23aef08b --- /dev/null +++ b/models/face_parsing/modules/misc.py @@ -0,0 +1,21 @@ +import torch.nn as nn +import torch +import torch.distributed as dist + +class GlobalAvgPool2d(nn.Module): + def __init__(self): + """Global average pooling over the input's spatial dimensions""" + super(GlobalAvgPool2d, self).__init__() + + def forward(self, inputs): + in_size = inputs.size() + return inputs.view((in_size[0], in_size[1], -1)).mean(dim=2) + +class SingleGPU(nn.Module): + def __init__(self, module): + super(SingleGPU, self).__init__() + self.module=module + + def forward(self, input): + return self.module(input.cuda(non_blocking=True)) + diff --git a/models/face_parsing/modules/residual.py b/models/face_parsing/modules/residual.py new file mode 100644 index 0000000000000000000000000000000000000000..b7d51ad274f3841813c1584a0ceb60ce58979d94 --- /dev/null +++ b/models/face_parsing/modules/residual.py @@ -0,0 +1,88 @@ +from collections import OrderedDict + +import torch.nn as nn + +from .bn import ABN + + +class IdentityResidualBlock(nn.Module): + def __init__(self, + in_channels, + channels, + stride=1, + dilation=1, + groups=1, + norm_act=ABN, + dropout=None): + """Configurable identity-mapping residual block + + Parameters + ---------- + in_channels : int + Number of input channels. + channels : list of int + Number of channels in the internal feature maps. Can either have two or three elements: if three construct + a residual block with two `3 x 3` convolutions, otherwise construct a bottleneck block with `1 x 1`, then + `3 x 3` then `1 x 1` convolutions. + stride : int + Stride of the first `3 x 3` convolution + dilation : int + Dilation to apply to the `3 x 3` convolutions. + groups : int + Number of convolution groups. This is used to create ResNeXt-style blocks and is only compatible with + bottleneck blocks. + norm_act : callable + Function to create normalization / activation Module. + dropout: callable + Function to create Dropout Module. + """ + super(IdentityResidualBlock, self).__init__() + + # Check parameters for inconsistencies + if len(channels) != 2 and len(channels) != 3: + raise ValueError("channels must contain either two or three values") + if len(channels) == 2 and groups != 1: + raise ValueError("groups > 1 are only valid if len(channels) == 3") + + is_bottleneck = len(channels) == 3 + need_proj_conv = stride != 1 or in_channels != channels[-1] + + self.bn1 = norm_act(in_channels) + if not is_bottleneck: + layers = [ + ("conv1", nn.Conv2d(in_channels, channels[0], 3, stride=stride, padding=dilation, bias=False, + dilation=dilation)), + ("bn2", norm_act(channels[0])), + ("conv2", nn.Conv2d(channels[0], channels[1], 3, stride=1, padding=dilation, bias=False, + dilation=dilation)) + ] + if dropout is not None: + layers = layers[0:2] + [("dropout", dropout())] + layers[2:] + else: + layers = [ + ("conv1", nn.Conv2d(in_channels, channels[0], 1, stride=stride, padding=0, bias=False)), + ("bn2", norm_act(channels[0])), + ("conv2", nn.Conv2d(channels[0], channels[1], 3, stride=1, padding=dilation, bias=False, + groups=groups, dilation=dilation)), + ("bn3", norm_act(channels[1])), + ("conv3", nn.Conv2d(channels[1], channels[2], 1, stride=1, padding=0, bias=False)) + ] + if dropout is not None: + layers = layers[0:4] + [("dropout", dropout())] + layers[4:] + self.convs = nn.Sequential(OrderedDict(layers)) + + if need_proj_conv: + self.proj_conv = nn.Conv2d(in_channels, channels[-1], 1, stride=stride, padding=0, bias=False) + + def forward(self, x): + if hasattr(self, "proj_conv"): + bn1 = self.bn1(x) + shortcut = self.proj_conv(bn1) + else: + shortcut = x.clone() + bn1 = self.bn1(x) + + out = self.convs(bn1) + out.add_(shortcut) + + return out diff --git a/models/face_parsing/modules/src/checks.h b/models/face_parsing/modules/src/checks.h new file mode 100644 index 0000000000000000000000000000000000000000..e761a6fe34d0789815b588eba7e3726026e0e868 --- /dev/null +++ b/models/face_parsing/modules/src/checks.h @@ -0,0 +1,15 @@ +#pragma once + +#include + +// Define AT_CHECK for old version of ATen where the same function was called AT_ASSERT +#ifndef AT_CHECK +#define AT_CHECK AT_ASSERT +#endif + +#define CHECK_CUDA(x) AT_CHECK((x).type().is_cuda(), #x " must be a CUDA tensor") +#define CHECK_CPU(x) AT_CHECK(!(x).type().is_cuda(), #x " must be a CPU tensor") +#define CHECK_CONTIGUOUS(x) AT_CHECK((x).is_contiguous(), #x " must be contiguous") + +#define CHECK_CUDA_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x) +#define CHECK_CPU_INPUT(x) CHECK_CPU(x); CHECK_CONTIGUOUS(x) \ No newline at end of file diff --git a/models/face_parsing/modules/src/inplace_abn.cpp b/models/face_parsing/modules/src/inplace_abn.cpp new file mode 100644 index 0000000000000000000000000000000000000000..0a6b1128cc20cbfc476134154e23e5869a92b856 --- /dev/null +++ b/models/face_parsing/modules/src/inplace_abn.cpp @@ -0,0 +1,95 @@ +#include + +#include + +#include "inplace_abn.h" + +std::vector mean_var(at::Tensor x) { + if (x.is_cuda()) { + if (x.type().scalarType() == at::ScalarType::Half) { + return mean_var_cuda_h(x); + } else { + return mean_var_cuda(x); + } + } else { + return mean_var_cpu(x); + } +} + +at::Tensor forward(at::Tensor x, at::Tensor mean, at::Tensor var, at::Tensor weight, at::Tensor bias, + bool affine, float eps) { + if (x.is_cuda()) { + if (x.type().scalarType() == at::ScalarType::Half) { + return forward_cuda_h(x, mean, var, weight, bias, affine, eps); + } else { + return forward_cuda(x, mean, var, weight, bias, affine, eps); + } + } else { + return forward_cpu(x, mean, var, weight, bias, affine, eps); + } +} + +std::vector edz_eydz(at::Tensor z, at::Tensor dz, at::Tensor weight, at::Tensor bias, + bool affine, float eps) { + if (z.is_cuda()) { + if (z.type().scalarType() == at::ScalarType::Half) { + return edz_eydz_cuda_h(z, dz, weight, bias, affine, eps); + } else { + return edz_eydz_cuda(z, dz, weight, bias, affine, eps); + } + } else { + return edz_eydz_cpu(z, dz, weight, bias, affine, eps); + } +} + +at::Tensor backward(at::Tensor z, at::Tensor dz, at::Tensor var, at::Tensor weight, at::Tensor bias, + at::Tensor edz, at::Tensor eydz, bool affine, float eps) { + if (z.is_cuda()) { + if (z.type().scalarType() == at::ScalarType::Half) { + return backward_cuda_h(z, dz, var, weight, bias, edz, eydz, affine, eps); + } else { + return backward_cuda(z, dz, var, weight, bias, edz, eydz, affine, eps); + } + } else { + return backward_cpu(z, dz, var, weight, bias, edz, eydz, affine, eps); + } +} + +void leaky_relu_forward(at::Tensor z, float slope) { + at::leaky_relu_(z, slope); +} + +void leaky_relu_backward(at::Tensor z, at::Tensor dz, float slope) { + if (z.is_cuda()) { + if (z.type().scalarType() == at::ScalarType::Half) { + return leaky_relu_backward_cuda_h(z, dz, slope); + } else { + return leaky_relu_backward_cuda(z, dz, slope); + } + } else { + return leaky_relu_backward_cpu(z, dz, slope); + } +} + +void elu_forward(at::Tensor z) { + at::elu_(z); +} + +void elu_backward(at::Tensor z, at::Tensor dz) { + if (z.is_cuda()) { + return elu_backward_cuda(z, dz); + } else { + return elu_backward_cpu(z, dz); + } +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("mean_var", &mean_var, "Mean and variance computation"); + m.def("forward", &forward, "In-place forward computation"); + m.def("edz_eydz", &edz_eydz, "First part of backward computation"); + m.def("backward", &backward, "Second part of backward computation"); + m.def("leaky_relu_forward", &leaky_relu_forward, "Leaky relu forward computation"); + m.def("leaky_relu_backward", &leaky_relu_backward, "Leaky relu backward computation and inversion"); + m.def("elu_forward", &elu_forward, "Elu forward computation"); + m.def("elu_backward", &elu_backward, "Elu backward computation and inversion"); +} diff --git a/models/face_parsing/modules/src/inplace_abn.h b/models/face_parsing/modules/src/inplace_abn.h new file mode 100644 index 0000000000000000000000000000000000000000..17afd1196449ecb6376f28961e54b55e1537492f --- /dev/null +++ b/models/face_parsing/modules/src/inplace_abn.h @@ -0,0 +1,88 @@ +#pragma once + +#include + +#include + +std::vector mean_var_cpu(at::Tensor x); +std::vector mean_var_cuda(at::Tensor x); +std::vector mean_var_cuda_h(at::Tensor x); + +at::Tensor forward_cpu(at::Tensor x, at::Tensor mean, at::Tensor var, at::Tensor weight, at::Tensor bias, + bool affine, float eps); +at::Tensor forward_cuda(at::Tensor x, at::Tensor mean, at::Tensor var, at::Tensor weight, at::Tensor bias, + bool affine, float eps); +at::Tensor forward_cuda_h(at::Tensor x, at::Tensor mean, at::Tensor var, at::Tensor weight, at::Tensor bias, + bool affine, float eps); + +std::vector edz_eydz_cpu(at::Tensor z, at::Tensor dz, at::Tensor weight, at::Tensor bias, + bool affine, float eps); +std::vector edz_eydz_cuda(at::Tensor z, at::Tensor dz, at::Tensor weight, at::Tensor bias, + bool affine, float eps); +std::vector edz_eydz_cuda_h(at::Tensor z, at::Tensor dz, at::Tensor weight, at::Tensor bias, + bool affine, float eps); + +at::Tensor backward_cpu(at::Tensor z, at::Tensor dz, at::Tensor var, at::Tensor weight, at::Tensor bias, + at::Tensor edz, at::Tensor eydz, bool affine, float eps); +at::Tensor backward_cuda(at::Tensor z, at::Tensor dz, at::Tensor var, at::Tensor weight, at::Tensor bias, + at::Tensor edz, at::Tensor eydz, bool affine, float eps); +at::Tensor backward_cuda_h(at::Tensor z, at::Tensor dz, at::Tensor var, at::Tensor weight, at::Tensor bias, + at::Tensor edz, at::Tensor eydz, bool affine, float eps); + +void leaky_relu_backward_cpu(at::Tensor z, at::Tensor dz, float slope); +void leaky_relu_backward_cuda(at::Tensor z, at::Tensor dz, float slope); +void leaky_relu_backward_cuda_h(at::Tensor z, at::Tensor dz, float slope); + +void elu_backward_cpu(at::Tensor z, at::Tensor dz); +void elu_backward_cuda(at::Tensor z, at::Tensor dz); + +static void get_dims(at::Tensor x, int64_t& num, int64_t& chn, int64_t& sp) { + num = x.size(0); + chn = x.size(1); + sp = 1; + for (int64_t i = 2; i < x.ndimension(); ++i) + sp *= x.size(i); +} + +/* + * Specialized CUDA reduction functions for BN + */ +#ifdef __CUDACC__ + +#include "utils/cuda.cuh" + +template +__device__ T reduce(Op op, int plane, int N, int S) { + T sum = (T)0; + for (int batch = 0; batch < N; ++batch) { + for (int x = threadIdx.x; x < S; x += blockDim.x) { + sum += op(batch, plane, x); + } + } + + // sum over NumThreads within a warp + sum = warpSum(sum); + + // 'transpose', and reduce within warp again + __shared__ T shared[32]; + __syncthreads(); + if (threadIdx.x % WARP_SIZE == 0) { + shared[threadIdx.x / WARP_SIZE] = sum; + } + if (threadIdx.x >= blockDim.x / WARP_SIZE && threadIdx.x < WARP_SIZE) { + // zero out the other entries in shared + shared[threadIdx.x] = (T)0; + } + __syncthreads(); + if (threadIdx.x / WARP_SIZE == 0) { + sum = warpSum(shared[threadIdx.x]); + if (threadIdx.x == 0) { + shared[0] = sum; + } + } + __syncthreads(); + + // Everyone picks it up, should be broadcast into the whole gradInput + return shared[0]; +} +#endif diff --git a/models/face_parsing/modules/src/inplace_abn_cpu.cpp b/models/face_parsing/modules/src/inplace_abn_cpu.cpp new file mode 100644 index 0000000000000000000000000000000000000000..ffc6d38c52ea31661b8dd438dc3fe1958f50b61e --- /dev/null +++ b/models/face_parsing/modules/src/inplace_abn_cpu.cpp @@ -0,0 +1,119 @@ +#include + +#include + +#include "utils/checks.h" +#include "inplace_abn.h" + +at::Tensor reduce_sum(at::Tensor x) { + if (x.ndimension() == 2) { + return x.sum(0); + } else { + auto x_view = x.view({x.size(0), x.size(1), -1}); + return x_view.sum(-1).sum(0); + } +} + +at::Tensor broadcast_to(at::Tensor v, at::Tensor x) { + if (x.ndimension() == 2) { + return v; + } else { + std::vector broadcast_size = {1, -1}; + for (int64_t i = 2; i < x.ndimension(); ++i) + broadcast_size.push_back(1); + + return v.view(broadcast_size); + } +} + +int64_t count(at::Tensor x) { + int64_t count = x.size(0); + for (int64_t i = 2; i < x.ndimension(); ++i) + count *= x.size(i); + + return count; +} + +at::Tensor invert_affine(at::Tensor z, at::Tensor weight, at::Tensor bias, bool affine, float eps) { + if (affine) { + return (z - broadcast_to(bias, z)) / broadcast_to(at::abs(weight) + eps, z); + } else { + return z; + } +} + +std::vector mean_var_cpu(at::Tensor x) { + auto num = count(x); + auto mean = reduce_sum(x) / num; + auto diff = x - broadcast_to(mean, x); + auto var = reduce_sum(diff.pow(2)) / num; + + return {mean, var}; +} + +at::Tensor forward_cpu(at::Tensor x, at::Tensor mean, at::Tensor var, at::Tensor weight, at::Tensor bias, + bool affine, float eps) { + auto gamma = affine ? at::abs(weight) + eps : at::ones_like(var); + auto mul = at::rsqrt(var + eps) * gamma; + + x.sub_(broadcast_to(mean, x)); + x.mul_(broadcast_to(mul, x)); + if (affine) x.add_(broadcast_to(bias, x)); + + return x; +} + +std::vector edz_eydz_cpu(at::Tensor z, at::Tensor dz, at::Tensor weight, at::Tensor bias, + bool affine, float eps) { + auto edz = reduce_sum(dz); + auto y = invert_affine(z, weight, bias, affine, eps); + auto eydz = reduce_sum(y * dz); + + return {edz, eydz}; +} + +at::Tensor backward_cpu(at::Tensor z, at::Tensor dz, at::Tensor var, at::Tensor weight, at::Tensor bias, + at::Tensor edz, at::Tensor eydz, bool affine, float eps) { + auto y = invert_affine(z, weight, bias, affine, eps); + auto mul = affine ? at::rsqrt(var + eps) * (at::abs(weight) + eps) : at::rsqrt(var + eps); + + auto num = count(z); + auto dx = (dz - broadcast_to(edz / num, dz) - y * broadcast_to(eydz / num, dz)) * broadcast_to(mul, dz); + return dx; +} + +void leaky_relu_backward_cpu(at::Tensor z, at::Tensor dz, float slope) { + CHECK_CPU_INPUT(z); + CHECK_CPU_INPUT(dz); + + AT_DISPATCH_FLOATING_TYPES(z.type(), "leaky_relu_backward_cpu", ([&] { + int64_t count = z.numel(); + auto *_z = z.data(); + auto *_dz = dz.data(); + + for (int64_t i = 0; i < count; ++i) { + if (_z[i] < 0) { + _z[i] *= 1 / slope; + _dz[i] *= slope; + } + } + })); +} + +void elu_backward_cpu(at::Tensor z, at::Tensor dz) { + CHECK_CPU_INPUT(z); + CHECK_CPU_INPUT(dz); + + AT_DISPATCH_FLOATING_TYPES(z.type(), "elu_backward_cpu", ([&] { + int64_t count = z.numel(); + auto *_z = z.data(); + auto *_dz = dz.data(); + + for (int64_t i = 0; i < count; ++i) { + if (_z[i] < 0) { + _z[i] = log1p(_z[i]); + _dz[i] *= (_z[i] + 1.f); + } + } + })); +} diff --git a/models/face_parsing/modules/src/inplace_abn_cuda.cu b/models/face_parsing/modules/src/inplace_abn_cuda.cu new file mode 100644 index 0000000000000000000000000000000000000000..b157b06d47173d1645c6a40c89f564b737e84d43 --- /dev/null +++ b/models/face_parsing/modules/src/inplace_abn_cuda.cu @@ -0,0 +1,333 @@ +#include + +#include +#include + +#include + +#include "utils/checks.h" +#include "utils/cuda.cuh" +#include "inplace_abn.h" + +#include + +// Operations for reduce +template +struct SumOp { + __device__ SumOp(const T *t, int c, int s) + : tensor(t), chn(c), sp(s) {} + __device__ __forceinline__ T operator()(int batch, int plane, int n) { + return tensor[(batch * chn + plane) * sp + n]; + } + const T *tensor; + const int chn; + const int sp; +}; + +template +struct VarOp { + __device__ VarOp(T m, const T *t, int c, int s) + : mean(m), tensor(t), chn(c), sp(s) {} + __device__ __forceinline__ T operator()(int batch, int plane, int n) { + T val = tensor[(batch * chn + plane) * sp + n]; + return (val - mean) * (val - mean); + } + const T mean; + const T *tensor; + const int chn; + const int sp; +}; + +template +struct GradOp { + __device__ GradOp(T _weight, T _bias, const T *_z, const T *_dz, int c, int s) + : weight(_weight), bias(_bias), z(_z), dz(_dz), chn(c), sp(s) {} + __device__ __forceinline__ Pair operator()(int batch, int plane, int n) { + T _y = (z[(batch * chn + plane) * sp + n] - bias) / weight; + T _dz = dz[(batch * chn + plane) * sp + n]; + return Pair(_dz, _y * _dz); + } + const T weight; + const T bias; + const T *z; + const T *dz; + const int chn; + const int sp; +}; + +/*********** + * mean_var + ***********/ + +template +__global__ void mean_var_kernel(const T *x, T *mean, T *var, int num, int chn, int sp) { + int plane = blockIdx.x; + T norm = T(1) / T(num * sp); + + T _mean = reduce>(SumOp(x, chn, sp), plane, num, sp) * norm; + __syncthreads(); + T _var = reduce>(VarOp(_mean, x, chn, sp), plane, num, sp) * norm; + + if (threadIdx.x == 0) { + mean[plane] = _mean; + var[plane] = _var; + } +} + +std::vector mean_var_cuda(at::Tensor x) { + CHECK_CUDA_INPUT(x); + + // Extract dimensions + int64_t num, chn, sp; + get_dims(x, num, chn, sp); + + // Prepare output tensors + auto mean = at::empty({chn}, x.options()); + auto var = at::empty({chn}, x.options()); + + // Run kernel + dim3 blocks(chn); + dim3 threads(getNumThreads(sp)); + auto stream = at::cuda::getCurrentCUDAStream(); + AT_DISPATCH_FLOATING_TYPES(x.type(), "mean_var_cuda", ([&] { + mean_var_kernel<<>>( + x.data(), + mean.data(), + var.data(), + num, chn, sp); + })); + + return {mean, var}; +} + +/********** + * forward + **********/ + +template +__global__ void forward_kernel(T *x, const T *mean, const T *var, const T *weight, const T *bias, + bool affine, float eps, int num, int chn, int sp) { + int plane = blockIdx.x; + + T _mean = mean[plane]; + T _var = var[plane]; + T _weight = affine ? abs(weight[plane]) + eps : T(1); + T _bias = affine ? bias[plane] : T(0); + + T mul = rsqrt(_var + eps) * _weight; + + for (int batch = 0; batch < num; ++batch) { + for (int n = threadIdx.x; n < sp; n += blockDim.x) { + T _x = x[(batch * chn + plane) * sp + n]; + T _y = (_x - _mean) * mul + _bias; + + x[(batch * chn + plane) * sp + n] = _y; + } + } +} + +at::Tensor forward_cuda(at::Tensor x, at::Tensor mean, at::Tensor var, at::Tensor weight, at::Tensor bias, + bool affine, float eps) { + CHECK_CUDA_INPUT(x); + CHECK_CUDA_INPUT(mean); + CHECK_CUDA_INPUT(var); + CHECK_CUDA_INPUT(weight); + CHECK_CUDA_INPUT(bias); + + // Extract dimensions + int64_t num, chn, sp; + get_dims(x, num, chn, sp); + + // Run kernel + dim3 blocks(chn); + dim3 threads(getNumThreads(sp)); + auto stream = at::cuda::getCurrentCUDAStream(); + AT_DISPATCH_FLOATING_TYPES(x.type(), "forward_cuda", ([&] { + forward_kernel<<>>( + x.data(), + mean.data(), + var.data(), + weight.data(), + bias.data(), + affine, eps, num, chn, sp); + })); + + return x; +} + +/*********** + * edz_eydz + ***********/ + +template +__global__ void edz_eydz_kernel(const T *z, const T *dz, const T *weight, const T *bias, + T *edz, T *eydz, bool affine, float eps, int num, int chn, int sp) { + int plane = blockIdx.x; + + T _weight = affine ? abs(weight[plane]) + eps : 1.f; + T _bias = affine ? bias[plane] : 0.f; + + Pair res = reduce, GradOp>(GradOp(_weight, _bias, z, dz, chn, sp), plane, num, sp); + __syncthreads(); + + if (threadIdx.x == 0) { + edz[plane] = res.v1; + eydz[plane] = res.v2; + } +} + +std::vector edz_eydz_cuda(at::Tensor z, at::Tensor dz, at::Tensor weight, at::Tensor bias, + bool affine, float eps) { + CHECK_CUDA_INPUT(z); + CHECK_CUDA_INPUT(dz); + CHECK_CUDA_INPUT(weight); + CHECK_CUDA_INPUT(bias); + + // Extract dimensions + int64_t num, chn, sp; + get_dims(z, num, chn, sp); + + auto edz = at::empty({chn}, z.options()); + auto eydz = at::empty({chn}, z.options()); + + // Run kernel + dim3 blocks(chn); + dim3 threads(getNumThreads(sp)); + auto stream = at::cuda::getCurrentCUDAStream(); + AT_DISPATCH_FLOATING_TYPES(z.type(), "edz_eydz_cuda", ([&] { + edz_eydz_kernel<<>>( + z.data(), + dz.data(), + weight.data(), + bias.data(), + edz.data(), + eydz.data(), + affine, eps, num, chn, sp); + })); + + return {edz, eydz}; +} + +/*********** + * backward + ***********/ + +template +__global__ void backward_kernel(const T *z, const T *dz, const T *var, const T *weight, const T *bias, const T *edz, + const T *eydz, T *dx, bool affine, float eps, int num, int chn, int sp) { + int plane = blockIdx.x; + + T _weight = affine ? abs(weight[plane]) + eps : 1.f; + T _bias = affine ? bias[plane] : 0.f; + T _var = var[plane]; + T _edz = edz[plane]; + T _eydz = eydz[plane]; + + T _mul = _weight * rsqrt(_var + eps); + T count = T(num * sp); + + for (int batch = 0; batch < num; ++batch) { + for (int n = threadIdx.x; n < sp; n += blockDim.x) { + T _dz = dz[(batch * chn + plane) * sp + n]; + T _y = (z[(batch * chn + plane) * sp + n] - _bias) / _weight; + + dx[(batch * chn + plane) * sp + n] = (_dz - _edz / count - _y * _eydz / count) * _mul; + } + } +} + +at::Tensor backward_cuda(at::Tensor z, at::Tensor dz, at::Tensor var, at::Tensor weight, at::Tensor bias, + at::Tensor edz, at::Tensor eydz, bool affine, float eps) { + CHECK_CUDA_INPUT(z); + CHECK_CUDA_INPUT(dz); + CHECK_CUDA_INPUT(var); + CHECK_CUDA_INPUT(weight); + CHECK_CUDA_INPUT(bias); + CHECK_CUDA_INPUT(edz); + CHECK_CUDA_INPUT(eydz); + + // Extract dimensions + int64_t num, chn, sp; + get_dims(z, num, chn, sp); + + auto dx = at::zeros_like(z); + + // Run kernel + dim3 blocks(chn); + dim3 threads(getNumThreads(sp)); + auto stream = at::cuda::getCurrentCUDAStream(); + AT_DISPATCH_FLOATING_TYPES(z.type(), "backward_cuda", ([&] { + backward_kernel<<>>( + z.data(), + dz.data(), + var.data(), + weight.data(), + bias.data(), + edz.data(), + eydz.data(), + dx.data(), + affine, eps, num, chn, sp); + })); + + return dx; +} + +/************** + * activations + **************/ + +template +inline void leaky_relu_backward_impl(T *z, T *dz, float slope, int64_t count) { + // Create thrust pointers + thrust::device_ptr th_z = thrust::device_pointer_cast(z); + thrust::device_ptr th_dz = thrust::device_pointer_cast(dz); + + auto stream = at::cuda::getCurrentCUDAStream(); + thrust::transform_if(thrust::cuda::par.on(stream), + th_dz, th_dz + count, th_z, th_dz, + [slope] __device__ (const T& dz) { return dz * slope; }, + [] __device__ (const T& z) { return z < 0; }); + thrust::transform_if(thrust::cuda::par.on(stream), + th_z, th_z + count, th_z, + [slope] __device__ (const T& z) { return z / slope; }, + [] __device__ (const T& z) { return z < 0; }); +} + +void leaky_relu_backward_cuda(at::Tensor z, at::Tensor dz, float slope) { + CHECK_CUDA_INPUT(z); + CHECK_CUDA_INPUT(dz); + + int64_t count = z.numel(); + + AT_DISPATCH_FLOATING_TYPES(z.type(), "leaky_relu_backward_cuda", ([&] { + leaky_relu_backward_impl(z.data(), dz.data(), slope, count); + })); +} + +template +inline void elu_backward_impl(T *z, T *dz, int64_t count) { + // Create thrust pointers + thrust::device_ptr th_z = thrust::device_pointer_cast(z); + thrust::device_ptr th_dz = thrust::device_pointer_cast(dz); + + auto stream = at::cuda::getCurrentCUDAStream(); + thrust::transform_if(thrust::cuda::par.on(stream), + th_dz, th_dz + count, th_z, th_z, th_dz, + [] __device__ (const T& dz, const T& z) { return dz * (z + 1.); }, + [] __device__ (const T& z) { return z < 0; }); + thrust::transform_if(thrust::cuda::par.on(stream), + th_z, th_z + count, th_z, + [] __device__ (const T& z) { return log1p(z); }, + [] __device__ (const T& z) { return z < 0; }); +} + +void elu_backward_cuda(at::Tensor z, at::Tensor dz) { + CHECK_CUDA_INPUT(z); + CHECK_CUDA_INPUT(dz); + + int64_t count = z.numel(); + + AT_DISPATCH_FLOATING_TYPES(z.type(), "leaky_relu_backward_cuda", ([&] { + elu_backward_impl(z.data(), dz.data(), count); + })); +} diff --git a/models/face_parsing/modules/src/inplace_abn_cuda_half.cu b/models/face_parsing/modules/src/inplace_abn_cuda_half.cu new file mode 100644 index 0000000000000000000000000000000000000000..bb63e73f9d90179e5bd5dae5579c4844da9c25e2 --- /dev/null +++ b/models/face_parsing/modules/src/inplace_abn_cuda_half.cu @@ -0,0 +1,275 @@ +#include + +#include + +#include + +#include "utils/checks.h" +#include "utils/cuda.cuh" +#include "inplace_abn.h" + +#include + +// Operations for reduce +struct SumOpH { + __device__ SumOpH(const half *t, int c, int s) + : tensor(t), chn(c), sp(s) {} + __device__ __forceinline__ float operator()(int batch, int plane, int n) { + return __half2float(tensor[(batch * chn + plane) * sp + n]); + } + const half *tensor; + const int chn; + const int sp; +}; + +struct VarOpH { + __device__ VarOpH(float m, const half *t, int c, int s) + : mean(m), tensor(t), chn(c), sp(s) {} + __device__ __forceinline__ float operator()(int batch, int plane, int n) { + const auto t = __half2float(tensor[(batch * chn + plane) * sp + n]); + return (t - mean) * (t - mean); + } + const float mean; + const half *tensor; + const int chn; + const int sp; +}; + +struct GradOpH { + __device__ GradOpH(float _weight, float _bias, const half *_z, const half *_dz, int c, int s) + : weight(_weight), bias(_bias), z(_z), dz(_dz), chn(c), sp(s) {} + __device__ __forceinline__ Pair operator()(int batch, int plane, int n) { + float _y = (__half2float(z[(batch * chn + plane) * sp + n]) - bias) / weight; + float _dz = __half2float(dz[(batch * chn + plane) * sp + n]); + return Pair(_dz, _y * _dz); + } + const float weight; + const float bias; + const half *z; + const half *dz; + const int chn; + const int sp; +}; + +/*********** + * mean_var + ***********/ + +__global__ void mean_var_kernel_h(const half *x, float *mean, float *var, int num, int chn, int sp) { + int plane = blockIdx.x; + float norm = 1.f / static_cast(num * sp); + + float _mean = reduce(SumOpH(x, chn, sp), plane, num, sp) * norm; + __syncthreads(); + float _var = reduce(VarOpH(_mean, x, chn, sp), plane, num, sp) * norm; + + if (threadIdx.x == 0) { + mean[plane] = _mean; + var[plane] = _var; + } +} + +std::vector mean_var_cuda_h(at::Tensor x) { + CHECK_CUDA_INPUT(x); + + // Extract dimensions + int64_t num, chn, sp; + get_dims(x, num, chn, sp); + + // Prepare output tensors + auto mean = at::empty({chn},x.options().dtype(at::kFloat)); + auto var = at::empty({chn},x.options().dtype(at::kFloat)); + + // Run kernel + dim3 blocks(chn); + dim3 threads(getNumThreads(sp)); + auto stream = at::cuda::getCurrentCUDAStream(); + mean_var_kernel_h<<>>( + reinterpret_cast(x.data()), + mean.data(), + var.data(), + num, chn, sp); + + return {mean, var}; +} + +/********** + * forward + **********/ + +__global__ void forward_kernel_h(half *x, const float *mean, const float *var, const float *weight, const float *bias, + bool affine, float eps, int num, int chn, int sp) { + int plane = blockIdx.x; + + const float _mean = mean[plane]; + const float _var = var[plane]; + const float _weight = affine ? abs(weight[plane]) + eps : 1.f; + const float _bias = affine ? bias[plane] : 0.f; + + const float mul = rsqrt(_var + eps) * _weight; + + for (int batch = 0; batch < num; ++batch) { + for (int n = threadIdx.x; n < sp; n += blockDim.x) { + half *x_ptr = x + (batch * chn + plane) * sp + n; + float _x = __half2float(*x_ptr); + float _y = (_x - _mean) * mul + _bias; + + *x_ptr = __float2half(_y); + } + } +} + +at::Tensor forward_cuda_h(at::Tensor x, at::Tensor mean, at::Tensor var, at::Tensor weight, at::Tensor bias, + bool affine, float eps) { + CHECK_CUDA_INPUT(x); + CHECK_CUDA_INPUT(mean); + CHECK_CUDA_INPUT(var); + CHECK_CUDA_INPUT(weight); + CHECK_CUDA_INPUT(bias); + + // Extract dimensions + int64_t num, chn, sp; + get_dims(x, num, chn, sp); + + // Run kernel + dim3 blocks(chn); + dim3 threads(getNumThreads(sp)); + auto stream = at::cuda::getCurrentCUDAStream(); + forward_kernel_h<<>>( + reinterpret_cast(x.data()), + mean.data(), + var.data(), + weight.data(), + bias.data(), + affine, eps, num, chn, sp); + + return x; +} + +__global__ void edz_eydz_kernel_h(const half *z, const half *dz, const float *weight, const float *bias, + float *edz, float *eydz, bool affine, float eps, int num, int chn, int sp) { + int plane = blockIdx.x; + + float _weight = affine ? abs(weight[plane]) + eps : 1.f; + float _bias = affine ? bias[plane] : 0.f; + + Pair res = reduce, GradOpH>(GradOpH(_weight, _bias, z, dz, chn, sp), plane, num, sp); + __syncthreads(); + + if (threadIdx.x == 0) { + edz[plane] = res.v1; + eydz[plane] = res.v2; + } +} + +std::vector edz_eydz_cuda_h(at::Tensor z, at::Tensor dz, at::Tensor weight, at::Tensor bias, + bool affine, float eps) { + CHECK_CUDA_INPUT(z); + CHECK_CUDA_INPUT(dz); + CHECK_CUDA_INPUT(weight); + CHECK_CUDA_INPUT(bias); + + // Extract dimensions + int64_t num, chn, sp; + get_dims(z, num, chn, sp); + + auto edz = at::empty({chn},z.options().dtype(at::kFloat)); + auto eydz = at::empty({chn},z.options().dtype(at::kFloat)); + + // Run kernel + dim3 blocks(chn); + dim3 threads(getNumThreads(sp)); + auto stream = at::cuda::getCurrentCUDAStream(); + edz_eydz_kernel_h<<>>( + reinterpret_cast(z.data()), + reinterpret_cast(dz.data()), + weight.data(), + bias.data(), + edz.data(), + eydz.data(), + affine, eps, num, chn, sp); + + return {edz, eydz}; +} + +__global__ void backward_kernel_h(const half *z, const half *dz, const float *var, const float *weight, const float *bias, const float *edz, + const float *eydz, half *dx, bool affine, float eps, int num, int chn, int sp) { + int plane = blockIdx.x; + + float _weight = affine ? abs(weight[plane]) + eps : 1.f; + float _bias = affine ? bias[plane] : 0.f; + float _var = var[plane]; + float _edz = edz[plane]; + float _eydz = eydz[plane]; + + float _mul = _weight * rsqrt(_var + eps); + float count = float(num * sp); + + for (int batch = 0; batch < num; ++batch) { + for (int n = threadIdx.x; n < sp; n += blockDim.x) { + float _dz = __half2float(dz[(batch * chn + plane) * sp + n]); + float _y = (__half2float(z[(batch * chn + plane) * sp + n]) - _bias) / _weight; + + dx[(batch * chn + plane) * sp + n] = __float2half((_dz - _edz / count - _y * _eydz / count) * _mul); + } + } +} + +at::Tensor backward_cuda_h(at::Tensor z, at::Tensor dz, at::Tensor var, at::Tensor weight, at::Tensor bias, + at::Tensor edz, at::Tensor eydz, bool affine, float eps) { + CHECK_CUDA_INPUT(z); + CHECK_CUDA_INPUT(dz); + CHECK_CUDA_INPUT(var); + CHECK_CUDA_INPUT(weight); + CHECK_CUDA_INPUT(bias); + CHECK_CUDA_INPUT(edz); + CHECK_CUDA_INPUT(eydz); + + // Extract dimensions + int64_t num, chn, sp; + get_dims(z, num, chn, sp); + + auto dx = at::zeros_like(z); + + // Run kernel + dim3 blocks(chn); + dim3 threads(getNumThreads(sp)); + auto stream = at::cuda::getCurrentCUDAStream(); + backward_kernel_h<<>>( + reinterpret_cast(z.data()), + reinterpret_cast(dz.data()), + var.data(), + weight.data(), + bias.data(), + edz.data(), + eydz.data(), + reinterpret_cast(dx.data()), + affine, eps, num, chn, sp); + + return dx; +} + +__global__ void leaky_relu_backward_impl_h(half *z, half *dz, float slope, int64_t count) { + for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < count; i += blockDim.x * gridDim.x){ + float _z = __half2float(z[i]); + if (_z < 0) { + dz[i] = __float2half(__half2float(dz[i]) * slope); + z[i] = __float2half(_z / slope); + } + } +} + +void leaky_relu_backward_cuda_h(at::Tensor z, at::Tensor dz, float slope) { + CHECK_CUDA_INPUT(z); + CHECK_CUDA_INPUT(dz); + + int64_t count = z.numel(); + dim3 threads(getNumThreads(count)); + dim3 blocks = (count + threads.x - 1) / threads.x; + auto stream = at::cuda::getCurrentCUDAStream(); + leaky_relu_backward_impl_h<<>>( + reinterpret_cast(z.data()), + reinterpret_cast(dz.data()), + slope, count); +} + diff --git a/models/face_parsing/modules/src/utils/checks.h b/models/face_parsing/modules/src/utils/checks.h new file mode 100644 index 0000000000000000000000000000000000000000..e761a6fe34d0789815b588eba7e3726026e0e868 --- /dev/null +++ b/models/face_parsing/modules/src/utils/checks.h @@ -0,0 +1,15 @@ +#pragma once + +#include + +// Define AT_CHECK for old version of ATen where the same function was called AT_ASSERT +#ifndef AT_CHECK +#define AT_CHECK AT_ASSERT +#endif + +#define CHECK_CUDA(x) AT_CHECK((x).type().is_cuda(), #x " must be a CUDA tensor") +#define CHECK_CPU(x) AT_CHECK(!(x).type().is_cuda(), #x " must be a CPU tensor") +#define CHECK_CONTIGUOUS(x) AT_CHECK((x).is_contiguous(), #x " must be contiguous") + +#define CHECK_CUDA_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x) +#define CHECK_CPU_INPUT(x) CHECK_CPU(x); CHECK_CONTIGUOUS(x) \ No newline at end of file diff --git a/models/face_parsing/modules/src/utils/common.h b/models/face_parsing/modules/src/utils/common.h new file mode 100644 index 0000000000000000000000000000000000000000..e8403eef8a233b75dd4bb353c16486fe1be2039a --- /dev/null +++ b/models/face_parsing/modules/src/utils/common.h @@ -0,0 +1,49 @@ +#pragma once + +#include + +/* + * Functions to share code between CPU and GPU + */ + +#ifdef __CUDACC__ +// CUDA versions + +#define HOST_DEVICE __host__ __device__ +#define INLINE_HOST_DEVICE __host__ __device__ inline +#define FLOOR(x) floor(x) + +#if __CUDA_ARCH__ >= 600 +// Recent compute capabilities have block-level atomicAdd for all data types, so we use that +#define ACCUM(x,y) atomicAdd_block(&(x),(y)) +#else +// Older architectures don't have block-level atomicAdd, nor atomicAdd for doubles, so we defer to atomicAdd for float +// and use the known atomicCAS-based implementation for double +template +__device__ inline data_t atomic_add(data_t *address, data_t val) { + return atomicAdd(address, val); +} + +template<> +__device__ inline double atomic_add(double *address, double val) { + unsigned long long int* address_as_ull = (unsigned long long int*)address; + unsigned long long int old = *address_as_ull, assumed; + do { + assumed = old; + old = atomicCAS(address_as_ull, assumed, __double_as_longlong(val + __longlong_as_double(assumed))); + } while (assumed != old); + return __longlong_as_double(old); +} + +#define ACCUM(x,y) atomic_add(&(x),(y)) +#endif // #if __CUDA_ARCH__ >= 600 + +#else +// CPU versions + +#define HOST_DEVICE +#define INLINE_HOST_DEVICE inline +#define FLOOR(x) std::floor(x) +#define ACCUM(x,y) (x) += (y) + +#endif // #ifdef __CUDACC__ \ No newline at end of file diff --git a/models/face_parsing/modules/src/utils/cuda.cuh b/models/face_parsing/modules/src/utils/cuda.cuh new file mode 100644 index 0000000000000000000000000000000000000000..60c0023835e02c5f7c539c28ac07b75b72df394b --- /dev/null +++ b/models/face_parsing/modules/src/utils/cuda.cuh @@ -0,0 +1,71 @@ +#pragma once + +/* + * General settings and functions + */ +const int WARP_SIZE = 32; +const int MAX_BLOCK_SIZE = 1024; + +static int getNumThreads(int nElem) { + int threadSizes[6] = {32, 64, 128, 256, 512, MAX_BLOCK_SIZE}; + for (int i = 0; i < 6; ++i) { + if (nElem <= threadSizes[i]) { + return threadSizes[i]; + } + } + return MAX_BLOCK_SIZE; +} + +/* + * Reduction utilities + */ +template +__device__ __forceinline__ T WARP_SHFL_XOR(T value, int laneMask, int width = warpSize, + unsigned int mask = 0xffffffff) { +#if CUDART_VERSION >= 9000 + return __shfl_xor_sync(mask, value, laneMask, width); +#else + return __shfl_xor(value, laneMask, width); +#endif +} + +__device__ __forceinline__ int getMSB(int val) { return 31 - __clz(val); } + +template +struct Pair { + T v1, v2; + __device__ Pair() {} + __device__ Pair(T _v1, T _v2) : v1(_v1), v2(_v2) {} + __device__ Pair(T v) : v1(v), v2(v) {} + __device__ Pair(int v) : v1(v), v2(v) {} + __device__ Pair &operator+=(const Pair &a) { + v1 += a.v1; + v2 += a.v2; + return *this; + } +}; + +template +static __device__ __forceinline__ T warpSum(T val) { +#if __CUDA_ARCH__ >= 300 + for (int i = 0; i < getMSB(WARP_SIZE); ++i) { + val += WARP_SHFL_XOR(val, 1 << i, WARP_SIZE); + } +#else + __shared__ T values[MAX_BLOCK_SIZE]; + values[threadIdx.x] = val; + __threadfence_block(); + const int base = (threadIdx.x / WARP_SIZE) * WARP_SIZE; + for (int i = 1; i < WARP_SIZE; i++) { + val += values[base + ((i + threadIdx.x) % WARP_SIZE)]; + } +#endif + return val; +} + +template +static __device__ __forceinline__ Pair warpSum(Pair value) { + value.v1 = warpSum(value.v1); + value.v2 = warpSum(value.v2); + return value; +} \ No newline at end of file diff --git a/models/face_parsing/prepropess_data.py b/models/face_parsing/prepropess_data.py new file mode 100644 index 0000000000000000000000000000000000000000..28ca5d767f5e7ba90a51c31e736a1dfc7fa69af8 --- /dev/null +++ b/models/face_parsing/prepropess_data.py @@ -0,0 +1,31 @@ +# import os +# import cv2 +# import glob +# import numpy as np +# +# +# label_list = ['skin', 'nose', 'eye_g', 'l_eye', 'r_eye', 'l_brow', 'r_brow', 'l_ear', 'r_ear', 'mouth', 'u_lip', +# 'l_lip', 'hair', 'hat', 'ear_r', 'neck_l', 'neck', 'cloth'] +# +# +# folder_base = '/media/zhup/Data/CelebAMask-HQ/CelebAMaskHQ-mask-anno' +# folder_save = '/media/zhup/Data/CelebAMask-HQ/Mask_less' +# img_num = 30000 +# +# correspond_list = [0, 1, 2, 3, 4, 4, 5, 5, 6, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15] +# +# for k in range(img_num): +# folder_num = int(k / 2000) +# im_base = np.zeros((512, 512)) +# for idx, label in enumerate(label_list, 1): +# filename = os.path.join(folder_base, str(folder_num), str(k).rjust(5, '0') + '_' + label + '.png') +# if (os.path.exists(filename)): +# # print(label, idx) +# im = cv2.imread(filename) +# im = im[:, :, 0] +# im_base[im != 0] = (correspond_list[idx]) +# +# filename_save = os.path.join(folder_save, str(k) + '.png') +# print(filename_save) +# cv2.imwrite(filename_save, im_base) +# diff --git a/models/face_parsing/resnet.py b/models/face_parsing/resnet.py new file mode 100644 index 0000000000000000000000000000000000000000..aa2bf95130e9815ba378cb6f73207068b81a04b9 --- /dev/null +++ b/models/face_parsing/resnet.py @@ -0,0 +1,109 @@ +#!/usr/bin/python +# -*- encoding: utf-8 -*- + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.utils.model_zoo as modelzoo + +# from modules.bn import InPlaceABNSync as BatchNorm2d + +resnet18_url = 'https://download.pytorch.org/models/resnet18-5c106cde.pth' + + +def conv3x3(in_planes, out_planes, stride=1): + """3x3 convolution with padding""" + return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, + padding=1, bias=False) + + +class BasicBlock(nn.Module): + def __init__(self, in_chan, out_chan, stride=1): + super(BasicBlock, self).__init__() + self.conv1 = conv3x3(in_chan, out_chan, stride) + self.bn1 = nn.BatchNorm2d(out_chan) + self.conv2 = conv3x3(out_chan, out_chan) + self.bn2 = nn.BatchNorm2d(out_chan) + self.relu = nn.ReLU(inplace=True) + self.downsample = None + if in_chan != out_chan or stride != 1: + self.downsample = nn.Sequential( + nn.Conv2d(in_chan, out_chan, + kernel_size=1, stride=stride, bias=False), + nn.BatchNorm2d(out_chan), + ) + + def forward(self, x): + residual = self.conv1(x) + residual = F.relu(self.bn1(residual)) + residual = self.conv2(residual) + residual = self.bn2(residual) + + shortcut = x + if self.downsample is not None: + shortcut = self.downsample(x) + + out = shortcut + residual + out = self.relu(out) + return out + + +def create_layer_basic(in_chan, out_chan, bnum, stride=1): + layers = [BasicBlock(in_chan, out_chan, stride=stride)] + for i in range(bnum-1): + layers.append(BasicBlock(out_chan, out_chan, stride=1)) + return nn.Sequential(*layers) + + +class Resnet18(nn.Module): + def __init__(self): + super(Resnet18, self).__init__() + self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, + bias=False) + self.bn1 = nn.BatchNorm2d(64) + self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) + self.layer1 = create_layer_basic(64, 64, bnum=2, stride=1) + self.layer2 = create_layer_basic(64, 128, bnum=2, stride=2) + self.layer3 = create_layer_basic(128, 256, bnum=2, stride=2) + self.layer4 = create_layer_basic(256, 512, bnum=2, stride=2) + self.init_weight() + + def forward(self, x): + x = self.conv1(x) + x = F.relu(self.bn1(x)) + x = self.maxpool(x) + + x = self.layer1(x) + feat8 = self.layer2(x) # 1/8 + feat16 = self.layer3(feat8) # 1/16 + feat32 = self.layer4(feat16) # 1/32 + return feat8, feat16, feat32 + + def init_weight(self): + state_dict = modelzoo.load_url(resnet18_url) + self_state_dict = self.state_dict() + for k, v in state_dict.items(): + if 'fc' in k: continue + self_state_dict.update({k: v}) + self.load_state_dict(self_state_dict) + + def get_params(self): + wd_params, nowd_params = [], [] + for name, module in self.named_modules(): + if isinstance(module, (nn.Linear, nn.Conv2d)): + wd_params.append(module.weight) + if not module.bias is None: + nowd_params.append(module.bias) + elif isinstance(module, nn.BatchNorm2d): + nowd_params += list(module.parameters()) + return wd_params, nowd_params + + +if __name__ == "__main__": + net = Resnet18() + x = torch.randn(16, 3, 224, 224) + out = net(x) + print(out[0].size()) + print(out[1].size()) + print(out[2].size()) + net.get_params() diff --git a/models/face_parsing/transform.py b/models/face_parsing/transform.py new file mode 100644 index 0000000000000000000000000000000000000000..a28e9d579ce2f57f230d6f0b52cd185e1525e516 --- /dev/null +++ b/models/face_parsing/transform.py @@ -0,0 +1,127 @@ +#!/usr/bin/python +# -*- encoding: utf-8 -*- + + +from PIL import Image +import PIL.ImageEnhance as ImageEnhance +import random +import numpy as np + +class RandomCrop(object): + def __init__(self, size, *args, **kwargs): + self.size = size + + def __call__(self, im_lb): + im = im_lb['im'] + lb = im_lb['lb'] + assert im.size == lb.size + W, H = self.size + w, h = im.size + + if (W, H) == (w, h): return dict(im=im, lb=lb) + if w < W or h < H: + scale = float(W) / w if w < h else float(H) / h + w, h = int(scale * w + 1), int(scale * h + 1) + im = im.resize((w, h), Image.BILINEAR) + lb = lb.resize((w, h), Image.NEAREST) + sw, sh = random.random() * (w - W), random.random() * (h - H) + crop = int(sw), int(sh), int(sw) + W, int(sh) + H + return dict( + im = im.crop(crop), + lb = lb.crop(crop) + ) + + +class HorizontalFlip(object): + def __init__(self, p=0.5, *args, **kwargs): + self.p = p + + def __call__(self, im_lb): + if random.random() > self.p: + return im_lb + else: + im = im_lb['im'] + lb = im_lb['lb'] + + + flip_lb = np.array(lb) + # flip_lb[lb == 2] = 3 + # flip_lb[lb == 3] = 2 + # flip_lb[lb == 4] = 5 + # flip_lb[lb == 5] = 4 + # flip_lb[lb == 7] = 8 + # flip_lb[lb == 8] = 7 + flip_lb = Image.fromarray(flip_lb) + return dict(im = im.transpose(Image.FLIP_LEFT_RIGHT), + lb = flip_lb.transpose(Image.FLIP_LEFT_RIGHT), + ) + + +class RandomScale(object): + def __init__(self, scales=(1, ), *args, **kwargs): + self.scales = scales + + def __call__(self, im_lb): + im = im_lb['im'] + lb = im_lb['lb'] + W, H = im.size + scale = random.choice(self.scales) + w, h = int(W * scale), int(H * scale) + return dict(im = im.resize((w, h), Image.BILINEAR), + lb = lb.resize((w, h), Image.NEAREST), + ) + + +class ColorJitter(object): + def __init__(self, brightness=None, contrast=None, saturation=None, *args, **kwargs): + if not brightness is None and brightness>0: + self.brightness = [max(1-brightness, 0), 1+brightness] + if not contrast is None and contrast>0: + self.contrast = [max(1-contrast, 0), 1+contrast] + if not saturation is None and saturation>0: + self.saturation = [max(1-saturation, 0), 1+saturation] + + def __call__(self, im_lb): + im = im_lb['im'] + lb = im_lb['lb'] + r_brightness = random.uniform(self.brightness[0], self.brightness[1]) + r_contrast = random.uniform(self.contrast[0], self.contrast[1]) + r_saturation = random.uniform(self.saturation[0], self.saturation[1]) + im = ImageEnhance.Brightness(im).enhance(r_brightness) + im = ImageEnhance.Contrast(im).enhance(r_contrast) + im = ImageEnhance.Color(im).enhance(r_saturation) + return dict(im = im, + lb = lb, + ) + + +class MultiScale(object): + def __init__(self, scales): + self.scales = scales + + def __call__(self, img): + W, H = img.size + sizes = [(int(W*ratio), int(H*ratio)) for ratio in self.scales] + imgs = [] + [imgs.append(img.resize(size, Image.BILINEAR)) for size in sizes] + return imgs + + +class Compose(object): + def __init__(self, do_list): + self.do_list = do_list + + def __call__(self, im_lb): + for comp in self.do_list: + im_lb = comp(im_lb) + return im_lb + + + + +if __name__ == '__main__': + flip = HorizontalFlip(p = 1) + crop = RandomCrop((321, 321)) + rscales = RandomScale((0.75, 1.0, 1.5, 1.75, 2.0)) + img = Image.open('data/img.jpg') + lb = Image.open('data/label.png') diff --git a/models/optimizer/ClampOptimizer.py b/models/optimizer/ClampOptimizer.py new file mode 100644 index 0000000000000000000000000000000000000000..cfcb907362057e9284f0e67b8a1845a39cdd7d41 --- /dev/null +++ b/models/optimizer/ClampOptimizer.py @@ -0,0 +1,30 @@ +import math +import torch +from torch.optim import Optimizer +import numpy as np + +class ClampOptimizer(Optimizer): + def __init__(self, optimizer, params, **kwargs): + self.opt = optimizer(params, **kwargs) + self.params = params + + + + + @torch.no_grad() + def step(self, closure=None): + loss = self.opt.step(closure) + + + for param in self.params: + tmp_latent_norm = torch.clamp(param.data, 0, 1) + param.data.add_(tmp_latent_norm - param.data) + + + return loss + + + def zero_grad(self): + self.opt.zero_grad() + + diff --git a/models/optimizer/__init__.py b/models/optimizer/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/models/sean_codes/.DS_Store b/models/sean_codes/.DS_Store new file mode 100644 index 0000000000000000000000000000000000000000..250b560c48bf8ebe2d4347df37704e90c6da5c94 Binary files /dev/null and b/models/sean_codes/.DS_Store differ diff --git a/models/sean_codes/__init__.py b/models/sean_codes/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3ee410e3c065678a2877bbb16cb2612c8c8a4c51 --- /dev/null +++ b/models/sean_codes/__init__.py @@ -0,0 +1,8 @@ +# -*- coding: utf-8 -*- + +""" +# File name: __init__.py.py +# Time : 2021/12/7 17:27 +# Author: xyguoo@163.com +# Description: +""" diff --git a/models/sean_codes/data/__init__.py b/models/sean_codes/data/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..7743f99330ba072b3a1773b7ba172c7c6540fb77 --- /dev/null +++ b/models/sean_codes/data/__init__.py @@ -0,0 +1,55 @@ +""" +Copyright (C) 2019 NVIDIA Corporation. All rights reserved. +Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode). +""" + +import importlib +import torch.utils.data +from .base_dataset import BaseDataset + + +def find_dataset_using_name(dataset_name): + # Given the option --dataset [datasetname], + # the file "datasets/datasetname_dataset.py" + # will be imported. + dataset_filename = "sean_codes.data." + dataset_name + "_dataset" + datasetlib = importlib.import_module(dataset_filename) + + # In the file, the class called DatasetNameDataset() will + # be instantiated. It has to be a subclass of BaseDataset, + # and it is case-insensitive. + dataset = None + target_dataset_name = dataset_name.replace('_', '') + 'dataset' + for name, cls in datasetlib.__dict__.items(): + if name.lower() == target_dataset_name.lower() \ + and issubclass(cls, BaseDataset): + dataset = cls + + if dataset is None: + raise ValueError("In %s.py, there should be a subclass of BaseDataset " + "with class name that matches %s in lowercase." % + (dataset_filename, target_dataset_name)) + + return dataset + + +def get_option_setter(dataset_name): + dataset_class = find_dataset_using_name(dataset_name) + return dataset_class.modify_commandline_options + + +def create_dataloader(opt): + dataset = find_dataset_using_name(opt.dataset_mode) + instance = dataset() + instance.initialize(opt) + print("dataset [%s] of size %d was created" % + (type(instance).__name__, len(instance))) + dataloader = torch.utils.data.DataLoader( + instance, + batch_size=opt.batchSize, + shuffle=not opt.serial_batches, + num_workers=min(opt.batchSize, int(opt.nThreads)), + drop_last=opt.isTrain, + pin_memory=True + ) + return dataloader diff --git a/models/sean_codes/data/base_dataset.py b/models/sean_codes/data/base_dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..2bb4b3cbf7d73e89eb53eeb14d8c5389a241fba2 --- /dev/null +++ b/models/sean_codes/data/base_dataset.py @@ -0,0 +1,128 @@ +""" +Copyright (C) 2019 NVIDIA Corporation. All rights reserved. +Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode). +""" + +import torch.utils.data as data +from PIL import Image +import torchvision.transforms as transforms +import numpy as np +import random + + +class BaseDataset(data.Dataset): + def __init__(self): + super(BaseDataset, self).__init__() + + @staticmethod + def modify_commandline_options(parser, is_train): + return parser + + def initialize(self, opt): + pass + + +def get_params(opt, size): + w, h = size + new_h = h + new_w = w + if opt.preprocess_mode == 'resize_and_crop': + new_h = new_w = opt.load_size + elif opt.preprocess_mode == 'scale_width_and_crop': + new_w = opt.load_size + new_h = opt.load_size * h // w + elif opt.preprocess_mode == 'scale_shortside_and_crop': + ss, ls = min(w, h), max(w, h) # shortside and longside + width_is_shorter = w == ss + ls = int(opt.load_size * ls / ss) + new_w, new_h = (ss, ls) if width_is_shorter else (ls, ss) + + x = random.randint(0, np.maximum(0, new_w - opt.crop_size)) + y = random.randint(0, np.maximum(0, new_h - opt.crop_size)) + + flip = random.random() > 0.5 + return {'crop_pos': (x, y), 'flip': flip} + + +def get_transform(opt, params, method=Image.BICUBIC, normalize=True, toTensor=True): + transform_list = [] + if 'resize' in opt.preprocess_mode: + osize = [opt.load_size, opt.load_size] + transform_list.append(transforms.Resize(osize, interpolation=method)) + elif 'scale_width' in opt.preprocess_mode: + transform_list.append(transforms.Lambda(lambda img: __scale_width(img, opt.load_size, method))) + elif 'scale_shortside' in opt.preprocess_mode: + transform_list.append(transforms.Lambda(lambda img: __scale_shortside(img, opt.load_size, method))) + + if 'crop' in opt.preprocess_mode: + transform_list.append(transforms.Lambda(lambda img: __crop(img, params['crop_pos'], opt.crop_size))) + + if opt.preprocess_mode == 'none': + base = 32 + transform_list.append(transforms.Lambda(lambda img: __make_power_2(img, base, method))) + + if opt.preprocess_mode == 'fixed': + w = opt.crop_size + h = round(opt.crop_size / opt.aspect_ratio) + transform_list.append(transforms.Lambda(lambda img: __resize(img, w, h, method))) + + if opt.isTrain and not opt.no_flip: + transform_list.append(transforms.Lambda(lambda img: __flip(img, params['flip']))) + + if toTensor: + transform_list += [transforms.ToTensor()] + + if normalize: + transform_list += [transforms.Normalize((0.5, 0.5, 0.5), + (0.5, 0.5, 0.5))] + return transforms.Compose(transform_list) + + +def normalize(): + return transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) + + +def __resize(img, w, h, method=Image.BICUBIC): + return img.resize((w, h), method) + + +def __make_power_2(img, base, method=Image.BICUBIC): + ow, oh = img.size + h = int(round(oh / base) * base) + w = int(round(ow / base) * base) + if (h == oh) and (w == ow): + return img + return img.resize((w, h), method) + + +def __scale_width(img, target_width, method=Image.BICUBIC): + ow, oh = img.size + if (ow == target_width): + return img + w = target_width + h = int(target_width * oh / ow) + return img.resize((w, h), method) + + +def __scale_shortside(img, target_width, method=Image.BICUBIC): + ow, oh = img.size + ss, ls = min(ow, oh), max(ow, oh) # shortside and longside + width_is_shorter = ow == ss + if (ss == target_width): + return img + ls = int(target_width * ls / ss) + nw, nh = (ss, ls) if width_is_shorter else (ls, ss) + return img.resize((nw, nh), method) + + +def __crop(img, pos, size): + ow, oh = img.size + x1, y1 = pos + tw = th = size + return img.crop((x1, y1, x1 + tw, y1 + th)) + + +def __flip(img, flip): + if flip: + return img.transpose(Image.FLIP_LEFT_RIGHT) + return img diff --git a/models/sean_codes/data/custom_dataset.py b/models/sean_codes/data/custom_dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..4b0b4eca3d1dd647088a7060d1c4fcc1c1d10782 --- /dev/null +++ b/models/sean_codes/data/custom_dataset.py @@ -0,0 +1,50 @@ +""" +Copyright (C) 2019 NVIDIA Corporation. All rights reserved. +Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode). +""" + +from .pix2pix_dataset import Pix2pixDataset +from .image_folder import make_dataset + + +class CustomDataset(Pix2pixDataset): + """ Dataset that loads images from directories + Use option --label_dir, --image_dir, --instance_dir to specify the directories. + The images in the directories are sorted in alphabetical order and paired in order. + """ + + @staticmethod + def modify_commandline_options(parser, is_train): + parser = Pix2pixDataset.modify_commandline_options(parser, is_train) + parser.set_defaults(preprocess_mode='resize_and_crop') + load_size = 286 if is_train else 256 + parser.set_defaults(load_size=load_size) + parser.set_defaults(crop_size=256) + parser.set_defaults(display_winsize=256) + parser.set_defaults(label_nc=19) + parser.set_defaults(contain_dontcare_label=False) + + parser.add_argument('--label_dir', type=str, default='datasets/CelebA-HQ/test/labels', + help='path to the directory that contains label images') + parser.add_argument('--image_dir', type=str, default='datasets/CelebA-HQ/test/images', + help='path to the directory that contains photo images') + parser.add_argument('--instance_dir', type=str, default='', + help='path to the directory that contains instance maps. Leave black if not exists') + return parser + + def get_paths(self, opt): + label_dir = opt.label_dir + label_paths = make_dataset(label_dir, recursive=False, read_cache=True) + + image_dir = opt.image_dir + image_paths = make_dataset(image_dir, recursive=False, read_cache=True) + + if len(opt.instance_dir) > 0: + instance_dir = opt.instance_dir + instance_paths = make_dataset(instance_dir, recursive=False, read_cache=True) + else: + instance_paths = [] + + assert len(label_paths) == len(image_paths), "The #images in %s and %s do not match. Is there something wrong?" + + return label_paths, image_paths, instance_paths diff --git a/models/sean_codes/data/image_folder.py b/models/sean_codes/data/image_folder.py new file mode 100644 index 0000000000000000000000000000000000000000..fd157935070526c823a831d4ad4e0b1970ccd14d --- /dev/null +++ b/models/sean_codes/data/image_folder.py @@ -0,0 +1,101 @@ +""" +Copyright (C) 2019 NVIDIA Corporation. All rights reserved. +Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode). +""" + +############################################################################### +# Code from +# https://github.com/pytorch/vision/blob/master/torchvision/datasets/folder.py +# Modified the original code so that it also loads images from the current +# directory as well as the subdirectories +############################################################################### +import torch.utils.data as data +from PIL import Image +import os + +IMG_EXTENSIONS = [ + '.jpg', '.JPG', '.jpeg', '.JPEG', + '.png', '.PNG', '.ppm', '.PPM', '.bmp', '.BMP', '.tiff', '.webp' +] + + +def is_image_file(filename): + return any(filename.endswith(extension) for extension in IMG_EXTENSIONS) + + +def make_dataset_rec(dir, images): + assert os.path.isdir(dir), '%s is not a valid directory' % dir + + for root, dnames, fnames in sorted(os.walk(dir, followlinks=True)): + for fname in fnames: + if is_image_file(fname): + path = os.path.join(root, fname) + images.append(path) + + +def make_dataset(dir, recursive=False, read_cache=False, write_cache=False): + """ + xuyang: could get image path list + """ + images = [] + + if read_cache: + possible_filelist = os.path.join(dir, 'files.list') + if os.path.isfile(possible_filelist): + with open(possible_filelist, 'r') as f: + images = f.read().splitlines() + return images + + if recursive: + make_dataset_rec(dir, images) + else: + assert os.path.isdir(dir) or os.path.islink(dir), '%s is not a valid directory' % dir + + for root, dnames, fnames in sorted(os.walk(dir)): + for fname in fnames: + if is_image_file(fname): + path = os.path.join(root, fname) + images.append(path) + + if write_cache: + filelist_cache = os.path.join(dir, 'files.list') + with open(filelist_cache, 'w') as f: + for path in images: + f.write("%s\n" % path) + print('wrote filelist cache at %s' % filelist_cache) + + return images + + +def default_loader(path): + return Image.open(path).convert('RGB') + + +class ImageFolder(data.Dataset): + + def __init__(self, root, transform=None, return_paths=False, + loader=default_loader): + imgs = make_dataset(root) + if len(imgs) == 0: + raise(RuntimeError("Found 0 images in: " + root + "\n" + "Supported image extensions are: " + + ",".join(IMG_EXTENSIONS))) + + self.root = root + self.imgs = imgs + self.transform = transform + self.return_paths = return_paths + self.loader = loader + + def __getitem__(self, index): + path = self.imgs[index] + img = self.loader(path) + if self.transform is not None: + img = self.transform(img) + if self.return_paths: + return img, path + else: + return img + + def __len__(self): + return len(self.imgs) diff --git a/models/sean_codes/data/pix2pix_dataset.py b/models/sean_codes/data/pix2pix_dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..7001596886d4f06554d926958f6de6ffe25155af --- /dev/null +++ b/models/sean_codes/data/pix2pix_dataset.py @@ -0,0 +1,139 @@ +""" +Copyright (C) 2019 NVIDIA Corporation. All rights reserved. +Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode). +""" + +from .base_dataset import BaseDataset, get_params, get_transform +from PIL import Image +import util.util as util +import os +import torch + + +class Pix2pixDataset(BaseDataset): + @staticmethod + def modify_commandline_options(parser, is_train): + parser.add_argument('--no_pairing_check', action='store_true', + help='If specified, skip sanity check of correct label-image file pairing') + return parser + + def initialize(self, opt): + self.opt = opt + + label_paths, image_paths, instance_paths = self.get_paths(opt) + + util.natural_sort(label_paths) + util.natural_sort(image_paths) + if not opt.no_instance: + util.natural_sort(instance_paths) + + label_paths = label_paths[:opt.max_dataset_size] + image_paths = image_paths[:opt.max_dataset_size] + instance_paths = instance_paths[:opt.max_dataset_size] + + if not opt.no_pairing_check: + for path1, path2 in zip(label_paths, image_paths): + assert self.paths_match(path1, path2), \ + "The label-image pair (%s, %s) do not look like the right pair because the filenames are quite different. Are you sure about the pairing? Please see data/pix2pix_dataset.py to see what is going on, and use --no_pairing_check to bypass this." % ( + path1, path2) + + self.label_paths = label_paths + self.image_paths = image_paths + self.instance_paths = instance_paths + + size = len(self.label_paths) + self.dataset_size = size + + def get_paths(self, opt): + label_paths = [] + image_paths = [] + instance_paths = [] + assert False, "A subclass of Pix2pixDataset must override self.get_paths(self, opt)" + return label_paths, image_paths, instance_paths + + def paths_match(self, path1, path2): + filename1_without_ext = os.path.splitext(os.path.basename(path1))[0] + filename2_without_ext = os.path.splitext(os.path.basename(path2))[0] + return filename1_without_ext == filename2_without_ext + + def __getitem__(self, index): + # Label Image + label_path = self.label_paths[index] + label = Image.open(label_path) + params = get_params(self.opt, label.size) + transform_label = get_transform(self.opt, params, method=Image.NEAREST, normalize=False) + label_tensor = transform_label(label) * 255.0 + label_tensor[label_tensor == 255] = self.opt.label_nc # 'unknown' is opt.label_nc + + # input image (real images) + image_path = self.image_paths[index] + assert self.paths_match(label_path, image_path), \ + "The label_path %s and image_path %s don't match." % \ + (label_path, image_path) + image = Image.open(image_path) + image = image.convert('RGB') + + transform_image = get_transform(self.opt, params) + image_tensor = transform_image(image) + + # if using instance maps + if self.opt.no_instance: + instance_tensor = 0 + else: + instance_path = self.instance_paths[index] + instance = Image.open(instance_path) + if instance.mode == 'L': + instance_tensor = transform_label(instance) * 255 + instance_tensor = instance_tensor.long() + else: + instance_tensor = transform_label(instance) + + input_dict = {'label': label_tensor, + 'instance': instance_tensor, + 'image': image_tensor, + 'path': image_path, + } + + # Give subclasses a chance to modify the final color_texture + + self.postprocess(input_dict) + + return input_dict + + def postprocess(self, input_dict): + return input_dict + + def __len__(self): + return self.dataset_size + + # Our codes get input images and labels + def get_input_by_names(self, image_path, image, label_img): + label = Image.fromarray(label_img) + params = get_params(self.opt, label.size) + transform_label = get_transform(self.opt, params, method=Image.NEAREST, normalize=False) + label_tensor = transform_label(label) * 255.0 + label_tensor[label_tensor == 255] = self.opt.label_nc # 'unknown' is opt.label_nc + label_tensor.unsqueeze_(0) + + # input image (real images)] + # image = Image.open(image_path) + # image = image.convert('RGB') + + transform_image = get_transform(self.opt, params) + image_tensor = transform_image(image) + image_tensor.unsqueeze_(0) + + # if using instance maps + if self.opt.no_instance: + instance_tensor = torch.Tensor([0]) + + input_dict = {'label': label_tensor, + 'instance': instance_tensor, + 'image': image_tensor, + 'path': image_path, + } + + # Give subclasses a chance to modify the final color_texture + self.postprocess(input_dict) + + return input_dict diff --git a/models/sean_codes/get_mean_code.py b/models/sean_codes/get_mean_code.py new file mode 100644 index 0000000000000000000000000000000000000000..31f5dae9633b5aa20a42722113878f7b16c06990 --- /dev/null +++ b/models/sean_codes/get_mean_code.py @@ -0,0 +1,45 @@ +# -*- coding: utf-8 -*- + +""" +# File name: get_mean_code.py +# Time : 2022/2/22 17:22 +# Author: xyguoo@163.com +# Description: +""" + +from glob import glob + +import numpy as np +import os + +layers_list = ['ACE.npy'] + +style_list = [] + +for cat_i in range(19): + for layer_j in layers_list: + tmp_list = glob('styles_test/style_codes/*/' + str(cat_i) + '/' + layer_j) + style_list = [] + + for k in tmp_list: + style_list.append(np.load(k)) + + if len(style_list) > 0: + style_list = np.array(style_list) + + style_list_norm2 = np.linalg.norm(style_list, axis=1, keepdims=True) ** 2 + dist_matrix = (style_list_norm2 + style_list_norm2.T -2 * style_list @ style_list.T) + dist_matrix[dist_matrix < 0] = 0 + dist_matrix = dist_matrix ** 0.5 + median_idx = dist_matrix.sum(axis=1).argmin() + feature = style_list[median_idx] + + save_folder = os.path.join('styles_test/mean_style_code/median', str(cat_i)) + + if not os.path.exists(save_folder): + os.makedirs(save_folder) + + save_name = os.path.join(save_folder, layer_j) + np.save(save_name, feature) + +print(100) diff --git a/models/sean_codes/models/__init__.py b/models/sean_codes/models/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..810d272b2e64dbeb6c7b941221dba59d1a62df8a --- /dev/null +++ b/models/sean_codes/models/__init__.py @@ -0,0 +1,44 @@ +""" +Copyright (C) 2019 NVIDIA Corporation. All rights reserved. +Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode). +""" + +import importlib +import torch + + +def find_model_using_name(model_name): + # Given the option --model [modelname], + # the file "models/modelname_model.py" + # will be imported. + model_filename = "sean_codes.models." + model_name + "_model" + modellib = importlib.import_module(model_filename) + + # In the file, the class called ModelNameModel() will + # be instantiated. It has to be a subclass of torch.nn.Module, + # and it is case-insensitive. + model = None + target_model_name = model_name.replace('_', '') + 'model' + for name, cls in modellib.__dict__.items(): + if name.lower() == target_model_name.lower() \ + and issubclass(cls, torch.nn.Module): + model = cls + + if model is None: + print("In %s.py, there should be a subclass of torch.nn.Module with class name that matches %s in lowercase." % (model_filename, target_model_name)) + exit(0) + + return model + + +def get_option_setter(model_name): + model_class = find_model_using_name(model_name) + return model_class.modify_commandline_options + + +def create_model(opt): + model = find_model_using_name(opt.sean_model) + instance = model(opt) + print("model [%s] was created" % (type(instance).__name__)) + + return instance diff --git a/models/sean_codes/models/networks/__init__.py b/models/sean_codes/models/networks/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..5b05d174064cb6fe736a8b091f0ed8e9c46e047b --- /dev/null +++ b/models/sean_codes/models/networks/__init__.py @@ -0,0 +1,62 @@ +""" +Copyright (C) 2019 NVIDIA Corporation. All rights reserved. +Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode). +""" + +from .base_network import BaseNetwork +from .loss import * +# from sean_codes.models.networks import * +from . import * +from .encoder import * +from models.sean_codes.util import util + + +def find_network_using_name(target_network_name, filename): + target_class_name = target_network_name + filename + module_name = 'sean_codes.models.networks.' + filename + network = util.find_class_in_module(target_class_name, module_name) + + assert issubclass(network, BaseNetwork), \ + "Class %s should be a subclass of BaseNetwork" % network + + return network + + +def modify_commandline_options(parser, is_train): + opt, _ = parser.parse_known_args() + + netG_cls = find_network_using_name(opt.netG, 'generator') + parser = netG_cls.modify_commandline_options(parser, is_train) + if is_train: + netD_cls = find_network_using_name(opt.netD, 'discriminator') + parser = netD_cls.modify_commandline_options(parser, is_train) + netE_cls = find_network_using_name('conv', 'encoder') + parser = netE_cls.modify_commandline_options(parser, is_train) + + return parser + + +def create_network(cls, opt): + net = cls(opt) + net.print_network() + if len(opt.gpu_ids) > 0: + assert(torch.cuda.is_available()) + net.cuda() + net.init_weights(opt.init_type, opt.init_variance) + return net + + +def define_G(opt): + netG_cls = find_network_using_name(opt.netG, 'generator') + return create_network(netG_cls, opt) + + +def define_D(opt): + netD_cls = find_network_using_name(opt.netD, 'discriminator') + return create_network(netD_cls, opt) + + +def define_E(opt): + # there exists only one encoder type + netE_cls = find_network_using_name('conv', 'encoder') + return create_network(netE_cls, opt) diff --git a/models/sean_codes/models/networks/architecture.py b/models/sean_codes/models/networks/architecture.py new file mode 100644 index 0000000000000000000000000000000000000000..c11afbcdb5e0eec9a91d52a68240578ba77ea306 --- /dev/null +++ b/models/sean_codes/models/networks/architecture.py @@ -0,0 +1,207 @@ +""" +Copyright (C) 2019 NVIDIA Corporation. All rights reserved. +Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode). +""" + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torchvision +import torch.nn.utils.spectral_norm as spectral_norm +from .normalization import ACE + + +# ResNet block that uses SPADE. +# It differs from the ResNet block of pix2pixHD in that +# it takes in the segmentation map as input, learns the skip connection if necessary, +# and applies normalization first and then convolution. +# This architecture seemed like a standard architecture for unconditional or +# class-conditional GAN architecture using residual block. +# The code was inspired from https://github.com/LMescheder/GAN_stability. +class SPADEResnetBlock(nn.Module): + def __init__(self, fin, fout, opt, Block_Name=None, use_rgb=True): + super().__init__() + + self.use_rgb = use_rgb + + self.Block_Name = Block_Name + self.status = opt.status + + # Attributes + self.learned_shortcut = (fin != fout) + fmiddle = min(fin, fout) + + # create conv layers + self.conv_0 = nn.Conv2d(fin, fmiddle, kernel_size=3, padding=1) + self.conv_1 = nn.Conv2d(fmiddle, fout, kernel_size=3, padding=1) + if self.learned_shortcut: + self.conv_s = nn.Conv2d(fin, fout, kernel_size=1, bias=False) + + # apply spectral norm if specified + if 'spectral' in opt.norm_G: + self.conv_0 = spectral_norm(self.conv_0) + self.conv_1 = spectral_norm(self.conv_1) + if self.learned_shortcut: + self.conv_s = spectral_norm(self.conv_s) + + # define normalization layers + spade_config_str = opt.norm_G.replace('spectral', '') + + ########### Modifications 1 + normtype_list = ['spadeinstance3x3', 'spadesyncbatch3x3', 'spadebatch3x3'] + our_norm_type = 'spadesyncbatch3x3' + + self.ace_0 = ACE(our_norm_type, fin, 3, ACE_Name=Block_Name + '_ACE_0', status=self.status, + spade_params=[spade_config_str, fin, opt.semantic_nc], use_rgb=use_rgb) + ########### Modifications 1 + + ########### Modifications 1 + self.ace_1 = ACE(our_norm_type, fmiddle, 3, ACE_Name=Block_Name + '_ACE_1', status=self.status, + spade_params=[spade_config_str, fmiddle, opt.semantic_nc], use_rgb=use_rgb) + ########### Modifications 1 + + if self.learned_shortcut: + self.ace_s = ACE(our_norm_type, fin, 3, ACE_Name=Block_Name + '_ACE_s', status=self.status, + spade_params=[spade_config_str, fin, opt.semantic_nc], use_rgb=use_rgb) + + # note the resnet block with SPADE also takes in |seg|, + # the semantic segmentation map as input + def forward(self, x, seg, style_codes, obj_dic=None): + + x_s = self.shortcut(x, seg, style_codes, obj_dic) + + ########### Modifications 1 + dx = self.ace_0(x, seg, style_codes, obj_dic) + + dx = self.conv_0(self.actvn(dx)) + + dx = self.ace_1(dx, seg, style_codes, obj_dic) + + dx = self.conv_1(self.actvn(dx)) + ########### Modifications 1 + + out = x_s + dx + return out + + def shortcut(self, x, seg, style_codes, obj_dic): + if self.learned_shortcut: + x_s = self.ace_s(x, seg, style_codes, obj_dic) + x_s = self.conv_s(x_s) + + else: + x_s = x + return x_s + + def actvn(self, x): + return F.leaky_relu(x, 2e-1) + + +# ResNet block used in pix2pixHD +# We keep the same architecture as pix2pixHD. +class ResnetBlock(nn.Module): + def __init__(self, dim, norm_layer, activation=nn.ReLU(False), kernel_size=3): + super().__init__() + + pw = (kernel_size - 1) // 2 + self.conv_block = nn.Sequential( + nn.ReflectionPad2d(pw), + norm_layer(nn.Conv2d(dim, dim, kernel_size=kernel_size)), + activation, + nn.ReflectionPad2d(pw), + norm_layer(nn.Conv2d(dim, dim, kernel_size=kernel_size)) + ) + + def forward(self, x): + y = self.conv_block(x) + out = x + y + return out + + +# VGG architecter, used for the perceptual loss using a pretrained VGG network +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 + + +class Zencoder(torch.nn.Module): + def __init__(self, input_nc, output_nc, ngf=32, n_downsampling=2, norm_layer=nn.InstanceNorm2d): + super(Zencoder, self).__init__() + self.output_nc = output_nc + + model = [nn.ReflectionPad2d(1), nn.Conv2d(input_nc, ngf, kernel_size=3, padding=0), + norm_layer(ngf), nn.LeakyReLU(0.2, False)] + ### downsample + for i in range(n_downsampling): + mult = 2 ** i + model += [nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3, stride=2, padding=1), + norm_layer(ngf * mult * 2), nn.LeakyReLU(0.2, False)] + + ### upsample + for i in range(1): + mult = 2 ** (n_downsampling - i) + model += [ + nn.ConvTranspose2d(ngf * mult, ngf * mult * 2, kernel_size=3, stride=2, padding=1, output_padding=1), + norm_layer(int(ngf * mult / 2)), nn.LeakyReLU(0.2, False)] + + model += [nn.ReflectionPad2d(1), nn.Conv2d(256, output_nc, kernel_size=3, padding=0), nn.Tanh()] + self.model = nn.Sequential(*model) + + def forward(self, input, segmap): + + codes = self.model(input) + # codes shape: [N, 512, 128, 128] + segmap = F.interpolate(segmap, size=codes.size()[2:], mode='nearest') + # segmap shape: [N, 19, 128, 128] + # print(segmap.shape) + # print(codes.shape) + + b_size = codes.shape[0] + # h_size = codes.shape[2] + # w_size = codes.shape[3] + f_size = codes.shape[1] + + s_size = segmap.shape[1] + + codes_vector = torch.zeros((b_size, s_size, f_size), dtype=codes.dtype, device=codes.device) + + for i in range(b_size): + for j in range(s_size): + component_mask_area = torch.sum(segmap.bool()[i, j]) + + if component_mask_area > 0: + codes_component_feature = codes[i].masked_select(segmap.bool()[i, j]).reshape(f_size, + component_mask_area).mean( + 1) + codes_vector[i][j] = codes_component_feature + + # codes_avg[i].masked_scatter_(segmap.bool()[i, j], codes_component_mu) + + return codes_vector diff --git a/models/sean_codes/models/networks/base_network.py b/models/sean_codes/models/networks/base_network.py new file mode 100644 index 0000000000000000000000000000000000000000..f79189f3c4c9b0032ea5e8ec28c8915a9112a754 --- /dev/null +++ b/models/sean_codes/models/networks/base_network.py @@ -0,0 +1,59 @@ +""" +Copyright (C) 2019 NVIDIA Corporation. All rights reserved. +Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode). +""" + +import torch.nn as nn +from torch.nn import init + + +class BaseNetwork(nn.Module): + def __init__(self): + super(BaseNetwork, self).__init__() + + @staticmethod + def modify_commandline_options(parser, is_train): + return parser + + def print_network(self): + if isinstance(self, list): + self = self[0] + num_params = 0 + for param in self.parameters(): + num_params += param.numel() + print('Network [%s] was created. Total number of parameters: %.1f million. ' + 'To see the architecture, do print(network).' + % (type(self).__name__, num_params / 1000000)) + + def init_weights(self, init_type='normal', gain=0.02): + def init_func(m): + classname = m.__class__.__name__ + if classname.find('BatchNorm2d') != -1: + if hasattr(m, 'weight') and m.weight is not None: + init.normal_(m.weight.data, 1.0, gain) + if hasattr(m, 'bias') and m.bias is not None: + init.constant_(m.bias.data, 0.0) + elif hasattr(m, 'weight') and (classname.find('Conv') != -1 or classname.find('Linear') != -1): + if init_type == 'normal': + init.normal_(m.weight.data, 0.0, gain) + elif init_type == 'xavier': + init.xavier_normal_(m.weight.data, gain=gain) + elif init_type == 'xavier_uniform': + init.xavier_uniform_(m.weight.data, gain=1.0) + 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=gain) + elif init_type == 'none': # uses pytorch's default init method + m.reset_parameters() + 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) + + self.apply(init_func) + + # propagate to children + for m in self.children(): + if hasattr(m, 'init_weights'): + m.init_weights(init_type, gain) diff --git a/models/sean_codes/models/networks/discriminator.py b/models/sean_codes/models/networks/discriminator.py new file mode 100644 index 0000000000000000000000000000000000000000..215f1bfae07d95d74c9c40799988e3ab5124c0bf --- /dev/null +++ b/models/sean_codes/models/networks/discriminator.py @@ -0,0 +1,119 @@ +""" +Copyright (C) 2019 NVIDIA Corporation. All rights reserved. +Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode). +""" + +import torch.nn as nn +import numpy as np +import torch.nn.functional as F +from sean_codes.models.networks.base_network import BaseNetwork +from sean_codes.models.networks.normalization import get_nonspade_norm_layer +import util.util as util + + +class MultiscaleDiscriminator(BaseNetwork): + @staticmethod + def modify_commandline_options(parser, is_train): + parser.add_argument('--netD_subarch', type=str, default='n_layer', + help='architecture of each discriminator') + parser.add_argument('--num_D', type=int, default=2, + help='number of discriminators to be used in multiscale') + opt, _ = parser.parse_known_args() + + # define properties of each discriminator of the multiscale discriminator + subnetD = util.find_class_in_module(opt.netD_subarch + 'discriminator', + 'models.networks.discriminator') + subnetD.modify_commandline_options(parser, is_train) + + return parser + + def __init__(self, opt): + super().__init__() + self.opt = opt + + for i in range(opt.num_D): + subnetD = self.create_single_discriminator(opt) + self.add_module('discriminator_%d' % i, subnetD) + + def create_single_discriminator(self, opt): + subarch = opt.netD_subarch + if subarch == 'n_layer': + netD = NLayerDiscriminator(opt) + else: + raise ValueError('unrecognized discriminator subarchitecture %s' % subarch) + return netD + + def downsample(self, input): + return F.avg_pool2d(input, kernel_size=3, + stride=2, padding=[1, 1], + count_include_pad=False) + + # Returns list of lists of discriminator outputs. + # The final result is of size opt.num_D x opt.n_layers_D + def forward(self, input): + result = [] + get_intermediate_features = not self.opt.no_ganFeat_loss + for name, D in self.named_children(): + out = D(input) + if not get_intermediate_features: + out = [out] + result.append(out) + input = self.downsample(input) + + return result + + +# Defines the PatchGAN discriminator with the specified arguments. +class NLayerDiscriminator(BaseNetwork): + @staticmethod + def modify_commandline_options(parser, is_train): + parser.add_argument('--n_layers_D', type=int, default=3, + help='# layers in each discriminator') + return parser + + def __init__(self, opt): + super().__init__() + self.opt = opt + + kw = 4 + padw = int(np.ceil((kw - 1.0) / 2)) + nf = opt.ndf + input_nc = self.compute_D_input_nc(opt) + + norm_layer = get_nonspade_norm_layer(opt, opt.norm_D) + sequence = [[nn.Conv2d(input_nc, nf, kernel_size=kw, stride=2, padding=padw), + nn.LeakyReLU(0.2, False)]] + + for n in range(1, opt.n_layers_D): + nf_prev = nf + nf = min(nf * 2, 512) + sequence += [[norm_layer(nn.Conv2d(nf_prev, nf, kernel_size=kw, + stride=2, padding=padw)), + nn.LeakyReLU(0.2, False) + ]] + + sequence += [[nn.Conv2d(nf, 1, kernel_size=kw, stride=1, padding=padw)]] + + # We divide the layers into groups to extract intermediate layer outputs + for n in range(len(sequence)): + self.add_module('model' + str(n), nn.Sequential(*sequence[n])) + + def compute_D_input_nc(self, opt): + input_nc = opt.label_nc + opt.output_nc + if opt.contain_dontcare_label: + input_nc += 1 + if not opt.no_instance: + input_nc += 1 + return input_nc + + def forward(self, input): + results = [input] + for submodel in self.children(): + intermediate_output = submodel(results[-1]) + results.append(intermediate_output) + + get_intermediate_features = not self.opt.no_ganFeat_loss + if get_intermediate_features: + return results[1:] + else: + return results[-1] diff --git a/models/sean_codes/models/networks/encoder.py b/models/sean_codes/models/networks/encoder.py new file mode 100644 index 0000000000000000000000000000000000000000..43c6a74a5e7f0c0326a4bd3bf124dd56965b4e78 --- /dev/null +++ b/models/sean_codes/models/networks/encoder.py @@ -0,0 +1,55 @@ +""" +Copyright (C) 2019 NVIDIA Corporation. All rights reserved. +Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode). +""" + +import torch.nn as nn +import numpy as np +import torch.nn.functional as F +from .base_network import BaseNetwork +from .normalization import get_nonspade_norm_layer + + +class ConvEncoder(BaseNetwork): + """ Same architecture as the image discriminator """ + + def __init__(self, opt): + super().__init__() + + kw = 3 + pw = int(np.ceil((kw - 1.0) / 2)) + ndf = opt.ngf + norm_layer = get_nonspade_norm_layer(opt, opt.norm_E) + self.layer1 = norm_layer(nn.Conv2d(3, ndf, kw, stride=2, padding=pw)) + self.layer2 = norm_layer(nn.Conv2d(ndf * 1, ndf * 2, kw, stride=2, padding=pw)) + self.layer3 = norm_layer(nn.Conv2d(ndf * 2, ndf * 4, kw, stride=2, padding=pw)) + self.layer4 = norm_layer(nn.Conv2d(ndf * 4, ndf * 8, kw, stride=2, padding=pw)) + self.layer5 = norm_layer(nn.Conv2d(ndf * 8, ndf * 8, kw, stride=2, padding=pw)) + if opt.crop_size >= 256: + self.layer6 = norm_layer(nn.Conv2d(ndf * 8, ndf * 8, kw, stride=2, padding=pw)) + + self.so = s0 = 4 + self.fc_mu = nn.Linear(ndf * 8 * s0 * s0, 256) + self.fc_var = nn.Linear(ndf * 8 * s0 * s0, 256) + + self.actvn = nn.LeakyReLU(0.2, False) + self.opt = opt + + def forward(self, x): + if x.size(2) != 256 or x.size(3) != 256: + x = F.interpolate(x, size=(256, 256), mode='bilinear') + + x = self.layer1(x) + x = self.layer2(self.actvn(x)) + x = self.layer3(self.actvn(x)) + x = self.layer4(self.actvn(x)) + x = self.layer5(self.actvn(x)) + if self.opt.crop_size >= 256: + x = self.layer6(self.actvn(x)) + x = self.actvn(x) + + x = x.view(x.size(0), -1) + mu = self.fc_mu(x) + logvar = self.fc_var(x) + + return mu, logvar diff --git a/models/sean_codes/models/networks/generator.py b/models/sean_codes/models/networks/generator.py new file mode 100644 index 0000000000000000000000000000000000000000..4600e6e15d352f183537756d737e6aac221de503 --- /dev/null +++ b/models/sean_codes/models/networks/generator.py @@ -0,0 +1,172 @@ +""" +Copyright (C) 2019 NVIDIA Corporation. All rights reserved. +Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode). +""" + +import torch.nn as nn +import torch.nn.functional as F +from .base_network import BaseNetwork +from .architecture import SPADEResnetBlock as SPADEResnetBlock +from .architecture import Zencoder +import torch + + +class SPADEGenerator(BaseNetwork): + @staticmethod + def modify_commandline_options(parser, is_train): + parser.set_defaults(norm_G='spectralspadesyncbatch3x3') + parser.add_argument('--num_upsampling_layers', + choices=('normal', 'more', 'most'), default='normal', + help="If 'more', adds upsampling layer between the two middle resnet blocks. If 'most', also add one more upsampling + resnet layer at the end of the generator") + + return parser + + def __init__(self, opt): + super().__init__() + self.opt = opt + nf = opt.ngf + + self.sw, self.sh = self.compute_latent_vector_size(opt) + + self.Zencoder = Zencoder(3, 512) + + self.fc = nn.Conv2d(self.opt.semantic_nc, 16 * nf, 3, padding=1) + + self.head_0 = SPADEResnetBlock(16 * nf, 16 * nf, opt, Block_Name='head_0') + + self.G_middle_0 = SPADEResnetBlock(16 * nf, 16 * nf, opt, Block_Name='G_middle_0') + self.G_middle_1 = SPADEResnetBlock(16 * nf, 16 * nf, opt, Block_Name='G_middle_1') + + self.up_0 = SPADEResnetBlock(16 * nf, 8 * nf, opt, Block_Name='up_0') + self.up_1 = SPADEResnetBlock(8 * nf, 4 * nf, opt, Block_Name='up_1') + self.up_2 = SPADEResnetBlock(4 * nf, 2 * nf, opt, Block_Name='up_2') + self.up_3 = SPADEResnetBlock(2 * nf, 1 * nf, opt, Block_Name='up_3', use_rgb=False) + + final_nc = nf + + if opt.num_upsampling_layers == 'most': + self.up_4 = SPADEResnetBlock(1 * nf, nf // 2, opt, Block_Name='up_4') + final_nc = nf // 2 + + self.conv_img = nn.Conv2d(final_nc, 3, 3, padding=1) + + self.up = nn.Upsample(scale_factor=2) + # self.up = nn.Upsample(scale_factor=2, mode='bilinear') + + def compute_latent_vector_size(self, opt): + if opt.num_upsampling_layers == 'normal': + num_up_layers = 5 + elif opt.num_upsampling_layers == 'more': + num_up_layers = 6 + elif opt.num_upsampling_layers == 'most': + num_up_layers = 7 + else: + raise ValueError('opt.num_upsampling_layers [%s] not recognized' % + opt.num_upsampling_layers) + + sw = opt.crop_size // (2 ** num_up_layers) + sh = round(sw / opt.aspect_ratio) + + return sw, sh + + def forward(self, input, rgb_img, obj_dic=None): + seg = input + + x = F.interpolate(seg, size=(self.sh, self.sw)) + x = self.fc(x) + + if rgb_img is None: + style_codes = None + else: + style_codes = self.Zencoder(input=rgb_img, segmap=seg) + + x = self.head_0(x, seg, style_codes, obj_dic=obj_dic) + + x = self.up(x) + x = self.G_middle_0(x, seg, style_codes, obj_dic=obj_dic) + + if self.opt.num_upsampling_layers == 'more' or \ + self.opt.num_upsampling_layers == 'most': + x = self.up(x) + + x = self.G_middle_1(x, seg, style_codes, obj_dic=obj_dic) + + x = self.up(x) + x = self.up_0(x, seg, style_codes, obj_dic=obj_dic) + x = self.up(x) + x = self.up_1(x, seg, style_codes, obj_dic=obj_dic) + x = self.up(x) + x = self.up_2(x, seg, style_codes, obj_dic=obj_dic) + x = self.up(x) + x = self.up_3(x, seg, style_codes, obj_dic=obj_dic) + + # if self.opt.num_upsampling_layers == 'most': + # x = self.up(x) + # x= self.up_4(x, seg, style_codes, obj_dic=obj_dic) + + x = self.conv_img(F.leaky_relu(x, 2e-1)) + x = torch.tanh(x) + return x + +# class Pix2PixHDGenerator(BaseNetwork): +# @staticmethod +# def modify_commandline_options(parser, is_train): +# parser.add_argument('--resnet_n_downsample', type=int, default=4, help='number of downsampling layers in netG') +# parser.add_argument('--resnet_n_blocks', type=int, default=9, help='number of residual blocks in the global generator network') +# parser.add_argument('--resnet_kernel_size', type=int, default=3, +# help='kernel size of the resnet block') +# parser.add_argument('--resnet_initial_kernel_size', type=int, default=7, +# help='kernel size of the first convolution') +# parser.set_defaults(norm_G='instance') +# return parser +# +# def __init__(self, opt): +# super().__init__() +# input_nc = opt.label_nc + (1 if opt.contain_dontcare_label else 0) + (0 if opt.no_instance else 1) +# +# norm_layer = get_nonspade_norm_layer(opt, opt.norm_G) +# activation = nn.ReLU(False) +# +# model = [] +# +# # initial conv +# model += [nn.ReflectionPad2d(opt.resnet_initial_kernel_size // 2), +# norm_layer(nn.Conv2d(input_nc, opt.ngf, +# kernel_size=opt.resnet_initial_kernel_size, +# padding=0)), +# activation] +# +# # downsample +# mult = 1 +# for i in range(opt.resnet_n_downsample): +# model += [norm_layer(nn.Conv2d(opt.ngf * mult, opt.ngf * mult * 2, +# kernel_size=3, stride=2, padding=1)), +# activation] +# mult *= 2 +# +# # resnet blocks +# for i in range(opt.resnet_n_blocks): +# model += [ResnetBlock(opt.ngf * mult, +# norm_layer=norm_layer, +# activation=activation, +# kernel_size=opt.resnet_kernel_size)] +# +# # upsample +# for i in range(opt.resnet_n_downsample): +# nc_in = int(opt.ngf * mult) +# nc_out = int((opt.ngf * mult) / 2) +# model += [norm_layer(nn.ConvTranspose2d(nc_in, nc_out, +# kernel_size=3, stride=2, +# padding=1, output_padding=1)), +# activation] +# mult = mult // 2 +# +# # final color_texture conv +# model += [nn.ReflectionPad2d(3), +# nn.Conv2d(nc_out, opt.output_nc, kernel_size=7, padding=0), +# nn.Tanh()] +# +# self.model = nn.Sequential(*model) +# +# def forward(self, input, z=None): +# return self.model(input) diff --git a/models/sean_codes/models/networks/loss.py b/models/sean_codes/models/networks/loss.py new file mode 100644 index 0000000000000000000000000000000000000000..f443cd71ca22a087949e11285996a153307f19cf --- /dev/null +++ b/models/sean_codes/models/networks/loss.py @@ -0,0 +1,120 @@ +""" +Copyright (C) 2019 NVIDIA Corporation. All rights reserved. +Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode). +""" + +import torch +import torch.nn as nn +import torch.nn.functional as F +from .architecture import VGG19 + + +# Defines the GAN loss which uses either LSGAN or the regular GAN. +# When LSGAN is used, it is basically same as MSELoss, +# but it abstracts away the need to create the target label tensor +# that has the same size as the input +class GANLoss(nn.Module): + def __init__(self, gan_mode, target_real_label=1.0, target_fake_label=0.0, + tensor=torch.FloatTensor, opt=None): + super(GANLoss, self).__init__() + self.real_label = target_real_label + self.fake_label = target_fake_label + self.real_label_tensor = None + self.fake_label_tensor = None + self.zero_tensor = None + self.Tensor = tensor + self.gan_mode = gan_mode + self.opt = opt + if gan_mode == 'ls': + pass + elif gan_mode == 'original': + pass + elif gan_mode == 'w': + pass + elif gan_mode == 'hinge': + pass + else: + raise ValueError('Unexpected gan_mode {}'.format(gan_mode)) + + def get_target_tensor(self, input, target_is_real): + if target_is_real: + if self.real_label_tensor is None: + self.real_label_tensor = self.Tensor(1).fill_(self.real_label) + self.real_label_tensor.requires_grad_(False) + return self.real_label_tensor.expand_as(input) + else: + if self.fake_label_tensor is None: + self.fake_label_tensor = self.Tensor(1).fill_(self.fake_label) + self.fake_label_tensor.requires_grad_(False) + return self.fake_label_tensor.expand_as(input) + + def get_zero_tensor(self, input): + if self.zero_tensor is None: + self.zero_tensor = self.Tensor(1).fill_(0) + self.zero_tensor.requires_grad_(False) + return self.zero_tensor.expand_as(input) + + def loss(self, input, target_is_real, for_discriminator=True): + if self.gan_mode == 'original': # cross entropy loss + target_tensor = self.get_target_tensor(input, target_is_real) + loss = F.binary_cross_entropy_with_logits(input, target_tensor) + return loss + elif self.gan_mode == 'ls': + target_tensor = self.get_target_tensor(input, target_is_real) + return F.mse_loss(input, target_tensor) + elif self.gan_mode == 'hinge': + if for_discriminator: + if target_is_real: + minval = torch.min(input - 1, self.get_zero_tensor(input)) + loss = -torch.mean(minval) + else: + minval = torch.min(-input - 1, self.get_zero_tensor(input)) + loss = -torch.mean(minval) + else: + assert target_is_real, "The generator's hinge loss must be aiming for real" + loss = -torch.mean(input) + return loss + else: + # wgan + if target_is_real: + return -input.mean() + else: + return input.mean() + + def __call__(self, input, target_is_real, for_discriminator=True): + # computing loss is a bit complicated because |input| may not be + # a tensor, but list of tensors in case of multiscale discriminator + if isinstance(input, list): + loss = 0 + for pred_i in input: + if isinstance(pred_i, list): + pred_i = pred_i[-1] + loss_tensor = self.loss(pred_i, target_is_real, for_discriminator) + bs = 1 if len(loss_tensor.size()) == 0 else loss_tensor.size(0) + new_loss = torch.mean(loss_tensor.view(bs, -1), dim=1) + loss += new_loss + return loss / len(input) + else: + return self.loss(input, target_is_real, for_discriminator) + + +# Perceptual loss that uses a pretrained VGG network +class VGGLoss(nn.Module): + def __init__(self, gpu_ids): + 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 + + +# KL Divergence loss used in VAE with an image encoder +# class KLDLoss(nn.Module): +# def forward(self, mu, logvar): +# return -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp()) diff --git a/models/sean_codes/models/networks/normalization.py b/models/sean_codes/models/networks/normalization.py new file mode 100644 index 0000000000000000000000000000000000000000..b6099805ed5ac8423291340322d039252d8bd7b3 --- /dev/null +++ b/models/sean_codes/models/networks/normalization.py @@ -0,0 +1,257 @@ +""" +Copyright (C) 2019 NVIDIA Corporation. All rights reserved. +Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode). +""" + +import re +import torch +import torch.nn as nn +import torch.nn.functional as F +from .sync_batchnorm import SynchronizedBatchNorm2d +import torch.nn.utils.spectral_norm as spectral_norm +import os +import numpy as np + + +# Returns a function that creates a normalization function +# that does not condition on semantic map +def get_nonspade_norm_layer(opt, norm_type='instance'): + # helper function to get # color_texture channels of the previous layer + def get_out_channel(layer): + if hasattr(layer, 'out_channels'): + return getattr(layer, 'out_channels') + return layer.weight.size(0) + + # this function will be returned + def add_norm_layer(layer): + nonlocal norm_type + if norm_type.startswith('spectral'): + layer = spectral_norm(layer) + subnorm_type = norm_type[len('spectral'):] + + if subnorm_type == 'none' or len(subnorm_type) == 0: + return layer + + # remove bias in the previous layer, which is meaningless + # since it has no effect after normalization + if getattr(layer, 'bias', None) is not None: + delattr(layer, 'bias') + layer.register_parameter('bias', None) + + if subnorm_type == 'batch': + norm_layer = nn.BatchNorm2d(get_out_channel(layer), affine=True) + elif subnorm_type == 'sync_batch': + norm_layer = SynchronizedBatchNorm2d(get_out_channel(layer), affine=True) + elif subnorm_type == 'instance': + norm_layer = nn.InstanceNorm2d(get_out_channel(layer), affine=False) + else: + raise ValueError('normalization layer %s is not recognized' % subnorm_type) + + return nn.Sequential(layer, norm_layer) + + return add_norm_layer + + +# Creates SPADE normalization layer based on the given configuration +# SPADE consists of two steps. First, it normalizes the activations using +# your favorite normalization method, such as Batch Norm or Instance Norm. +# Second, it applies scale and bias to the normalized color_texture, conditioned on +# the segmentation map. +# The format of |config_text| is spade(norm)(ks), where +# (norm) specifies the type of parameter-free normalization. +# (e.g. syncbatch, batch, instance) +# (ks) specifies the size of kernel in the SPADE module (e.g. 3x3) +# Example |config_text| will be spadesyncbatch3x3, or spadeinstance5x5. +# Also, the other arguments are +# |norm_nc|: the #channels of the normalized activations, hence the color_texture dim of SPADE +# |label_nc|: the #channels of the input semantic map, hence the input dim of SPADE + + +class ACE(nn.Module): + def __init__(self, config_text, norm_nc, label_nc, ACE_Name=None, status='scripts', spade_params=None, use_rgb=True): + super().__init__() + + self.ACE_Name = ACE_Name + self.status = status + self.save_npy = True + self.Spade = SPADE(*spade_params) + self.use_rgb = use_rgb + self.style_length = 512 + self.blending_gamma = nn.Parameter(torch.zeros(1), requires_grad=True) + self.blending_beta = nn.Parameter(torch.zeros(1), requires_grad=True) + self.noise_var = nn.Parameter(torch.zeros(norm_nc), requires_grad=True) + + assert config_text.startswith('spade') + parsed = re.search('spade(\D+)(\d)x\d', config_text) + param_free_norm_type = str(parsed.group(1)) + ks = int(parsed.group(2)) + pw = ks // 2 + + if param_free_norm_type == 'instance': + self.param_free_norm = nn.InstanceNorm2d(norm_nc, affine=False) + elif param_free_norm_type == 'syncbatch': + self.param_free_norm = SynchronizedBatchNorm2d(norm_nc, affine=False) + elif param_free_norm_type == 'batch': + self.param_free_norm = nn.BatchNorm2d(norm_nc, affine=False) + else: + raise ValueError('%s is not a recognized param-free norm type in SPADE' + % param_free_norm_type) + + # The dimension of the intermediate embedding space. Yes, hardcoded. + + if self.use_rgb: + self.create_gamma_beta_fc_layers() + + self.conv_gamma = nn.Conv2d(self.style_length, norm_nc, kernel_size=ks, padding=pw) + self.conv_beta = nn.Conv2d(self.style_length, norm_nc, kernel_size=ks, padding=pw) + + def forward(self, x, segmap, style_codes=None, obj_dic=None): + + # Part 1. generate parameter-free normalized activations + added_noise = (torch.randn(x.shape[0], x.shape[3], x.shape[2], 1, device='cuda') * self.noise_var).transpose(1, 3) + normalized = self.param_free_norm(x + added_noise) + + # Part 2. produce scaling and bias conditioned on semantic map + segmap = F.interpolate(segmap, size=x.size()[2:], mode='nearest') + + if self.use_rgb: + [b_size, f_size, h_size, w_size] = normalized.shape + middle_avg = torch.zeros((b_size, self.style_length, h_size, w_size), device=normalized.device) + + if self.status == 'UI_mode': + ############## hard coding + + for i in range(1): + for j in range(segmap.shape[1]): + + component_mask_area = torch.sum(segmap.bool()[i, j]) + + if component_mask_area > 0: + if obj_dic is None: + print('wrong even it is the first input') + else: + style_code_tmp = obj_dic[str(j)]['ACE'] + + middle_mu = F.relu(self.__getattr__('fc_mu' + str(j))(style_code_tmp)) + component_mu = middle_mu.reshape(self.style_length, 1).expand(self.style_length, + component_mask_area) + + middle_avg[i].masked_scatter_(segmap.bool()[i, j], component_mu) + + else: + + for i in range(b_size): + for j in range(segmap.shape[1]): + component_mask_area = torch.sum(segmap.bool()[i, j]) + + if component_mask_area > 0: + + middle_mu = F.relu(self.__getattr__('fc_mu' + str(j))(style_codes[i][j])) + component_mu = middle_mu.reshape(self.style_length, 1).expand(self.style_length, + component_mask_area) + + middle_avg[i].masked_scatter_(segmap.bool()[i, j], component_mu) + + if self.status == 'test' and self.save_npy and self.ACE_Name == 'up_2_ACE_0': + tmp = style_codes[i][j].cpu().numpy() + # gil + if obj_dic[i].find('temp_path') != -1: + dir_path = os.path.split(obj_dic[i])[0] + else: + dir_path = 'styles_test' + im_name = os.path.basename(obj_dic[i]) + folder_path = os.path.join(dir_path, 'style_codes', im_name, str(j)) + ############### some problem with obj_dic[i] + + if not os.path.exists(folder_path): + os.makedirs(folder_path) + + style_code_path = os.path.join(folder_path, 'ACE.npy') + np.save(style_code_path, tmp) + + gamma_avg = self.conv_gamma(middle_avg) + beta_avg = self.conv_beta(middle_avg) + + gamma_spade, beta_spade = self.Spade(segmap) + + gamma_alpha = torch.sigmoid(self.blending_gamma) + beta_alpha = torch.sigmoid(self.blending_beta) + + gamma_final = gamma_alpha * gamma_avg + (1 - gamma_alpha) * gamma_spade + beta_final = beta_alpha * beta_avg + (1 - beta_alpha) * beta_spade + out = normalized * (1 + gamma_final) + beta_final + else: + gamma_spade, beta_spade = self.Spade(segmap) + gamma_final = gamma_spade + beta_final = beta_spade + out = normalized * (1 + gamma_final) + beta_final + + return out + + def create_gamma_beta_fc_layers(self): + + ################### These codes should be replaced with torch.nn.ModuleList + + style_length = self.style_length + + self.fc_mu0 = nn.Linear(style_length, style_length) + self.fc_mu1 = nn.Linear(style_length, style_length) + self.fc_mu2 = nn.Linear(style_length, style_length) + self.fc_mu3 = nn.Linear(style_length, style_length) + self.fc_mu4 = nn.Linear(style_length, style_length) + self.fc_mu5 = nn.Linear(style_length, style_length) + self.fc_mu6 = nn.Linear(style_length, style_length) + self.fc_mu7 = nn.Linear(style_length, style_length) + self.fc_mu8 = nn.Linear(style_length, style_length) + self.fc_mu9 = nn.Linear(style_length, style_length) + self.fc_mu10 = nn.Linear(style_length, style_length) + self.fc_mu11 = nn.Linear(style_length, style_length) + self.fc_mu12 = nn.Linear(style_length, style_length) + self.fc_mu13 = nn.Linear(style_length, style_length) + self.fc_mu14 = nn.Linear(style_length, style_length) + self.fc_mu15 = nn.Linear(style_length, style_length) + self.fc_mu16 = nn.Linear(style_length, style_length) + self.fc_mu17 = nn.Linear(style_length, style_length) + self.fc_mu18 = nn.Linear(style_length, style_length) + + +class SPADE(nn.Module): + def __init__(self, config_text, norm_nc, label_nc): + super().__init__() + + assert config_text.startswith('spade') + parsed = re.search('spade(\D+)(\d)x\d', config_text) + param_free_norm_type = str(parsed.group(1)) + ks = int(parsed.group(2)) + + if param_free_norm_type == 'instance': + self.param_free_norm = nn.InstanceNorm2d(norm_nc, affine=False) + elif param_free_norm_type == 'syncbatch': + self.param_free_norm = SynchronizedBatchNorm2d(norm_nc, affine=False) + elif param_free_norm_type == 'batch': + self.param_free_norm = nn.BatchNorm2d(norm_nc, affine=False) + else: + raise ValueError('%s is not a recognized param-free norm type in SPADE' + % param_free_norm_type) + + # The dimension of the intermediate embedding space. Yes, hardcoded. + nhidden = 128 + + pw = ks // 2 + self.mlp_shared = nn.Sequential( + nn.Conv2d(label_nc, nhidden, kernel_size=ks, padding=pw), + nn.ReLU() + ) + + self.mlp_gamma = nn.Conv2d(nhidden, norm_nc, kernel_size=ks, padding=pw) + self.mlp_beta = nn.Conv2d(nhidden, norm_nc, kernel_size=ks, padding=pw) + + def forward(self, segmap): + + inputmap = segmap + + actv = self.mlp_shared(inputmap) + gamma = self.mlp_gamma(actv) + beta = self.mlp_beta(actv) + + return gamma, beta diff --git a/models/sean_codes/models/networks/sync_batchnorm/__init__.py b/models/sean_codes/models/networks/sync_batchnorm/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..84ef0a02ec3d1649a62052c65ef1c75e2eaeb5bb --- /dev/null +++ b/models/sean_codes/models/networks/sync_batchnorm/__init__.py @@ -0,0 +1,13 @@ +# -*- coding: utf-8 -*- +# File : __init__.py +# Author : Jiayuan Mao +# Email : maojiayuan@gmail.com +# Date : 27/01/2018 +# +# This file is part of Synchronized-BatchNorm-PyTorch. +# https://github.com/vacancy/Synchronized-BatchNorm-PyTorch +# Distributed under MIT License. + +from .batchnorm import SynchronizedBatchNorm1d, SynchronizedBatchNorm2d, SynchronizedBatchNorm3d +from .batchnorm import convert_model +from .replicate import DataParallelWithCallback, patch_replication_callback diff --git a/models/sean_codes/models/networks/sync_batchnorm/batchnorm.py b/models/sean_codes/models/networks/sync_batchnorm/batchnorm.py new file mode 100644 index 0000000000000000000000000000000000000000..27f4b8bc7a24d19d5ab75e2d2bc2f5b23d2a95d2 --- /dev/null +++ b/models/sean_codes/models/networks/sync_batchnorm/batchnorm.py @@ -0,0 +1,361 @@ +# -*- coding: utf-8 -*- +# File : batchnorm.py +# Author : Jiayuan Mao +# Email : maojiayuan@gmail.com +# Date : 27/01/2018 +# +# This file is part of Synchronized-BatchNorm-PyTorch. +# https://github.com/vacancy/Synchronized-BatchNorm-PyTorch +# Distributed under MIT License. + +import collections + +import torch +import torch.nn.functional as F + +from torch.nn.modules.batchnorm import _BatchNorm +from torch.nn.parallel._functions import ReduceAddCoalesced, Broadcast + +from .comm import SyncMaster +from .replicate import DataParallelWithCallback + +__all__ = ['SynchronizedBatchNorm1d', 'SynchronizedBatchNorm2d', + 'SynchronizedBatchNorm3d', 'convert_model'] + + +def _sum_ft(tensor): + """sum over the first and last dimention""" + return tensor.sum(dim=0).sum(dim=-1) + + +def _unsqueeze_ft(tensor): + """add new dementions at the front and the tail""" + return tensor.unsqueeze(0).unsqueeze(-1) + + +_ChildMessage = collections.namedtuple('_ChildMessage', ['sum', 'ssum', 'sum_size']) +_MasterMessage = collections.namedtuple('_MasterMessage', ['sum', 'inv_std']) + + +class _SynchronizedBatchNorm(_BatchNorm): + def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=True): + super(_SynchronizedBatchNorm, self).__init__(num_features, eps=eps, momentum=momentum, affine=affine) + + self._sync_master = SyncMaster(self._data_parallel_master) + + self._is_parallel = False + self._parallel_id = None + self._slave_pipe = None + + def forward(self, input): + # If it is not parallel computation or is in evaluation mode, use PyTorch's implementation. + if not (self._is_parallel and self.training): + return F.batch_norm( + input, self.running_mean, self.running_var, self.weight, self.bias, + self.training, self.momentum, self.eps) + + # Resize the input to (B, C, -1). + input_shape = input.size() + input = input.view(input.size(0), self.num_features, -1) + + # Compute the sum and square-sum. + sum_size = input.size(0) * input.size(2) + input_sum = _sum_ft(input) + input_ssum = _sum_ft(input ** 2) + + # Reduce-and-broadcast the statistics. + if self._parallel_id == 0: + mean, inv_std = self._sync_master.run_master(_ChildMessage(input_sum, input_ssum, sum_size)) + else: + mean, inv_std = self._slave_pipe.run_slave(_ChildMessage(input_sum, input_ssum, sum_size)) + + # Compute the color_texture. + if self.affine: + # MJY:: Fuse the multiplication for speed. + output = (input - _unsqueeze_ft(mean)) * _unsqueeze_ft(inv_std * self.weight) + _unsqueeze_ft(self.bias) + else: + output = (input - _unsqueeze_ft(mean)) * _unsqueeze_ft(inv_std) + + # Reshape it. + return output.view(input_shape) + + def __data_parallel_replicate__(self, ctx, copy_id): + self._is_parallel = True + self._parallel_id = copy_id + + # parallel_id == 0 means master device. + if self._parallel_id == 0: + ctx.sync_master = self._sync_master + else: + self._slave_pipe = ctx.sync_master.register_slave(copy_id) + + def _data_parallel_master(self, intermediates): + """Reduce the sum and square-sum, compute the statistics, and broadcast it.""" + + # Always using same "device order" makes the ReduceAdd operation faster. + # Thanks to:: Tete Xiao (http://tetexiao.com/) + intermediates = sorted(intermediates, key=lambda i: i[1].sum.get_device()) + + to_reduce = [i[1][:2] for i in intermediates] + to_reduce = [j for i in to_reduce for j in i] # flatten + target_gpus = [i[1].sum.get_device() for i in intermediates] + + sum_size = sum([i[1].sum_size for i in intermediates]) + sum_, ssum = ReduceAddCoalesced.apply(target_gpus[0], 2, *to_reduce) + mean, inv_std = self._compute_mean_std(sum_, ssum, sum_size) + + broadcasted = Broadcast.apply(target_gpus, mean, inv_std) + + outputs = [] + for i, rec in enumerate(intermediates): + outputs.append((rec[0], _MasterMessage(*broadcasted[i*2:i*2+2]))) + + return outputs + + def _compute_mean_std(self, sum_, ssum, size): + """Compute the mean and standard-deviation with sum and square-sum. This method + also maintains the moving average on the master device.""" + assert size > 1, 'BatchNorm computes unbiased standard-deviation, which requires size > 1.' + mean = sum_ / size + sumvar = ssum - sum_ * mean + unbias_var = sumvar / (size - 1) + bias_var = sumvar / size + + self.running_mean = (1 - self.momentum) * self.running_mean + self.momentum * mean.data + self.running_var = (1 - self.momentum) * self.running_var + self.momentum * unbias_var.data + + return mean, bias_var.clamp(self.eps) ** -0.5 + + +class SynchronizedBatchNorm1d(_SynchronizedBatchNorm): + r"""Applies Synchronized Batch Normalization over a 2d or 3d input that is seen as a + mini-batch. + + .. math:: + + y = \frac{x - mean[x]}{ \sqrt{Var[x] + \epsilon}} * gamma + beta + + This module differs from the built-in PyTorch BatchNorm1d as the mean and + standard-deviation are reduced across all devices during training. + + For example, when one uses `nn.DataParallel` to wrap the network during + training, PyTorch's implementation normalize the tensor on each device using + the statistics only on that device, which accelerated the computation and + is also easy to implement, but the statistics might be inaccurate. + Instead, in this synchronized version, the statistics will be computed + over all training samples distributed on multiple devices. + + Note that, for one-GPU or CPU-only case, this module behaves exactly same + as the built-in PyTorch implementation. + + The mean and standard-deviation are calculated per-dimension over + the mini-batches and gamma and beta are learnable parameter vectors + of size C (where C is the input size). + + During training, this layer keeps a running estimate of its computed mean + and variance. The running sum is kept with a default momentum of 0.1. + + During evaluation, this running mean/variance is used for normalization. + + Because the BatchNorm is done over the `C` dimension, computing statistics + on `(N, L)` slices, it's common terminology to call this Temporal BatchNorm + + Args: + num_features: num_features from an expected input of size + `batch_size x num_features [x width]` + eps: a value added to the denominator for numerical stability. + Default: 1e-5 + momentum: the value used for the running_mean and running_var + computation. Default: 0.1 + affine: a boolean value that when set to ``True``, gives the layer learnable + affine parameters. Default: ``True`` + + Shape: + - Input: :math:`(N, C)` or :math:`(N, C, L)` + - Output: :math:`(N, C)` or :math:`(N, C, L)` (same shape as input) + + Examples: + >>> # With Learnable Parameters + >>> m = SynchronizedBatchNorm1d(100) + >>> # Without Learnable Parameters + >>> m = SynchronizedBatchNorm1d(100, affine=False) + >>> input = torch.autograd.Variable(torch.randn(20, 100)) + >>> color_texture = m(input) + """ + + def _check_input_dim(self, input): + if input.dim() != 2 and input.dim() != 3: + raise ValueError('expected 2D or 3D input (got {}D input)' + .format(input.dim())) + super(SynchronizedBatchNorm1d, self)._check_input_dim(input) + + +class SynchronizedBatchNorm2d(_SynchronizedBatchNorm): + r"""Applies Batch Normalization over a 4d input that is seen as a mini-batch + of 3d inputs + + .. math:: + + y = \frac{x - mean[x]}{ \sqrt{Var[x] + \epsilon}} * gamma + beta + + This module differs from the built-in PyTorch BatchNorm2d as the mean and + standard-deviation are reduced across all devices during training. + + For example, when one uses `nn.DataParallel` to wrap the network during + training, PyTorch's implementation normalize the tensor on each device using + the statistics only on that device, which accelerated the computation and + is also easy to implement, but the statistics might be inaccurate. + Instead, in this synchronized version, the statistics will be computed + over all training samples distributed on multiple devices. + + Note that, for one-GPU or CPU-only case, this module behaves exactly same + as the built-in PyTorch implementation. + + The mean and standard-deviation are calculated per-dimension over + the mini-batches and gamma and beta are learnable parameter vectors + of size C (where C is the input size). + + During training, this layer keeps a running estimate of its computed mean + and variance. The running sum is kept with a default momentum of 0.1. + + During evaluation, this running mean/variance is used for normalization. + + Because the BatchNorm is done over the `C` dimension, computing statistics + on `(N, H, W)` slices, it's common terminology to call this Spatial BatchNorm + + Args: + num_features: num_features from an expected input of + size batch_size x num_features x height x width + eps: a value added to the denominator for numerical stability. + Default: 1e-5 + momentum: the value used for the running_mean and running_var + computation. Default: 0.1 + affine: a boolean value that when set to ``True``, gives the layer learnable + affine parameters. Default: ``True`` + + Shape: + - Input: :math:`(N, C, H, W)` + - Output: :math:`(N, C, H, W)` (same shape as input) + + Examples: + >>> # With Learnable Parameters + >>> m = SynchronizedBatchNorm2d(100) + >>> # Without Learnable Parameters + >>> m = SynchronizedBatchNorm2d(100, affine=False) + >>> input = torch.autograd.Variable(torch.randn(20, 100, 35, 45)) + >>> color_texture = m(input) + """ + + def _check_input_dim(self, input): + if input.dim() != 4: + raise ValueError('expected 4D input (got {}D input)' + .format(input.dim())) + super(SynchronizedBatchNorm2d, self)._check_input_dim(input) + + +class SynchronizedBatchNorm3d(_SynchronizedBatchNorm): + r"""Applies Batch Normalization over a 5d input that is seen as a mini-batch + of 4d inputs + + .. math:: + + y = \frac{x - mean[x]}{ \sqrt{Var[x] + \epsilon}} * gamma + beta + + This module differs from the built-in PyTorch BatchNorm3d as the mean and + standard-deviation are reduced across all devices during training. + + For example, when one uses `nn.DataParallel` to wrap the network during + training, PyTorch's implementation normalize the tensor on each device using + the statistics only on that device, which accelerated the computation and + is also easy to implement, but the statistics might be inaccurate. + Instead, in this synchronized version, the statistics will be computed + over all training samples distributed on multiple devices. + + Note that, for one-GPU or CPU-only case, this module behaves exactly same + as the built-in PyTorch implementation. + + The mean and standard-deviation are calculated per-dimension over + the mini-batches and gamma and beta are learnable parameter vectors + of size C (where C is the input size). + + During training, this layer keeps a running estimate of its computed mean + and variance. The running sum is kept with a default momentum of 0.1. + + During evaluation, this running mean/variance is used for normalization. + + Because the BatchNorm is done over the `C` dimension, computing statistics + on `(N, D, H, W)` slices, it's common terminology to call this Volumetric BatchNorm + or Spatio-temporal BatchNorm + + Args: + num_features: num_features from an expected input of + size batch_size x num_features x depth x height x width + eps: a value added to the denominator for numerical stability. + Default: 1e-5 + momentum: the value used for the running_mean and running_var + computation. Default: 0.1 + affine: a boolean value that when set to ``True``, gives the layer learnable + affine parameters. Default: ``True`` + + Shape: + - Input: :math:`(N, C, D, H, W)` + - Output: :math:`(N, C, D, H, W)` (same shape as input) + + Examples: + >>> # With Learnable Parameters + >>> m = SynchronizedBatchNorm3d(100) + >>> # Without Learnable Parameters + >>> m = SynchronizedBatchNorm3d(100, affine=False) + >>> input = torch.autograd.Variable(torch.randn(20, 100, 35, 45, 10)) + >>> color_texture = m(input) + """ + + def _check_input_dim(self, input): + if input.dim() != 5: + raise ValueError('expected 5D input (got {}D input)' + .format(input.dim())) + super(SynchronizedBatchNorm3d, self)._check_input_dim(input) + + +def convert_model(module): + """Traverse the input module and its child recursively + and replace all instance of torch.nn.modules.batchnorm.BatchNorm*N*d + to SynchronizedBatchNorm*N*d + + Args: + module: the input module needs to be convert to SyncBN model + + Examples: + >>> import torch.nn as nn + >>> import torchvision + >>> # m is a standard pytorch model + >>> m = torchvision.models.resnet18(True) + >>> m = nn.DataParallel(m) + >>> # after convert, m is using SyncBN + >>> m = convert_model(m) + """ + if isinstance(module, torch.nn.DataParallel): + mod = module.module + mod = convert_model(mod) + mod = DataParallelWithCallback(mod) + return mod + + mod = module + for pth_module, sync_module in zip([torch.nn.modules.batchnorm.BatchNorm1d, + torch.nn.modules.batchnorm.BatchNorm2d, + torch.nn.modules.batchnorm.BatchNorm3d], + [SynchronizedBatchNorm1d, + SynchronizedBatchNorm2d, + SynchronizedBatchNorm3d]): + if isinstance(module, pth_module): + mod = sync_module(module.num_features, module.eps, module.momentum, module.affine) + mod.running_mean = module.running_mean + mod.running_var = module.running_var + if module.affine: + mod.weight.data = module.weight.data.clone().detach() + mod.bias.data = module.bias.data.clone().detach() + + for name, child in module.named_children(): + mod.add_module(name, convert_model(child)) + + return mod diff --git a/models/sean_codes/models/networks/sync_batchnorm/batchnorm_reimpl.py b/models/sean_codes/models/networks/sync_batchnorm/batchnorm_reimpl.py new file mode 100644 index 0000000000000000000000000000000000000000..7afcdaff9c56d7ac9c487f2dbe61fe6cb9c353a0 --- /dev/null +++ b/models/sean_codes/models/networks/sync_batchnorm/batchnorm_reimpl.py @@ -0,0 +1,74 @@ +#! /usr/bin/env python3 +# -*- coding: utf-8 -*- +# File : batchnorm_reimpl.py +# Author : acgtyrant +# Date : 11/01/2018 +# +# This file is part of Synchronized-BatchNorm-PyTorch. +# https://github.com/vacancy/Synchronized-BatchNorm-PyTorch +# Distributed under MIT License. + +import torch +import torch.nn as nn +import torch.nn.init as init + +__all__ = ['BatchNormReimpl'] + + +class BatchNorm2dReimpl(nn.Module): + """ + A re-implementation of batch normalization, used for testing the numerical + stability. + + Author: acgtyrant + See also: + https://github.com/vacancy/Synchronized-BatchNorm-PyTorch/issues/14 + """ + def __init__(self, num_features, eps=1e-5, momentum=0.1): + super().__init__() + + self.num_features = num_features + self.eps = eps + self.momentum = momentum + self.weight = nn.Parameter(torch.empty(num_features)) + self.bias = nn.Parameter(torch.empty(num_features)) + self.register_buffer('running_mean', torch.zeros(num_features)) + self.register_buffer('running_var', torch.ones(num_features)) + self.reset_parameters() + + def reset_running_stats(self): + self.running_mean.zero_() + self.running_var.fill_(1) + + def reset_parameters(self): + self.reset_running_stats() + init.uniform_(self.weight) + init.zeros_(self.bias) + + def forward(self, input_): + batchsize, channels, height, width = input_.size() + numel = batchsize * height * width + input_ = input_.permute(1, 0, 2, 3).contiguous().view(channels, numel) + sum_ = input_.sum(1) + sum_of_square = input_.pow(2).sum(1) + mean = sum_ / numel + sumvar = sum_of_square - sum_ * mean + + self.running_mean = ( + (1 - self.momentum) * self.running_mean + + self.momentum * mean.detach() + ) + unbias_var = sumvar / (numel - 1) + self.running_var = ( + (1 - self.momentum) * self.running_var + + self.momentum * unbias_var.detach() + ) + + bias_var = sumvar / numel + inv_std = 1 / (bias_var + self.eps).pow(0.5) + output = ( + (input_ - mean.unsqueeze(1)) * inv_std.unsqueeze(1) * + self.weight.unsqueeze(1) + self.bias.unsqueeze(1)) + + return output.view(channels, batchsize, height, width).permute(1, 0, 2, 3).contiguous() + diff --git a/models/sean_codes/models/networks/sync_batchnorm/comm.py b/models/sean_codes/models/networks/sync_batchnorm/comm.py new file mode 100644 index 0000000000000000000000000000000000000000..922f8c4a3adaa9b32fdcaef09583be03b0d7eb2b --- /dev/null +++ b/models/sean_codes/models/networks/sync_batchnorm/comm.py @@ -0,0 +1,137 @@ +# -*- coding: utf-8 -*- +# File : comm.py +# Author : Jiayuan Mao +# Email : maojiayuan@gmail.com +# Date : 27/01/2018 +# +# This file is part of Synchronized-BatchNorm-PyTorch. +# https://github.com/vacancy/Synchronized-BatchNorm-PyTorch +# Distributed under MIT License. + +import queue +import collections +import threading + +__all__ = ['FutureResult', 'SlavePipe', 'SyncMaster'] + + +class FutureResult(object): + """A thread-safe future implementation. Used only as one-to-one pipe.""" + + def __init__(self): + self._result = None + self._lock = threading.Lock() + self._cond = threading.Condition(self._lock) + + def put(self, result): + with self._lock: + assert self._result is None, 'Previous result has\'t been fetched.' + self._result = result + self._cond.notify() + + def get(self): + with self._lock: + if self._result is None: + self._cond.wait() + + res = self._result + self._result = None + return res + + +_MasterRegistry = collections.namedtuple('MasterRegistry', ['result']) +_SlavePipeBase = collections.namedtuple('_SlavePipeBase', ['identifier', 'queue', 'result']) + + +class SlavePipe(_SlavePipeBase): + """Pipe for master-slave communication.""" + + def run_slave(self, msg): + self.queue.put((self.identifier, msg)) + ret = self.result.get() + self.queue.put(True) + return ret + + +class SyncMaster(object): + """An abstract `SyncMaster` object. + + - During the replication, as the data parallel will trigger an callback of each module, all slave devices should + call `register(id)` and obtain an `SlavePipe` to communicate with the master. + - During the forward pass, master device invokes `run_master`, all messages from slave devices will be collected, + and passed to a registered callback. + - After receiving the messages, the master device should gather the information and determine to message passed + back to each slave devices. + """ + + def __init__(self, master_callback): + """ + + Args: + master_callback: a callback to be invoked after having collected messages from slave devices. + """ + self._master_callback = master_callback + self._queue = queue.Queue() + self._registry = collections.OrderedDict() + self._activated = False + + def __getstate__(self): + return {'master_callback': self._master_callback} + + def __setstate__(self, state): + self.__init__(state['master_callback']) + + def register_slave(self, identifier): + """ + Register an slave device. + + Args: + identifier: an identifier, usually is the device id. + + Returns: a `SlavePipe` object which can be used to communicate with the master device. + + """ + if self._activated: + assert self._queue.empty(), 'Queue is not clean before next initialization.' + self._activated = False + self._registry.clear() + future = FutureResult() + self._registry[identifier] = _MasterRegistry(future) + return SlavePipe(identifier, self._queue, future) + + def run_master(self, master_msg): + """ + Main entry for the master device in each forward pass. + The messages were first collected from each devices (including the master device), and then + an callback will be invoked to compute the message to be sent back to each devices + (including the master device). + + Args: + master_msg: the message that the master want to send to itself. This will be placed as the first + message when calling `master_callback`. For detailed usage, see `_SynchronizedBatchNorm` for an example. + + Returns: the message to be sent back to the master device. + + """ + self._activated = True + + intermediates = [(0, master_msg)] + for i in range(self.nr_slaves): + intermediates.append(self._queue.get()) + + results = self._master_callback(intermediates) + assert results[0][0] == 0, 'The first result should belongs to the master.' + + for i, res in results: + if i == 0: + continue + self._registry[i].result.put(res) + + for i in range(self.nr_slaves): + assert self._queue.get() is True + + return results[0][1] + + @property + def nr_slaves(self): + return len(self._registry) diff --git a/models/sean_codes/models/networks/sync_batchnorm/replicate.py b/models/sean_codes/models/networks/sync_batchnorm/replicate.py new file mode 100644 index 0000000000000000000000000000000000000000..b71c7b8ed51a1d6c55b1f753bdd8d90bad79bd06 --- /dev/null +++ b/models/sean_codes/models/networks/sync_batchnorm/replicate.py @@ -0,0 +1,94 @@ +# -*- coding: utf-8 -*- +# File : replicate.py +# Author : Jiayuan Mao +# Email : maojiayuan@gmail.com +# Date : 27/01/2018 +# +# This file is part of Synchronized-BatchNorm-PyTorch. +# https://github.com/vacancy/Synchronized-BatchNorm-PyTorch +# Distributed under MIT License. + +import functools + +from torch.nn.parallel.data_parallel import DataParallel + +__all__ = [ + 'CallbackContext', + 'execute_replication_callbacks', + 'DataParallelWithCallback', + 'patch_replication_callback' +] + + +class CallbackContext(object): + pass + + +def execute_replication_callbacks(modules): + """ + Execute an replication callback `__data_parallel_replicate__` on each module created by original replication. + + The callback will be invoked with arguments `__data_parallel_replicate__(ctx, copy_id)` + + Note that, as all modules are isomorphism, we assign each sub-module with a context + (shared among multiple copies of this module on different devices). + Through this context, different copies can share some information. + + We guarantee that the callback on the master copy (the first copy) will be called ahead of calling the callback + of any slave copies. + """ + master_copy = modules[0] + nr_modules = len(list(master_copy.modules())) + ctxs = [CallbackContext() for _ in range(nr_modules)] + + for i, module in enumerate(modules): + for j, m in enumerate(module.modules()): + if hasattr(m, '__data_parallel_replicate__'): + m.__data_parallel_replicate__(ctxs[j], i) + + +class DataParallelWithCallback(DataParallel): + """ + Data Parallel with a replication callback. + + An replication callback `__data_parallel_replicate__` of each module will be invoked after being created by + original `replicate` function. + The callback will be invoked with arguments `__data_parallel_replicate__(ctx, copy_id)` + + Examples: + > sync_bn = SynchronizedBatchNorm1d(10, eps=1e-5, affine=False) + > sync_bn = DataParallelWithCallback(sync_bn, device_ids=[0, 1]) + # sync_bn.__data_parallel_replicate__ will be invoked. + """ + + def replicate(self, module, device_ids): + modules = super(DataParallelWithCallback, self).replicate(module, device_ids) + execute_replication_callbacks(modules) + return modules + + +def patch_replication_callback(data_parallel): + """ + Monkey-patch an existing `DataParallel` object. Add the replication callback. + Useful when you have customized `DataParallel` implementation. + + Examples: + > sync_bn = SynchronizedBatchNorm1d(10, eps=1e-5, affine=False) + > sync_bn = DataParallel(sync_bn, device_ids=[0, 1]) + > patch_replication_callback(sync_bn) + # this is equivalent to + > sync_bn = SynchronizedBatchNorm1d(10, eps=1e-5, affine=False) + > sync_bn = DataParallelWithCallback(sync_bn, device_ids=[0, 1]) + """ + + assert isinstance(data_parallel, DataParallel) + + old_replicate = data_parallel.replicate + + @functools.wraps(old_replicate) + def new_replicate(module, device_ids): + modules = old_replicate(module, device_ids) + execute_replication_callbacks(modules) + return modules + + data_parallel.replicate = new_replicate diff --git a/models/sean_codes/models/networks/sync_batchnorm/unittest.py b/models/sean_codes/models/networks/sync_batchnorm/unittest.py new file mode 100644 index 0000000000000000000000000000000000000000..bed56f1caa929ac3e9a57c583f8d3e42624f58be --- /dev/null +++ b/models/sean_codes/models/networks/sync_batchnorm/unittest.py @@ -0,0 +1,29 @@ +# -*- coding: utf-8 -*- +# File : unittest.py +# Author : Jiayuan Mao +# Email : maojiayuan@gmail.com +# Date : 27/01/2018 +# +# This file is part of Synchronized-BatchNorm-PyTorch. +# https://github.com/vacancy/Synchronized-BatchNorm-PyTorch +# Distributed under MIT License. + +import unittest +import torch + + +class TorchTestCase(unittest.TestCase): + def assertTensorClose(self, x, y): + adiff = float((x - y).abs().max()) + if (y == 0).all(): + rdiff = 'NaN' + else: + rdiff = float((adiff / y).abs().max()) + + message = ( + 'Tensor close check failed\n' + 'adiff={}\n' + 'rdiff={}\n' + ).format(adiff, rdiff) + self.assertTrue(torch.allclose(x, y), message) + diff --git a/models/sean_codes/models/pix2pix_model.py b/models/sean_codes/models/pix2pix_model.py new file mode 100644 index 0000000000000000000000000000000000000000..ba1df7da165bb2b669adad2a2057c9af31c3a9b6 --- /dev/null +++ b/models/sean_codes/models/pix2pix_model.py @@ -0,0 +1,339 @@ +""" +Copyright (C) 2019 NVIDIA Corporation. All rights reserved. +Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode). +""" +import os +from argparse import Namespace +from glob import glob + +import numpy as np +import torch + +from models.sean_codes.util import util +from . import networks + + +class Pix2PixModel(torch.nn.Module): + @staticmethod + def modify_commandline_options(parser, is_train): + networks.modify_commandline_options(parser, is_train) + return parser + + def __init__(self, opt): + super().__init__() + self.opt = opt + self.FloatTensor = torch.cuda.FloatTensor if self.use_gpu() \ + else torch.FloatTensor + self.ByteTensor = torch.cuda.ByteTensor if self.use_gpu() \ + else torch.ByteTensor + + self.netG, self.netD, self.netE = self.initialize_networks(opt) + + # set loss functions + if opt.isTrain: + self.criterionGAN = networks.GANLoss( + opt.gan_mode, tensor=self.FloatTensor, opt=self.opt) + self.criterionFeat = torch.nn.L1Loss() + if not opt.no_vgg_loss: + self.criterionVGG = networks.VGGLoss(self.opt.gpu_ids) + + # Entry point for all calls involving forward pass + # of deep networks. We used this approach since DataParallel module + # can't parallelize custom functions, we branch to different + # routines based on |mode|. + def forward(self, data, mode): + input_semantics, real_image = self.preprocess_input(data) + + if mode == 'generator': + g_loss, generated = self.compute_generator_loss( + input_semantics, real_image) + return g_loss, generated + elif mode == 'discriminator': + d_loss = self.compute_discriminator_loss( + input_semantics, real_image) + return d_loss + elif mode == 'encode_only': + z, mu, logvar = self.encode_z(real_image) + return mu, logvar + elif mode == 'inference': + with torch.no_grad(): + # fake_image, _ = self.generate_fake(input_semantics, real_image) + obj_dic = data['path'] + fake_image = self.save_style_codes(input_semantics, real_image, obj_dic) + return fake_image + elif mode == 'UI_mode': + with torch.no_grad(): + # fake_image, _ = self.generate_fake(input_semantics, real_image) + + ################### some problems here + obj_dic = data['obj_dic'] + # if isinstance(obj_dic, str): + # obj_dic = [obj_dic] + fake_image = self.use_style_codes(input_semantics, real_image, obj_dic) + return fake_image + elif mode == 'style_code': + with torch.no_grad(): + style_codes = self.netG.Zencoder(input=real_image, segmap=input_semantics) + return style_codes + else: + raise ValueError("|mode| is invalid") + + def create_optimizers(self, opt): + G_params = list(self.netG.parameters()) + if opt.use_vae: + G_params += list(self.netE.parameters()) + if opt.isTrain: + D_params = list(self.netD.parameters()) + + if opt.no_TTUR: + beta1, beta2 = opt.beta1, opt.beta2 + G_lr, D_lr = opt.lr, opt.lr + else: + beta1, beta2 = 0, 0.9 + G_lr, D_lr = opt.lr / 2, opt.lr * 2 + + optimizer_G = torch.optim.Adam(G_params, lr=G_lr, betas=(beta1, beta2)) + optimizer_D = torch.optim.Adam(D_params, lr=D_lr, betas=(beta1, beta2)) + + return optimizer_G, optimizer_D + + def save(self, epoch): + util.save_network(self.netG, 'G', epoch, self.opt) + util.save_network(self.netD, 'D', epoch, self.opt) + + ############################################################################ + # Private helper methods + ############################################################################ + + def initialize_networks(self, opt): + netG = networks.define_G(opt) + netD = networks.define_D(opt) if opt.isTrain else None + netE = networks.define_E(opt) if opt.use_vae else None + + if not opt.isTrain or opt.continue_train: + netG = util.load_network(netG, 'G', opt.which_epoch, opt) + if opt.isTrain: + netD = util.load_network(netD, 'D', opt.which_epoch, opt) + + return netG, netD, netE + + # preprocess the input, such as moving the tensors to GPUs and + # transforming the label map to one-hot encoding + # |data|: dictionary of the input data + + def preprocess_input(self, data): + # move to GPU and change data types + data['label'] = data['label'].long() + if self.use_gpu(): + for param in ['label', 'instance', 'image']: + if param in data and data[param] is not None: + data[param] = data[param].cuda(non_blocking=True) + if 'obj_dic' in data: + for idx in range(19): + if data['obj_dic'][str(idx)]['ACE'].device.type == 'cpu': + data['obj_dic'][str(idx)]['ACE'] = data['obj_dic'][str(idx)]['ACE'].cuda(non_blocking=True) + # create one-hot label map + label_map = data['label'] + bs, _, h, w = label_map.size() + nc = self.opt.label_nc + 1 if self.opt.contain_dontcare_label \ + else self.opt.label_nc + input_label = self.FloatTensor(bs, nc, h, w).zero_() + input_semantics = input_label.scatter_(1, label_map, 1.0) + + # concatenate instance map if it exists + if not self.opt.no_instance: + inst_map = data['instance'] + instance_edge_map = self.get_edges(inst_map) + input_semantics = torch.cat((input_semantics, instance_edge_map), dim=1) + + return input_semantics, data['image'] + + def compute_generator_loss(self, input_semantics, real_image): + G_losses = {} + + fake_image = self.generate_fake( + input_semantics, real_image, compute_kld_loss=self.opt.use_vae) + + pred_fake, pred_real = self.discriminate( + input_semantics, fake_image, real_image) + + G_losses['GAN'] = self.criterionGAN(pred_fake, True, + for_discriminator=False) + + if not self.opt.no_ganFeat_loss: + num_D = len(pred_fake) + GAN_Feat_loss = self.FloatTensor(1).fill_(0) + for i in range(num_D): # for each discriminator + # last color_texture is the final prediction, so we exclude it + num_intermediate_outputs = len(pred_fake[i]) - 1 + for j in range(num_intermediate_outputs): # for each layer color_texture + unweighted_loss = self.criterionFeat( + pred_fake[i][j], pred_real[i][j].detach()) + GAN_Feat_loss += unweighted_loss * self.opt.lambda_feat / num_D + G_losses['GAN_Feat'] = GAN_Feat_loss + + if not self.opt.no_vgg_loss: + G_losses['VGG'] = self.criterionVGG(fake_image, real_image) \ + * self.opt.lambda_vgg + + return G_losses, fake_image + + def compute_discriminator_loss(self, input_semantics, real_image): + D_losses = {} + with torch.no_grad(): + fake_image = self.generate_fake(input_semantics, real_image) + fake_image = fake_image.detach() + fake_image.requires_grad_() + + pred_fake, pred_real = self.discriminate( + input_semantics, fake_image, real_image) + + D_losses['D_Fake'] = self.criterionGAN(pred_fake, False, + for_discriminator=True) + D_losses['D_real'] = self.criterionGAN(pred_real, True, + for_discriminator=True) + + return D_losses + + def encode_z(self, real_image): + mu, logvar = self.netE(real_image) + z = self.reparameterize(mu, logvar) + return z, mu, logvar + + def generate_fake(self, input_semantics, real_image, compute_kld_loss=False): + + fake_image = self.netG(input_semantics, real_image) + + return fake_image + + ############################################################### + + def save_style_codes(self, input_semantics, real_image, obj_dic): + + fake_image = self.netG(input_semantics, real_image, obj_dic=obj_dic) + + return fake_image + + def use_style_codes(self, input_semantics, real_image, obj_dic): + fake_image = self.netG(input_semantics, real_image, obj_dic=obj_dic) + + return fake_image + + # Given fake and real image, return the prediction of discriminator + # for each fake and real image. + + def discriminate(self, input_semantics, fake_image, real_image): + fake_concat = torch.cat([input_semantics, fake_image], dim=1) + real_concat = torch.cat([input_semantics, real_image], dim=1) + + # In Batch Normalization, the fake and real images are + # recommended to be in the same batch to avoid disparate + # statistics in fake and real images. + # So both fake and real images are fed to D all at once. + fake_and_real = torch.cat([fake_concat, real_concat], dim=0) + + discriminator_out = self.netD(fake_and_real) + + pred_fake, pred_real = self.divide_pred(discriminator_out) + + return pred_fake, pred_real + + # Take the prediction of fake and real images from the combined batch + def divide_pred(self, pred): + # the prediction contains the intermediate outputs of multiscale GAN, + # so it's usually a list + if type(pred) == list: + fake = [] + real = [] + for p in pred: + fake.append([tensor[:tensor.size(0) // 2] for tensor in p]) + real.append([tensor[tensor.size(0) // 2:] for tensor in p]) + else: + fake = pred[:pred.size(0) // 2] + real = pred[pred.size(0) // 2:] + + return fake, real + + def get_edges(self, t): + edge = self.ByteTensor(t.size()).zero_() + edge[:, :, :, 1:] = edge[:, :, :, 1:] | (t[:, :, :, 1:] != t[:, :, :, :-1]) + edge[:, :, :, :-1] = edge[:, :, :, :-1] | (t[:, :, :, 1:] != t[:, :, :, :-1]) + edge[:, :, 1:, :] = edge[:, :, 1:, :] | (t[:, :, 1:, :] != t[:, :, :-1, :]) + edge[:, :, :-1, :] = edge[:, :, :-1, :] | (t[:, :, 1:, :] != t[:, :, :-1, :]) + return edge.float() + + def reparameterize(self, mu, logvar): + std = torch.exp(0.5 * logvar) + eps = torch.randn_like(std) + return eps.mul(std) + mu + + def use_gpu(self): + return len(self.opt.gpu_ids) > 0 + + +def load_average_feature(): + ############### load average features + # average_style_code_folder = 'styles_test/mean_style_code/mean/' + average_style_code_folder = 'models/sean_codes/styles_test/mean_style_code/median/' + input_style_dic = {} + + ############### hard coding for categories + for i in range(19): + input_style_dic[str(i)] = {} + average_category_folder_list = glob(os.path.join(average_style_code_folder, str(i), '*.npy')) + average_category_list = [os.path.splitext(os.path.basename(name))[0] for name in + average_category_folder_list] + + for style_code_path in average_category_list: + input_style_dic[str(i)][style_code_path] = torch.from_numpy( + np.load(os.path.join(average_style_code_folder, str(i), style_code_path + '.npy'))).cuda() + return input_style_dic + + +def change_status(model, new_status): + for m in model.modules(): + if hasattr(m, 'status'): + m.status = new_status + + +def encode_sean(sean_model, images, labels): + data = {'label': labels, + 'instance': torch.tensor(0), + 'image': images, + 'path': ['temp/temp_npy']} + change_status(sean_model, 'test') + img_codes = sean_model(data, mode='style_code') # [2, 19, 512] + return img_codes + + +def decode_sean(sean_model, image_code, target_mask): + obj_dic = load_average_feature() + + for idx in range(19): + cur_code = image_code[0, idx] + if not torch.all(cur_code == 0): + obj_dic[str(idx)]['ACE'] = cur_code + + temp_face_image = torch.zeros((0, 3, 256, 256)) # place holder + + data = {'label': target_mask, + 'instance': torch.tensor(0), + 'image': temp_face_image.clone().detach(), + 'obj_dic': obj_dic} + change_status(sean_model, 'UI_mode') + generated = sean_model(data, mode='UI_mode')[0] # [3, 256, 256] + return generated + + +SEAN_OPT = Namespace(name='CelebA-HQ_pretrained', gpu_ids=[0], checkpoints_dir='pretrained_models/sean_checkpoints', + model='pix2pix', norm_G='spectralspadesyncbatch3x3', norm_D='spectralinstance', + norm_E='spectralinstance', phase='test', batchSize=1, preprocess_mode='resize_and_crop', + load_size=256, crop_size=256, aspect_ratio=1.0, label_nc=19, contain_dontcare_label=False, + output_nc=3, dataroot='./datasets/cityscapes/', dataset_mode='custom', serial_batches=True, + no_flip=True, nThreads=28, max_dataset_size=9223372036854775807, load_from_opt_file=False, + cache_filelist_write=False, cache_filelist_read=False, display_winsize=256, netG='spade', ngf=64, + init_type='xavier', init_variance=0.02, z_dim=256, no_instance=True, nef=16, use_vae=False, + results_dir='./results/', which_epoch='latest', how_many=float("inf"), status='test', config='001', + gpu='0', need_crop=True, no_blending=False, num_upsampling_layers='normal', no_pairing_check=False, + label_dir='datasets/CelebA-HQ/test/labels', image_dir='datasets/CelebA-HQ/test/images', + instance_dir='', isTrain=False, semantic_nc=19) diff --git a/models/sean_codes/options/__init__.py b/models/sean_codes/options/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..ba4ddb2d78be7bbf79bb0671eefd9296d3a168b7 --- /dev/null +++ b/models/sean_codes/options/__init__.py @@ -0,0 +1,4 @@ +""" +Copyright (C) 2019 NVIDIA Corporation. All rights reserved. +Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode). +""" \ No newline at end of file diff --git a/models/sean_codes/options/base_options.py b/models/sean_codes/options/base_options.py new file mode 100644 index 0000000000000000000000000000000000000000..fd171f25c97c57236363fc34380241096814bb65 --- /dev/null +++ b/models/sean_codes/options/base_options.py @@ -0,0 +1,180 @@ +""" +Copyright (C) 2019 NVIDIA Corporation. All rights reserved. +Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode). +""" + +import sys +import argparse +import os +from util import util +import torch +from sean_codes import models, data +import pickle + + +class BaseOptions(): + def __init__(self): + self.initialized = False + + def initialize(self, parser): + # experiment specifics + parser.add_argument('--name', type=str, default='CelebA-HQ_pretrained', + help='name of the experiment. It decides where to store samples and models') + + parser.add_argument('--gpu_ids', type=str, default='0', help='gpu ids: e.g. 0 0,1,2, 0,2. use -1 for CPU') + parser.add_argument('--checkpoints_dir', type=str, default='pretrained_models/sean_checkpoints', + help='models are saved here') + parser.add_argument('--model', type=str, default='pix2pix', help='which model to use') + parser.add_argument('--norm_G', type=str, default='spectralinstance', help='instance normalization or batch normalization') + parser.add_argument('--norm_D', type=str, default='spectralinstance', help='instance normalization or batch normalization') + parser.add_argument('--norm_E', type=str, default='spectralinstance', help='instance normalization or batch normalization') + parser.add_argument('--phase', type=str, default='scripts', help='scripts, val, test, etc') + + # input/color_texture sizes + parser.add_argument('--batchSize', type=int, default=1, help='input batch size') + parser.add_argument('--preprocess_mode', type=str, default='scale_width_and_crop', help='scaling and cropping of images at load time.', choices=("resize_and_crop", "crop", "scale_width", "scale_width_and_crop", "scale_shortside", "scale_shortside_and_crop", "fixed", "none")) + parser.add_argument('--load_size', type=int, default=256, help='Scale images to this size. The final image will be cropped to --crop_size.') + parser.add_argument('--crop_size', type=int, default=256, help='Crop to the width of crop_size (after initially scaling the images to load_size.)') + parser.add_argument('--aspect_ratio', type=float, default=1.0, help='The ratio width/height. The final height of the load image will be crop_size/aspect_ratio') + parser.add_argument('--label_nc', type=int, default=19, help='# of input label classes without unknown class. If you have unknown class as class label, specify --contain_dopntcare_label.') + parser.add_argument('--contain_dontcare_label', action='store_true', help='if the label map contains dontcare label (dontcare=255)') + parser.add_argument('--output_nc', type=int, default=3, help='# of color_texture image channels') + + # for setting inputs + parser.add_argument('--dataroot', type=str, default='./datasets/cityscapes/') + parser.add_argument('--dataset_mode', type=str, default='custom') + parser.add_argument('--serial_batches', action='store_true', help='if true, takes images in order to make batches, otherwise takes them randomly') + parser.add_argument('--no_flip', action='store_true', help='if specified, do not flip the images for data argumentation') + parser.add_argument('--nThreads', default=28, type=int, help='# threads for loading data') + parser.add_argument('--max_dataset_size', type=int, default=sys.maxsize, help='Maximum number of samples allowed per dataset. If the dataset directory contains more than max_dataset_size, only a subset is loaded.') + parser.add_argument('--load_from_opt_file', action='store_true', help='load the options from sean_checkpoints and use that as default') + parser.add_argument('--cache_filelist_write', action='store_true', help='saves the current filelist into a text file, so that it loads faster') + parser.add_argument('--cache_filelist_read', action='store_true', help='reads from the file list cache') + + # for displays + parser.add_argument('--display_winsize', type=int, default=400, help='display window size') + + # for generator + parser.add_argument('--netG', type=str, default='spade', help='selects model to use for netG (pix2pixhd | spade)') + parser.add_argument('--ngf', type=int, default=64, help='# of gen filters in first conv layer') + parser.add_argument('--init_type', type=str, default='xavier', help='network initialization [normal|xavier|kaiming|orthogonal]') + parser.add_argument('--init_variance', type=float, default=0.02, help='variance of the initialization distribution') + parser.add_argument('--z_dim', type=int, default=256, + help="dimension of the latent z vector") + + # for instance-wise features + parser.add_argument('--no_instance', type=bool, default=True, + help='if specified, do *not* add instance map as input') + parser.add_argument('--nef', type=int, default=16, help='# of encoder filters in the first conv layer') + parser.add_argument('--use_vae', action='store_true', help='enable training with an image encoder.') + + self.initialized = True + return parser + + def gather_options(self): + # initialize parser with basic options + if not self.initialized: + parser = argparse.ArgumentParser( + formatter_class=argparse.ArgumentDefaultsHelpFormatter) + parser = self.initialize(parser) + + # get the basic options + opt, unknown = parser.parse_known_args() + + # modify model-related parser options + model_name = opt.model + model_option_setter = models.get_option_setter(model_name) + parser = model_option_setter(parser, self.isTrain) + + # modify dataset-related parser options + dataset_mode = opt.dataset_mode + dataset_option_setter = data.get_option_setter(dataset_mode) + parser = dataset_option_setter(parser, self.isTrain) + + opt, unknown = parser.parse_known_args() + + # if there is opt_file, load it. + # The previous default options will be overwritten + if opt.load_from_opt_file: + parser = self.update_options_from_file(parser, opt) + + opt = parser.parse_args() + self.parser = parser + return opt + + def print_options(self, opt): + message = '' + message += '----------------- Options ---------------\n' + for k, v in sorted(vars(opt).items()): + comment = '' + default = self.parser.get_default(k) + if v != default: + comment = '\t[default: %s]' % str(default) + message += '{:>25}: {:<30}{}\n'.format(str(k), str(v), comment) + message += '----------------- End -------------------' + print(message) + + def option_file_path(self, opt, makedir=False): + expr_dir = os.path.join(opt.checkpoints_dir, opt.name) + if makedir: + util.mkdirs(expr_dir) + file_name = os.path.join(expr_dir, 'opt') + return file_name + + def save_options(self, opt): + file_name = self.option_file_path(opt, makedir=True) + with open(file_name + '.txt', 'wt') as opt_file: + for k, v in sorted(vars(opt).items()): + comment = '' + default = self.parser.get_default(k) + if v != default: + comment = '\t[default: %s]' % str(default) + opt_file.write('{:>25}: {:<30}{}\n'.format(str(k), str(v), comment)) + + with open(file_name + '.pkl', 'wb') as opt_file: + pickle.dump(opt, opt_file) + + def update_options_from_file(self, parser, opt): + new_opt = self.load_options(opt) + for k, v in sorted(vars(opt).items()): + if hasattr(new_opt, k) and v != getattr(new_opt, k): + new_val = getattr(new_opt, k) + parser.set_defaults(**{k: new_val}) + return parser + + def load_options(self, opt): + file_name = self.option_file_path(opt, makedir=False) + new_opt = pickle.load(open(file_name + '.pkl', 'rb')) + return new_opt + + def parse(self, save=False): + + opt = self.gather_options() + opt.isTrain = self.isTrain # scripts or test + + # self.print_options(opt) + if opt.isTrain: + self.save_options(opt) + + # Set semantic_nc based on the option. + # This will be convenient in many places + opt.semantic_nc = opt.label_nc + \ + (1 if opt.contain_dontcare_label else 0) + \ + (0 if opt.no_instance else 1) + + # set gpu ids + str_ids = opt.gpu_ids.split(',') + opt.gpu_ids = [] + for str_id in str_ids: + id = int(str_id) + if id >= 0: + opt.gpu_ids.append(id) + if len(opt.gpu_ids) > 0: + torch.cuda.set_device(opt.gpu_ids[0]) + + assert len(opt.gpu_ids) == 0 or opt.batchSize % len(opt.gpu_ids) == 0, \ + "Batch size %d is wrong. It must be a multiple of # GPUs %d." \ + % (opt.batchSize, len(opt.gpu_ids)) + + self.opt = opt + return self.opt diff --git a/models/sean_codes/options/test_options.py b/models/sean_codes/options/test_options.py new file mode 100644 index 0000000000000000000000000000000000000000..d7e06df9ac5e17c875623c9c54cfeb1d6c54616e --- /dev/null +++ b/models/sean_codes/options/test_options.py @@ -0,0 +1,27 @@ +""" +Copyright (C) 2019 NVIDIA Corporation. All rights reserved. +Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode). +""" +from util.common_options import ctrl_hair_parser_options +from .base_options import BaseOptions + + +class TestOptions(BaseOptions): + def initialize(self, parser): + BaseOptions.initialize(self, parser) + parser.add_argument('--results_dir', type=str, default='./results/', help='saves results here.') + parser.add_argument('--which_epoch', type=str, default='latest', + help='which epoch to load? set to latest to use latest cached model') + parser.add_argument('--how_many', type=int, default=float("inf"), help='how many test images to run') + + parser.set_defaults(preprocess_mode='scale_width_and_crop', crop_size=256, load_size=256, display_winsize=256) + parser.set_defaults(serial_batches=True) + parser.set_defaults(no_flip=True) + parser.set_defaults(phase='test') + + parser.add_argument('--status', type=str, default='test') + + ctrl_hair_parser_options(parser) + + self.isTrain = False + return parser diff --git a/models/sean_codes/options/train_options.py b/models/sean_codes/options/train_options.py new file mode 100644 index 0000000000000000000000000000000000000000..1eb1496ab13f5d57eac9460583d82a931fd9a68c --- /dev/null +++ b/models/sean_codes/options/train_options.py @@ -0,0 +1,46 @@ +""" +Copyright (C) 2019 NVIDIA Corporation. All rights reserved. +Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode). +""" + +from .base_options import BaseOptions + + +class TrainOptions(BaseOptions): + def initialize(self, parser): + BaseOptions.initialize(self, parser) + # for displays + parser.add_argument('--display_freq', type=int, default=100, help='frequency of showing training results on screen') + parser.add_argument('--print_freq', type=int, default=100, help='frequency of showing training results on console') + parser.add_argument('--save_latest_freq', type=int, default=5000, help='frequency of saving the latest results') + parser.add_argument('--save_epoch_freq', type=int, default=10, help='frequency of saving sean_checkpoints at the end of epochs') + parser.add_argument('--no_html', default=False, help='do not save intermediate training results to [opt.checkpoints_dir]/[opt.name]/web/') + parser.add_argument('--debug', action='store_true', help='only do one epoch and displays at each iteration') + parser.add_argument('--tf_log', action='store_true', help='if specified, use tensorboard logging. Requires tensorflow installed') + + # for training + parser.add_argument('--continue_train', action='store_true', help='continue training: load the latest model') + parser.add_argument('--which_epoch', type=str, default='latest', help='which epoch to load? set to latest to use latest cached model') + parser.add_argument('--niter', type=int, default=50, help='# of iter at starting learning rate. This is NOT the total #epochs. Totla #epochs is niter + niter_decay') + parser.add_argument('--niter_decay', type=int, default=50, help='# of iter to linearly decay learning rate to zero') + parser.add_argument('--optimizer', type=str, default='adam') + parser.add_argument('--beta1', type=float, default=0.5, help='momentum term of adam') + parser.add_argument('--beta2', type=float, default=0.999, help='momentum term of adam') + parser.add_argument('--lr', type=float, default=0.0002, help='initial learning rate for adam') + parser.add_argument('--D_steps_per_G', type=int, default=1, help='number of discriminator iterations per generator iterations.') + + # for discriminators + parser.add_argument('--ndf', type=int, default=64, help='# of discrim filters in first conv layer') + parser.add_argument('--lambda_feat', type=float, default=10.0, help='weight for feature matching loss') + parser.add_argument('--lambda_vgg', type=float, default=10.0, help='weight for vgg loss') + parser.add_argument('--no_ganFeat_loss', action='store_true', help='if specified, do *not* use discriminator feature matching loss') + parser.add_argument('--no_vgg_loss', action='store_true', help='if specified, do *not* use VGG feature matching loss') + parser.add_argument('--gan_mode', type=str, default='hinge', help='(ls|original|hinge)') + parser.add_argument('--netD', type=str, default='multiscale', help='(n_layers|multiscale|image)') + parser.add_argument('--no_TTUR', action='store_true', help='Use TTUR training scheme') + parser.add_argument('--lambda_kld', type=float, default=0.005) + + parser.add_argument('--status', type=str, default='scripts') + + self.isTrain = True + return parser diff --git a/models/sean_codes/options/util/__init__.py b/models/sean_codes/options/util/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/models/sean_codes/options/util/canvas_grid.py b/models/sean_codes/options/util/canvas_grid.py new file mode 100644 index 0000000000000000000000000000000000000000..65a200854f842ce8d83876be58bc401c2c24bcc9 --- /dev/null +++ b/models/sean_codes/options/util/canvas_grid.py @@ -0,0 +1,34 @@ +# -*- coding: utf-8 -*- + +""" +# File name: canvas_grid.py +# Time : 2021/11/18 16:40 +# Author: xyguoo@163.com +# Description: This is the util that generating image in a big grid picture +""" + +import numpy as np + +from util.imutil import write_rgb + + +class Canvas: + def __init__(self, row, col, img_size=256, margin=0): + self.row = row + self.col = col + self.img_size = img_size + self.margin = margin + self.canvas = np.ones((row * img_size, col * img_size + margin * (col - 1), 3), dtype='uint8') * 255 + + def process_draw_image(self, img, i, j): + if img.dtype in [np.float32, np.float, np.float64]: + if img.min() < 0: + img = img * 127.5 + 127.5 + elif img.max() <= 1: + img = img * 255 + img = img.astype('uint8') + i_start, j_start = int(i * self.img_size), int(j * self.img_size) + int(j * self.margin) + self.canvas[i_start: i_start + img.shape[0], j_start: j_start + img.shape[1], :] = img + + def write_(self, file): + write_rgb(file, self.canvas) diff --git a/models/sean_codes/options/util/color_from_hsv_to_gaussian.py b/models/sean_codes/options/util/color_from_hsv_to_gaussian.py new file mode 100644 index 0000000000000000000000000000000000000000..5a45b8bacc231d6430eaf971288bd859a5a8d55d --- /dev/null +++ b/models/sean_codes/options/util/color_from_hsv_to_gaussian.py @@ -0,0 +1,40 @@ +# -*- coding: utf-8 -*- + +""" +# File name: color_from_hsv_to_gaussian.py +# Time : 2022/3/6 16:00 +# Author: xyguoo@163.com +# Description: transfer HSV value to Gaussian latent value according to the distribution of the dataset +""" +import os +import pickle as pkl +from bisect import bisect_left, bisect_right + +import scipy.stats as st + + +class DistTranslation: + def __init__(self): + hair_root = 'dataset_info_ctrlhair' + with open(os.path.join(hair_root, 'hsv_stat_dict_ordered.pkl'), 'rb') as f: + self.cols_hsv = pkl.load(f) + + def gaussian_to_val(self, dim, val): + # if dim == 0: + # return (val + 2.) / 4 * 179 + return self.cols_hsv[int((st.norm.cdf(val)) * self.cols_hsv.shape[0])][dim] + + def val_to_gaussian(self, dim, val): + # if dim == 0: + # return val / 179 * 2 * 2. - 2. + + left_v = bisect_left(self.cols_hsv[:, dim], val) + right_v = bisect_right(self.cols_hsv[:, dim], val) + return st.norm.ppf((left_v + right_v) / 2 / self.cols_hsv.shape[0]) + +# +# if __name__ == '__main__': +# dt = DistTranslation() +# +# with open('hsv_stat_dict_ordered.pkl', 'wb') as f: +# pkl.dump(dt.cols_hsv, f) diff --git a/models/sean_codes/options/util/common_options.py b/models/sean_codes/options/util/common_options.py new file mode 100644 index 0000000000000000000000000000000000000000..7f54b7f53ca743e0da96d2a2c05aad5cd234d569 --- /dev/null +++ b/models/sean_codes/options/util/common_options.py @@ -0,0 +1,15 @@ +# -*- coding: utf-8 -*- + +""" +# File name: common_options.py +# Time : 2022/7/12 +# Author: xyguoo@163.com +# Description: +""" + +def ctrl_hair_parser_options(parser): + parser.add_argument('-c', '--config', type=str, default='001') + parser.add_argument('-g', '--gpu', type=str, help='Specify GPU number', default='0') + parser.add_argument('-n', '--need_crop', type=bool, help='whether images need crop', default=True) + parser.add_argument('--no_blending', action='store_true', + help='whether using poisson blending as post processing', default=False) \ No newline at end of file diff --git a/models/sean_codes/options/util/find_semantic_direction.py b/models/sean_codes/options/util/find_semantic_direction.py new file mode 100644 index 0000000000000000000000000000000000000000..e2f8b06abc4cd0906825404cfeb2639a3ac82b31 --- /dev/null +++ b/models/sean_codes/options/util/find_semantic_direction.py @@ -0,0 +1,21 @@ +# -*- coding: utf-8 -*- + +""" +# File name: find_semantic_direction.py +# Time : 2022/07/17 +# Author: xyguoo@163.com +# Description: +""" +import torch + + +def get_random_direction(dim, existing_dirs): + dir = torch.randn(dim) + + for dd in existing_dirs: + dir = dir - torch.dot(dir, dd) * dd + + if dir[0] < 0: + dir = -dir + dir = dir / dir.norm() + return dir diff --git a/models/sean_codes/options/util/imutil.py b/models/sean_codes/options/util/imutil.py new file mode 100644 index 0000000000000000000000000000000000000000..539fa8c449eb192a8c84d5c01588f029a5ecd93d --- /dev/null +++ b/models/sean_codes/options/util/imutil.py @@ -0,0 +1,24 @@ +# -*- coding: utf-8 -*- + +""" +# File name: imutil.py +# Time : 2021/12/7 14:55 +# Author: xyguoo@163.com +# Description: +""" +import cv2 +import numpy as np + + +def read_rgb(img_path): + im = cv2.imread(img_path) + im = cv2.cvtColor(im, cv2.COLOR_BGR2RGB) + return im + + +def write_rgb(file_name, img): + if len(img.shape) == 2: + img = np.stack([img] * 3, axis=2) + elif img.shape[2] == 1: + img = np.tile(img, [1, 1, 3]) + cv2.imwrite(file_name, cv2.cvtColor(img, cv2.COLOR_RGB2BGR)) diff --git a/models/sean_codes/options/util/mask_color_util.py b/models/sean_codes/options/util/mask_color_util.py new file mode 100644 index 0000000000000000000000000000000000000000..c496255368d29cf4cd511387306aa066dffbc7a0 --- /dev/null +++ b/models/sean_codes/options/util/mask_color_util.py @@ -0,0 +1,64 @@ +# -*- coding: utf-8 -*- + +""" +# File name: mask_color_util.py +# Time : 2021/12/7 11:13 +# Author: xyguoo@163.com +# Description: +""" + +import numpy as np + +from global_value_utils import HAIR_IDX + + +def mask_to_rgb(pred, draw_type=2): + """ + generate visual mask image + :param pred: pred is the label image with pixel vale in {0, 1, 2, ..., 20} + :param draw_type: 0: all part; 1: {bg, face, hair}; 2: {hair, others} + :return: + """ + if len(pred.shape) == 3 and pred.shape[0] == 1: + pred = pred[0] + num_labels = 19 + color = np.array([[0, 128, 64], + [204, 0, 0], + [76, 153, 0], + [204, 204, 0], ## + [51, 51, 255], ## + [204, 0, 204], ## + [0, 255, 255], ## + [51, 255, 255], ## + [102, 51, 0], ## + [255, 0, 0], ## + [102, 204, 0], ## + [255, 255, 0], ## + [0, 0, 153], ## + [0, 0, 204], ## + [255, 51, 153], ## + [0, 204, 204], ## + [0, 51, 0], ## + [255, 153, 51], + [0, 204, 0], + ]) + + for cc in range(len(color)): + if draw_type == 2: + if cc != HAIR_IDX: + color[cc] = [255, 255, 255] + elif draw_type == 1: + if cc != HAIR_IDX and cc != 0: + color[cc] = [237, 28, 36] + + h, w = np.shape(pred) + rgb = np.zeros((h, w, 3), dtype=np.uint8) + # print(color.shape) + for ii in range(num_labels): + # print(ii) + mask = pred == ii + rgb[mask, None] = color[ii, :] + # Correct unk + unk = pred == 255 + rgb[unk, None] = 255 + return rgb diff --git a/models/sean_codes/options/util/util.py b/models/sean_codes/options/util/util.py new file mode 100644 index 0000000000000000000000000000000000000000..d3eb65e92ea5963f3df4e6ed4c9018e870cd0181 --- /dev/null +++ b/models/sean_codes/options/util/util.py @@ -0,0 +1,281 @@ +""" +These interfaces are most based on SEAN. But we modify some of them, and add some new interface for our CtrlHair. +""" + +import argparse +import importlib +import os +import re + +import cv2 +import dill as pickle +import numpy as np +import torch +from PIL import Image + + +def save_obj(obj, name): + with open(name, 'wb') as f: + pickle.dump(obj, f, pickle.HIGHEST_PROTOCOL) + + +def load_obj(name): + with open(name, 'rb') as f: + return pickle.load(f) + +# returns a configuration for creating a generator +# |default_opt| should be the opt of the current experiment +# |**kwargs|: if any configuration should be overriden, it can be specified here + + +def copyconf(default_opt, **kwargs): + conf = argparse.Namespace(**vars(default_opt)) + for key in kwargs: + print(key, kwargs[key]) + setattr(conf, key, kwargs[key]) + return conf + + +def tile_images(imgs, picturesPerRow=4): + """ Code borrowed from + https://stackoverflow.com/questions/26521365/cleanly-tile-numpy-array-of-images-stored-in-a-flattened-1d-format/26521997 + """ + + # Padding + if imgs.shape[0] % picturesPerRow == 0: + rowPadding = 0 + else: + rowPadding = picturesPerRow - imgs.shape[0] % picturesPerRow + if rowPadding > 0: + imgs = np.concatenate([imgs, np.zeros((rowPadding, *imgs.shape[1:]), dtype=imgs.dtype)], axis=0) + + # Tiling Loop (The conditionals are not necessary anymore) + tiled = [] + for i in range(0, imgs.shape[0], picturesPerRow): + tiled.append(np.concatenate([imgs[j] for j in range(i, i + picturesPerRow)], axis=1)) + + tiled = np.concatenate(tiled, axis=0) + return tiled + + +# Converts a Tensor into a Numpy array +# |imtype|: the desired type of the converted numpy array +def tensor2im(image_tensor, imtype=np.uint8, normalize=True, tile=False): + if isinstance(image_tensor, list): + image_numpy = [] + for i in range(len(image_tensor)): + image_numpy.append(tensor2im(image_tensor[i], imtype, normalize)) + return image_numpy + + if image_tensor.dim() == 4: + # transform each image in the batch + images_np = [] + for b in range(image_tensor.size(0)): + one_image = image_tensor[b] + one_image_np = tensor2im(one_image) + images_np.append(one_image_np.reshape(1, *one_image_np.shape)) + images_np = np.concatenate(images_np, axis=0) + if tile: + images_tiled = tile_images(images_np) + return images_tiled + else: + return images_np + + if image_tensor.dim() == 2: + image_tensor = image_tensor.unsqueeze(0) + image_numpy = image_tensor.detach().cpu().float().numpy() + if normalize: + image_numpy = (np.transpose(image_numpy, (1, 2, 0)) + 1) / 2.0 * 255.0 + else: + image_numpy = np.transpose(image_numpy, (1, 2, 0)) * 255.0 + image_numpy = np.clip(image_numpy, 0, 255) + if image_numpy.shape[2] == 1: + image_numpy = image_numpy[:, :, 0] + return image_numpy.astype(imtype) + + +# Converts a one-hot tensor into a colorful label map +def tensor2label(label_tensor, n_label, imtype=np.uint8, tile=False): + if label_tensor.dim() == 4: + # transform each image in the batch + images_np = [] + for b in range(label_tensor.size(0)): + one_image = label_tensor[b] + one_image_np = tensor2label(one_image, n_label, imtype) + images_np.append(one_image_np.reshape(1, *one_image_np.shape)) + images_np = np.concatenate(images_np, axis=0) + if tile: + images_tiled = tile_images(images_np) + return images_tiled + else: + images_np = images_np[0] + return images_np + + if label_tensor.dim() == 1: + return np.zeros((64, 64, 3), dtype=np.uint8) + if n_label == 0: + return tensor2im(label_tensor, imtype) + label_tensor = label_tensor.cpu().float() + if label_tensor.size()[0] > 1: + label_tensor = label_tensor.max(0, keepdim=True)[1] + label_tensor = Colorize(n_label)(label_tensor) + label_numpy = np.transpose(label_tensor.numpy(), (1, 2, 0)) + result = label_numpy.astype(imtype) + return result + + +def save_image(image_numpy, image_path, create_dir=False): + if create_dir: + os.makedirs(os.path.dirname(image_path), exist_ok=True) + if len(image_numpy.shape) == 2: + image_numpy = np.expand_dims(image_numpy, axis=2) + if image_numpy.shape[2] == 1: + image_numpy = np.repeat(image_numpy, 3, 2) + image_pil = Image.fromarray(image_numpy) + + # save to png + image_pil.save(image_path.replace('.jpg', '.png')) + + +def mkdirs(paths): + if isinstance(paths, list) and not isinstance(paths, str): + for path in paths: + mkdir(path) + else: + mkdir(paths) + + +def mkdir(path): + if not os.path.exists(path): + os.makedirs(path) + + +def atoi(text): + return int(text) if text.isdigit() else text + + +def natural_keys(text): + ''' + alist.sort(key=natural_keys) sorts in human order + http://nedbatchelder.com/blog/200712/human_sorting.html + (See Toothy's implementation in the comments) + ''' + return [atoi(c) for c in re.split('(\d+)', text)] + + +def natural_sort(items): + items.sort(key=natural_keys) + + +def str2bool(v): + if v.lower() in ('yes', 'true', 't', 'y', '1'): + return True + elif v.lower() in ('no', 'false', 'f', 'n', '0'): + return False + else: + raise argparse.ArgumentTypeError('Boolean value expected.') + + +def find_class_in_module(target_cls_name, module): + target_cls_name = target_cls_name.replace('_', '').lower() + clslib = importlib.import_module(module) + cls = None + for name, clsobj in clslib.__dict__.items(): + if name.lower() == target_cls_name: + cls = clsobj + + if cls is None: + print("In %s, there should be a class whose name matches %s in lowercase without underscore(_)" % (module, target_cls_name)) + exit(0) + + return cls + + +def save_network(net, label, epoch, opt): + save_filename = '%s_net_%s.pth' % (epoch, label) + save_path = os.path.join(opt.checkpoints_dir, opt.name, save_filename) + torch.save(net.cpu().state_dict(), save_path) + if len(opt.gpu_ids) and torch.cuda.is_available(): + net.cuda() + + +def load_network(net, label, epoch, opt): + save_filename = '%s_net_%s.pth' % (epoch, label) + save_dir = os.path.join(opt.checkpoints_dir, opt.name) + save_path = os.path.join(save_dir, save_filename) + weights = torch.load(save_path) + net.load_state_dict(weights) + return net + + +############################################################################### +# Code from +# https://github.com/ycszen/pytorch-seg/blob/master/transform.py +# Modified so it complies with the Citscape label map colors +############################################################################### +def uint82bin(n, count=8): + """returns the binary of integer n, count refers to amount of bits""" + return ''.join([str((n >> y) & 1) for y in range(count - 1, -1, -1)]) + + +def labelcolormap(N): + if N == 35: # cityscape + cmap = np.array([(0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), (111, 74, 0), (81, 0, 81), + (128, 64, 128), (244, 35, 232), (250, 170, 160), (230, 150, 140), (70, 70, 70), (102, 102, 156), (190, 153, 153), + (180, 165, 180), (150, 100, 100), (150, 120, 90), (153, 153, 153), (153, 153, 153), (250, 170, 30), (220, 220, 0), + (107, 142, 35), (152, 251, 152), (70, 130, 180), (220, 20, 60), (255, 0, 0), (0, 0, 142), (0, 0, 70), + (0, 60, 100), (0, 0, 90), (0, 0, 110), (0, 80, 100), (0, 0, 230), (119, 11, 32), (0, 0, 142)], + dtype=np.uint8) + else: + cmap = np.zeros((N, 3), dtype=np.uint8) + for i in range(N): + r, g, b = 0, 0, 0 + id = i + 1 # let's give 0 a color + for j in range(7): + str_id = uint82bin(id) + r = r ^ (np.uint8(str_id[-1]) << (7 - j)) + g = g ^ (np.uint8(str_id[-2]) << (7 - j)) + b = b ^ (np.uint8(str_id[-3]) << (7 - j)) + id = id >> 3 + cmap[i, 0] = r + cmap[i, 1] = g + cmap[i, 2] = b + + return cmap + + +class Colorize(object): + def __init__(self, n=35): + self.cmap = labelcolormap(n) + self.cmap = torch.from_numpy(self.cmap[:n]) + + def __call__(self, gray_image): + size = gray_image.size() + color_image = torch.ByteTensor(3, size[1], size[2]).fill_(0) + + for label in range(0, len(self.cmap)): + mask = (label == gray_image[0]).cpu() + color_image[0][mask] = self.cmap[label][0] + color_image[1][mask] = self.cmap[label][1] + color_image[2][mask] = self.cmap[label][2] + + return color_image + + +def make_folder(path): + if not os.path.exists(os.path.join(path)): + os.makedirs(os.path.join(path)) + + +def path_join_abs(*paths): + return os.path.abspath(os.path.join(*paths)) + + +def draw_landmark(landmarks, img_rd): + landmarks = landmarks.astype('int') + img_rd = img_rd.copy() + for idx, point in enumerate(landmarks): + pos = (point[0], point[1]) + cv2.circle(img_rd, pos, 2, color=(139, 0, 0)) + cv2.putText(img_rd, str(idx + 1), pos, cv2.FONT_HERSHEY_SIMPLEX, 0.3, (0, 0, 255), 1, cv2.LINE_AA) + return img_rd \ No newline at end of file diff --git a/models/sean_codes/styles_test/mean_style_code/mean/0/ACE.npy b/models/sean_codes/styles_test/mean_style_code/mean/0/ACE.npy new file mode 100644 index 0000000000000000000000000000000000000000..6d2197b1b46e84db722266cd12c04d886bdaa7e8 --- /dev/null +++ b/models/sean_codes/styles_test/mean_style_code/mean/0/ACE.npy @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:6df6679e50cd9b9fa367bfb7b4bc8583e4284f4c3f20b1dca99018bd93535d12 +size 2176 diff --git a/models/sean_codes/styles_test/mean_style_code/mean/1/ACE.npy b/models/sean_codes/styles_test/mean_style_code/mean/1/ACE.npy new file mode 100644 index 0000000000000000000000000000000000000000..3bc1d2a718cbdcaec856386edb1a6b304d762b38 --- /dev/null +++ b/models/sean_codes/styles_test/mean_style_code/mean/1/ACE.npy @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:069998c6759dec39bd83c86ec997c82488c838a8a2734b65f8996d0a61521a48 +size 2176 diff --git a/models/sean_codes/styles_test/mean_style_code/mean/10/ACE.npy b/models/sean_codes/styles_test/mean_style_code/mean/10/ACE.npy new file mode 100644 index 0000000000000000000000000000000000000000..930dfc1981b1e320bf15c7962205b1fc9f3f54f5 --- /dev/null +++ b/models/sean_codes/styles_test/mean_style_code/mean/10/ACE.npy @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:888ea376f03463d6973108233a83f17629b0231a69e0af0c73f836033d25fe1d +size 2176 diff --git a/models/sean_codes/styles_test/mean_style_code/mean/11/ACE.npy b/models/sean_codes/styles_test/mean_style_code/mean/11/ACE.npy new file mode 100644 index 0000000000000000000000000000000000000000..47525a27b3a210c43b7cdf5e7138ee6f63dd162f --- /dev/null +++ b/models/sean_codes/styles_test/mean_style_code/mean/11/ACE.npy @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:46781acb7d092581799ff668db98e0ca50c4a3138013aa3a120e74285cc182ae +size 2176 diff --git a/models/sean_codes/styles_test/mean_style_code/mean/12/ACE.npy b/models/sean_codes/styles_test/mean_style_code/mean/12/ACE.npy new file mode 100644 index 0000000000000000000000000000000000000000..d868eac0675a372f223dc053dde11852e6b43733 --- /dev/null +++ b/models/sean_codes/styles_test/mean_style_code/mean/12/ACE.npy @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:1c6c95dc6b510d1bf4ff9e7a536f79511aae6d98906427b20e4fe3309d94f291 +size 2176 diff --git a/models/sean_codes/styles_test/mean_style_code/mean/13/ACE.npy b/models/sean_codes/styles_test/mean_style_code/mean/13/ACE.npy new file mode 100644 index 0000000000000000000000000000000000000000..3a9d827baea9f27bd2c4ea0bb35aa7f11c4d8eb1 --- /dev/null +++ b/models/sean_codes/styles_test/mean_style_code/mean/13/ACE.npy @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:95f916d1570754b9d3c698a5982260094f1c7bdea4c955298015cb3289f48774 +size 2176 diff --git a/models/sean_codes/styles_test/mean_style_code/mean/14/ACE.npy b/models/sean_codes/styles_test/mean_style_code/mean/14/ACE.npy new file mode 100644 index 0000000000000000000000000000000000000000..99db5a4604b5f9bf0422ae98fb8145beed107845 --- /dev/null +++ b/models/sean_codes/styles_test/mean_style_code/mean/14/ACE.npy @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:9f0b1488d88ae4b393e709db67aeb04f52e73dba0c7d63c978ec489c053dac30 +size 2176 diff --git a/models/sean_codes/styles_test/mean_style_code/mean/15/ACE.npy b/models/sean_codes/styles_test/mean_style_code/mean/15/ACE.npy new file mode 100644 index 0000000000000000000000000000000000000000..0f61fae7063f91345782ebad22bc4243cc3688c0 --- /dev/null +++ b/models/sean_codes/styles_test/mean_style_code/mean/15/ACE.npy @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:0ff1e5bd7cd1a00ca1b065ff358e82a0cc56744ffb544d72a8a4b130a3081ea6 +size 2176 diff --git a/models/sean_codes/styles_test/mean_style_code/mean/16/ACE.npy b/models/sean_codes/styles_test/mean_style_code/mean/16/ACE.npy new file mode 100644 index 0000000000000000000000000000000000000000..4b9212d1aa487e3e62e3ed1c79362f174fa53a83 --- /dev/null +++ b/models/sean_codes/styles_test/mean_style_code/mean/16/ACE.npy @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:ceebf902fc78355df4d4a3b95b57c48dbce61e166c0233c73464e5c07fce4166 +size 2176 diff --git a/models/sean_codes/styles_test/mean_style_code/mean/17/ACE.npy b/models/sean_codes/styles_test/mean_style_code/mean/17/ACE.npy new file mode 100644 index 0000000000000000000000000000000000000000..ce96f5be3916536c9b6d8f98be8ca8660a09d77d --- /dev/null +++ b/models/sean_codes/styles_test/mean_style_code/mean/17/ACE.npy @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:00faed73adf5726d0d2fd6f9f763e3651be602c948245704d55b1d35adf5e38c +size 2176 diff --git a/models/sean_codes/styles_test/mean_style_code/mean/18/ACE.npy b/models/sean_codes/styles_test/mean_style_code/mean/18/ACE.npy new file mode 100644 index 0000000000000000000000000000000000000000..e727d2ca8050058a1ede22c2ab543e28737b30c3 --- /dev/null +++ b/models/sean_codes/styles_test/mean_style_code/mean/18/ACE.npy @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:cd49b67d5f6e468b3df2c340007722feb9490a899c506b862327d41d7536225b +size 2176 diff --git a/models/sean_codes/styles_test/mean_style_code/mean/2/ACE.npy b/models/sean_codes/styles_test/mean_style_code/mean/2/ACE.npy new file mode 100644 index 0000000000000000000000000000000000000000..b8e09628e2dc42ccd10f801504dd1679dcae0e55 --- /dev/null +++ b/models/sean_codes/styles_test/mean_style_code/mean/2/ACE.npy @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:56a2dfee3f66cdbf1ab860824a52b06974a2a2275b010849761c4c86ab9c583b +size 2176 diff --git a/models/sean_codes/styles_test/mean_style_code/mean/3/ACE.npy b/models/sean_codes/styles_test/mean_style_code/mean/3/ACE.npy new file mode 100644 index 0000000000000000000000000000000000000000..c8c148ec7ffb78eee20467089b5bc1bf559a2401 --- /dev/null +++ b/models/sean_codes/styles_test/mean_style_code/mean/3/ACE.npy @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:55e6cfe4053ea221774f6f9cbff403d9092953d5c80375db821d0d6d08175551 +size 2176 diff --git a/models/sean_codes/styles_test/mean_style_code/mean/4/ACE.npy b/models/sean_codes/styles_test/mean_style_code/mean/4/ACE.npy new file mode 100644 index 0000000000000000000000000000000000000000..32768d886efc9f0499c5b99ddab970058f1bb9f7 --- /dev/null +++ b/models/sean_codes/styles_test/mean_style_code/mean/4/ACE.npy @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:acd24c3dfd93f4114b1ec12194196580890e300b9ac10cd4181ccfbabfcb6718 +size 2176 diff --git a/models/sean_codes/styles_test/mean_style_code/mean/5/ACE.npy b/models/sean_codes/styles_test/mean_style_code/mean/5/ACE.npy new file mode 100644 index 0000000000000000000000000000000000000000..4af0b1e470756c4bfdba2859c54622722d1c38e0 --- /dev/null +++ b/models/sean_codes/styles_test/mean_style_code/mean/5/ACE.npy @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:8c0ce9398bdeb4a90d69ba99980813e3c637aea6f841b6111c36fd87441badde +size 2176 diff --git a/models/sean_codes/styles_test/mean_style_code/mean/6/ACE.npy b/models/sean_codes/styles_test/mean_style_code/mean/6/ACE.npy new file mode 100644 index 0000000000000000000000000000000000000000..7701e924191a1014e84811f41c2ac05b8994ad22 --- /dev/null +++ b/models/sean_codes/styles_test/mean_style_code/mean/6/ACE.npy @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:5ae2e169fb73a5bdb79bc4e0e858d8bb79cba5ee011f42826e4599245c149eae +size 2176 diff --git a/models/sean_codes/styles_test/mean_style_code/mean/7/ACE.npy b/models/sean_codes/styles_test/mean_style_code/mean/7/ACE.npy new file mode 100644 index 0000000000000000000000000000000000000000..b81d76e8d031a049535b700a183f6a7cb218c052 --- /dev/null +++ b/models/sean_codes/styles_test/mean_style_code/mean/7/ACE.npy @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:484433038cd434b8dfe590bd659c52ac685b0b4669b848897bd7416787a042ef +size 2176 diff --git a/models/sean_codes/styles_test/mean_style_code/mean/8/ACE.npy b/models/sean_codes/styles_test/mean_style_code/mean/8/ACE.npy new file mode 100644 index 0000000000000000000000000000000000000000..8ab5e27a848e7797a10d75b073930921c1cb95ed --- /dev/null +++ b/models/sean_codes/styles_test/mean_style_code/mean/8/ACE.npy @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:92e06adecd89b948d136b5628974973a843675e9193cf7e854a91af542c94179 +size 2176 diff --git a/models/sean_codes/styles_test/mean_style_code/mean/9/ACE.npy b/models/sean_codes/styles_test/mean_style_code/mean/9/ACE.npy new file mode 100644 index 0000000000000000000000000000000000000000..2cbd13efaec0219568cc814028e00572ad7d683f --- /dev/null +++ b/models/sean_codes/styles_test/mean_style_code/mean/9/ACE.npy @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:90f8e4b9306cd55bb241bce0ec919296f7106b123fa9216e68478736e9188e78 +size 2176 diff --git a/models/sean_codes/styles_test/mean_style_code/median/0/ACE.npy b/models/sean_codes/styles_test/mean_style_code/median/0/ACE.npy new file mode 100644 index 0000000000000000000000000000000000000000..0d65e4559f7bd806f06909f0ed1d3ccc908a15e6 --- /dev/null +++ b/models/sean_codes/styles_test/mean_style_code/median/0/ACE.npy @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:01f260cb73bb826fe09af14d57aeb29ef5c614c62718a642fff45d86b99546b3 +size 2176 diff --git a/models/sean_codes/styles_test/mean_style_code/median/1/ACE.npy b/models/sean_codes/styles_test/mean_style_code/median/1/ACE.npy new file mode 100644 index 0000000000000000000000000000000000000000..feaf10805b7b3292a0e1f32c0adbd184e63180b2 --- /dev/null +++ b/models/sean_codes/styles_test/mean_style_code/median/1/ACE.npy @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:5e2b44c1187f089f598621a6ec09bf3630c184e0274513812c3d642c4b10a7c9 +size 2176 diff --git a/models/sean_codes/styles_test/mean_style_code/median/10/ACE.npy b/models/sean_codes/styles_test/mean_style_code/median/10/ACE.npy new file mode 100644 index 0000000000000000000000000000000000000000..ec1571f57a8191cc3551e83e880c67d28685de3e --- /dev/null +++ b/models/sean_codes/styles_test/mean_style_code/median/10/ACE.npy @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:9686a06242ed8ddd899334ee74225ae6266942e3dae624304ea9929d112d6ebb +size 2176 diff --git a/models/sean_codes/styles_test/mean_style_code/median/11/ACE.npy b/models/sean_codes/styles_test/mean_style_code/median/11/ACE.npy new file mode 100644 index 0000000000000000000000000000000000000000..ed88c2c321046b3236c9111b89502e043465568c --- /dev/null +++ b/models/sean_codes/styles_test/mean_style_code/median/11/ACE.npy @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:b5b8a75e9b60c770ddcf2920d61093db81f5c725936451e873ab3a93cbec4918 +size 2176 diff --git a/models/sean_codes/styles_test/mean_style_code/median/12/ACE.npy b/models/sean_codes/styles_test/mean_style_code/median/12/ACE.npy new file mode 100644 index 0000000000000000000000000000000000000000..0b032ba2b7764a842384708e166c50bd2c74ab27 --- /dev/null +++ b/models/sean_codes/styles_test/mean_style_code/median/12/ACE.npy @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:6be4d2262f3f7444f8aa2dbe9459b5a80b44f48b5eb8fd61f599f52c5830d441 +size 2176 diff --git a/models/sean_codes/styles_test/mean_style_code/median/13/ACE.npy b/models/sean_codes/styles_test/mean_style_code/median/13/ACE.npy new file mode 100644 index 0000000000000000000000000000000000000000..6a08f63a39d192fdb8ef9cddc627570d732211fb --- /dev/null +++ b/models/sean_codes/styles_test/mean_style_code/median/13/ACE.npy @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:2b81b64eded777babb6cb4cc77348487c1274ccf46211826a3dedb620b3c5318 +size 2176 diff --git a/models/sean_codes/styles_test/mean_style_code/median/14/ACE.npy b/models/sean_codes/styles_test/mean_style_code/median/14/ACE.npy new file mode 100644 index 0000000000000000000000000000000000000000..4542adebb22eb106231d89a537b238b3312a2f66 --- /dev/null +++ b/models/sean_codes/styles_test/mean_style_code/median/14/ACE.npy @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:8b1bda532440da30c1eae7873d5610fdf6bbc2628223a2eba6b4908c1476bb8d +size 2176 diff --git a/models/sean_codes/styles_test/mean_style_code/median/15/ACE.npy b/models/sean_codes/styles_test/mean_style_code/median/15/ACE.npy new file mode 100644 index 0000000000000000000000000000000000000000..6e933de7025533c6438002de05c79988afa716c6 --- /dev/null +++ b/models/sean_codes/styles_test/mean_style_code/median/15/ACE.npy @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a5b4db7ea6cdf99cb0e86609ac4cd996103af7aa2d6dcb13652ad26f170d1a56 +size 2176 diff --git a/models/sean_codes/styles_test/mean_style_code/median/16/ACE.npy b/models/sean_codes/styles_test/mean_style_code/median/16/ACE.npy new file mode 100644 index 0000000000000000000000000000000000000000..9850c8af2f253f45f920aa6db028cae6c0bf4669 --- /dev/null +++ b/models/sean_codes/styles_test/mean_style_code/median/16/ACE.npy @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:e5ae251680f4afe303c88342ab27f98d75fdf3428a4b69e0e4de5c42160b5245 +size 2176 diff --git a/models/sean_codes/styles_test/mean_style_code/median/17/ACE.npy b/models/sean_codes/styles_test/mean_style_code/median/17/ACE.npy new file mode 100644 index 0000000000000000000000000000000000000000..93bc7838406b09cf3b4d3a2061f70a8f9492efd5 --- /dev/null +++ b/models/sean_codes/styles_test/mean_style_code/median/17/ACE.npy @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:0bdb5b8f1407fe0ee0059032fc801fe7b34ca4f8195a33d8e409682697f7ae04 +size 2176 diff --git a/models/sean_codes/styles_test/mean_style_code/median/18/ACE.npy b/models/sean_codes/styles_test/mean_style_code/median/18/ACE.npy new file mode 100644 index 0000000000000000000000000000000000000000..ef139a1b4208faa54c53796e6ab08a49e1332b5b --- /dev/null +++ b/models/sean_codes/styles_test/mean_style_code/median/18/ACE.npy @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:e8f5c975a9b60ec10bea335e663eb4feebeb4108f8f76955203beefd1dd3db60 +size 2176 diff --git a/models/sean_codes/styles_test/mean_style_code/median/2/ACE.npy b/models/sean_codes/styles_test/mean_style_code/median/2/ACE.npy new file mode 100644 index 0000000000000000000000000000000000000000..d212f43f5047139326acda225089a8ebdef59a48 --- /dev/null +++ b/models/sean_codes/styles_test/mean_style_code/median/2/ACE.npy @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:71c1cd3d3d49ce097447f5789f93ba3928b54af7d7fd0340896331a66d8a7067 +size 2176 diff --git a/models/sean_codes/styles_test/mean_style_code/median/3/ACE.npy b/models/sean_codes/styles_test/mean_style_code/median/3/ACE.npy new file mode 100644 index 0000000000000000000000000000000000000000..a804404f4f3ff72734414000c7b607ac6c7235dc --- /dev/null +++ b/models/sean_codes/styles_test/mean_style_code/median/3/ACE.npy @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:8ef5824b3f4ae142107567caedeed38e6dbf243ef7b748f484813a6d61258bdd +size 2176 diff --git a/models/sean_codes/styles_test/mean_style_code/median/4/ACE.npy b/models/sean_codes/styles_test/mean_style_code/median/4/ACE.npy new file mode 100644 index 0000000000000000000000000000000000000000..adf59ab10af468d3553ca81a0027e7abd633da98 --- /dev/null +++ b/models/sean_codes/styles_test/mean_style_code/median/4/ACE.npy @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:e7eb947b107ac66aa984b7daa032f95cc5ce9d2199bc964ad165e63536060b42 +size 2176 diff --git a/models/sean_codes/styles_test/mean_style_code/median/5/ACE.npy b/models/sean_codes/styles_test/mean_style_code/median/5/ACE.npy new file mode 100644 index 0000000000000000000000000000000000000000..88b7642157ca0daab09fc2a2fdc1fa15c1d645f2 --- /dev/null +++ b/models/sean_codes/styles_test/mean_style_code/median/5/ACE.npy @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:d1f3ed64ac47ee99ea6521683f66f3f6624ac4e7506f70eaddbefab43891a7ed +size 2176 diff --git a/models/sean_codes/styles_test/mean_style_code/median/6/ACE.npy b/models/sean_codes/styles_test/mean_style_code/median/6/ACE.npy new file mode 100644 index 0000000000000000000000000000000000000000..75d90174e480137445a1579a06c1023357667102 --- /dev/null +++ b/models/sean_codes/styles_test/mean_style_code/median/6/ACE.npy @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:238b505caec470e82ed0d375678de303e5d21b1af0a1fc44a1a0ab5331edfcf0 +size 2176 diff --git a/models/sean_codes/styles_test/mean_style_code/median/7/ACE.npy b/models/sean_codes/styles_test/mean_style_code/median/7/ACE.npy new file mode 100644 index 0000000000000000000000000000000000000000..a8330f0583ca3b07fa2759a5a10879903bab3d20 --- /dev/null +++ b/models/sean_codes/styles_test/mean_style_code/median/7/ACE.npy @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:24258a2a07bfbddd751c63a9011b57b74b0cc199a010850378c7a9bff0c65084 +size 2176 diff --git a/models/sean_codes/styles_test/mean_style_code/median/8/ACE.npy b/models/sean_codes/styles_test/mean_style_code/median/8/ACE.npy new file mode 100644 index 0000000000000000000000000000000000000000..18811d022f9f22a742c0639a1729c2022e7b492c --- /dev/null +++ b/models/sean_codes/styles_test/mean_style_code/median/8/ACE.npy @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:fd62e3dda70f74301a1fe1d54cb6efb34807d9877584873dcbe18941a90df40a +size 2176 diff --git a/models/sean_codes/styles_test/mean_style_code/median/9/ACE.npy b/models/sean_codes/styles_test/mean_style_code/median/9/ACE.npy new file mode 100644 index 0000000000000000000000000000000000000000..2fca73059c48338708fe657c71a1b46d8e7f3168 --- /dev/null +++ b/models/sean_codes/styles_test/mean_style_code/median/9/ACE.npy @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:1e528a9e65f84d8329532ef6ba227ef2b0bd95d844880078e3868143e18ff9bc +size 2176 diff --git a/models/sean_codes/trainers/__init__.py b/models/sean_codes/trainers/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..7c23f2831952e1a35fd8f170760c627f7fb20ed9 --- /dev/null +++ b/models/sean_codes/trainers/__init__.py @@ -0,0 +1,4 @@ +""" +Copyright (C) 2019 NVIDIA Corporation. All rights reserved. +Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode). +""" diff --git a/models/sean_codes/trainers/pix2pix_trainer.py b/models/sean_codes/trainers/pix2pix_trainer.py new file mode 100644 index 0000000000000000000000000000000000000000..988f3244fe6d5152b0e51561635066432dfd0b77 --- /dev/null +++ b/models/sean_codes/trainers/pix2pix_trainer.py @@ -0,0 +1,86 @@ +""" +Copyright (C) 2019 NVIDIA Corporation. All rights reserved. +Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode). +""" + +from sean_codes.models.networks.sync_batchnorm import DataParallelWithCallback +from sean_codes.models.pix2pix_model import Pix2PixModel + + +class Pix2PixTrainer(): + """ + Trainer creates the model and optimizers, and uses them to + updates the weights of the network while reporting losses + and the latest visuals to visualize the progress in training. + """ + + def __init__(self, opt): + self.opt = opt + self.pix2pix_model = Pix2PixModel(opt) + if len(opt.gpu_ids) > 0: + self.pix2pix_model = DataParallelWithCallback(self.pix2pix_model, + device_ids=opt.gpu_ids) + self.pix2pix_model_on_one_gpu = self.pix2pix_model.module + else: + self.pix2pix_model_on_one_gpu = self.pix2pix_model + + self.generated = None + if opt.isTrain: + self.optimizer_G, self.optimizer_D = \ + self.pix2pix_model_on_one_gpu.create_optimizers(opt) + self.old_lr = opt.lr + + def run_generator_one_step(self, data): + self.optimizer_G.zero_grad() + g_losses, generated = self.pix2pix_model(data, mode='generator') + g_loss = sum(g_losses.values()).mean() + g_loss.backward() + self.optimizer_G.step() + self.g_losses = g_losses + self.generated = generated + + def run_discriminator_one_step(self, data): + self.optimizer_D.zero_grad() + d_losses = self.pix2pix_model(data, mode='discriminator') + d_loss = sum(d_losses.values()).mean() + d_loss.backward() + self.optimizer_D.step() + self.d_losses = d_losses + + def get_latest_losses(self): + return {**self.g_losses, **self.d_losses} + + def get_latest_generated(self): + return self.generated + + def update_learning_rate(self, epoch): + self.update_learning_rate(epoch) + + def save(self, epoch): + self.pix2pix_model_on_one_gpu.save(epoch) + + ################################################################## + # Helper functions + ################################################################## + + def update_learning_rate(self, epoch): + if epoch > self.opt.niter: + lrd = self.opt.lr / self.opt.niter_decay + new_lr = self.old_lr - lrd + else: + new_lr = self.old_lr + + if new_lr != self.old_lr: + if self.opt.no_TTUR: + new_lr_G = new_lr + new_lr_D = new_lr + else: + new_lr_G = new_lr / 2 + new_lr_D = new_lr * 2 + + for param_group in self.optimizer_D.param_groups: + param_group['lr'] = new_lr_D + for param_group in self.optimizer_G.param_groups: + param_group['lr'] = new_lr_G + print('update learning rate: %f -> %f' % (self.old_lr, new_lr)) + self.old_lr = new_lr diff --git a/models/sean_codes/util/__init__.py b/models/sean_codes/util/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/models/sean_codes/util/canvas_grid.py b/models/sean_codes/util/canvas_grid.py new file mode 100644 index 0000000000000000000000000000000000000000..65a200854f842ce8d83876be58bc401c2c24bcc9 --- /dev/null +++ b/models/sean_codes/util/canvas_grid.py @@ -0,0 +1,34 @@ +# -*- coding: utf-8 -*- + +""" +# File name: canvas_grid.py +# Time : 2021/11/18 16:40 +# Author: xyguoo@163.com +# Description: This is the util that generating image in a big grid picture +""" + +import numpy as np + +from util.imutil import write_rgb + + +class Canvas: + def __init__(self, row, col, img_size=256, margin=0): + self.row = row + self.col = col + self.img_size = img_size + self.margin = margin + self.canvas = np.ones((row * img_size, col * img_size + margin * (col - 1), 3), dtype='uint8') * 255 + + def process_draw_image(self, img, i, j): + if img.dtype in [np.float32, np.float, np.float64]: + if img.min() < 0: + img = img * 127.5 + 127.5 + elif img.max() <= 1: + img = img * 255 + img = img.astype('uint8') + i_start, j_start = int(i * self.img_size), int(j * self.img_size) + int(j * self.margin) + self.canvas[i_start: i_start + img.shape[0], j_start: j_start + img.shape[1], :] = img + + def write_(self, file): + write_rgb(file, self.canvas) diff --git a/models/sean_codes/util/color_from_hsv_to_gaussian.py b/models/sean_codes/util/color_from_hsv_to_gaussian.py new file mode 100644 index 0000000000000000000000000000000000000000..5a45b8bacc231d6430eaf971288bd859a5a8d55d --- /dev/null +++ b/models/sean_codes/util/color_from_hsv_to_gaussian.py @@ -0,0 +1,40 @@ +# -*- coding: utf-8 -*- + +""" +# File name: color_from_hsv_to_gaussian.py +# Time : 2022/3/6 16:00 +# Author: xyguoo@163.com +# Description: transfer HSV value to Gaussian latent value according to the distribution of the dataset +""" +import os +import pickle as pkl +from bisect import bisect_left, bisect_right + +import scipy.stats as st + + +class DistTranslation: + def __init__(self): + hair_root = 'dataset_info_ctrlhair' + with open(os.path.join(hair_root, 'hsv_stat_dict_ordered.pkl'), 'rb') as f: + self.cols_hsv = pkl.load(f) + + def gaussian_to_val(self, dim, val): + # if dim == 0: + # return (val + 2.) / 4 * 179 + return self.cols_hsv[int((st.norm.cdf(val)) * self.cols_hsv.shape[0])][dim] + + def val_to_gaussian(self, dim, val): + # if dim == 0: + # return val / 179 * 2 * 2. - 2. + + left_v = bisect_left(self.cols_hsv[:, dim], val) + right_v = bisect_right(self.cols_hsv[:, dim], val) + return st.norm.ppf((left_v + right_v) / 2 / self.cols_hsv.shape[0]) + +# +# if __name__ == '__main__': +# dt = DistTranslation() +# +# with open('hsv_stat_dict_ordered.pkl', 'wb') as f: +# pkl.dump(dt.cols_hsv, f) diff --git a/models/sean_codes/util/common_options.py b/models/sean_codes/util/common_options.py new file mode 100644 index 0000000000000000000000000000000000000000..7f54b7f53ca743e0da96d2a2c05aad5cd234d569 --- /dev/null +++ b/models/sean_codes/util/common_options.py @@ -0,0 +1,15 @@ +# -*- coding: utf-8 -*- + +""" +# File name: common_options.py +# Time : 2022/7/12 +# Author: xyguoo@163.com +# Description: +""" + +def ctrl_hair_parser_options(parser): + parser.add_argument('-c', '--config', type=str, default='001') + parser.add_argument('-g', '--gpu', type=str, help='Specify GPU number', default='0') + parser.add_argument('-n', '--need_crop', type=bool, help='whether images need crop', default=True) + parser.add_argument('--no_blending', action='store_true', + help='whether using poisson blending as post processing', default=False) \ No newline at end of file diff --git a/models/sean_codes/util/find_semantic_direction.py b/models/sean_codes/util/find_semantic_direction.py new file mode 100644 index 0000000000000000000000000000000000000000..e2f8b06abc4cd0906825404cfeb2639a3ac82b31 --- /dev/null +++ b/models/sean_codes/util/find_semantic_direction.py @@ -0,0 +1,21 @@ +# -*- coding: utf-8 -*- + +""" +# File name: find_semantic_direction.py +# Time : 2022/07/17 +# Author: xyguoo@163.com +# Description: +""" +import torch + + +def get_random_direction(dim, existing_dirs): + dir = torch.randn(dim) + + for dd in existing_dirs: + dir = dir - torch.dot(dir, dd) * dd + + if dir[0] < 0: + dir = -dir + dir = dir / dir.norm() + return dir diff --git a/models/sean_codes/util/imutil.py b/models/sean_codes/util/imutil.py new file mode 100644 index 0000000000000000000000000000000000000000..539fa8c449eb192a8c84d5c01588f029a5ecd93d --- /dev/null +++ b/models/sean_codes/util/imutil.py @@ -0,0 +1,24 @@ +# -*- coding: utf-8 -*- + +""" +# File name: imutil.py +# Time : 2021/12/7 14:55 +# Author: xyguoo@163.com +# Description: +""" +import cv2 +import numpy as np + + +def read_rgb(img_path): + im = cv2.imread(img_path) + im = cv2.cvtColor(im, cv2.COLOR_BGR2RGB) + return im + + +def write_rgb(file_name, img): + if len(img.shape) == 2: + img = np.stack([img] * 3, axis=2) + elif img.shape[2] == 1: + img = np.tile(img, [1, 1, 3]) + cv2.imwrite(file_name, cv2.cvtColor(img, cv2.COLOR_RGB2BGR)) diff --git a/models/sean_codes/util/mask_color_util.py b/models/sean_codes/util/mask_color_util.py new file mode 100644 index 0000000000000000000000000000000000000000..c496255368d29cf4cd511387306aa066dffbc7a0 --- /dev/null +++ b/models/sean_codes/util/mask_color_util.py @@ -0,0 +1,64 @@ +# -*- coding: utf-8 -*- + +""" +# File name: mask_color_util.py +# Time : 2021/12/7 11:13 +# Author: xyguoo@163.com +# Description: +""" + +import numpy as np + +from global_value_utils import HAIR_IDX + + +def mask_to_rgb(pred, draw_type=2): + """ + generate visual mask image + :param pred: pred is the label image with pixel vale in {0, 1, 2, ..., 20} + :param draw_type: 0: all part; 1: {bg, face, hair}; 2: {hair, others} + :return: + """ + if len(pred.shape) == 3 and pred.shape[0] == 1: + pred = pred[0] + num_labels = 19 + color = np.array([[0, 128, 64], + [204, 0, 0], + [76, 153, 0], + [204, 204, 0], ## + [51, 51, 255], ## + [204, 0, 204], ## + [0, 255, 255], ## + [51, 255, 255], ## + [102, 51, 0], ## + [255, 0, 0], ## + [102, 204, 0], ## + [255, 255, 0], ## + [0, 0, 153], ## + [0, 0, 204], ## + [255, 51, 153], ## + [0, 204, 204], ## + [0, 51, 0], ## + [255, 153, 51], + [0, 204, 0], + ]) + + for cc in range(len(color)): + if draw_type == 2: + if cc != HAIR_IDX: + color[cc] = [255, 255, 255] + elif draw_type == 1: + if cc != HAIR_IDX and cc != 0: + color[cc] = [237, 28, 36] + + h, w = np.shape(pred) + rgb = np.zeros((h, w, 3), dtype=np.uint8) + # print(color.shape) + for ii in range(num_labels): + # print(ii) + mask = pred == ii + rgb[mask, None] = color[ii, :] + # Correct unk + unk = pred == 255 + rgb[unk, None] = 255 + return rgb diff --git a/models/sean_codes/util/util.py b/models/sean_codes/util/util.py new file mode 100644 index 0000000000000000000000000000000000000000..057647cf579c0e8c7a432c4e6abb4b78af2700b7 --- /dev/null +++ b/models/sean_codes/util/util.py @@ -0,0 +1,283 @@ +""" +These interfaces are most based on SEAN. But we modify some of them, and add some new interface for our CtrlHair. +""" + +import argparse +import importlib +import os +import re +import sys + +import cv2 +import dill as pickle +import numpy as np +import torch +from PIL import Image + + +def save_obj(obj, name): + with open(name, 'wb') as f: + pickle.dump(obj, f, pickle.HIGHEST_PROTOCOL) + + +def load_obj(name): + with open(name, 'rb') as f: + return pickle.load(f) + +# returns a configuration for creating a generator +# |default_opt| should be the opt of the current experiment +# |**kwargs|: if any configuration should be overriden, it can be specified here + + +def copyconf(default_opt, **kwargs): + conf = argparse.Namespace(**vars(default_opt)) + for key in kwargs: + print(key, kwargs[key]) + setattr(conf, key, kwargs[key]) + return conf + + +def tile_images(imgs, picturesPerRow=4): + """ Code borrowed from + https://stackoverflow.com/questions/26521365/cleanly-tile-numpy-array-of-images-stored-in-a-flattened-1d-format/26521997 + """ + + # Padding + if imgs.shape[0] % picturesPerRow == 0: + rowPadding = 0 + else: + rowPadding = picturesPerRow - imgs.shape[0] % picturesPerRow + if rowPadding > 0: + imgs = np.concatenate([imgs, np.zeros((rowPadding, *imgs.shape[1:]), dtype=imgs.dtype)], axis=0) + + # Tiling Loop (The conditionals are not necessary anymore) + tiled = [] + for i in range(0, imgs.shape[0], picturesPerRow): + tiled.append(np.concatenate([imgs[j] for j in range(i, i + picturesPerRow)], axis=1)) + + tiled = np.concatenate(tiled, axis=0) + return tiled + + +# Converts a Tensor into a Numpy array +# |imtype|: the desired type of the converted numpy array +def tensor2im(image_tensor, imtype=np.uint8, normalize=True, tile=False): + if isinstance(image_tensor, list): + image_numpy = [] + for i in range(len(image_tensor)): + image_numpy.append(tensor2im(image_tensor[i], imtype, normalize)) + return image_numpy + + if image_tensor.dim() == 4: + # transform each image in the batch + images_np = [] + for b in range(image_tensor.size(0)): + one_image = image_tensor[b] + one_image_np = tensor2im(one_image) + images_np.append(one_image_np.reshape(1, *one_image_np.shape)) + images_np = np.concatenate(images_np, axis=0) + if tile: + images_tiled = tile_images(images_np) + return images_tiled + else: + return images_np + + if image_tensor.dim() == 2: + image_tensor = image_tensor.unsqueeze(0) + image_numpy = image_tensor.detach().cpu().float().numpy() + if normalize: + image_numpy = (np.transpose(image_numpy, (1, 2, 0)) + 1) / 2.0 * 255.0 + else: + image_numpy = np.transpose(image_numpy, (1, 2, 0)) * 255.0 + image_numpy = np.clip(image_numpy, 0, 255) + if image_numpy.shape[2] == 1: + image_numpy = image_numpy[:, :, 0] + return image_numpy.astype(imtype) + + +# Converts a one-hot tensor into a colorful label map +def tensor2label(label_tensor, n_label, imtype=np.uint8, tile=False): + if label_tensor.dim() == 4: + # transform each image in the batch + images_np = [] + for b in range(label_tensor.size(0)): + one_image = label_tensor[b] + one_image_np = tensor2label(one_image, n_label, imtype) + images_np.append(one_image_np.reshape(1, *one_image_np.shape)) + images_np = np.concatenate(images_np, axis=0) + if tile: + images_tiled = tile_images(images_np) + return images_tiled + else: + images_np = images_np[0] + return images_np + + if label_tensor.dim() == 1: + return np.zeros((64, 64, 3), dtype=np.uint8) + if n_label == 0: + return tensor2im(label_tensor, imtype) + label_tensor = label_tensor.cpu().float() + if label_tensor.size()[0] > 1: + label_tensor = label_tensor.max(0, keepdim=True)[1] + label_tensor = Colorize(n_label)(label_tensor) + label_numpy = np.transpose(label_tensor.numpy(), (1, 2, 0)) + result = label_numpy.astype(imtype) + return result + + +def save_image(image_numpy, image_path, create_dir=False): + if create_dir: + os.makedirs(os.path.dirname(image_path), exist_ok=True) + if len(image_numpy.shape) == 2: + image_numpy = np.expand_dims(image_numpy, axis=2) + if image_numpy.shape[2] == 1: + image_numpy = np.repeat(image_numpy, 3, 2) + image_pil = Image.fromarray(image_numpy) + + # save to png + image_pil.save(image_path.replace('.jpg', '.png')) + + +def mkdirs(paths): + if isinstance(paths, list) and not isinstance(paths, str): + for path in paths: + mkdir(path) + else: + mkdir(paths) + + +def mkdir(path): + if not os.path.exists(path): + os.makedirs(path) + + +def atoi(text): + return int(text) if text.isdigit() else text + + +def natural_keys(text): + ''' + alist.sort(key=natural_keys) sorts in human order + http://nedbatchelder.com/blog/200712/human_sorting.html + (See Toothy's implementation in the comments) + ''' + return [atoi(c) for c in re.split('(\d+)', text)] + + +def natural_sort(items): + items.sort(key=natural_keys) + + +def str2bool(v): + if v.lower() in ('yes', 'true', 't', 'y', '1'): + return True + elif v.lower() in ('no', 'false', 'f', 'n', '0'): + return False + else: + raise argparse.ArgumentTypeError('Boolean value expected.') + + +def find_class_in_module(target_cls_name, module): + module = f'models.{module}' + target_cls_name = target_cls_name.replace('_', '').lower() + clslib = importlib.import_module(module) + cls = None + for name, clsobj in clslib.__dict__.items(): + if name.lower() == target_cls_name: + cls = clsobj + + if cls is None: + print("In %s, there should be a class whose name matches %s in lowercase without underscore(_)" % (module, target_cls_name)) + exit(0) + + return cls + + +def save_network(net, label, epoch, opt): + save_filename = '%s_net_%s.pth' % (epoch, label) + save_path = os.path.join(opt.checkpoints_dir, opt.name, save_filename) + torch.save(net.cpu().state_dict(), save_path) + if len(opt.gpu_ids) and torch.cuda.is_available(): + net.cuda() + + +def load_network(net, label, epoch, opt): + save_filename = '%s_net_%s.pth' % (epoch, label) + save_dir = os.path.join(opt.checkpoints_dir, opt.name) + save_path = os.path.join(save_dir, save_filename) + weights = torch.load(save_path) + net.load_state_dict(weights) + return net + + +############################################################################### +# Code from +# https://github.com/ycszen/pytorch-seg/blob/master/transform.py +# Modified so it complies with the Citscape label map colors +############################################################################### +def uint82bin(n, count=8): + """returns the binary of integer n, count refers to amount of bits""" + return ''.join([str((n >> y) & 1) for y in range(count - 1, -1, -1)]) + + +def labelcolormap(N): + if N == 35: # cityscape + cmap = np.array([(0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), (111, 74, 0), (81, 0, 81), + (128, 64, 128), (244, 35, 232), (250, 170, 160), (230, 150, 140), (70, 70, 70), (102, 102, 156), (190, 153, 153), + (180, 165, 180), (150, 100, 100), (150, 120, 90), (153, 153, 153), (153, 153, 153), (250, 170, 30), (220, 220, 0), + (107, 142, 35), (152, 251, 152), (70, 130, 180), (220, 20, 60), (255, 0, 0), (0, 0, 142), (0, 0, 70), + (0, 60, 100), (0, 0, 90), (0, 0, 110), (0, 80, 100), (0, 0, 230), (119, 11, 32), (0, 0, 142)], + dtype=np.uint8) + else: + cmap = np.zeros((N, 3), dtype=np.uint8) + for i in range(N): + r, g, b = 0, 0, 0 + id = i + 1 # let's give 0 a color + for j in range(7): + str_id = uint82bin(id) + r = r ^ (np.uint8(str_id[-1]) << (7 - j)) + g = g ^ (np.uint8(str_id[-2]) << (7 - j)) + b = b ^ (np.uint8(str_id[-3]) << (7 - j)) + id = id >> 3 + cmap[i, 0] = r + cmap[i, 1] = g + cmap[i, 2] = b + + return cmap + + +class Colorize(object): + def __init__(self, n=35): + self.cmap = labelcolormap(n) + self.cmap = torch.from_numpy(self.cmap[:n]) + + def __call__(self, gray_image): + size = gray_image.size() + color_image = torch.ByteTensor(3, size[1], size[2]).fill_(0) + + for label in range(0, len(self.cmap)): + mask = (label == gray_image[0]).cpu() + color_image[0][mask] = self.cmap[label][0] + color_image[1][mask] = self.cmap[label][1] + color_image[2][mask] = self.cmap[label][2] + + return color_image + + +def make_folder(path): + if not os.path.exists(os.path.join(path)): + os.makedirs(os.path.join(path)) + + +def path_join_abs(*paths): + return os.path.abspath(os.path.join(*paths)) + + +def draw_landmark(landmarks, img_rd): + landmarks = landmarks.astype('int') + img_rd = img_rd.copy() + for idx, point in enumerate(landmarks): + pos = (point[0], point[1]) + cv2.circle(img_rd, pos, 2, color=(139, 0, 0)) + cv2.putText(img_rd, str(idx + 1), pos, cv2.FONT_HERSHEY_SIMPLEX, 0.3, (0, 0, 255), 1, cv2.LINE_AA) + return img_rd \ No newline at end of file diff --git a/models/stylegan2/.DS_Store b/models/stylegan2/.DS_Store new file mode 100644 index 0000000000000000000000000000000000000000..06a8775e69b7505ac9488a82a0c67712a9c3d8ca Binary files /dev/null and b/models/stylegan2/.DS_Store differ diff --git a/models/stylegan2/__init__.py b/models/stylegan2/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/models/stylegan2/dnnlib/__init__.py b/models/stylegan2/dnnlib/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..0a2cd19cb09ef6cd5a9f74d3b97a91c6aa080558 --- /dev/null +++ b/models/stylegan2/dnnlib/__init__.py @@ -0,0 +1,9 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +# +# NVIDIA CORPORATION and its licensors retain all intellectual property +# and proprietary rights in and to this software, related documentation +# and any modifications thereto. Any use, reproduction, disclosure or +# distribution of this software and related documentation without an express +# license agreement from NVIDIA CORPORATION is strictly prohibited. + +from .util import EasyDict, make_cache_dir_path diff --git a/models/stylegan2/dnnlib/util.py b/models/stylegan2/dnnlib/util.py new file mode 100644 index 0000000000000000000000000000000000000000..4771462e2a75387cfc2b9eec793f8238cb5c345c --- /dev/null +++ b/models/stylegan2/dnnlib/util.py @@ -0,0 +1,477 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +# +# NVIDIA CORPORATION and its licensors retain all intellectual property +# and proprietary rights in and to this software, related documentation +# and any modifications thereto. Any use, reproduction, disclosure or +# distribution of this software and related documentation without an express +# license agreement from NVIDIA CORPORATION is strictly prohibited. + +"""Miscellaneous utility classes and functions.""" + +import ctypes +import fnmatch +import importlib +import inspect +import numpy as np +import os +import shutil +import sys +import types +import io +import pickle +import re +import requests +import html +import hashlib +import glob +import tempfile +import urllib +import urllib.request +import uuid + +from distutils.util import strtobool +from typing import Any, List, Tuple, Union + + +# Util classes +# ------------------------------------------------------------------------------------------ + + +class EasyDict(dict): + """Convenience class that behaves like a dict but allows access with the attribute syntax.""" + + def __getattr__(self, name: str) -> Any: + try: + return self[name] + except KeyError: + raise AttributeError(name) + + def __setattr__(self, name: str, value: Any) -> None: + self[name] = value + + def __delattr__(self, name: str) -> None: + del self[name] + + +class Logger(object): + """Redirect stderr to stdout, optionally print stdout to a file, and optionally force flushing on both stdout and the file.""" + + def __init__(self, file_name: str = None, file_mode: str = "w", should_flush: bool = True): + self.file = None + + if file_name is not None: + self.file = open(file_name, file_mode) + + self.should_flush = should_flush + self.stdout = sys.stdout + self.stderr = sys.stderr + + sys.stdout = self + sys.stderr = self + + def __enter__(self) -> "Logger": + return self + + def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None: + self.close() + + def write(self, text: Union[str, bytes]) -> None: + """Write text to stdout (and a file) and optionally flush.""" + if isinstance(text, bytes): + text = text.decode() + if len(text) == 0: # workaround for a bug in VSCode debugger: sys.stdout.write(''); sys.stdout.flush() => crash + return + + if self.file is not None: + self.file.write(text) + + self.stdout.write(text) + + if self.should_flush: + self.flush() + + def flush(self) -> None: + """Flush written text to both stdout and a file, if open.""" + if self.file is not None: + self.file.flush() + + self.stdout.flush() + + def close(self) -> None: + """Flush, close possible files, and remove stdout/stderr mirroring.""" + self.flush() + + # if using multiple loggers, prevent closing in wrong order + if sys.stdout is self: + sys.stdout = self.stdout + if sys.stderr is self: + sys.stderr = self.stderr + + if self.file is not None: + self.file.close() + self.file = None + + +# Cache directories +# ------------------------------------------------------------------------------------------ + +_dnnlib_cache_dir = None + +def set_cache_dir(path: str) -> None: + global _dnnlib_cache_dir + _dnnlib_cache_dir = path + +def make_cache_dir_path(*paths: str) -> str: + if _dnnlib_cache_dir is not None: + return os.path.join(_dnnlib_cache_dir, *paths) + if 'DNNLIB_CACHE_DIR' in os.environ: + return os.path.join(os.environ['DNNLIB_CACHE_DIR'], *paths) + if 'HOME' in os.environ: + return os.path.join(os.environ['HOME'], '.cache', 'dnnlib', *paths) + if 'USERPROFILE' in os.environ: + return os.path.join(os.environ['USERPROFILE'], '.cache', 'dnnlib', *paths) + return os.path.join(tempfile.gettempdir(), '.cache', 'dnnlib', *paths) + +# Small util functions +# ------------------------------------------------------------------------------------------ + + +def format_time(seconds: Union[int, float]) -> str: + """Convert the seconds to human readable string with days, hours, minutes and seconds.""" + s = int(np.rint(seconds)) + + if s < 60: + return "{0}s".format(s) + elif s < 60 * 60: + return "{0}m {1:02}s".format(s // 60, s % 60) + elif s < 24 * 60 * 60: + return "{0}h {1:02}m {2:02}s".format(s // (60 * 60), (s // 60) % 60, s % 60) + else: + return "{0}d {1:02}h {2:02}m".format(s // (24 * 60 * 60), (s // (60 * 60)) % 24, (s // 60) % 60) + + +def ask_yes_no(question: str) -> bool: + """Ask the user the question until the user inputs a valid answer.""" + while True: + try: + print("{0} [y/n]".format(question)) + return strtobool(input().lower()) + except ValueError: + pass + + +def tuple_product(t: Tuple) -> Any: + """Calculate the product of the tuple elements.""" + result = 1 + + for v in t: + result *= v + + return result + + +_str_to_ctype = { + "uint8": ctypes.c_ubyte, + "uint16": ctypes.c_uint16, + "uint32": ctypes.c_uint32, + "uint64": ctypes.c_uint64, + "int8": ctypes.c_byte, + "int16": ctypes.c_int16, + "int32": ctypes.c_int32, + "int64": ctypes.c_int64, + "float32": ctypes.c_float, + "float64": ctypes.c_double +} + + +def get_dtype_and_ctype(type_obj: Any) -> Tuple[np.dtype, Any]: + """Given a type name string (or an object having a __name__ attribute), return matching Numpy and ctypes types that have the same size in bytes.""" + type_str = None + + if isinstance(type_obj, str): + type_str = type_obj + elif hasattr(type_obj, "__name__"): + type_str = type_obj.__name__ + elif hasattr(type_obj, "name"): + type_str = type_obj.name + else: + raise RuntimeError("Cannot infer type name from input") + + assert type_str in _str_to_ctype.keys() + + my_dtype = np.dtype(type_str) + my_ctype = _str_to_ctype[type_str] + + assert my_dtype.itemsize == ctypes.sizeof(my_ctype) + + return my_dtype, my_ctype + + +def is_pickleable(obj: Any) -> bool: + try: + with io.BytesIO() as stream: + pickle.dump(obj, stream) + return True + except: + return False + + +# Functionality to import modules/objects by name, and call functions by name +# ------------------------------------------------------------------------------------------ + +def get_module_from_obj_name(obj_name: str) -> Tuple[types.ModuleType, str]: + """Searches for the underlying module behind the name to some python object. + Returns the module and the object name (original name with module part removed).""" + + # allow convenience shorthands, substitute them by full names + obj_name = re.sub("^np.", "numpy.", obj_name) + obj_name = re.sub("^tf.", "tensorflow.", obj_name) + + # list alternatives for (module_name, local_obj_name) + parts = obj_name.split(".") + name_pairs = [(".".join(parts[:i]), ".".join(parts[i:])) for i in range(len(parts), 0, -1)] + + # try each alternative in turn + for module_name, local_obj_name in name_pairs: + try: + module = importlib.import_module(module_name) # may raise ImportError + get_obj_from_module(module, local_obj_name) # may raise AttributeError + return module, local_obj_name + except: + pass + + # maybe some of the modules themselves contain errors? + for module_name, _local_obj_name in name_pairs: + try: + importlib.import_module(module_name) # may raise ImportError + except ImportError: + if not str(sys.exc_info()[1]).startswith("No module named '" + module_name + "'"): + raise + + # maybe the requested attribute is missing? + for module_name, local_obj_name in name_pairs: + try: + module = importlib.import_module(module_name) # may raise ImportError + get_obj_from_module(module, local_obj_name) # may raise AttributeError + except ImportError: + pass + + # we are out of luck, but we have no idea why + raise ImportError(obj_name) + + +def get_obj_from_module(module: types.ModuleType, obj_name: str) -> Any: + """Traverses the object name and returns the last (rightmost) python object.""" + if obj_name == '': + return module + obj = module + for part in obj_name.split("."): + obj = getattr(obj, part) + return obj + + +def get_obj_by_name(name: str) -> Any: + """Finds the python object with the given name.""" + module, obj_name = get_module_from_obj_name(name) + return get_obj_from_module(module, obj_name) + + +def call_func_by_name(*args, func_name: str = None, **kwargs) -> Any: + """Finds the python object with the given name and calls it as a function.""" + assert func_name is not None + func_obj = get_obj_by_name(func_name) + assert callable(func_obj) + return func_obj(*args, **kwargs) + + +def construct_class_by_name(*args, class_name: str = None, **kwargs) -> Any: + """Finds the python class with the given name and constructs it with the given arguments.""" + return call_func_by_name(*args, func_name=class_name, **kwargs) + + +def get_module_dir_by_obj_name(obj_name: str) -> str: + """Get the directory path of the module containing the given object name.""" + module, _ = get_module_from_obj_name(obj_name) + return os.path.dirname(inspect.getfile(module)) + + +def is_top_level_function(obj: Any) -> bool: + """Determine whether the given object is a top-level function, i.e., defined at module scope using 'def'.""" + return callable(obj) and obj.__name__ in sys.modules[obj.__module__].__dict__ + + +def get_top_level_function_name(obj: Any) -> str: + """Return the fully-qualified name of a top-level function.""" + assert is_top_level_function(obj) + module = obj.__module__ + if module == '__main__': + module = os.path.splitext(os.path.basename(sys.modules[module].__file__))[0] + return module + "." + obj.__name__ + + +# File system helpers +# ------------------------------------------------------------------------------------------ + +def list_dir_recursively_with_ignore(dir_path: str, ignores: List[str] = None, add_base_to_relative: bool = False) -> List[Tuple[str, str]]: + """List all files recursively in a given directory while ignoring given file and directory names. + Returns list of tuples containing both absolute and relative paths.""" + assert os.path.isdir(dir_path) + base_name = os.path.basename(os.path.normpath(dir_path)) + + if ignores is None: + ignores = [] + + result = [] + + for root, dirs, files in os.walk(dir_path, topdown=True): + for ignore_ in ignores: + dirs_to_remove = [d for d in dirs if fnmatch.fnmatch(d, ignore_)] + + # dirs need to be edited in-place + for d in dirs_to_remove: + dirs.remove(d) + + files = [f for f in files if not fnmatch.fnmatch(f, ignore_)] + + absolute_paths = [os.path.join(root, f) for f in files] + relative_paths = [os.path.relpath(p, dir_path) for p in absolute_paths] + + if add_base_to_relative: + relative_paths = [os.path.join(base_name, p) for p in relative_paths] + + assert len(absolute_paths) == len(relative_paths) + result += zip(absolute_paths, relative_paths) + + return result + + +def copy_files_and_create_dirs(files: List[Tuple[str, str]]) -> None: + """Takes in a list of tuples of (src, dst) paths and copies files. + Will create all necessary directories.""" + for file in files: + target_dir_name = os.path.dirname(file[1]) + + # will create all intermediate-level directories + if not os.path.exists(target_dir_name): + os.makedirs(target_dir_name) + + shutil.copyfile(file[0], file[1]) + + +# URL helpers +# ------------------------------------------------------------------------------------------ + +def is_url(obj: Any, allow_file_urls: bool = False) -> bool: + """Determine whether the given object is a valid URL string.""" + if not isinstance(obj, str) or not "://" in obj: + return False + if allow_file_urls and obj.startswith('file://'): + return True + try: + res = requests.compat.urlparse(obj) + if not res.scheme or not res.netloc or not "." in res.netloc: + return False + res = requests.compat.urlparse(requests.compat.urljoin(obj, "/")) + if not res.scheme or not res.netloc or not "." in res.netloc: + return False + except: + return False + return True + + +def open_url(url: str, cache_dir: str = None, num_attempts: int = 10, verbose: bool = True, return_filename: bool = False, cache: bool = True) -> Any: + """Download the given URL and return a binary-mode file object to access the data.""" + assert num_attempts >= 1 + assert not (return_filename and (not cache)) + + # Doesn't look like an URL scheme so interpret it as a local filename. + if not re.match('^[a-z]+://', url): + return url if return_filename else open(url, "rb") + + # Handle file URLs. This code handles unusual file:// patterns that + # arise on Windows: + # + # file:///c:/foo.txt + # + # which would translate to a local '/c:/foo.txt' filename that's + # invalid. Drop the forward slash for such pathnames. + # + # If you touch this code path, you should test it on both Linux and + # Windows. + # + # Some internet resources suggest using urllib.request.url2pathname() but + # but that converts forward slashes to backslashes and this causes + # its own set of problems. + if url.startswith('file://'): + filename = urllib.parse.urlparse(url).path + if re.match(r'^/[a-zA-Z]:', filename): + filename = filename[1:] + return filename if return_filename else open(filename, "rb") + + assert is_url(url) + + # Lookup from cache. + if cache_dir is None: + cache_dir = make_cache_dir_path('downloads') + + url_md5 = hashlib.md5(url.encode("utf-8")).hexdigest() + if cache: + cache_files = glob.glob(os.path.join(cache_dir, url_md5 + "_*")) + if len(cache_files) == 1: + filename = cache_files[0] + return filename if return_filename else open(filename, "rb") + + # Download. + url_name = None + url_data = None + with requests.Session() as session: + if verbose: + print("Downloading %s ..." % url, end="", flush=True) + for attempts_left in reversed(range(num_attempts)): + try: + with session.get(url) as res: + res.raise_for_status() + if len(res.content) == 0: + raise IOError("No data received") + + if len(res.content) < 8192: + content_str = res.content.decode("utf-8") + if "download_warning" in res.headers.get("Set-Cookie", ""): + links = [html.unescape(link) for link in content_str.split('"') if "export=download" in link] + if len(links) == 1: + url = requests.compat.urljoin(url, links[0]) + raise IOError("Google Drive virus checker nag") + if "Google Drive - Quota exceeded" in content_str: + raise IOError("Google Drive download quota exceeded -- please try again later") + + match = re.search(r'filename="([^"]*)"', res.headers.get("Content-Disposition", "")) + url_name = match[1] if match else url + url_data = res.content + if verbose: + print(" done") + break + except KeyboardInterrupt: + raise + except: + if not attempts_left: + if verbose: + print(" failed") + raise + if verbose: + print(".", end="", flush=True) + + # Save to cache. + if cache: + safe_name = re.sub(r"[^0-9a-zA-Z-._]", "_", url_name) + cache_file = os.path.join(cache_dir, url_md5 + "_" + safe_name) + temp_file = os.path.join(cache_dir, "tmp_" + uuid.uuid4().hex + "_" + url_md5 + "_" + safe_name) + os.makedirs(cache_dir, exist_ok=True) + with open(temp_file, "wb") as f: + f.write(url_data) + os.replace(temp_file, cache_file) # atomic + if return_filename: + return cache_file + + # Return data as file object. + assert not return_filename + return io.BytesIO(url_data) diff --git a/models/stylegan2/model.py b/models/stylegan2/model.py new file mode 100644 index 0000000000000000000000000000000000000000..2fdd5ac7b5cd49826f173d115b8c159b040e8bc9 --- /dev/null +++ b/models/stylegan2/model.py @@ -0,0 +1,726 @@ +import math +import random +import functools +import operator + +import torch +from torch import nn +from torch.nn import functional as F +from torch.autograd import Function + +from models.stylegan2.op import FusedLeakyReLU, fused_leaky_relu, upfirdn2d +import torchvision +toPIL = torchvision.transforms.ToPILImage() +import numpy as np + +class PixelNorm(nn.Module): + def __init__(self): + super().__init__() + + def forward(self, input): + return input * torch.rsqrt(torch.mean(input ** 2, dim=1, keepdim=True) + 1e-8) + + +def make_kernel(k): + k = torch.tensor(k, dtype=torch.float32) + + if k.ndim == 1: + k = k[None, :] * k[:, None] + + k /= k.sum() + + return k + + +class Upsample(nn.Module): + def __init__(self, kernel, factor=2): + super().__init__() + + self.factor = factor + kernel = make_kernel(kernel) * (factor ** 2) + self.register_buffer('kernel', kernel) + + p = kernel.shape[0] - factor + + pad0 = (p + 1) // 2 + factor - 1 + pad1 = p // 2 + + self.pad = (pad0, pad1) + + def forward(self, input): + out = upfirdn2d(input, self.kernel, up=self.factor, down=1, pad=self.pad) + + return out + + +class Downsample(nn.Module): + def __init__(self, kernel, factor=2): + super().__init__() + + self.factor = factor + kernel = make_kernel(kernel) + self.register_buffer('kernel', kernel) + + p = kernel.shape[0] - factor + + pad0 = (p + 1) // 2 + pad1 = p // 2 + + self.pad = (pad0, pad1) + + def forward(self, input): + out = upfirdn2d(input, self.kernel, up=1, down=self.factor, pad=self.pad) + + return out + + +class Blur(nn.Module): + def __init__(self, kernel, pad, upsample_factor=1): + super().__init__() + + kernel = make_kernel(kernel) + + if upsample_factor > 1: + kernel = kernel * (upsample_factor ** 2) + + self.register_buffer('kernel', kernel) + + self.pad = pad + + def forward(self, input): + out = upfirdn2d(input, self.kernel, pad=self.pad) + + return out + + +class EqualConv2d(nn.Module): + def __init__( + self, in_channel, out_channel, kernel_size, stride=1, padding=0, bias=True + ): + super().__init__() + + self.weight = nn.Parameter( + torch.randn(out_channel, in_channel, kernel_size, kernel_size) + ) + self.scale = 1 / math.sqrt(in_channel * kernel_size ** 2) + + self.stride = stride + self.padding = padding + + if bias: + self.bias = nn.Parameter(torch.zeros(out_channel)) + + else: + self.bias = None + + def forward(self, input): + out = F.conv2d( + input, + self.weight * self.scale, + bias=self.bias, + stride=self.stride, + padding=self.padding, + ) + + return out + + def __repr__(self): + return ( + f'{self.__class__.__name__}({self.weight.shape[1]}, {self.weight.shape[0]},' + f' {self.weight.shape[2]}, stride={self.stride}, padding={self.padding})' + ) + + +class EqualLinear(nn.Module): + def __init__( + self, in_dim, out_dim, bias=True, bias_init=0, lr_mul=1, activation=None + ): + super().__init__() + + self.weight = nn.Parameter(torch.randn(out_dim, in_dim).div_(lr_mul)) + + if bias: + self.bias = nn.Parameter(torch.zeros(out_dim).fill_(bias_init)) + + else: + self.bias = None + + self.activation = activation + + self.scale = (1 / math.sqrt(in_dim)) * lr_mul + self.lr_mul = lr_mul + + def forward(self, input): + if self.activation: + out = F.linear(input, self.weight * self.scale) + out = fused_leaky_relu(out, self.bias * self.lr_mul) + + else: + out = F.linear( + input, self.weight * self.scale, bias=self.bias * self.lr_mul + ) + + return out + + def __repr__(self): + return ( + f'{self.__class__.__name__}({self.weight.shape[1]}, {self.weight.shape[0]})' + ) + + +class ScaledLeakyReLU(nn.Module): + def __init__(self, negative_slope=0.2): + super().__init__() + + self.negative_slope = negative_slope + + def forward(self, input): + out = F.leaky_relu(input, negative_slope=self.negative_slope) + + return out * math.sqrt(2) + + +class ModulatedConv2d(nn.Module): + def __init__( + self, + in_channel, + out_channel, + kernel_size, + style_dim, + demodulate=True, + upsample=False, + downsample=False, + blur_kernel=[1, 3, 3, 1], + ): + super().__init__() + + self.eps = 1e-8 + self.kernel_size = kernel_size + self.in_channel = in_channel + self.out_channel = out_channel + self.upsample = upsample + self.downsample = downsample + + if upsample: + factor = 2 + p = (len(blur_kernel) - factor) - (kernel_size - 1) + pad0 = (p + 1) // 2 + factor - 1 + pad1 = p // 2 + 1 + + self.blur = Blur(blur_kernel, pad=(pad0, pad1), upsample_factor=factor) + + if downsample: + factor = 2 + p = (len(blur_kernel) - factor) + (kernel_size - 1) + pad0 = (p + 1) // 2 + pad1 = p // 2 + + self.blur = Blur(blur_kernel, pad=(pad0, pad1)) + + fan_in = in_channel * kernel_size ** 2 + self.scale = 1 / math.sqrt(fan_in) + self.padding = kernel_size // 2 + + self.weight = nn.Parameter( + torch.randn(1, out_channel, in_channel, kernel_size, kernel_size) + ) + + self.modulation = EqualLinear(style_dim, in_channel, bias_init=1) + + self.demodulate = demodulate + + def __repr__(self): + return ( + f'{self.__class__.__name__}({self.in_channel}, {self.out_channel}, {self.kernel_size}, ' + f'upsample={self.upsample}, downsample={self.downsample})' + ) + + def forward(self, input, style): + batch, in_channel, height, width = input.shape + + style = self.modulation(style).view(batch, 1, in_channel, 1, 1) + weight = self.scale * self.weight * style + + if self.demodulate: + demod = torch.rsqrt(weight.pow(2).sum([2, 3, 4]) + 1e-8) + weight = weight * demod.view(batch, self.out_channel, 1, 1, 1) + + weight = weight.view( + batch * self.out_channel, in_channel, self.kernel_size, self.kernel_size + ) + + if self.upsample: + input = input.view(1, batch * in_channel, height, width) + weight = weight.view( + batch, self.out_channel, in_channel, self.kernel_size, self.kernel_size + ) + weight = weight.transpose(1, 2).reshape( + batch * in_channel, self.out_channel, self.kernel_size, self.kernel_size + ) + out = F.conv_transpose2d(input, weight, padding=0, stride=2, groups=batch) + _, _, height, width = out.shape + out = out.view(batch, self.out_channel, height, width) + out = self.blur(out) + + elif self.downsample: + input = self.blur(input) + _, _, height, width = input.shape + input = input.view(1, batch * in_channel, height, width) + out = F.conv2d(input, weight, padding=0, stride=2, groups=batch) + _, _, height, width = out.shape + out = out.view(batch, self.out_channel, height, width) + + else: + input = input.view(1, batch * in_channel, height, width) + out = F.conv2d(input, weight, padding=self.padding, groups=batch) + _, _, height, width = out.shape + out = out.view(batch, self.out_channel, height, width) + + return out + + +class NoiseInjection(nn.Module): + def __init__(self): + super().__init__() + + self.weight = nn.Parameter(torch.zeros(1)) + + def forward(self, image, noise=None): + if noise is None: + batch, _, height, width = image.shape + noise = image.new_empty(batch, 1, height, width).normal_() + + return image + self.weight * noise + + +class ConstantInput(nn.Module): + def __init__(self, channel, size=4): + super().__init__() + + self.input = nn.Parameter(torch.randn(1, channel, size, size)) + + def forward(self, input): + batch = input.shape[0] + out = self.input.repeat(batch, 1, 1, 1) + + return out + + +class StyledConv(nn.Module): + def __init__( + self, + in_channel, + out_channel, + kernel_size, + style_dim, + upsample=False, + blur_kernel=[1, 3, 3, 1], + demodulate=True, + ): + super().__init__() + + self.conv = ModulatedConv2d( + in_channel, + out_channel, + kernel_size, + style_dim, + upsample=upsample, + blur_kernel=blur_kernel, + demodulate=demodulate, + ) + + self.noise = NoiseInjection() + # self.bias = nn.Parameter(torch.zeros(1, out_channel, 1, 1)) + # self.activate = ScaledLeakyReLU(0.2) + self.activate = FusedLeakyReLU(out_channel) + + def forward(self, input, style, noise=None): + out = self.conv(input, style) + out = self.noise(out, noise=noise) + # out = out + self.bias + out = self.activate(out) + + return out + + +class ToRGB(nn.Module): + def __init__(self, in_channel, style_dim, upsample=True, blur_kernel=[1, 3, 3, 1]): + super().__init__() + + if upsample: + self.upsample = Upsample(blur_kernel) + + self.conv = ModulatedConv2d(in_channel, 3, 1, style_dim, demodulate=False) + self.bias = nn.Parameter(torch.zeros(1, 3, 1, 1)) + + def forward(self, input, style, skip=None): + out = self.conv(input, style) + out = out + self.bias + + if skip is not None: + skip = self.upsample(skip) + + out = out + skip + + return out + + +class Generator(nn.Module): + def __init__( + self, + size, + style_dim, + n_mlp, + channel_multiplier=2, + blur_kernel=[1, 3, 3, 1], + lr_mlp=0.01, + ): + super().__init__() + + self.size = size + + self.style_dim = style_dim + + layers = [PixelNorm()] + + for i in range(n_mlp): + layers.append( + EqualLinear( + style_dim, style_dim, lr_mul=lr_mlp, activation='fused_lrelu' + ) + ) + + self.style = nn.Sequential(*layers) + + self.channels = { + 4: 512, + 8: 512, + 16: 512, + 32: 512, + 64: 256 * channel_multiplier, + 128: 128 * channel_multiplier, + 256: 64 * channel_multiplier, + 512: 32 * channel_multiplier, + 1024: 16 * channel_multiplier, + } + + + self.input = ConstantInput(self.channels[4]) + self.conv1 = StyledConv( + self.channels[4], self.channels[4], 3, style_dim, blur_kernel=blur_kernel + ) + self.to_rgb1 = ToRGB(self.channels[4], style_dim, upsample=False) + + self.log_size = int(math.log(size, 2)) + self.num_layers = (self.log_size - 2) * 2 + 1 + + self.convs = nn.ModuleList() + self.upsamples = nn.ModuleList() + self.to_rgbs = nn.ModuleList() + self.noises = nn.Module() + + in_channel = self.channels[4] + + for layer_idx in range(self.num_layers): + res = (layer_idx + 5) // 2 + shape = [1, 1, 2 ** res, 2 ** res] + self.noises.register_buffer(f'noise_{layer_idx}', torch.randn(*shape)) + + for i in range(3, self.log_size + 1): + out_channel = self.channels[2 ** i] + + self.convs.append( + StyledConv( + in_channel, + out_channel, + 3, + style_dim, + upsample=True, + blur_kernel=blur_kernel, + ) + ) + + self.convs.append( + StyledConv( + out_channel, out_channel, 3, style_dim, blur_kernel=blur_kernel + ) + ) + + self.to_rgbs.append(ToRGB(out_channel, style_dim)) + + in_channel = out_channel + + self.n_latent = self.log_size * 2 - 2 + + def make_noise(self): + device = self.input.input.device + + noises = [torch.randn(1, 1, 2 ** 2, 2 ** 2, device=device)] + + for i in range(3, self.log_size + 1): + for _ in range(2): + noises.append(torch.randn(1, 1, 2 ** i, 2 ** i, device=device)) + + return noises + + def mean_latent(self, n_latent): + latent_in = torch.randn( + n_latent, self.style_dim, device=self.input.input.device + ) + latent = self.style(latent_in).mean(0, keepdim=True) + + return latent + + def get_latent(self, input): + return self.style(input) + + def forward( + self, + styles, + return_latents=False, + inject_index=None, + truncation=1, + truncation_latent=None, + input_is_latent=False, + noise=None, + randomize_noise=True, + layer_in=None, + skip=None, + start_layer=0, + end_layer=8, + return_rgb=False, + + ): + if not input_is_latent: + styles = [self.style(s) for s in styles] + + if noise is None: + if randomize_noise: + noise = [None] * self.num_layers + else: + noise = [ + getattr(self.noises, f'noise_{i}') for i in range(self.num_layers) + ] + + if truncation < 1: + style_t = [] + + for style in styles: + style_t.append( + truncation_latent + truncation * (style - truncation_latent) + ) + + styles = style_t + + if len(styles) < 2: + inject_index = self.n_latent + + if styles[0].ndim < 3: + latent = styles[0].unsqueeze(1).repeat(1, inject_index, 1) + + else: + latent = styles[0] + + else: + if inject_index is None: + inject_index = random.randint(1, self.n_latent - 1) + + latent = styles[0].unsqueeze(1).repeat(1, inject_index, 1) + latent2 = styles[1].unsqueeze(1).repeat(1, self.n_latent - inject_index, 1) + + latent = torch.cat([latent, latent2], 1) + out = self.input(latent) + + if start_layer == 0: + out = self.conv1(out, latent[:, 0], noise=noise[0]) # 0th layer + skip = self.to_rgb1(out, latent[:, 1]) + if end_layer == 0: + return out, skip + i = 1 + current_layer = 1 + for conv1, conv2, noise1, noise2, to_rgb in zip( + self.convs[::2], self.convs[1::2], noise[1::2], noise[2::2], self.to_rgbs + ): + if current_layer < start_layer: + pass + elif current_layer == start_layer: + out = conv1(layer_in, latent[:, i], noise=noise1) + out = conv2(out, latent[:, i + 1], noise=noise2) + skip = to_rgb(out, latent[:, i + 2], skip) + elif current_layer > end_layer: + return out, skip + else: + out = conv1(out, latent[:, i], noise=noise1) + out = conv2(out, latent[:, i + 1], noise=noise2) + skip = to_rgb(out, latent[:, i + 2], skip) + current_layer += 1 + i += 2 + + image = skip + + if return_latents: + return image, latent + + else: + return image, None + + def generate_im_from_w_space(self, code, noises=None): + latent = torch.from_numpy(code).cuda() + I_G, _ = self([latent], input_is_latent=True, return_latents=False, noise=noises, start_layer=0, + end_layer=8) + I_G_0_1 = (I_G + 1) / 2 + im = np.array(toPIL(I_G_0_1[0].cpu().detach().clamp(0, 1))) + return im + + def generate_initial_intermediate(self, code, noises=None): + latent = torch.from_numpy(code).cuda() + intermediate, _ = self([latent], input_is_latent=True, return_latents=False, noise=noises, + start_layer=0, end_layer=3) + return intermediate + + + def update_on_FS(self, code, initial_intermediate, initial_F, initial_S, noises=None): + latent = torch.from_numpy(code).cuda() + intermediate, _ = self([latent], input_is_latent=True, return_latents=False, noise=noises, + start_layer=0, end_layer=3) + + difference = initial_F - initial_intermediate + new_intermediate = intermediate + difference + + I_G, _ = self([initial_S], input_is_latent=True, return_latents=False, noise=noises, start_layer=4, + end_layer=8, layer_in=new_intermediate) + I_G_0_1 = (I_G + 1) / 2 + im = np.array(toPIL(I_G_0_1[0].cpu().detach().clamp(0, 1))) + return im + + +class ConvLayer(nn.Sequential): + def __init__( + self, + in_channel, + out_channel, + kernel_size, + downsample=False, + blur_kernel=[1, 3, 3, 1], + bias=True, + activate=True, + ): + layers = [] + + if downsample: + factor = 2 + p = (len(blur_kernel) - factor) + (kernel_size - 1) + pad0 = (p + 1) // 2 + pad1 = p // 2 + + layers.append(Blur(blur_kernel, pad=(pad0, pad1))) + + stride = 2 + self.padding = 0 + + else: + stride = 1 + self.padding = kernel_size // 2 + + layers.append( + EqualConv2d( + in_channel, + out_channel, + kernel_size, + padding=self.padding, + stride=stride, + bias=bias and not activate, + ) + ) + + if activate: + if bias: + layers.append(FusedLeakyReLU(out_channel)) + + else: + layers.append(ScaledLeakyReLU(0.2)) + + super().__init__(*layers) + + +class ResBlock(nn.Module): + def __init__(self, in_channel, out_channel, blur_kernel=[1, 3, 3, 1]): + super().__init__() + + self.conv1 = ConvLayer(in_channel, in_channel, 3) + self.conv2 = ConvLayer(in_channel, out_channel, 3, downsample=True) + + self.skip = ConvLayer( + in_channel, out_channel, 1, downsample=True, activate=False, bias=False + ) + + def forward(self, input): + out = self.conv1(input) + out = self.conv2(out) + + skip = self.skip(input) + out = (out + skip) / math.sqrt(2) + + return out + + +class Discriminator(nn.Module): + def __init__(self, size, channel_multiplier=2, blur_kernel=[1, 3, 3, 1]): + super().__init__() + + channels = { + 4: 512, + 8: 512, + 16: 512, + 32: 512, + 64: 256 * channel_multiplier, + 128: 128 * channel_multiplier, + 256: 64 * channel_multiplier, + 512: 32 * channel_multiplier, + 1024: 16 * channel_multiplier, + } + + convs = [ConvLayer(3, channels[size], 1)] + + log_size = int(math.log(size, 2)) + + in_channel = channels[size] + + for i in range(log_size, 2, -1): + out_channel = channels[2 ** (i - 1)] + + convs.append(ResBlock(in_channel, out_channel, blur_kernel)) + + in_channel = out_channel + + self.convs = nn.Sequential(*convs) + + self.stddev_group = 4 + self.stddev_feat = 1 + + self.final_conv = ConvLayer(in_channel + 1, channels[4], 3) + self.final_linear = nn.Sequential( + EqualLinear(channels[4] * 4 * 4, channels[4], activation='fused_lrelu'), + EqualLinear(channels[4], 1), + ) + + def forward(self, input): + out = self.convs(input) + + batch, channel, height, width = out.shape + group = min(batch, self.stddev_group) + stddev = out.view( + group, -1, self.stddev_feat, channel // self.stddev_feat, height, width + ) + stddev = torch.sqrt(stddev.var(0, unbiased=False) + 1e-8) + stddev = stddev.mean([2, 3, 4], keepdims=True).squeeze(2) + stddev = stddev.repeat(group, 1, height, width) + out = torch.cat([out, stddev], 1) + + out = self.final_conv(out) + + out = out.view(batch, -1) + out = self.final_linear(out) + + return out + diff --git a/models/stylegan2/op/__init__.py b/models/stylegan2/op/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d0918d92285955855be89f00096b888ee5597ce3 --- /dev/null +++ b/models/stylegan2/op/__init__.py @@ -0,0 +1,2 @@ +from .fused_act import FusedLeakyReLU, fused_leaky_relu +from .upfirdn2d import upfirdn2d diff --git a/models/stylegan2/op/fused_act.py b/models/stylegan2/op/fused_act.py new file mode 100644 index 0000000000000000000000000000000000000000..80c9519ccfd212c799380bc0d7bce6232fd2753b --- /dev/null +++ b/models/stylegan2/op/fused_act.py @@ -0,0 +1,96 @@ +import os + +import torch +from torch import nn +from torch.nn import functional as F +from torch.autograd import Function +from torch.utils.cpp_extension import load + +module_path = os.path.dirname(__file__) +fused = load( + "fused", + sources=[ + os.path.join(module_path, "fused_bias_act.cpp"), + os.path.join(module_path, "fused_bias_act_kernel.cu"), + ], +) + + +class FusedLeakyReLUFunctionBackward(Function): + @staticmethod + def forward(ctx, grad_output, out, negative_slope, scale): + ctx.save_for_backward(out) + ctx.negative_slope = negative_slope + ctx.scale = scale + + empty = grad_output.new_empty(0) + + grad_input = fused.fused_bias_act( + grad_output, empty, out, 3, 1, negative_slope, scale + ) + + dim = [0] + + if grad_input.ndim > 2: + dim += list(range(2, grad_input.ndim)) + + grad_bias = grad_input.sum(dim).detach() + + return grad_input, grad_bias + + @staticmethod + def backward(ctx, gradgrad_input, gradgrad_bias): + (out,) = ctx.saved_tensors + gradgrad_out = fused.fused_bias_act( + gradgrad_input, gradgrad_bias, out, 3, 1, ctx.negative_slope, ctx.scale + ) + + return gradgrad_out, None, None, None + + +class FusedLeakyReLUFunction(Function): + @staticmethod + def forward(ctx, input, bias, negative_slope, scale): + empty = input.new_empty(0) + out = fused.fused_bias_act(input, bias, empty, 3, 0, negative_slope, scale) + ctx.save_for_backward(out) + ctx.negative_slope = negative_slope + ctx.scale = scale + + return out + + @staticmethod + def backward(ctx, grad_output): + (out,) = ctx.saved_tensors + + grad_input, grad_bias = FusedLeakyReLUFunctionBackward.apply( + grad_output, out, ctx.negative_slope, ctx.scale + ) + + return grad_input, grad_bias, None, None + + +class FusedLeakyReLU(nn.Module): + def __init__(self, channel, negative_slope=0.2, scale=2 ** 0.5): + super().__init__() + + self.bias = nn.Parameter(torch.zeros(channel)) + self.negative_slope = negative_slope + self.scale = scale + + def forward(self, input): + return fused_leaky_relu(input, self.bias, self.negative_slope, self.scale) + + +def fused_leaky_relu(input, bias, negative_slope=0.2, scale=2 ** 0.5): + if input.device.type == "cpu": + rest_dim = [1] * (input.ndim - bias.ndim - 1) + return ( + F.leaky_relu( + input + bias.view(1, bias.shape[0], *rest_dim), negative_slope=0.2 + ) + * scale + ) + + else: + return FusedLeakyReLUFunction.apply(input, bias, negative_slope, scale) diff --git a/models/stylegan2/op/fused_bias_act.cpp b/models/stylegan2/op/fused_bias_act.cpp new file mode 100644 index 0000000000000000000000000000000000000000..f00f8255e053cc5354c2f213c18c0ef332a70344 --- /dev/null +++ b/models/stylegan2/op/fused_bias_act.cpp @@ -0,0 +1,21 @@ +#include + + +torch::Tensor fused_bias_act_op(const torch::Tensor& input, const torch::Tensor& bias, const torch::Tensor& refer, + int act, int grad, float alpha, float scale); + +#define CHECK_CUDA(x) TORCH_CHECK(x.type().is_cuda(), #x " must be a CUDA tensor") +#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous") +#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x) + +torch::Tensor fused_bias_act(const torch::Tensor& input, const torch::Tensor& bias, const torch::Tensor& refer, + int act, int grad, float alpha, float scale) { + CHECK_CUDA(input); + CHECK_CUDA(bias); + + return fused_bias_act_op(input, bias, refer, act, grad, alpha, scale); +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("fused_bias_act", &fused_bias_act, "fused bias act (CUDA)"); +} diff --git a/models/stylegan2/op/fused_bias_act_kernel.cu b/models/stylegan2/op/fused_bias_act_kernel.cu new file mode 100644 index 0000000000000000000000000000000000000000..57e109bdb712eb439ca0b36b52790c2e1379c5d5 --- /dev/null +++ b/models/stylegan2/op/fused_bias_act_kernel.cu @@ -0,0 +1,99 @@ +// Copyright (c) 2019, NVIDIA Corporation. All rights reserved. +// +// This work is made available under the Nvidia Source Code License-NC. +// To view a copy of this license, visit +// https://nvlabs.github.io/stylegan2/license.html + +#include + +#include +#include +#include +#include + +#include +#include + + +template +static __global__ void fused_bias_act_kernel(scalar_t* out, const scalar_t* p_x, const scalar_t* p_b, const scalar_t* p_ref, + int act, int grad, scalar_t alpha, scalar_t scale, int loop_x, int size_x, int step_b, int size_b, int use_bias, int use_ref) { + int xi = blockIdx.x * loop_x * blockDim.x + threadIdx.x; + + scalar_t zero = 0.0; + + for (int loop_idx = 0; loop_idx < loop_x && xi < size_x; loop_idx++, xi += blockDim.x) { + scalar_t x = p_x[xi]; + + if (use_bias) { + x += p_b[(xi / step_b) % size_b]; + } + + scalar_t ref = use_ref ? p_ref[xi] : zero; + + scalar_t y; + + switch (act * 10 + grad) { + default: + case 10: y = x; break; + case 11: y = x; break; + case 12: y = 0.0; break; + + case 30: y = (x > 0.0) ? x : x * alpha; break; + case 31: y = (ref > 0.0) ? x : x * alpha; break; + case 32: y = 0.0; break; + } + + out[xi] = y * scale; + } +} + + +torch::Tensor fused_bias_act_op(const torch::Tensor& input, const torch::Tensor& bias, const torch::Tensor& refer, + int act, int grad, float alpha, float scale) { + int curDevice = -1; + cudaGetDevice(&curDevice); + cudaStream_t stream = at::cuda::getCurrentCUDAStream(curDevice); + + auto x = input.contiguous(); + auto b = bias.contiguous(); + auto ref = refer.contiguous(); + + int use_bias = b.numel() ? 1 : 0; + int use_ref = ref.numel() ? 1 : 0; + + int size_x = x.numel(); + int size_b = b.numel(); + int step_b = 1; + + for (int i = 1 + 1; i < x.dim(); i++) { + step_b *= x.size(i); + } + + int loop_x = 4; + int block_size = 4 * 32; + int grid_size = (size_x - 1) / (loop_x * block_size) + 1; + + auto y = torch::empty_like(x); + + AT_DISPATCH_FLOATING_TYPES_AND_HALF(x.scalar_type(), "fused_bias_act_kernel", [&] { + fused_bias_act_kernel<<>>( + y.data_ptr(), + x.data_ptr(), + b.data_ptr(), + ref.data_ptr(), + act, + grad, + alpha, + scale, + loop_x, + size_x, + step_b, + size_b, + use_bias, + use_ref + ); + }); + + return y; +} diff --git a/models/stylegan2/op/slurm-935635.out b/models/stylegan2/op/slurm-935635.out new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/models/stylegan2/op/upfirdn2d.cpp b/models/stylegan2/op/upfirdn2d.cpp new file mode 100644 index 0000000000000000000000000000000000000000..d2e633dc896433c205e18bc3e455539192ff968e --- /dev/null +++ b/models/stylegan2/op/upfirdn2d.cpp @@ -0,0 +1,23 @@ +#include + + +torch::Tensor upfirdn2d_op(const torch::Tensor& input, const torch::Tensor& kernel, + int up_x, int up_y, int down_x, int down_y, + int pad_x0, int pad_x1, int pad_y0, int pad_y1); + +#define CHECK_CUDA(x) TORCH_CHECK(x.type().is_cuda(), #x " must be a CUDA tensor") +#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous") +#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x) + +torch::Tensor upfirdn2d(const torch::Tensor& input, const torch::Tensor& kernel, + int up_x, int up_y, int down_x, int down_y, + int pad_x0, int pad_x1, int pad_y0, int pad_y1) { + CHECK_CUDA(input); + CHECK_CUDA(kernel); + + return upfirdn2d_op(input, kernel, up_x, up_y, down_x, down_y, pad_x0, pad_x1, pad_y0, pad_y1); +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("upfirdn2d", &upfirdn2d, "upfirdn2d (CUDA)"); +} \ No newline at end of file diff --git a/models/stylegan2/op/upfirdn2d.py b/models/stylegan2/op/upfirdn2d.py new file mode 100644 index 0000000000000000000000000000000000000000..39d2953949c410880fd8e563496273e9feefc9b9 --- /dev/null +++ b/models/stylegan2/op/upfirdn2d.py @@ -0,0 +1,200 @@ +import os + +import torch +from torch.nn import functional as F +from torch.autograd import Function +from torch.utils.cpp_extension import load + + +module_path = os.path.dirname(__file__) +upfirdn2d_op = load( + "upfirdn2d", + sources=[ + os.path.join(module_path, "upfirdn2d.cpp"), + os.path.join(module_path, "upfirdn2d_kernel.cu"), + ], +) + + +class UpFirDn2dBackward(Function): + @staticmethod + def forward( + ctx, grad_output, kernel, grad_kernel, up, down, pad, g_pad, in_size, out_size + ): + + up_x, up_y = up + down_x, down_y = down + g_pad_x0, g_pad_x1, g_pad_y0, g_pad_y1 = g_pad + + grad_output = grad_output.reshape(-1, out_size[0], out_size[1], 1) + + grad_input = upfirdn2d_op.upfirdn2d( + grad_output, + grad_kernel, + down_x, + down_y, + up_x, + up_y, + g_pad_x0, + g_pad_x1, + g_pad_y0, + g_pad_y1, + ) + grad_input = grad_input.view(in_size[0], in_size[1], in_size[2], in_size[3]) + + ctx.save_for_backward(kernel) + + pad_x0, pad_x1, pad_y0, pad_y1 = pad + + ctx.up_x = up_x + ctx.up_y = up_y + ctx.down_x = down_x + ctx.down_y = down_y + ctx.pad_x0 = pad_x0 + ctx.pad_x1 = pad_x1 + ctx.pad_y0 = pad_y0 + ctx.pad_y1 = pad_y1 + ctx.in_size = in_size + ctx.out_size = out_size + + return grad_input + + @staticmethod + def backward(ctx, gradgrad_input): + (kernel,) = ctx.saved_tensors + + gradgrad_input = gradgrad_input.reshape(-1, ctx.in_size[2], ctx.in_size[3], 1) + + gradgrad_out = upfirdn2d_op.upfirdn2d( + gradgrad_input, + kernel, + ctx.up_x, + ctx.up_y, + ctx.down_x, + ctx.down_y, + ctx.pad_x0, + ctx.pad_x1, + ctx.pad_y0, + ctx.pad_y1, + ) + # gradgrad_out = gradgrad_out.view(ctx.in_size[0], ctx.out_size[0], ctx.out_size[1], ctx.in_size[3]) + gradgrad_out = gradgrad_out.view( + ctx.in_size[0], ctx.in_size[1], ctx.out_size[0], ctx.out_size[1] + ) + + return gradgrad_out, None, None, None, None, None, None, None, None + + +class UpFirDn2d(Function): + @staticmethod + def forward(ctx, input, kernel, up, down, pad): + up_x, up_y = up + down_x, down_y = down + pad_x0, pad_x1, pad_y0, pad_y1 = pad + + kernel_h, kernel_w = kernel.shape + batch, channel, in_h, in_w = input.shape + ctx.in_size = input.shape + + input = input.reshape(-1, in_h, in_w, 1) + + ctx.save_for_backward(kernel, torch.flip(kernel, [0, 1])) + + out_h = (in_h * up_y + pad_y0 + pad_y1 - kernel_h) // down_y + 1 + out_w = (in_w * up_x + pad_x0 + pad_x1 - kernel_w) // down_x + 1 + ctx.out_size = (out_h, out_w) + + ctx.up = (up_x, up_y) + ctx.down = (down_x, down_y) + ctx.pad = (pad_x0, pad_x1, pad_y0, pad_y1) + + g_pad_x0 = kernel_w - pad_x0 - 1 + g_pad_y0 = kernel_h - pad_y0 - 1 + g_pad_x1 = in_w * up_x - out_w * down_x + pad_x0 - up_x + 1 + g_pad_y1 = in_h * up_y - out_h * down_y + pad_y0 - up_y + 1 + + ctx.g_pad = (g_pad_x0, g_pad_x1, g_pad_y0, g_pad_y1) + + out = upfirdn2d_op.upfirdn2d( + input, kernel, up_x, up_y, down_x, down_y, pad_x0, pad_x1, pad_y0, pad_y1 + ) + # out = out.view(major, out_h, out_w, minor) + out = out.view(-1, channel, out_h, out_w) + + return out + + @staticmethod + def backward(ctx, grad_output): + kernel, grad_kernel = ctx.saved_tensors + + grad_input = UpFirDn2dBackward.apply( + grad_output, + kernel, + grad_kernel, + ctx.up, + ctx.down, + ctx.pad, + ctx.g_pad, + ctx.in_size, + ctx.out_size, + ) + + return grad_input, None, None, None, None + + +def upfirdn2d(input, kernel, up=1, down=1, pad=(0, 0)): + if input.device.type == "cpu": + out = upfirdn2d_native( + input, kernel, up, up, down, down, pad[0], pad[1], pad[0], pad[1] + ) + + else: + out = UpFirDn2d.apply( + input, kernel, (up, up), (down, down), (pad[0], pad[1], pad[0], pad[1]) + ) + + return out + + +def upfirdn2d_native( + input, kernel, up_x, up_y, down_x, down_y, pad_x0, pad_x1, pad_y0, pad_y1 +): + _, channel, in_h, in_w = input.shape + input = input.reshape(-1, in_h, in_w, 1) + + _, in_h, in_w, minor = input.shape + kernel_h, kernel_w = kernel.shape + + out = input.view(-1, in_h, 1, in_w, 1, minor) + out = F.pad(out, [0, 0, 0, up_x - 1, 0, 0, 0, up_y - 1]) + out = out.view(-1, in_h * up_y, in_w * up_x, minor) + + out = F.pad( + out, [0, 0, max(pad_x0, 0), max(pad_x1, 0), max(pad_y0, 0), max(pad_y1, 0)] + ) + out = out[ + :, + max(-pad_y0, 0) : out.shape[1] - max(-pad_y1, 0), + max(-pad_x0, 0) : out.shape[2] - max(-pad_x1, 0), + :, + ] + + out = out.permute(0, 3, 1, 2) + out = out.reshape( + [-1, 1, in_h * up_y + pad_y0 + pad_y1, in_w * up_x + pad_x0 + pad_x1] + ) + w = torch.flip(kernel, [0, 1]).view(1, 1, kernel_h, kernel_w) + out = F.conv2d(out, w) + out = out.reshape( + -1, + minor, + in_h * up_y + pad_y0 + pad_y1 - kernel_h + 1, + in_w * up_x + pad_x0 + pad_x1 - kernel_w + 1, + ) + out = out.permute(0, 2, 3, 1) + out = out[:, ::down_y, ::down_x, :] + + out_h = (in_h * up_y + pad_y0 + pad_y1 - kernel_h) // down_y + 1 + out_w = (in_w * up_x + pad_x0 + pad_x1 - kernel_w) // down_x + 1 + + return out.view(-1, channel, out_h, out_w) diff --git a/models/stylegan2/op/upfirdn2d_kernel.cu b/models/stylegan2/op/upfirdn2d_kernel.cu new file mode 100644 index 0000000000000000000000000000000000000000..a88bc7720da6cd54fccd0c4a03dd20fde85c063d --- /dev/null +++ b/models/stylegan2/op/upfirdn2d_kernel.cu @@ -0,0 +1,369 @@ +// Copyright (c) 2019, NVIDIA Corporation. All rights reserved. +// +// This work is made available under the Nvidia Source Code License-NC. +// To view a copy of this license, visit +// https://nvlabs.github.io/stylegan2/license.html + +#include + +#include +#include +#include +#include + +#include +#include + +static __host__ __device__ __forceinline__ int floor_div(int a, int b) { + int c = a / b; + + if (c * b > a) { + c--; + } + + return c; +} + +struct UpFirDn2DKernelParams { + int up_x; + int up_y; + int down_x; + int down_y; + int pad_x0; + int pad_x1; + int pad_y0; + int pad_y1; + + int major_dim; + int in_h; + int in_w; + int minor_dim; + int kernel_h; + int kernel_w; + int out_h; + int out_w; + int loop_major; + int loop_x; +}; + +template +__global__ void upfirdn2d_kernel_large(scalar_t *out, const scalar_t *input, + const scalar_t *kernel, + const UpFirDn2DKernelParams p) { + int minor_idx = blockIdx.x * blockDim.x + threadIdx.x; + int out_y = minor_idx / p.minor_dim; + minor_idx -= out_y * p.minor_dim; + int out_x_base = blockIdx.y * p.loop_x * blockDim.y + threadIdx.y; + int major_idx_base = blockIdx.z * p.loop_major; + + if (out_x_base >= p.out_w || out_y >= p.out_h || + major_idx_base >= p.major_dim) { + return; + } + + int mid_y = out_y * p.down_y + p.up_y - 1 - p.pad_y0; + int in_y = min(max(floor_div(mid_y, p.up_y), 0), p.in_h); + int h = min(max(floor_div(mid_y + p.kernel_h, p.up_y), 0), p.in_h) - in_y; + int kernel_y = mid_y + p.kernel_h - (in_y + 1) * p.up_y; + + for (int loop_major = 0, major_idx = major_idx_base; + loop_major < p.loop_major && major_idx < p.major_dim; + loop_major++, major_idx++) { + for (int loop_x = 0, out_x = out_x_base; + loop_x < p.loop_x && out_x < p.out_w; loop_x++, out_x += blockDim.y) { + int mid_x = out_x * p.down_x + p.up_x - 1 - p.pad_x0; + int in_x = min(max(floor_div(mid_x, p.up_x), 0), p.in_w); + int w = min(max(floor_div(mid_x + p.kernel_w, p.up_x), 0), p.in_w) - in_x; + int kernel_x = mid_x + p.kernel_w - (in_x + 1) * p.up_x; + + const scalar_t *x_p = + &input[((major_idx * p.in_h + in_y) * p.in_w + in_x) * p.minor_dim + + minor_idx]; + const scalar_t *k_p = &kernel[kernel_y * p.kernel_w + kernel_x]; + int x_px = p.minor_dim; + int k_px = -p.up_x; + int x_py = p.in_w * p.minor_dim; + int k_py = -p.up_y * p.kernel_w; + + scalar_t v = 0.0f; + + for (int y = 0; y < h; y++) { + for (int x = 0; x < w; x++) { + v += static_cast(*x_p) * static_cast(*k_p); + x_p += x_px; + k_p += k_px; + } + + x_p += x_py - w * x_px; + k_p += k_py - w * k_px; + } + + out[((major_idx * p.out_h + out_y) * p.out_w + out_x) * p.minor_dim + + minor_idx] = v; + } + } +} + +template +__global__ void upfirdn2d_kernel(scalar_t *out, const scalar_t *input, + const scalar_t *kernel, + const UpFirDn2DKernelParams p) { + const int tile_in_h = ((tile_out_h - 1) * down_y + kernel_h - 1) / up_y + 1; + const int tile_in_w = ((tile_out_w - 1) * down_x + kernel_w - 1) / up_x + 1; + + __shared__ volatile float sk[kernel_h][kernel_w]; + __shared__ volatile float sx[tile_in_h][tile_in_w]; + + int minor_idx = blockIdx.x; + int tile_out_y = minor_idx / p.minor_dim; + minor_idx -= tile_out_y * p.minor_dim; + tile_out_y *= tile_out_h; + int tile_out_x_base = blockIdx.y * p.loop_x * tile_out_w; + int major_idx_base = blockIdx.z * p.loop_major; + + if (tile_out_x_base >= p.out_w | tile_out_y >= p.out_h | + major_idx_base >= p.major_dim) { + return; + } + + for (int tap_idx = threadIdx.x; tap_idx < kernel_h * kernel_w; + tap_idx += blockDim.x) { + int ky = tap_idx / kernel_w; + int kx = tap_idx - ky * kernel_w; + scalar_t v = 0.0; + + if (kx < p.kernel_w & ky < p.kernel_h) { + v = kernel[(p.kernel_h - 1 - ky) * p.kernel_w + (p.kernel_w - 1 - kx)]; + } + + sk[ky][kx] = v; + } + + for (int loop_major = 0, major_idx = major_idx_base; + loop_major < p.loop_major & major_idx < p.major_dim; + loop_major++, major_idx++) { + for (int loop_x = 0, tile_out_x = tile_out_x_base; + loop_x < p.loop_x & tile_out_x < p.out_w; + loop_x++, tile_out_x += tile_out_w) { + int tile_mid_x = tile_out_x * down_x + up_x - 1 - p.pad_x0; + int tile_mid_y = tile_out_y * down_y + up_y - 1 - p.pad_y0; + int tile_in_x = floor_div(tile_mid_x, up_x); + int tile_in_y = floor_div(tile_mid_y, up_y); + + __syncthreads(); + + for (int in_idx = threadIdx.x; in_idx < tile_in_h * tile_in_w; + in_idx += blockDim.x) { + int rel_in_y = in_idx / tile_in_w; + int rel_in_x = in_idx - rel_in_y * tile_in_w; + int in_x = rel_in_x + tile_in_x; + int in_y = rel_in_y + tile_in_y; + + scalar_t v = 0.0; + + if (in_x >= 0 & in_y >= 0 & in_x < p.in_w & in_y < p.in_h) { + v = input[((major_idx * p.in_h + in_y) * p.in_w + in_x) * + p.minor_dim + + minor_idx]; + } + + sx[rel_in_y][rel_in_x] = v; + } + + __syncthreads(); + for (int out_idx = threadIdx.x; out_idx < tile_out_h * tile_out_w; + out_idx += blockDim.x) { + int rel_out_y = out_idx / tile_out_w; + int rel_out_x = out_idx - rel_out_y * tile_out_w; + int out_x = rel_out_x + tile_out_x; + int out_y = rel_out_y + tile_out_y; + + int mid_x = tile_mid_x + rel_out_x * down_x; + int mid_y = tile_mid_y + rel_out_y * down_y; + int in_x = floor_div(mid_x, up_x); + int in_y = floor_div(mid_y, up_y); + int rel_in_x = in_x - tile_in_x; + int rel_in_y = in_y - tile_in_y; + int kernel_x = (in_x + 1) * up_x - mid_x - 1; + int kernel_y = (in_y + 1) * up_y - mid_y - 1; + + scalar_t v = 0.0; + +#pragma unroll + for (int y = 0; y < kernel_h / up_y; y++) +#pragma unroll + for (int x = 0; x < kernel_w / up_x; x++) + v += sx[rel_in_y + y][rel_in_x + x] * + sk[kernel_y + y * up_y][kernel_x + x * up_x]; + + if (out_x < p.out_w & out_y < p.out_h) { + out[((major_idx * p.out_h + out_y) * p.out_w + out_x) * p.minor_dim + + minor_idx] = v; + } + } + } + } +} + +torch::Tensor upfirdn2d_op(const torch::Tensor &input, + const torch::Tensor &kernel, int up_x, int up_y, + int down_x, int down_y, int pad_x0, int pad_x1, + int pad_y0, int pad_y1) { + int curDevice = -1; + cudaGetDevice(&curDevice); + cudaStream_t stream = at::cuda::getCurrentCUDAStream(curDevice); + + UpFirDn2DKernelParams p; + + auto x = input.contiguous(); + auto k = kernel.contiguous(); + + p.major_dim = x.size(0); + p.in_h = x.size(1); + p.in_w = x.size(2); + p.minor_dim = x.size(3); + p.kernel_h = k.size(0); + p.kernel_w = k.size(1); + p.up_x = up_x; + p.up_y = up_y; + p.down_x = down_x; + p.down_y = down_y; + p.pad_x0 = pad_x0; + p.pad_x1 = pad_x1; + p.pad_y0 = pad_y0; + p.pad_y1 = pad_y1; + + p.out_h = (p.in_h * p.up_y + p.pad_y0 + p.pad_y1 - p.kernel_h + p.down_y) / + p.down_y; + p.out_w = (p.in_w * p.up_x + p.pad_x0 + p.pad_x1 - p.kernel_w + p.down_x) / + p.down_x; + + auto out = + at::empty({p.major_dim, p.out_h, p.out_w, p.minor_dim}, x.options()); + + int mode = -1; + + int tile_out_h = -1; + int tile_out_w = -1; + + if (p.up_x == 1 && p.up_y == 1 && p.down_x == 1 && p.down_y == 1 && + p.kernel_h <= 4 && p.kernel_w <= 4) { + mode = 1; + tile_out_h = 16; + tile_out_w = 64; + } + + if (p.up_x == 1 && p.up_y == 1 && p.down_x == 1 && p.down_y == 1 && + p.kernel_h <= 3 && p.kernel_w <= 3) { + mode = 2; + tile_out_h = 16; + tile_out_w = 64; + } + + if (p.up_x == 2 && p.up_y == 2 && p.down_x == 1 && p.down_y == 1 && + p.kernel_h <= 4 && p.kernel_w <= 4) { + mode = 3; + tile_out_h = 16; + tile_out_w = 64; + } + + if (p.up_x == 2 && p.up_y == 2 && p.down_x == 1 && p.down_y == 1 && + p.kernel_h <= 2 && p.kernel_w <= 2) { + mode = 4; + tile_out_h = 16; + tile_out_w = 64; + } + + if (p.up_x == 1 && p.up_y == 1 && p.down_x == 2 && p.down_y == 2 && + p.kernel_h <= 4 && p.kernel_w <= 4) { + mode = 5; + tile_out_h = 8; + tile_out_w = 32; + } + + if (p.up_x == 1 && p.up_y == 1 && p.down_x == 2 && p.down_y == 2 && + p.kernel_h <= 2 && p.kernel_w <= 2) { + mode = 6; + tile_out_h = 8; + tile_out_w = 32; + } + + dim3 block_size; + dim3 grid_size; + + if (tile_out_h > 0 && tile_out_w > 0) { + p.loop_major = (p.major_dim - 1) / 16384 + 1; + p.loop_x = 1; + block_size = dim3(32 * 8, 1, 1); + grid_size = dim3(((p.out_h - 1) / tile_out_h + 1) * p.minor_dim, + (p.out_w - 1) / (p.loop_x * tile_out_w) + 1, + (p.major_dim - 1) / p.loop_major + 1); + } else { + p.loop_major = (p.major_dim - 1) / 16384 + 1; + p.loop_x = 4; + block_size = dim3(4, 32, 1); + grid_size = dim3((p.out_h * p.minor_dim - 1) / block_size.x + 1, + (p.out_w - 1) / (p.loop_x * block_size.y) + 1, + (p.major_dim - 1) / p.loop_major + 1); + } + + AT_DISPATCH_FLOATING_TYPES_AND_HALF(x.scalar_type(), "upfirdn2d_cuda", [&] { + switch (mode) { + case 1: + upfirdn2d_kernel + <<>>(out.data_ptr(), + x.data_ptr(), + k.data_ptr(), p); + + break; + + case 2: + upfirdn2d_kernel + <<>>(out.data_ptr(), + x.data_ptr(), + k.data_ptr(), p); + + break; + + case 3: + upfirdn2d_kernel + <<>>(out.data_ptr(), + x.data_ptr(), + k.data_ptr(), p); + + break; + + case 4: + upfirdn2d_kernel + <<>>(out.data_ptr(), + x.data_ptr(), + k.data_ptr(), p); + + break; + + case 5: + upfirdn2d_kernel + <<>>(out.data_ptr(), + x.data_ptr(), + k.data_ptr(), p); + + break; + + case 6: + upfirdn2d_kernel + <<>>(out.data_ptr(), + x.data_ptr(), + k.data_ptr(), p); + + break; + + default: + upfirdn2d_kernel_large<<>>( + out.data_ptr(), x.data_ptr(), + k.data_ptr(), p); + } + }); + + return out; +} \ No newline at end of file diff --git a/models/stylegan2/torch_utils/__init__.py b/models/stylegan2/torch_utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e1e1a5ba99e56a56ecaa14f7d4fa41777789c0cf --- /dev/null +++ b/models/stylegan2/torch_utils/__init__.py @@ -0,0 +1,9 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +# +# NVIDIA CORPORATION and its licensors retain all intellectual property +# and proprietary rights in and to this software, related documentation +# and any modifications thereto. Any use, reproduction, disclosure or +# distribution of this software and related documentation without an express +# license agreement from NVIDIA CORPORATION is strictly prohibited. + +# empty diff --git a/models/stylegan2/torch_utils/custom_ops.py b/models/stylegan2/torch_utils/custom_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..4cc4e43fc6f6ce79f2bd68a44ba87990b9b8564e --- /dev/null +++ b/models/stylegan2/torch_utils/custom_ops.py @@ -0,0 +1,126 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +# +# NVIDIA CORPORATION and its licensors retain all intellectual property +# and proprietary rights in and to this software, related documentation +# and any modifications thereto. Any use, reproduction, disclosure or +# distribution of this software and related documentation without an express +# license agreement from NVIDIA CORPORATION is strictly prohibited. + +import os +import glob +import torch +import torch.utils.cpp_extension +import importlib +import hashlib +import shutil +from pathlib import Path + +from torch.utils.file_baton import FileBaton + +#---------------------------------------------------------------------------- +# Global options. + +verbosity = 'brief' # Verbosity level: 'none', 'brief', 'full' + +#---------------------------------------------------------------------------- +# Internal helper funcs. + +def _find_compiler_bindir(): + patterns = [ + 'C:/Program Files (x86)/Microsoft Visual Studio/*/Professional/VC/Tools/MSVC/*/bin/Hostx64/x64', + 'C:/Program Files (x86)/Microsoft Visual Studio/*/BuildTools/VC/Tools/MSVC/*/bin/Hostx64/x64', + 'C:/Program Files (x86)/Microsoft Visual Studio/*/Community/VC/Tools/MSVC/*/bin/Hostx64/x64', + 'C:/Program Files (x86)/Microsoft Visual Studio */vc/bin', + ] + for pattern in patterns: + matches = sorted(glob.glob(pattern)) + if len(matches): + return matches[-1] + return None + +#---------------------------------------------------------------------------- +# Main entry point for compiling and loading C++/CUDA plugins. + +_cached_plugins = dict() + +def get_plugin(module_name, sources, **build_kwargs): + assert verbosity in ['none', 'brief', 'full'] + + # Already cached? + if module_name in _cached_plugins: + return _cached_plugins[module_name] + + # Print status. + if verbosity == 'full': + print(f'Setting up PyTorch plugin "{module_name}"...') + elif verbosity == 'brief': + print(f'Setting up PyTorch plugin "{module_name}"... ', end='', flush=True) + + try: # pylint: disable=too-many-nested-blocks + # Make sure we can find the necessary compiler binaries. + if os.name == 'nt' and os.system("where cl.exe >nul 2>nul") != 0: + compiler_bindir = _find_compiler_bindir() + if compiler_bindir is None: + raise RuntimeError(f'Could not find MSVC/GCC/CLANG installation on this computer. Check _find_compiler_bindir() in "{__file__}".') + os.environ['PATH'] += ';' + compiler_bindir + + # Compile and load. + verbose_build = (verbosity == 'full') + + # Incremental build md5sum trickery. Copies all the input source files + # into a cached build directory under a combined md5 digest of the input + # source files. Copying is done only if the combined digest has changed. + # This keeps input file timestamps and filenames the same as in previous + # extension builds, allowing for fast incremental rebuilds. + # + # This optimization is done only in case all the source files reside in + # a single directory (just for simplicity) and if the TORCH_EXTENSIONS_DIR + # environment variable is set (we take this as a signal that the user + # actually cares about this.) + source_dirs_set = set(os.path.dirname(source) for source in sources) + if len(source_dirs_set) == 1 and ('TORCH_EXTENSIONS_DIR' in os.environ): + all_source_files = sorted(list(x for x in Path(list(source_dirs_set)[0]).iterdir() if x.is_file())) + + # Compute a combined hash digest for all source files in the same + # custom op directory (usually .cu, .cpp, .py and .h files). + hash_md5 = hashlib.md5() + for src in all_source_files: + with open(src, 'rb') as f: + hash_md5.update(f.read()) + build_dir = torch.utils.cpp_extension._get_build_directory(module_name, verbose=verbose_build) # pylint: disable=protected-access + digest_build_dir = os.path.join(build_dir, hash_md5.hexdigest()) + + if not os.path.isdir(digest_build_dir): + os.makedirs(digest_build_dir, exist_ok=True) + baton = FileBaton(os.path.join(digest_build_dir, 'lock')) + if baton.try_acquire(): + try: + for src in all_source_files: + shutil.copyfile(src, os.path.join(digest_build_dir, os.path.basename(src))) + finally: + baton.release() + else: + # Someone else is copying source files under the digest dir, + # wait until done and continue. + baton.wait() + digest_sources = [os.path.join(digest_build_dir, os.path.basename(x)) for x in sources] + torch.utils.cpp_extension.load(name=module_name, build_directory=build_dir, + verbose=verbose_build, sources=digest_sources, **build_kwargs) + else: + torch.utils.cpp_extension.load(name=module_name, verbose=verbose_build, sources=sources, **build_kwargs) + module = importlib.import_module(module_name) + + except: + if verbosity == 'brief': + print('Failed!') + raise + + # Print status and add to cache. + if verbosity == 'full': + print(f'Done setting up PyTorch plugin "{module_name}".') + elif verbosity == 'brief': + print('Done.') + _cached_plugins[module_name] = module + return module + +#---------------------------------------------------------------------------- diff --git a/models/stylegan2/torch_utils/misc.py b/models/stylegan2/torch_utils/misc.py new file mode 100644 index 0000000000000000000000000000000000000000..97cc1325c811849a95a74dfce76d3ef6d711dccd --- /dev/null +++ b/models/stylegan2/torch_utils/misc.py @@ -0,0 +1,262 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +# +# NVIDIA CORPORATION and its licensors retain all intellectual property +# and proprietary rights in and to this software, related documentation +# and any modifications thereto. Any use, reproduction, disclosure or +# distribution of this software and related documentation without an express +# license agreement from NVIDIA CORPORATION is strictly prohibited. + +import re +import contextlib +import numpy as np +import torch +import warnings +import models.stylegan2.dnnlib as dnnlib + +#---------------------------------------------------------------------------- +# Cached construction of constant tensors. Avoids CPU=>GPU copy when the +# same constant is used multiple times. + +_constant_cache = dict() + +def constant(value, shape=None, dtype=None, device=None, memory_format=None): + value = np.asarray(value) + if shape is not None: + shape = tuple(shape) + if dtype is None: + dtype = torch.get_default_dtype() + if device is None: + device = torch.device('cpu') + if memory_format is None: + memory_format = torch.contiguous_format + + key = (value.shape, value.dtype, value.tobytes(), shape, dtype, device, memory_format) + tensor = _constant_cache.get(key, None) + if tensor is None: + tensor = torch.as_tensor(value.copy(), dtype=dtype, device=device) + if shape is not None: + tensor, _ = torch.broadcast_tensors(tensor, torch.empty(shape)) + tensor = tensor.contiguous(memory_format=memory_format) + _constant_cache[key] = tensor + return tensor + +#---------------------------------------------------------------------------- +# Replace NaN/Inf with specified numerical values. + +try: + nan_to_num = torch.nan_to_num # 1.8.0a0 +except AttributeError: + def nan_to_num(input, nan=0.0, posinf=None, neginf=None, *, out=None): # pylint: disable=redefined-builtin + assert isinstance(input, torch.Tensor) + if posinf is None: + posinf = torch.finfo(input.dtype).max + if neginf is None: + neginf = torch.finfo(input.dtype).min + assert nan == 0 + return torch.clamp(input.unsqueeze(0).nansum(0), min=neginf, max=posinf, out=out) + +#---------------------------------------------------------------------------- +# Symbolic assert. + +try: + symbolic_assert = torch._assert # 1.8.0a0 # pylint: disable=protected-access +except AttributeError: + symbolic_assert = torch.Assert # 1.7.0 + +#---------------------------------------------------------------------------- +# Context manager to suppress known warnings in torch.jit.trace(). + +class suppress_tracer_warnings(warnings.catch_warnings): + def __enter__(self): + super().__enter__() + warnings.simplefilter('ignore', category=torch.jit.TracerWarning) + return self + +#---------------------------------------------------------------------------- +# Assert that the shape of a tensor matches the given list of integers. +# None indicates that the size of a dimension is allowed to vary. +# Performs symbolic assertion when used in torch.jit.trace(). + +def assert_shape(tensor, ref_shape): + if tensor.ndim != len(ref_shape): + raise AssertionError(f'Wrong number of dimensions: got {tensor.ndim}, expected {len(ref_shape)}') + for idx, (size, ref_size) in enumerate(zip(tensor.shape, ref_shape)): + if ref_size is None: + pass + elif isinstance(ref_size, torch.Tensor): + with suppress_tracer_warnings(): # as_tensor results are registered as constants + symbolic_assert(torch.equal(torch.as_tensor(size), ref_size), f'Wrong size for dimension {idx}') + elif isinstance(size, torch.Tensor): + with suppress_tracer_warnings(): # as_tensor results are registered as constants + symbolic_assert(torch.equal(size, torch.as_tensor(ref_size)), f'Wrong size for dimension {idx}: expected {ref_size}') + elif size != ref_size: + raise AssertionError(f'Wrong size for dimension {idx}: got {size}, expected {ref_size}') + +#---------------------------------------------------------------------------- +# Function decorator that calls torch.autograd.profiler.record_function(). + +def profiled_function(fn): + def decorator(*args, **kwargs): + with torch.autograd.profiler.record_function(fn.__name__): + return fn(*args, **kwargs) + decorator.__name__ = fn.__name__ + return decorator + +#---------------------------------------------------------------------------- +# Sampler for torch.utils.data.DataLoader that loops over the dataset +# indefinitely, shuffling items as it goes. + +class InfiniteSampler(torch.utils.data.Sampler): + def __init__(self, dataset, rank=0, num_replicas=1, shuffle=True, seed=0, window_size=0.5): + assert len(dataset) > 0 + assert num_replicas > 0 + assert 0 <= rank < num_replicas + assert 0 <= window_size <= 1 + super().__init__(dataset) + self.dataset = dataset + self.rank = rank + self.num_replicas = num_replicas + self.shuffle = shuffle + self.seed = seed + self.window_size = window_size + + def __iter__(self): + order = np.arange(len(self.dataset)) + rnd = None + window = 0 + if self.shuffle: + rnd = np.random.RandomState(self.seed) + rnd.shuffle(order) + window = int(np.rint(order.size * self.window_size)) + + idx = 0 + while True: + i = idx % order.size + if idx % self.num_replicas == self.rank: + yield order[i] + if window >= 2: + j = (i - rnd.randint(window)) % order.size + order[i], order[j] = order[j], order[i] + idx += 1 + +#---------------------------------------------------------------------------- +# Utilities for operating with torch.nn.Module parameters and buffers. + +def params_and_buffers(module): + assert isinstance(module, torch.nn.Module) + return list(module.parameters()) + list(module.buffers()) + +def named_params_and_buffers(module): + assert isinstance(module, torch.nn.Module) + return list(module.named_parameters()) + list(module.named_buffers()) + +def copy_params_and_buffers(src_module, dst_module, require_all=False): + assert isinstance(src_module, torch.nn.Module) + assert isinstance(dst_module, torch.nn.Module) + src_tensors = {name: tensor for name, tensor in named_params_and_buffers(src_module)} + for name, tensor in named_params_and_buffers(dst_module): + assert (name in src_tensors) or (not require_all) + if name in src_tensors: + tensor.copy_(src_tensors[name].detach()).requires_grad_(tensor.requires_grad) + +#---------------------------------------------------------------------------- +# Context manager for easily enabling/disabling DistributedDataParallel +# synchronization. + +@contextlib.contextmanager +def ddp_sync(module, sync): + assert isinstance(module, torch.nn.Module) + if sync or not isinstance(module, torch.nn.parallel.DistributedDataParallel): + yield + else: + with module.no_sync(): + yield + +#---------------------------------------------------------------------------- +# Check DistributedDataParallel consistency across processes. + +def check_ddp_consistency(module, ignore_regex=None): + assert isinstance(module, torch.nn.Module) + for name, tensor in named_params_and_buffers(module): + fullname = type(module).__name__ + '.' + name + if ignore_regex is not None and re.fullmatch(ignore_regex, fullname): + continue + tensor = tensor.detach() + other = tensor.clone() + torch.distributed.broadcast(tensor=other, src=0) + assert (nan_to_num(tensor) == nan_to_num(other)).all(), fullname + +#---------------------------------------------------------------------------- +# Print summary table of module hierarchy. + +def print_module_summary(module, inputs, max_nesting=3, skip_redundant=True): + assert isinstance(module, torch.nn.Module) + assert not isinstance(module, torch.jit.ScriptModule) + assert isinstance(inputs, (tuple, list)) + + # Register hooks. + entries = [] + nesting = [0] + def pre_hook(_mod, _inputs): + nesting[0] += 1 + def post_hook(mod, _inputs, outputs): + nesting[0] -= 1 + if nesting[0] <= max_nesting: + outputs = list(outputs) if isinstance(outputs, (tuple, list)) else [outputs] + outputs = [t for t in outputs if isinstance(t, torch.Tensor)] + entries.append(dnnlib.EasyDict(mod=mod, outputs=outputs)) + hooks = [mod.register_forward_pre_hook(pre_hook) for mod in module.modules()] + hooks += [mod.register_forward_hook(post_hook) for mod in module.modules()] + + # Run module. + outputs = module(*inputs) + for hook in hooks: + hook.remove() + + # Identify unique outputs, parameters, and buffers. + tensors_seen = set() + for e in entries: + e.unique_params = [t for t in e.mod.parameters() if id(t) not in tensors_seen] + e.unique_buffers = [t for t in e.mod.buffers() if id(t) not in tensors_seen] + e.unique_outputs = [t for t in e.outputs if id(t) not in tensors_seen] + tensors_seen |= {id(t) for t in e.unique_params + e.unique_buffers + e.unique_outputs} + + # Filter out redundant entries. + if skip_redundant: + entries = [e for e in entries if len(e.unique_params) or len(e.unique_buffers) or len(e.unique_outputs)] + + # Construct table. + rows = [[type(module).__name__, 'Parameters', 'Buffers', 'Output shape', 'Datatype']] + rows += [['---'] * len(rows[0])] + param_total = 0 + buffer_total = 0 + submodule_names = {mod: name for name, mod in module.named_modules()} + for e in entries: + name = '' if e.mod is module else submodule_names[e.mod] + param_size = sum(t.numel() for t in e.unique_params) + buffer_size = sum(t.numel() for t in e.unique_buffers) + output_shapes = [str(list(e.outputs[0].shape)) for t in e.outputs] + output_dtypes = [str(t.dtype).split('.')[-1] for t in e.outputs] + rows += [[ + name + (':0' if len(e.outputs) >= 2 else ''), + str(param_size) if param_size else '-', + str(buffer_size) if buffer_size else '-', + (output_shapes + ['-'])[0], + (output_dtypes + ['-'])[0], + ]] + for idx in range(1, len(e.outputs)): + rows += [[name + f':{idx}', '-', '-', output_shapes[idx], output_dtypes[idx]]] + param_total += param_size + buffer_total += buffer_size + rows += [['---'] * len(rows[0])] + rows += [['Total', str(param_total), str(buffer_total), '-', '-']] + + # Print table. + widths = [max(len(cell) for cell in column) for column in zip(*rows)] + print() + for row in rows: + print(' '.join(cell + ' ' * (width - len(cell)) for cell, width in zip(row, widths))) + print() + return outputs + +#---------------------------------------------------------------------------- diff --git a/models/stylegan2/torch_utils/ops/__init__.py b/models/stylegan2/torch_utils/ops/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e1e1a5ba99e56a56ecaa14f7d4fa41777789c0cf --- /dev/null +++ b/models/stylegan2/torch_utils/ops/__init__.py @@ -0,0 +1,9 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +# +# NVIDIA CORPORATION and its licensors retain all intellectual property +# and proprietary rights in and to this software, related documentation +# and any modifications thereto. Any use, reproduction, disclosure or +# distribution of this software and related documentation without an express +# license agreement from NVIDIA CORPORATION is strictly prohibited. + +# empty diff --git a/models/stylegan2/torch_utils/ops/bias_act.cpp b/models/stylegan2/torch_utils/ops/bias_act.cpp new file mode 100644 index 0000000000000000000000000000000000000000..5d2425d8054991a8e8b6f7a940fd0ff7fa0bb330 --- /dev/null +++ b/models/stylegan2/torch_utils/ops/bias_act.cpp @@ -0,0 +1,99 @@ +// Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +// +// NVIDIA CORPORATION and its licensors retain all intellectual property +// and proprietary rights in and to this software, related documentation +// and any modifications thereto. Any use, reproduction, disclosure or +// distribution of this software and related documentation without an express +// license agreement from NVIDIA CORPORATION is strictly prohibited. + +#include +#include +#include +#include "bias_act.h" + +//------------------------------------------------------------------------ + +static bool has_same_layout(torch::Tensor x, torch::Tensor y) +{ + if (x.dim() != y.dim()) + return false; + for (int64_t i = 0; i < x.dim(); i++) + { + if (x.size(i) != y.size(i)) + return false; + if (x.size(i) >= 2 && x.stride(i) != y.stride(i)) + return false; + } + return true; +} + +//------------------------------------------------------------------------ + +static torch::Tensor bias_act(torch::Tensor x, torch::Tensor b, torch::Tensor xref, torch::Tensor yref, torch::Tensor dy, int grad, int dim, int act, float alpha, float gain, float clamp) +{ + // Validate arguments. + TORCH_CHECK(x.is_cuda(), "x must reside on CUDA device"); + TORCH_CHECK(b.numel() == 0 || (b.dtype() == x.dtype() && b.device() == x.device()), "b must have the same dtype and device as x"); + TORCH_CHECK(xref.numel() == 0 || (xref.sizes() == x.sizes() && xref.dtype() == x.dtype() && xref.device() == x.device()), "xref must have the same shape, dtype, and device as x"); + TORCH_CHECK(yref.numel() == 0 || (yref.sizes() == x.sizes() && yref.dtype() == x.dtype() && yref.device() == x.device()), "yref must have the same shape, dtype, and device as x"); + TORCH_CHECK(dy.numel() == 0 || (dy.sizes() == x.sizes() && dy.dtype() == x.dtype() && dy.device() == x.device()), "dy must have the same dtype and device as x"); + TORCH_CHECK(x.numel() <= INT_MAX, "x is too large"); + TORCH_CHECK(b.dim() == 1, "b must have rank 1"); + TORCH_CHECK(b.numel() == 0 || (dim >= 0 && dim < x.dim()), "dim is out of bounds"); + TORCH_CHECK(b.numel() == 0 || b.numel() == x.size(dim), "b has wrong number of elements"); + TORCH_CHECK(grad >= 0, "grad must be non-negative"); + + // Validate layout. + TORCH_CHECK(x.is_non_overlapping_and_dense(), "x must be non-overlapping and dense"); + TORCH_CHECK(b.is_contiguous(), "b must be contiguous"); + TORCH_CHECK(xref.numel() == 0 || has_same_layout(xref, x), "xref must have the same layout as x"); + TORCH_CHECK(yref.numel() == 0 || has_same_layout(yref, x), "yref must have the same layout as x"); + TORCH_CHECK(dy.numel() == 0 || has_same_layout(dy, x), "dy must have the same layout as x"); + + // Create output tensor. + const at::cuda::OptionalCUDAGuard device_guard(device_of(x)); + torch::Tensor y = torch::empty_like(x); + TORCH_CHECK(has_same_layout(y, x), "y must have the same layout as x"); + + // Initialize CUDA kernel parameters. + bias_act_kernel_params p; + p.x = x.data_ptr(); + p.b = (b.numel()) ? b.data_ptr() : NULL; + p.xref = (xref.numel()) ? xref.data_ptr() : NULL; + p.yref = (yref.numel()) ? yref.data_ptr() : NULL; + p.dy = (dy.numel()) ? dy.data_ptr() : NULL; + p.y = y.data_ptr(); + p.grad = grad; + p.act = act; + p.alpha = alpha; + p.gain = gain; + p.clamp = clamp; + p.sizeX = (int)x.numel(); + p.sizeB = (int)b.numel(); + p.stepB = (b.numel()) ? (int)x.stride(dim) : 1; + + // Choose CUDA kernel. + void* kernel; + AT_DISPATCH_FLOATING_TYPES_AND_HALF(x.scalar_type(), "upfirdn2d_cuda", [&] + { + kernel = choose_bias_act_kernel(p); + }); + TORCH_CHECK(kernel, "no CUDA kernel found for the specified activation func"); + + // Launch CUDA kernel. + p.loopX = 4; + int blockSize = 4 * 32; + int gridSize = (p.sizeX - 1) / (p.loopX * blockSize) + 1; + void* args[] = {&p}; + AT_CUDA_CHECK(cudaLaunchKernel(kernel, gridSize, blockSize, args, 0, at::cuda::getCurrentCUDAStream())); + return y; +} + +//------------------------------------------------------------------------ + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) +{ + m.def("bias_act", &bias_act); +} + +//------------------------------------------------------------------------ diff --git a/models/stylegan2/torch_utils/ops/bias_act.cu b/models/stylegan2/torch_utils/ops/bias_act.cu new file mode 100644 index 0000000000000000000000000000000000000000..dd8fc4756d7d94727f94af738665b68d9c518880 --- /dev/null +++ b/models/stylegan2/torch_utils/ops/bias_act.cu @@ -0,0 +1,173 @@ +// Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +// +// NVIDIA CORPORATION and its licensors retain all intellectual property +// and proprietary rights in and to this software, related documentation +// and any modifications thereto. Any use, reproduction, disclosure or +// distribution of this software and related documentation without an express +// license agreement from NVIDIA CORPORATION is strictly prohibited. + +#include +#include "bias_act.h" + +//------------------------------------------------------------------------ +// Helpers. + +template struct InternalType; +template <> struct InternalType { typedef double scalar_t; }; +template <> struct InternalType { typedef float scalar_t; }; +template <> struct InternalType { typedef float scalar_t; }; + +//------------------------------------------------------------------------ +// CUDA kernel. + +template +__global__ void bias_act_kernel(bias_act_kernel_params p) +{ + typedef typename InternalType::scalar_t scalar_t; + int G = p.grad; + scalar_t alpha = (scalar_t)p.alpha; + scalar_t gain = (scalar_t)p.gain; + scalar_t clamp = (scalar_t)p.clamp; + scalar_t one = (scalar_t)1; + scalar_t two = (scalar_t)2; + scalar_t expRange = (scalar_t)80; + scalar_t halfExpRange = (scalar_t)40; + scalar_t seluScale = (scalar_t)1.0507009873554804934193349852946; + scalar_t seluAlpha = (scalar_t)1.6732632423543772848170429916717; + + // Loop over elements. + int xi = blockIdx.x * p.loopX * blockDim.x + threadIdx.x; + for (int loopIdx = 0; loopIdx < p.loopX && xi < p.sizeX; loopIdx++, xi += blockDim.x) + { + // Load. + scalar_t x = (scalar_t)((const T*)p.x)[xi]; + scalar_t b = (p.b) ? (scalar_t)((const T*)p.b)[(xi / p.stepB) % p.sizeB] : 0; + scalar_t xref = (p.xref) ? (scalar_t)((const T*)p.xref)[xi] : 0; + scalar_t yref = (p.yref) ? (scalar_t)((const T*)p.yref)[xi] : 0; + scalar_t dy = (p.dy) ? (scalar_t)((const T*)p.dy)[xi] : one; + scalar_t yy = (gain != 0) ? yref / gain : 0; + scalar_t y = 0; + + // Apply bias. + ((G == 0) ? x : xref) += b; + + // linear + if (A == 1) + { + if (G == 0) y = x; + if (G == 1) y = x; + } + + // relu + if (A == 2) + { + if (G == 0) y = (x > 0) ? x : 0; + if (G == 1) y = (yy > 0) ? x : 0; + } + + // lrelu + if (A == 3) + { + if (G == 0) y = (x > 0) ? x : x * alpha; + if (G == 1) y = (yy > 0) ? x : x * alpha; + } + + // tanh + if (A == 4) + { + if (G == 0) { scalar_t c = exp(x); scalar_t d = one / c; y = (x < -expRange) ? -one : (x > expRange) ? one : (c - d) / (c + d); } + if (G == 1) y = x * (one - yy * yy); + if (G == 2) y = x * (one - yy * yy) * (-two * yy); + } + + // sigmoid + if (A == 5) + { + if (G == 0) y = (x < -expRange) ? 0 : one / (exp(-x) + one); + if (G == 1) y = x * yy * (one - yy); + if (G == 2) y = x * yy * (one - yy) * (one - two * yy); + } + + // elu + if (A == 6) + { + if (G == 0) y = (x >= 0) ? x : exp(x) - one; + if (G == 1) y = (yy >= 0) ? x : x * (yy + one); + if (G == 2) y = (yy >= 0) ? 0 : x * (yy + one); + } + + // selu + if (A == 7) + { + if (G == 0) y = (x >= 0) ? seluScale * x : (seluScale * seluAlpha) * (exp(x) - one); + if (G == 1) y = (yy >= 0) ? x * seluScale : x * (yy + seluScale * seluAlpha); + if (G == 2) y = (yy >= 0) ? 0 : x * (yy + seluScale * seluAlpha); + } + + // softplus + if (A == 8) + { + if (G == 0) y = (x > expRange) ? x : log(exp(x) + one); + if (G == 1) y = x * (one - exp(-yy)); + if (G == 2) { scalar_t c = exp(-yy); y = x * c * (one - c); } + } + + // swish + if (A == 9) + { + if (G == 0) + y = (x < -expRange) ? 0 : x / (exp(-x) + one); + else + { + scalar_t c = exp(xref); + scalar_t d = c + one; + if (G == 1) + y = (xref > halfExpRange) ? x : x * c * (xref + d) / (d * d); + else + y = (xref > halfExpRange) ? 0 : x * c * (xref * (two - d) + two * d) / (d * d * d); + yref = (xref < -expRange) ? 0 : xref / (exp(-xref) + one) * gain; + } + } + + // Apply gain. + y *= gain * dy; + + // Clamp. + if (clamp >= 0) + { + if (G == 0) + y = (y > -clamp & y < clamp) ? y : (y >= 0) ? clamp : -clamp; + else + y = (yref > -clamp & yref < clamp) ? y : 0; + } + + // Store. + ((T*)p.y)[xi] = (T)y; + } +} + +//------------------------------------------------------------------------ +// CUDA kernel selection. + +template void* choose_bias_act_kernel(const bias_act_kernel_params& p) +{ + if (p.act == 1) return (void*)bias_act_kernel; + if (p.act == 2) return (void*)bias_act_kernel; + if (p.act == 3) return (void*)bias_act_kernel; + if (p.act == 4) return (void*)bias_act_kernel; + if (p.act == 5) return (void*)bias_act_kernel; + if (p.act == 6) return (void*)bias_act_kernel; + if (p.act == 7) return (void*)bias_act_kernel; + if (p.act == 8) return (void*)bias_act_kernel; + if (p.act == 9) return (void*)bias_act_kernel; + return NULL; +} + +//------------------------------------------------------------------------ +// Template specializations. + +template void* choose_bias_act_kernel (const bias_act_kernel_params& p); +template void* choose_bias_act_kernel (const bias_act_kernel_params& p); +template void* choose_bias_act_kernel (const bias_act_kernel_params& p); + +//------------------------------------------------------------------------ diff --git a/models/stylegan2/torch_utils/ops/bias_act.h b/models/stylegan2/torch_utils/ops/bias_act.h new file mode 100644 index 0000000000000000000000000000000000000000..a32187e1fb7e3bae509d4eceaf900866866875a4 --- /dev/null +++ b/models/stylegan2/torch_utils/ops/bias_act.h @@ -0,0 +1,38 @@ +// Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +// +// NVIDIA CORPORATION and its licensors retain all intellectual property +// and proprietary rights in and to this software, related documentation +// and any modifications thereto. Any use, reproduction, disclosure or +// distribution of this software and related documentation without an express +// license agreement from NVIDIA CORPORATION is strictly prohibited. + +//------------------------------------------------------------------------ +// CUDA kernel parameters. + +struct bias_act_kernel_params +{ + const void* x; // [sizeX] + const void* b; // [sizeB] or NULL + const void* xref; // [sizeX] or NULL + const void* yref; // [sizeX] or NULL + const void* dy; // [sizeX] or NULL + void* y; // [sizeX] + + int grad; + int act; + float alpha; + float gain; + float clamp; + + int sizeX; + int sizeB; + int stepB; + int loopX; +}; + +//------------------------------------------------------------------------ +// CUDA kernel selection. + +template void* choose_bias_act_kernel(const bias_act_kernel_params& p); + +//------------------------------------------------------------------------ diff --git a/models/stylegan2/torch_utils/ops/bias_act.py b/models/stylegan2/torch_utils/ops/bias_act.py new file mode 100644 index 0000000000000000000000000000000000000000..a03673e66f277be9c4a051ac739b711e6a427618 --- /dev/null +++ b/models/stylegan2/torch_utils/ops/bias_act.py @@ -0,0 +1,212 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +# +# NVIDIA CORPORATION and its licensors retain all intellectual property +# and proprietary rights in and to this software, related documentation +# and any modifications thereto. Any use, reproduction, disclosure or +# distribution of this software and related documentation without an express +# license agreement from NVIDIA CORPORATION is strictly prohibited. + +"""Custom PyTorch ops for efficient bias and activation.""" + +import os +import warnings +import numpy as np +import torch +import models.stylegan2.dnnlib as dnnlib +import traceback + +from .. import custom_ops +from .. import misc + +#---------------------------------------------------------------------------- + +activation_funcs = { + 'linear': dnnlib.EasyDict(func=lambda x, **_: x, def_alpha=0, def_gain=1, cuda_idx=1, ref='', has_2nd_grad=False), + 'relu': dnnlib.EasyDict(func=lambda x, **_: torch.nn.functional.relu(x), def_alpha=0, def_gain=np.sqrt(2), cuda_idx=2, ref='y', has_2nd_grad=False), + 'lrelu': dnnlib.EasyDict(func=lambda x, alpha, **_: torch.nn.functional.leaky_relu(x, alpha), def_alpha=0.2, def_gain=np.sqrt(2), cuda_idx=3, ref='y', has_2nd_grad=False), + 'tanh': dnnlib.EasyDict(func=lambda x, **_: torch.tanh(x), def_alpha=0, def_gain=1, cuda_idx=4, ref='y', has_2nd_grad=True), + 'sigmoid': dnnlib.EasyDict(func=lambda x, **_: torch.sigmoid(x), def_alpha=0, def_gain=1, cuda_idx=5, ref='y', has_2nd_grad=True), + 'elu': dnnlib.EasyDict(func=lambda x, **_: torch.nn.functional.elu(x), def_alpha=0, def_gain=1, cuda_idx=6, ref='y', has_2nd_grad=True), + 'selu': dnnlib.EasyDict(func=lambda x, **_: torch.nn.functional.selu(x), def_alpha=0, def_gain=1, cuda_idx=7, ref='y', has_2nd_grad=True), + 'softplus': dnnlib.EasyDict(func=lambda x, **_: torch.nn.functional.softplus(x), def_alpha=0, def_gain=1, cuda_idx=8, ref='y', has_2nd_grad=True), + 'swish': dnnlib.EasyDict(func=lambda x, **_: torch.sigmoid(x) * x, def_alpha=0, def_gain=np.sqrt(2), cuda_idx=9, ref='x', has_2nd_grad=True), +} + +#---------------------------------------------------------------------------- + +_inited = False +_plugin = None +_null_tensor = torch.empty([0]) + +def _init(): + global _inited, _plugin + if not _inited: + _inited = True + sources = ['bias_act.cpp', 'bias_act.cu'] + sources = [os.path.join(os.path.dirname(__file__), s) for s in sources] + try: + _plugin = custom_ops.get_plugin('bias_act_plugin', sources=sources, extra_cuda_cflags=['--use_fast_math']) + except: + warnings.warn('Failed to build CUDA kernels for bias_act. Falling back to slow reference implementation. Details:\n\n' + traceback.format_exc()) + return _plugin is not None + +#---------------------------------------------------------------------------- + +def bias_act(x, b=None, dim=1, act='linear', alpha=None, gain=None, clamp=None, impl='cuda'): + r"""Fused bias and activation function. + + Adds bias `b` to activation tensor `x`, evaluates activation function `act`, + and scales the result by `gain`. Each of the steps is optional. In most cases, + the fused op is considerably more efficient than performing the same calculation + using standard PyTorch ops. It supports first and second order gradients, + but not third order gradients. + + Args: + x: Input activation tensor. Can be of any shape. + b: Bias vector, or `None` to disable. Must be a 1D tensor of the same type + as `x`. The shape must be known, and it must match the dimension of `x` + corresponding to `dim`. + dim: The dimension in `x` corresponding to the elements of `b`. + The value of `dim` is ignored if `b` is not specified. + act: Name of the activation function to evaluate, or `"linear"` to disable. + Can be e.g. `"relu"`, `"lrelu"`, `"tanh"`, `"sigmoid"`, `"swish"`, etc. + See `activation_funcs` for a full list. `None` is not allowed. + alpha: Shape parameter for the activation function, or `None` to use the default. + gain: Scaling factor for the output tensor, or `None` to use default. + See `activation_funcs` for the default scaling of each activation function. + If unsure, consider specifying 1. + clamp: Clamp the output values to `[-clamp, +clamp]`, or `None` to disable + the clamping (default). + impl: Name of the implementation to use. Can be `"ref"` or `"cuda"` (default). + + Returns: + Tensor of the same shape and datatype as `x`. + """ + assert isinstance(x, torch.Tensor) + assert impl in ['ref', 'cuda'] + if impl == 'cuda' and x.device.type == 'cuda' and _init(): + return _bias_act_cuda(dim=dim, act=act, alpha=alpha, gain=gain, clamp=clamp).apply(x, b) + return _bias_act_ref(x=x, b=b, dim=dim, act=act, alpha=alpha, gain=gain, clamp=clamp) + +#---------------------------------------------------------------------------- + +@misc.profiled_function +def _bias_act_ref(x, b=None, dim=1, act='linear', alpha=None, gain=None, clamp=None): + """Slow reference implementation of `bias_act()` using standard TensorFlow ops. + """ + assert isinstance(x, torch.Tensor) + assert clamp is None or clamp >= 0 + spec = activation_funcs[act] + alpha = float(alpha if alpha is not None else spec.def_alpha) + gain = float(gain if gain is not None else spec.def_gain) + clamp = float(clamp if clamp is not None else -1) + + # Add bias. + if b is not None: + assert isinstance(b, torch.Tensor) and b.ndim == 1 + assert 0 <= dim < x.ndim + assert b.shape[0] == x.shape[dim] + x = x + b.reshape([-1 if i == dim else 1 for i in range(x.ndim)]) + + # Evaluate activation function. + alpha = float(alpha) + x = spec.func(x, alpha=alpha) + + # Scale by gain. + gain = float(gain) + if gain != 1: + x = x * gain + + # Clamp. + if clamp >= 0: + x = x.clamp(-clamp, clamp) # pylint: disable=invalid-unary-operand-type + return x + +#---------------------------------------------------------------------------- + +_bias_act_cuda_cache = dict() + +def _bias_act_cuda(dim=1, act='linear', alpha=None, gain=None, clamp=None): + """Fast CUDA implementation of `bias_act()` using custom ops. + """ + # Parse arguments. + assert clamp is None or clamp >= 0 + spec = activation_funcs[act] + alpha = float(alpha if alpha is not None else spec.def_alpha) + gain = float(gain if gain is not None else spec.def_gain) + clamp = float(clamp if clamp is not None else -1) + + # Lookup from cache. + key = (dim, act, alpha, gain, clamp) + if key in _bias_act_cuda_cache: + return _bias_act_cuda_cache[key] + + # Forward op. + class BiasActCuda(torch.autograd.Function): + @staticmethod + def forward(ctx, x, b): # pylint: disable=arguments-differ + ctx.memory_format = torch.channels_last if x.ndim > 2 and x.stride()[1] == 1 else torch.contiguous_format + x = x.contiguous(memory_format=ctx.memory_format) + b = b.contiguous() if b is not None else _null_tensor + y = x + if act != 'linear' or gain != 1 or clamp >= 0 or b is not _null_tensor: + y = _plugin.bias_act(x, b, _null_tensor, _null_tensor, _null_tensor, 0, dim, spec.cuda_idx, alpha, gain, clamp) + ctx.save_for_backward( + x if 'x' in spec.ref or spec.has_2nd_grad else _null_tensor, + b if 'x' in spec.ref or spec.has_2nd_grad else _null_tensor, + y if 'y' in spec.ref else _null_tensor) + return y + + @staticmethod + def backward(ctx, dy): # pylint: disable=arguments-differ + dy = dy.contiguous(memory_format=ctx.memory_format) + x, b, y = ctx.saved_tensors + dx = None + db = None + + if ctx.needs_input_grad[0] or ctx.needs_input_grad[1]: + dx = dy + if act != 'linear' or gain != 1 or clamp >= 0: + dx = BiasActCudaGrad.apply(dy, x, b, y) + + if ctx.needs_input_grad[1]: + db = dx.sum([i for i in range(dx.ndim) if i != dim]) + + return dx, db + + # Backward op. + class BiasActCudaGrad(torch.autograd.Function): + @staticmethod + def forward(ctx, dy, x, b, y): # pylint: disable=arguments-differ + ctx.memory_format = torch.channels_last if dy.ndim > 2 and dy.stride()[1] == 1 else torch.contiguous_format + dx = _plugin.bias_act(dy, b, x, y, _null_tensor, 1, dim, spec.cuda_idx, alpha, gain, clamp) + ctx.save_for_backward( + dy if spec.has_2nd_grad else _null_tensor, + x, b, y) + return dx + + @staticmethod + def backward(ctx, d_dx): # pylint: disable=arguments-differ + d_dx = d_dx.contiguous(memory_format=ctx.memory_format) + dy, x, b, y = ctx.saved_tensors + d_dy = None + d_x = None + d_b = None + d_y = None + + if ctx.needs_input_grad[0]: + d_dy = BiasActCudaGrad.apply(d_dx, x, b, y) + + if spec.has_2nd_grad and (ctx.needs_input_grad[1] or ctx.needs_input_grad[2]): + d_x = _plugin.bias_act(d_dx, b, x, y, dy, 2, dim, spec.cuda_idx, alpha, gain, clamp) + + if spec.has_2nd_grad and ctx.needs_input_grad[2]: + d_b = d_x.sum([i for i in range(d_x.ndim) if i != dim]) + + return d_dy, d_x, d_b, d_y + + # Add to cache. + _bias_act_cuda_cache[key] = BiasActCuda + return BiasActCuda + +#---------------------------------------------------------------------------- diff --git a/models/stylegan2/torch_utils/ops/conv2d_gradfix.py b/models/stylegan2/torch_utils/ops/conv2d_gradfix.py new file mode 100644 index 0000000000000000000000000000000000000000..e95e10d0b1d0315a63a76446fd4c5c293c8bbc6d --- /dev/null +++ b/models/stylegan2/torch_utils/ops/conv2d_gradfix.py @@ -0,0 +1,170 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +# +# NVIDIA CORPORATION and its licensors retain all intellectual property +# and proprietary rights in and to this software, related documentation +# and any modifications thereto. Any use, reproduction, disclosure or +# distribution of this software and related documentation without an express +# license agreement from NVIDIA CORPORATION is strictly prohibited. + +"""Custom replacement for `torch.nn.functional.conv2d` that supports +arbitrarily high order gradients with zero performance penalty.""" + +import warnings +import contextlib +import torch + +# pylint: disable=redefined-builtin +# pylint: disable=arguments-differ +# pylint: disable=protected-access + +#---------------------------------------------------------------------------- + +enabled = False # Enable the custom op by setting this to true. +weight_gradients_disabled = False # Forcefully disable computation of gradients with respect to the weights. + +@contextlib.contextmanager +def no_weight_gradients(): + global weight_gradients_disabled + old = weight_gradients_disabled + weight_gradients_disabled = True + yield + weight_gradients_disabled = old + +#---------------------------------------------------------------------------- + +def conv2d(input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1): + if _should_use_custom_op(input): + return _conv2d_gradfix(transpose=False, weight_shape=weight.shape, stride=stride, padding=padding, output_padding=0, dilation=dilation, groups=groups).apply(input, weight, bias) + return torch.nn.functional.conv2d(input=input, weight=weight, bias=bias, stride=stride, padding=padding, dilation=dilation, groups=groups) + +def conv_transpose2d(input, weight, bias=None, stride=1, padding=0, output_padding=0, groups=1, dilation=1): + if _should_use_custom_op(input): + return _conv2d_gradfix(transpose=True, weight_shape=weight.shape, stride=stride, padding=padding, output_padding=output_padding, groups=groups, dilation=dilation).apply(input, weight, bias) + return torch.nn.functional.conv_transpose2d(input=input, weight=weight, bias=bias, stride=stride, padding=padding, output_padding=output_padding, groups=groups, dilation=dilation) + +#---------------------------------------------------------------------------- + +def _should_use_custom_op(input): + assert isinstance(input, torch.Tensor) + if (not enabled) or (not torch.backends.cudnn.enabled): + return False + if input.device.type != 'cuda': + return False + if any(torch.__version__.startswith(x) for x in ['1.7.', '1.8.', '1.9']): + return True + warnings.warn(f'conv2d_gradfix not supported on PyTorch {torch.__version__}. Falling back to torch.nn.functional.conv2d().') + return False + +def _tuple_of_ints(xs, ndim): + xs = tuple(xs) if isinstance(xs, (tuple, list)) else (xs,) * ndim + assert len(xs) == ndim + assert all(isinstance(x, int) for x in xs) + return xs + +#---------------------------------------------------------------------------- + +_conv2d_gradfix_cache = dict() + +def _conv2d_gradfix(transpose, weight_shape, stride, padding, output_padding, dilation, groups): + # Parse arguments. + ndim = 2 + weight_shape = tuple(weight_shape) + stride = _tuple_of_ints(stride, ndim) + padding = _tuple_of_ints(padding, ndim) + output_padding = _tuple_of_ints(output_padding, ndim) + dilation = _tuple_of_ints(dilation, ndim) + + # Lookup from cache. + key = (transpose, weight_shape, stride, padding, output_padding, dilation, groups) + if key in _conv2d_gradfix_cache: + return _conv2d_gradfix_cache[key] + + # Validate arguments. + assert groups >= 1 + assert len(weight_shape) == ndim + 2 + assert all(stride[i] >= 1 for i in range(ndim)) + assert all(padding[i] >= 0 for i in range(ndim)) + assert all(dilation[i] >= 0 for i in range(ndim)) + if not transpose: + assert all(output_padding[i] == 0 for i in range(ndim)) + else: # transpose + assert all(0 <= output_padding[i] < max(stride[i], dilation[i]) for i in range(ndim)) + + # Helpers. + common_kwargs = dict(stride=stride, padding=padding, dilation=dilation, groups=groups) + def calc_output_padding(input_shape, output_shape): + if transpose: + return [0, 0] + return [ + input_shape[i + 2] + - (output_shape[i + 2] - 1) * stride[i] + - (1 - 2 * padding[i]) + - dilation[i] * (weight_shape[i + 2] - 1) + for i in range(ndim) + ] + + # Forward & backward. + class Conv2d(torch.autograd.Function): + @staticmethod + def forward(ctx, input, weight, bias): + assert weight.shape == weight_shape + if not transpose: + output = torch.nn.functional.conv2d(input=input, weight=weight, bias=bias, **common_kwargs) + else: # transpose + output = torch.nn.functional.conv_transpose2d(input=input, weight=weight, bias=bias, output_padding=output_padding, **common_kwargs) + ctx.save_for_backward(input, weight) + return output + + @staticmethod + def backward(ctx, grad_output): + input, weight = ctx.saved_tensors + grad_input = None + grad_weight = None + grad_bias = None + + if ctx.needs_input_grad[0]: + p = calc_output_padding(input_shape=input.shape, output_shape=grad_output.shape) + grad_input = _conv2d_gradfix(transpose=(not transpose), weight_shape=weight_shape, output_padding=p, **common_kwargs).apply(grad_output, weight, None) + assert grad_input.shape == input.shape + + if ctx.needs_input_grad[1] and not weight_gradients_disabled: + grad_weight = Conv2dGradWeight.apply(grad_output, input) + assert grad_weight.shape == weight_shape + + if ctx.needs_input_grad[2]: + grad_bias = grad_output.sum([0, 2, 3]) + + return grad_input, grad_weight, grad_bias + + # Gradient with respect to the weights. + class Conv2dGradWeight(torch.autograd.Function): + @staticmethod + def forward(ctx, grad_output, input): + op = torch._C._jit_get_operation('aten::cudnn_convolution_backward_weight' if not transpose else 'aten::cudnn_convolution_transpose_backward_weight') + flags = [torch.backends.cudnn.benchmark, torch.backends.cudnn.deterministic, torch.backends.cudnn.allow_tf32] + grad_weight = op(weight_shape, grad_output, input, padding, stride, dilation, groups, *flags) + assert grad_weight.shape == weight_shape + ctx.save_for_backward(grad_output, input) + return grad_weight + + @staticmethod + def backward(ctx, grad2_grad_weight): + grad_output, input = ctx.saved_tensors + grad2_grad_output = None + grad2_input = None + + if ctx.needs_input_grad[0]: + grad2_grad_output = Conv2d.apply(input, grad2_grad_weight, None) + assert grad2_grad_output.shape == grad_output.shape + + if ctx.needs_input_grad[1]: + p = calc_output_padding(input_shape=input.shape, output_shape=grad_output.shape) + grad2_input = _conv2d_gradfix(transpose=(not transpose), weight_shape=weight_shape, output_padding=p, **common_kwargs).apply(grad_output, grad2_grad_weight, None) + assert grad2_input.shape == input.shape + + return grad2_grad_output, grad2_input + + _conv2d_gradfix_cache[key] = Conv2d + return Conv2d + +#---------------------------------------------------------------------------- diff --git a/models/stylegan2/torch_utils/ops/conv2d_resample.py b/models/stylegan2/torch_utils/ops/conv2d_resample.py new file mode 100644 index 0000000000000000000000000000000000000000..cd4750744c83354bab78704d4ef51ad1070fcc4a --- /dev/null +++ b/models/stylegan2/torch_utils/ops/conv2d_resample.py @@ -0,0 +1,156 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +# +# NVIDIA CORPORATION and its licensors retain all intellectual property +# and proprietary rights in and to this software, related documentation +# and any modifications thereto. Any use, reproduction, disclosure or +# distribution of this software and related documentation without an express +# license agreement from NVIDIA CORPORATION is strictly prohibited. + +"""2D convolution with optional up/downsampling.""" + +import torch + +from .. import misc +from . import conv2d_gradfix +from . import upfirdn2d +from .upfirdn2d import _parse_padding +from .upfirdn2d import _get_filter_size + +#---------------------------------------------------------------------------- + +def _get_weight_shape(w): + with misc.suppress_tracer_warnings(): # this value will be treated as a constant + shape = [int(sz) for sz in w.shape] + misc.assert_shape(w, shape) + return shape + +#---------------------------------------------------------------------------- + +def _conv2d_wrapper(x, w, stride=1, padding=0, groups=1, transpose=False, flip_weight=True): + """Wrapper for the underlying `conv2d()` and `conv_transpose2d()` implementations. + """ + out_channels, in_channels_per_group, kh, kw = _get_weight_shape(w) + + # Flip weight if requested. + if not flip_weight: # conv2d() actually performs correlation (flip_weight=True) not convolution (flip_weight=False). + w = w.flip([2, 3]) + + # Workaround performance pitfall in cuDNN 8.0.5, triggered when using + # 1x1 kernel + memory_format=channels_last + less than 64 channels. + if kw == 1 and kh == 1 and stride == 1 and padding in [0, [0, 0], (0, 0)] and not transpose: + if x.stride()[1] == 1 and min(out_channels, in_channels_per_group) < 64: + if out_channels <= 4 and groups == 1: + in_shape = x.shape + x = w.squeeze(3).squeeze(2) @ x.reshape([in_shape[0], in_channels_per_group, -1]) + x = x.reshape([in_shape[0], out_channels, in_shape[2], in_shape[3]]) + else: + x = x.to(memory_format=torch.contiguous_format) + w = w.to(memory_format=torch.contiguous_format) + x = conv2d_gradfix.conv2d(x, w, groups=groups) + return x.to(memory_format=torch.channels_last) + + # Otherwise => execute using conv2d_gradfix. + op = conv2d_gradfix.conv_transpose2d if transpose else conv2d_gradfix.conv2d + return op(x, w, stride=stride, padding=padding, groups=groups) + +#---------------------------------------------------------------------------- + +@misc.profiled_function +def conv2d_resample(x, w, f=None, up=1, down=1, padding=0, groups=1, flip_weight=True, flip_filter=False): + r"""2D convolution with optional up/downsampling. + + Padding is performed only once at the beginning, not between the operations. + + Args: + x: Input tensor of shape + `[batch_size, in_channels, in_height, in_width]`. + w: Weight tensor of shape + `[out_channels, in_channels//groups, kernel_height, kernel_width]`. + f: Low-pass filter for up/downsampling. Must be prepared beforehand by + calling upfirdn2d.setup_filter(). None = identity (default). + up: Integer upsampling factor (default: 1). + down: Integer downsampling factor (default: 1). + padding: Padding with respect to the upsampled image. Can be a single number + or a list/tuple `[x, y]` or `[x_before, x_after, y_before, y_after]` + (default: 0). + groups: Split input channels into N groups (default: 1). + flip_weight: False = convolution, True = correlation (default: True). + flip_filter: False = convolution, True = correlation (default: False). + + Returns: + Tensor of the shape `[batch_size, num_channels, out_height, out_width]`. + """ + # Validate arguments. + assert isinstance(x, torch.Tensor) and (x.ndim == 4) + assert isinstance(w, torch.Tensor) and (w.ndim == 4) and (w.dtype == x.dtype) + assert f is None or (isinstance(f, torch.Tensor) and f.ndim in [1, 2] and f.dtype == torch.float32) + assert isinstance(up, int) and (up >= 1) + assert isinstance(down, int) and (down >= 1) + assert isinstance(groups, int) and (groups >= 1) + out_channels, in_channels_per_group, kh, kw = _get_weight_shape(w) + fw, fh = _get_filter_size(f) + px0, px1, py0, py1 = _parse_padding(padding) + + # Adjust padding to account for up/downsampling. + if up > 1: + px0 += (fw + up - 1) // 2 + px1 += (fw - up) // 2 + py0 += (fh + up - 1) // 2 + py1 += (fh - up) // 2 + if down > 1: + px0 += (fw - down + 1) // 2 + px1 += (fw - down) // 2 + py0 += (fh - down + 1) // 2 + py1 += (fh - down) // 2 + + # Fast path: 1x1 convolution with downsampling only => downsample first, then convolve. + if kw == 1 and kh == 1 and (down > 1 and up == 1): + x = upfirdn2d.upfirdn2d(x=x, f=f, down=down, padding=[px0,px1,py0,py1], flip_filter=flip_filter) + x = _conv2d_wrapper(x=x, w=w, groups=groups, flip_weight=flip_weight) + return x + + # Fast path: 1x1 convolution with upsampling only => convolve first, then upsample. + if kw == 1 and kh == 1 and (up > 1 and down == 1): + x = _conv2d_wrapper(x=x, w=w, groups=groups, flip_weight=flip_weight) + x = upfirdn2d.upfirdn2d(x=x, f=f, up=up, padding=[px0,px1,py0,py1], gain=up**2, flip_filter=flip_filter) + return x + + # Fast path: downsampling only => use strided convolution. + if down > 1 and up == 1: + x = upfirdn2d.upfirdn2d(x=x, f=f, padding=[px0,px1,py0,py1], flip_filter=flip_filter) + x = _conv2d_wrapper(x=x, w=w, stride=down, groups=groups, flip_weight=flip_weight) + return x + + # Fast path: upsampling with optional downsampling => use transpose strided convolution. + if up > 1: + if groups == 1: + w = w.transpose(0, 1) + else: + w = w.reshape(groups, out_channels // groups, in_channels_per_group, kh, kw) + w = w.transpose(1, 2) + w = w.reshape(groups * in_channels_per_group, out_channels // groups, kh, kw) + px0 -= kw - 1 + px1 -= kw - up + py0 -= kh - 1 + py1 -= kh - up + pxt = max(min(-px0, -px1), 0) + pyt = max(min(-py0, -py1), 0) + x = _conv2d_wrapper(x=x, w=w, stride=up, padding=[pyt,pxt], groups=groups, transpose=True, flip_weight=(not flip_weight)) + x = upfirdn2d.upfirdn2d(x=x, f=f, padding=[px0+pxt,px1+pxt,py0+pyt,py1+pyt], gain=up**2, flip_filter=flip_filter) + if down > 1: + x = upfirdn2d.upfirdn2d(x=x, f=f, down=down, flip_filter=flip_filter) + return x + + # Fast path: no up/downsampling, padding supported by the underlying implementation => use plain conv2d. + if up == 1 and down == 1: + if px0 == px1 and py0 == py1 and px0 >= 0 and py0 >= 0: + return _conv2d_wrapper(x=x, w=w, padding=[py0,px0], groups=groups, flip_weight=flip_weight) + + # Fallback: Generic reference implementation. + x = upfirdn2d.upfirdn2d(x=x, f=(f if up > 1 else None), up=up, padding=[px0,px1,py0,py1], gain=up**2, flip_filter=flip_filter) + x = _conv2d_wrapper(x=x, w=w, groups=groups, flip_weight=flip_weight) + if down > 1: + x = upfirdn2d.upfirdn2d(x=x, f=f, down=down, flip_filter=flip_filter) + return x + +#---------------------------------------------------------------------------- diff --git a/models/stylegan2/torch_utils/ops/fma.py b/models/stylegan2/torch_utils/ops/fma.py new file mode 100644 index 0000000000000000000000000000000000000000..2eeac58a626c49231e04122b93e321ada954c5d3 --- /dev/null +++ b/models/stylegan2/torch_utils/ops/fma.py @@ -0,0 +1,60 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +# +# NVIDIA CORPORATION and its licensors retain all intellectual property +# and proprietary rights in and to this software, related documentation +# and any modifications thereto. Any use, reproduction, disclosure or +# distribution of this software and related documentation without an express +# license agreement from NVIDIA CORPORATION is strictly prohibited. + +"""Fused multiply-add, with slightly faster gradients than `torch.addcmul()`.""" + +import torch + +#---------------------------------------------------------------------------- + +def fma(a, b, c): # => a * b + c + return _FusedMultiplyAdd.apply(a, b, c) + +#---------------------------------------------------------------------------- + +class _FusedMultiplyAdd(torch.autograd.Function): # a * b + c + @staticmethod + def forward(ctx, a, b, c): # pylint: disable=arguments-differ + out = torch.addcmul(c, a, b) + ctx.save_for_backward(a, b) + ctx.c_shape = c.shape + return out + + @staticmethod + def backward(ctx, dout): # pylint: disable=arguments-differ + a, b = ctx.saved_tensors + c_shape = ctx.c_shape + da = None + db = None + dc = None + + if ctx.needs_input_grad[0]: + da = _unbroadcast(dout * b, a.shape) + + if ctx.needs_input_grad[1]: + db = _unbroadcast(dout * a, b.shape) + + if ctx.needs_input_grad[2]: + dc = _unbroadcast(dout, c_shape) + + return da, db, dc + +#---------------------------------------------------------------------------- + +def _unbroadcast(x, shape): + extra_dims = x.ndim - len(shape) + assert extra_dims >= 0 + dim = [i for i in range(x.ndim) if x.shape[i] > 1 and (i < extra_dims or shape[i - extra_dims] == 1)] + if len(dim): + x = x.sum(dim=dim, keepdim=True) + if extra_dims: + x = x.reshape(-1, *x.shape[extra_dims+1:]) + assert x.shape == shape + return x + +#---------------------------------------------------------------------------- diff --git a/models/stylegan2/torch_utils/ops/grid_sample_gradfix.py b/models/stylegan2/torch_utils/ops/grid_sample_gradfix.py new file mode 100644 index 0000000000000000000000000000000000000000..ca6b3413ea72a734703c34382c023b84523601fd --- /dev/null +++ b/models/stylegan2/torch_utils/ops/grid_sample_gradfix.py @@ -0,0 +1,83 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +# +# NVIDIA CORPORATION and its licensors retain all intellectual property +# and proprietary rights in and to this software, related documentation +# and any modifications thereto. Any use, reproduction, disclosure or +# distribution of this software and related documentation without an express +# license agreement from NVIDIA CORPORATION is strictly prohibited. + +"""Custom replacement for `torch.nn.functional.grid_sample` that +supports arbitrarily high order gradients between the input and output. +Only works on 2D images and assumes +`mode='bilinear'`, `padding_mode='zeros'`, `align_corners=False`.""" + +import warnings +import torch + +# pylint: disable=redefined-builtin +# pylint: disable=arguments-differ +# pylint: disable=protected-access + +#---------------------------------------------------------------------------- + +enabled = False # Enable the custom op by setting this to true. + +#---------------------------------------------------------------------------- + +def grid_sample(input, grid): + if _should_use_custom_op(): + return _GridSample2dForward.apply(input, grid) + return torch.nn.functional.grid_sample(input=input, grid=grid, mode='bilinear', padding_mode='zeros', align_corners=False) + +#---------------------------------------------------------------------------- + +def _should_use_custom_op(): + if not enabled: + return False + if any(torch.__version__.startswith(x) for x in ['1.7.', '1.8.', '1.9']): + return True + warnings.warn(f'grid_sample_gradfix not supported on PyTorch {torch.__version__}. Falling back to torch.nn.functional.grid_sample().') + return False + +#---------------------------------------------------------------------------- + +class _GridSample2dForward(torch.autograd.Function): + @staticmethod + def forward(ctx, input, grid): + assert input.ndim == 4 + assert grid.ndim == 4 + output = torch.nn.functional.grid_sample(input=input, grid=grid, mode='bilinear', padding_mode='zeros', align_corners=False) + ctx.save_for_backward(input, grid) + return output + + @staticmethod + def backward(ctx, grad_output): + input, grid = ctx.saved_tensors + grad_input, grad_grid = _GridSample2dBackward.apply(grad_output, input, grid) + return grad_input, grad_grid + +#---------------------------------------------------------------------------- + +class _GridSample2dBackward(torch.autograd.Function): + @staticmethod + def forward(ctx, grad_output, input, grid): + op = torch._C._jit_get_operation('aten::grid_sampler_2d_backward') + grad_input, grad_grid = op(grad_output, input, grid, 0, 0, False) + ctx.save_for_backward(grid) + return grad_input, grad_grid + + @staticmethod + def backward(ctx, grad2_grad_input, grad2_grad_grid): + _ = grad2_grad_grid # unused + grid, = ctx.saved_tensors + grad2_grad_output = None + grad2_input = None + grad2_grid = None + + if ctx.needs_input_grad[0]: + grad2_grad_output = _GridSample2dForward.apply(grad2_grad_input, grid) + + assert not ctx.needs_input_grad[2] + return grad2_grad_output, grad2_input, grad2_grid + +#---------------------------------------------------------------------------- diff --git a/models/stylegan2/torch_utils/ops/upfirdn2d.cpp b/models/stylegan2/torch_utils/ops/upfirdn2d.cpp new file mode 100644 index 0000000000000000000000000000000000000000..2d7177fc60040751d20e9a8da0301fa3ab64968a --- /dev/null +++ b/models/stylegan2/torch_utils/ops/upfirdn2d.cpp @@ -0,0 +1,103 @@ +// Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +// +// NVIDIA CORPORATION and its licensors retain all intellectual property +// and proprietary rights in and to this software, related documentation +// and any modifications thereto. Any use, reproduction, disclosure or +// distribution of this software and related documentation without an express +// license agreement from NVIDIA CORPORATION is strictly prohibited. + +#include +#include +#include +#include "upfirdn2d.h" + +//------------------------------------------------------------------------ + +static torch::Tensor upfirdn2d(torch::Tensor x, torch::Tensor f, int upx, int upy, int downx, int downy, int padx0, int padx1, int pady0, int pady1, bool flip, float gain) +{ + // Validate arguments. + TORCH_CHECK(x.is_cuda(), "x must reside on CUDA device"); + TORCH_CHECK(f.device() == x.device(), "f must reside on the same device as x"); + TORCH_CHECK(f.dtype() == torch::kFloat, "f must be float32"); + TORCH_CHECK(x.numel() <= INT_MAX, "x is too large"); + TORCH_CHECK(f.numel() <= INT_MAX, "f is too large"); + TORCH_CHECK(x.dim() == 4, "x must be rank 4"); + TORCH_CHECK(f.dim() == 2, "f must be rank 2"); + TORCH_CHECK(f.size(0) >= 1 && f.size(1) >= 1, "f must be at least 1x1"); + TORCH_CHECK(upx >= 1 && upy >= 1, "upsampling factor must be at least 1"); + TORCH_CHECK(downx >= 1 && downy >= 1, "downsampling factor must be at least 1"); + + // Create output tensor. + const at::cuda::OptionalCUDAGuard device_guard(device_of(x)); + int outW = ((int)x.size(3) * upx + padx0 + padx1 - (int)f.size(1) + downx) / downx; + int outH = ((int)x.size(2) * upy + pady0 + pady1 - (int)f.size(0) + downy) / downy; + TORCH_CHECK(outW >= 1 && outH >= 1, "output must be at least 1x1"); + torch::Tensor y = torch::empty({x.size(0), x.size(1), outH, outW}, x.options(), x.suggest_memory_format()); + TORCH_CHECK(y.numel() <= INT_MAX, "output is too large"); + + // Initialize CUDA kernel parameters. + upfirdn2d_kernel_params p; + p.x = x.data_ptr(); + p.f = f.data_ptr(); + p.y = y.data_ptr(); + p.up = make_int2(upx, upy); + p.down = make_int2(downx, downy); + p.pad0 = make_int2(padx0, pady0); + p.flip = (flip) ? 1 : 0; + p.gain = gain; + p.inSize = make_int4((int)x.size(3), (int)x.size(2), (int)x.size(1), (int)x.size(0)); + p.inStride = make_int4((int)x.stride(3), (int)x.stride(2), (int)x.stride(1), (int)x.stride(0)); + p.filterSize = make_int2((int)f.size(1), (int)f.size(0)); + p.filterStride = make_int2((int)f.stride(1), (int)f.stride(0)); + p.outSize = make_int4((int)y.size(3), (int)y.size(2), (int)y.size(1), (int)y.size(0)); + p.outStride = make_int4((int)y.stride(3), (int)y.stride(2), (int)y.stride(1), (int)y.stride(0)); + p.sizeMajor = (p.inStride.z == 1) ? p.inSize.w : p.inSize.w * p.inSize.z; + p.sizeMinor = (p.inStride.z == 1) ? p.inSize.z : 1; + + // Choose CUDA kernel. + upfirdn2d_kernel_spec spec; + AT_DISPATCH_FLOATING_TYPES_AND_HALF(x.scalar_type(), "upfirdn2d_cuda", [&] + { + spec = choose_upfirdn2d_kernel(p); + }); + + // Set looping options. + p.loopMajor = (p.sizeMajor - 1) / 16384 + 1; + p.loopMinor = spec.loopMinor; + p.loopX = spec.loopX; + p.launchMinor = (p.sizeMinor - 1) / p.loopMinor + 1; + p.launchMajor = (p.sizeMajor - 1) / p.loopMajor + 1; + + // Compute grid size. + dim3 blockSize, gridSize; + if (spec.tileOutW < 0) // large + { + blockSize = dim3(4, 32, 1); + gridSize = dim3( + ((p.outSize.y - 1) / blockSize.x + 1) * p.launchMinor, + (p.outSize.x - 1) / (blockSize.y * p.loopX) + 1, + p.launchMajor); + } + else // small + { + blockSize = dim3(256, 1, 1); + gridSize = dim3( + ((p.outSize.y - 1) / spec.tileOutH + 1) * p.launchMinor, + (p.outSize.x - 1) / (spec.tileOutW * p.loopX) + 1, + p.launchMajor); + } + + // Launch CUDA kernel. + void* args[] = {&p}; + AT_CUDA_CHECK(cudaLaunchKernel(spec.kernel, gridSize, blockSize, args, 0, at::cuda::getCurrentCUDAStream())); + return y; +} + +//------------------------------------------------------------------------ + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) +{ + m.def("upfirdn2d", &upfirdn2d); +} + +//------------------------------------------------------------------------ diff --git a/models/stylegan2/torch_utils/ops/upfirdn2d.cu b/models/stylegan2/torch_utils/ops/upfirdn2d.cu new file mode 100644 index 0000000000000000000000000000000000000000..ebdd9879f4bb16fc57a23cbc81f9de8ef54e4916 --- /dev/null +++ b/models/stylegan2/torch_utils/ops/upfirdn2d.cu @@ -0,0 +1,350 @@ +// Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +// +// NVIDIA CORPORATION and its licensors retain all intellectual property +// and proprietary rights in and to this software, related documentation +// and any modifications thereto. Any use, reproduction, disclosure or +// distribution of this software and related documentation without an express +// license agreement from NVIDIA CORPORATION is strictly prohibited. + +#include +#include "upfirdn2d.h" + +//------------------------------------------------------------------------ +// Helpers. + +template struct InternalType; +template <> struct InternalType { typedef double scalar_t; }; +template <> struct InternalType { typedef float scalar_t; }; +template <> struct InternalType { typedef float scalar_t; }; + +static __device__ __forceinline__ int floor_div(int a, int b) +{ + int t = 1 - a / b; + return (a + t * b) / b - t; +} + +//------------------------------------------------------------------------ +// Generic CUDA implementation for large filters. + +template static __global__ void upfirdn2d_kernel_large(upfirdn2d_kernel_params p) +{ + typedef typename InternalType::scalar_t scalar_t; + + // Calculate thread index. + int minorBase = blockIdx.x * blockDim.x + threadIdx.x; + int outY = minorBase / p.launchMinor; + minorBase -= outY * p.launchMinor; + int outXBase = blockIdx.y * p.loopX * blockDim.y + threadIdx.y; + int majorBase = blockIdx.z * p.loopMajor; + if (outXBase >= p.outSize.x | outY >= p.outSize.y | majorBase >= p.sizeMajor) + return; + + // Setup Y receptive field. + int midY = outY * p.down.y + p.up.y - 1 - p.pad0.y; + int inY = min(max(floor_div(midY, p.up.y), 0), p.inSize.y); + int h = min(max(floor_div(midY + p.filterSize.y, p.up.y), 0), p.inSize.y) - inY; + int filterY = midY + p.filterSize.y - (inY + 1) * p.up.y; + if (p.flip) + filterY = p.filterSize.y - 1 - filterY; + + // Loop over major, minor, and X. + for (int majorIdx = 0, major = majorBase; majorIdx < p.loopMajor & major < p.sizeMajor; majorIdx++, major++) + for (int minorIdx = 0, minor = minorBase; minorIdx < p.loopMinor & minor < p.sizeMinor; minorIdx++, minor += p.launchMinor) + { + int nc = major * p.sizeMinor + minor; + int n = nc / p.inSize.z; + int c = nc - n * p.inSize.z; + for (int loopX = 0, outX = outXBase; loopX < p.loopX & outX < p.outSize.x; loopX++, outX += blockDim.y) + { + // Setup X receptive field. + int midX = outX * p.down.x + p.up.x - 1 - p.pad0.x; + int inX = min(max(floor_div(midX, p.up.x), 0), p.inSize.x); + int w = min(max(floor_div(midX + p.filterSize.x, p.up.x), 0), p.inSize.x) - inX; + int filterX = midX + p.filterSize.x - (inX + 1) * p.up.x; + if (p.flip) + filterX = p.filterSize.x - 1 - filterX; + + // Initialize pointers. + const T* xp = &((const T*)p.x)[inX * p.inStride.x + inY * p.inStride.y + c * p.inStride.z + n * p.inStride.w]; + const float* fp = &p.f[filterX * p.filterStride.x + filterY * p.filterStride.y]; + int filterStepX = ((p.flip) ? p.up.x : -p.up.x) * p.filterStride.x; + int filterStepY = ((p.flip) ? p.up.y : -p.up.y) * p.filterStride.y; + + // Inner loop. + scalar_t v = 0; + for (int y = 0; y < h; y++) + { + for (int x = 0; x < w; x++) + { + v += (scalar_t)(*xp) * (scalar_t)(*fp); + xp += p.inStride.x; + fp += filterStepX; + } + xp += p.inStride.y - w * p.inStride.x; + fp += filterStepY - w * filterStepX; + } + + // Store result. + v *= p.gain; + ((T*)p.y)[outX * p.outStride.x + outY * p.outStride.y + c * p.outStride.z + n * p.outStride.w] = (T)v; + } + } +} + +//------------------------------------------------------------------------ +// Specialized CUDA implementation for small filters. + +template +static __global__ void upfirdn2d_kernel_small(upfirdn2d_kernel_params p) +{ + typedef typename InternalType::scalar_t scalar_t; + const int tileInW = ((tileOutW - 1) * downx + filterW - 1) / upx + 1; + const int tileInH = ((tileOutH - 1) * downy + filterH - 1) / upy + 1; + __shared__ volatile scalar_t sf[filterH][filterW]; + __shared__ volatile scalar_t sx[tileInH][tileInW][loopMinor]; + + // Calculate tile index. + int minorBase = blockIdx.x; + int tileOutY = minorBase / p.launchMinor; + minorBase -= tileOutY * p.launchMinor; + minorBase *= loopMinor; + tileOutY *= tileOutH; + int tileOutXBase = blockIdx.y * p.loopX * tileOutW; + int majorBase = blockIdx.z * p.loopMajor; + if (tileOutXBase >= p.outSize.x | tileOutY >= p.outSize.y | majorBase >= p.sizeMajor) + return; + + // Load filter (flipped). + for (int tapIdx = threadIdx.x; tapIdx < filterH * filterW; tapIdx += blockDim.x) + { + int fy = tapIdx / filterW; + int fx = tapIdx - fy * filterW; + scalar_t v = 0; + if (fx < p.filterSize.x & fy < p.filterSize.y) + { + int ffx = (p.flip) ? fx : p.filterSize.x - 1 - fx; + int ffy = (p.flip) ? fy : p.filterSize.y - 1 - fy; + v = (scalar_t)p.f[ffx * p.filterStride.x + ffy * p.filterStride.y]; + } + sf[fy][fx] = v; + } + + // Loop over major and X. + for (int majorIdx = 0, major = majorBase; majorIdx < p.loopMajor & major < p.sizeMajor; majorIdx++, major++) + { + int baseNC = major * p.sizeMinor + minorBase; + int n = baseNC / p.inSize.z; + int baseC = baseNC - n * p.inSize.z; + for (int loopX = 0, tileOutX = tileOutXBase; loopX < p.loopX & tileOutX < p.outSize.x; loopX++, tileOutX += tileOutW) + { + // Load input pixels. + int tileMidX = tileOutX * downx + upx - 1 - p.pad0.x; + int tileMidY = tileOutY * downy + upy - 1 - p.pad0.y; + int tileInX = floor_div(tileMidX, upx); + int tileInY = floor_div(tileMidY, upy); + __syncthreads(); + for (int inIdx = threadIdx.x; inIdx < tileInH * tileInW * loopMinor; inIdx += blockDim.x) + { + int relC = inIdx; + int relInX = relC / loopMinor; + int relInY = relInX / tileInW; + relC -= relInX * loopMinor; + relInX -= relInY * tileInW; + int c = baseC + relC; + int inX = tileInX + relInX; + int inY = tileInY + relInY; + scalar_t v = 0; + if (inX >= 0 & inY >= 0 & inX < p.inSize.x & inY < p.inSize.y & c < p.inSize.z) + v = (scalar_t)((const T*)p.x)[inX * p.inStride.x + inY * p.inStride.y + c * p.inStride.z + n * p.inStride.w]; + sx[relInY][relInX][relC] = v; + } + + // Loop over output pixels. + __syncthreads(); + for (int outIdx = threadIdx.x; outIdx < tileOutH * tileOutW * loopMinor; outIdx += blockDim.x) + { + int relC = outIdx; + int relOutX = relC / loopMinor; + int relOutY = relOutX / tileOutW; + relC -= relOutX * loopMinor; + relOutX -= relOutY * tileOutW; + int c = baseC + relC; + int outX = tileOutX + relOutX; + int outY = tileOutY + relOutY; + + // Setup receptive field. + int midX = tileMidX + relOutX * downx; + int midY = tileMidY + relOutY * downy; + int inX = floor_div(midX, upx); + int inY = floor_div(midY, upy); + int relInX = inX - tileInX; + int relInY = inY - tileInY; + int filterX = (inX + 1) * upx - midX - 1; // flipped + int filterY = (inY + 1) * upy - midY - 1; // flipped + + // Inner loop. + if (outX < p.outSize.x & outY < p.outSize.y & c < p.outSize.z) + { + scalar_t v = 0; + #pragma unroll + for (int y = 0; y < filterH / upy; y++) + #pragma unroll + for (int x = 0; x < filterW / upx; x++) + v += sx[relInY + y][relInX + x][relC] * sf[filterY + y * upy][filterX + x * upx]; + v *= p.gain; + ((T*)p.y)[outX * p.outStride.x + outY * p.outStride.y + c * p.outStride.z + n * p.outStride.w] = (T)v; + } + } + } + } +} + +//------------------------------------------------------------------------ +// CUDA kernel selection. + +template upfirdn2d_kernel_spec choose_upfirdn2d_kernel(const upfirdn2d_kernel_params& p) +{ + int s = p.inStride.z, fx = p.filterSize.x, fy = p.filterSize.y; + + upfirdn2d_kernel_spec spec = {(void*)upfirdn2d_kernel_large, -1,-1,1, 4}; // contiguous + if (s == 1) spec = {(void*)upfirdn2d_kernel_large, -1,-1,4, 1}; // channels_last + + if (s != 1 && p.up.x == 1 && p.up.y == 1 && p.down.x == 1 && p.down.y == 1) // contiguous + { + if (fx <= 7 && fy <= 7 ) spec = {(void*)upfirdn2d_kernel_small, 64,16,1, 1}; + if (fx <= 6 && fy <= 6 ) spec = {(void*)upfirdn2d_kernel_small, 64,16,1, 1}; + if (fx <= 5 && fy <= 5 ) spec = {(void*)upfirdn2d_kernel_small, 64,16,1, 1}; + if (fx <= 4 && fy <= 4 ) spec = {(void*)upfirdn2d_kernel_small, 64,16,1, 1}; + if (fx <= 3 && fy <= 3 ) spec = {(void*)upfirdn2d_kernel_small, 64,16,1, 1}; + if (fx <= 24 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small, 128,8,1, 1}; + if (fx <= 20 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small, 128,8,1, 1}; + if (fx <= 16 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small, 128,8,1, 1}; + if (fx <= 12 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small, 128,8,1, 1}; + if (fx <= 8 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small, 128,8,1, 1}; + if (fx <= 1 && fy <= 24) spec = {(void*)upfirdn2d_kernel_small, 32,32,1, 1}; + if (fx <= 1 && fy <= 20) spec = {(void*)upfirdn2d_kernel_small, 32,32,1, 1}; + if (fx <= 1 && fy <= 16) spec = {(void*)upfirdn2d_kernel_small, 32,32,1, 1}; + if (fx <= 1 && fy <= 12) spec = {(void*)upfirdn2d_kernel_small, 32,32,1, 1}; + if (fx <= 1 && fy <= 8 ) spec = {(void*)upfirdn2d_kernel_small, 32,32,1, 1}; + } + if (s == 1 && p.up.x == 1 && p.up.y == 1 && p.down.x == 1 && p.down.y == 1) // channels_last + { + if (fx <= 7 && fy <= 7 ) spec = {(void*)upfirdn2d_kernel_small, 16,16,8, 1}; + if (fx <= 6 && fy <= 6 ) spec = {(void*)upfirdn2d_kernel_small, 16,16,8, 1}; + if (fx <= 5 && fy <= 5 ) spec = {(void*)upfirdn2d_kernel_small, 16,16,8, 1}; + if (fx <= 4 && fy <= 4 ) spec = {(void*)upfirdn2d_kernel_small, 16,16,8, 1}; + if (fx <= 3 && fy <= 3 ) spec = {(void*)upfirdn2d_kernel_small, 16,16,8, 1}; + if (fx <= 24 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small, 128,1,16, 1}; + if (fx <= 20 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small, 128,1,16, 1}; + if (fx <= 16 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small, 128,1,16, 1}; + if (fx <= 12 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small, 128,1,16, 1}; + if (fx <= 8 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small, 128,1,16, 1}; + if (fx <= 1 && fy <= 24) spec = {(void*)upfirdn2d_kernel_small, 1,128,16, 1}; + if (fx <= 1 && fy <= 20) spec = {(void*)upfirdn2d_kernel_small, 1,128,16, 1}; + if (fx <= 1 && fy <= 16) spec = {(void*)upfirdn2d_kernel_small, 1,128,16, 1}; + if (fx <= 1 && fy <= 12) spec = {(void*)upfirdn2d_kernel_small, 1,128,16, 1}; + if (fx <= 1 && fy <= 8 ) spec = {(void*)upfirdn2d_kernel_small, 1,128,16, 1}; + } + if (s != 1 && p.up.x == 2 && p.up.y == 2 && p.down.x == 1 && p.down.y == 1) // contiguous + { + if (fx <= 8 && fy <= 8 ) spec = {(void*)upfirdn2d_kernel_small, 64,16,1, 1}; + if (fx <= 6 && fy <= 6 ) spec = {(void*)upfirdn2d_kernel_small, 64,16,1, 1}; + if (fx <= 4 && fy <= 4 ) spec = {(void*)upfirdn2d_kernel_small, 64,16,1, 1}; + if (fx <= 2 && fy <= 2 ) spec = {(void*)upfirdn2d_kernel_small, 64,16,1, 1}; + } + if (s == 1 && p.up.x == 2 && p.up.y == 2 && p.down.x == 1 && p.down.y == 1) // channels_last + { + if (fx <= 8 && fy <= 8 ) spec = {(void*)upfirdn2d_kernel_small, 16,16,8, 1}; + if (fx <= 6 && fy <= 6 ) spec = {(void*)upfirdn2d_kernel_small, 16,16,8, 1}; + if (fx <= 4 && fy <= 4 ) spec = {(void*)upfirdn2d_kernel_small, 16,16,8, 1}; + if (fx <= 2 && fy <= 2 ) spec = {(void*)upfirdn2d_kernel_small, 16,16,8, 1}; + } + if (s != 1 && p.up.x == 2 && p.up.y == 1 && p.down.x == 1 && p.down.y == 1) // contiguous + { + if (fx <= 24 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small, 128,8,1, 1}; + if (fx <= 20 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small, 128,8,1, 1}; + if (fx <= 16 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small, 128,8,1, 1}; + if (fx <= 12 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small, 128,8,1, 1}; + if (fx <= 8 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small, 128,8,1, 1}; + } + if (s == 1 && p.up.x == 2 && p.up.y == 1 && p.down.x == 1 && p.down.y == 1) // channels_last + { + if (fx <= 24 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small, 128,1,16, 1}; + if (fx <= 20 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small, 128,1,16, 1}; + if (fx <= 16 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small, 128,1,16, 1}; + if (fx <= 12 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small, 128,1,16, 1}; + if (fx <= 8 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small, 128,1,16, 1}; + } + if (s != 1 && p.up.x == 1 && p.up.y == 2 && p.down.x == 1 && p.down.y == 1) // contiguous + { + if (fx <= 1 && fy <= 24) spec = {(void*)upfirdn2d_kernel_small, 32,32,1, 1}; + if (fx <= 1 && fy <= 20) spec = {(void*)upfirdn2d_kernel_small, 32,32,1, 1}; + if (fx <= 1 && fy <= 16) spec = {(void*)upfirdn2d_kernel_small, 32,32,1, 1}; + if (fx <= 1 && fy <= 12) spec = {(void*)upfirdn2d_kernel_small, 32,32,1, 1}; + if (fx <= 1 && fy <= 8 ) spec = {(void*)upfirdn2d_kernel_small, 32,32,1, 1}; + } + if (s == 1 && p.up.x == 1 && p.up.y == 2 && p.down.x == 1 && p.down.y == 1) // channels_last + { + if (fx <= 1 && fy <= 24) spec = {(void*)upfirdn2d_kernel_small, 1,128,16, 1}; + if (fx <= 1 && fy <= 20) spec = {(void*)upfirdn2d_kernel_small, 1,128,16, 1}; + if (fx <= 1 && fy <= 16) spec = {(void*)upfirdn2d_kernel_small, 1,128,16, 1}; + if (fx <= 1 && fy <= 12) spec = {(void*)upfirdn2d_kernel_small, 1,128,16, 1}; + if (fx <= 1 && fy <= 8 ) spec = {(void*)upfirdn2d_kernel_small, 1,128,16, 1}; + } + if (s != 1 && p.up.x == 1 && p.up.y == 1 && p.down.x == 2 && p.down.y == 2) // contiguous + { + if (fx <= 8 && fy <= 8 ) spec = {(void*)upfirdn2d_kernel_small, 32,8,1, 1}; + if (fx <= 6 && fy <= 6 ) spec = {(void*)upfirdn2d_kernel_small, 32,8,1, 1}; + if (fx <= 4 && fy <= 4 ) spec = {(void*)upfirdn2d_kernel_small, 32,8,1, 1}; + if (fx <= 2 && fy <= 2 ) spec = {(void*)upfirdn2d_kernel_small, 32,8,1, 1}; + } + if (s == 1 && p.up.x == 1 && p.up.y == 1 && p.down.x == 2 && p.down.y == 2) // channels_last + { + if (fx <= 8 && fy <= 8 ) spec = {(void*)upfirdn2d_kernel_small, 8,8,8, 1}; + if (fx <= 6 && fy <= 6 ) spec = {(void*)upfirdn2d_kernel_small, 8,8,8, 1}; + if (fx <= 4 && fy <= 4 ) spec = {(void*)upfirdn2d_kernel_small, 8,8,8, 1}; + if (fx <= 2 && fy <= 2 ) spec = {(void*)upfirdn2d_kernel_small, 8,8,8, 1}; + } + if (s != 1 && p.up.x == 1 && p.up.y == 1 && p.down.x == 2 && p.down.y == 1) // contiguous + { + if (fx <= 24 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small, 64,8,1, 1}; + if (fx <= 20 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small, 64,8,1, 1}; + if (fx <= 16 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small, 64,8,1, 1}; + if (fx <= 12 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small, 64,8,1, 1}; + if (fx <= 8 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small, 64,8,1, 1}; + } + if (s == 1 && p.up.x == 1 && p.up.y == 1 && p.down.x == 2 && p.down.y == 1) // channels_last + { + if (fx <= 24 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small, 64,1,8, 1}; + if (fx <= 20 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small, 64,1,8, 1}; + if (fx <= 16 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small, 64,1,8, 1}; + if (fx <= 12 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small, 64,1,8, 1}; + if (fx <= 8 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small, 64,1,8, 1}; + } + if (s != 1 && p.up.x == 1 && p.up.y == 1 && p.down.x == 1 && p.down.y == 2) // contiguous + { + if (fx <= 1 && fy <= 24) spec = {(void*)upfirdn2d_kernel_small, 32,16,1, 1}; + if (fx <= 1 && fy <= 20) spec = {(void*)upfirdn2d_kernel_small, 32,16,1, 1}; + if (fx <= 1 && fy <= 16) spec = {(void*)upfirdn2d_kernel_small, 32,16,1, 1}; + if (fx <= 1 && fy <= 12) spec = {(void*)upfirdn2d_kernel_small, 32,16,1, 1}; + if (fx <= 1 && fy <= 8 ) spec = {(void*)upfirdn2d_kernel_small, 32,16,1, 1}; + } + if (s == 1 && p.up.x == 1 && p.up.y == 1 && p.down.x == 1 && p.down.y == 2) // channels_last + { + if (fx <= 1 && fy <= 24) spec = {(void*)upfirdn2d_kernel_small, 1,64,8, 1}; + if (fx <= 1 && fy <= 20) spec = {(void*)upfirdn2d_kernel_small, 1,64,8, 1}; + if (fx <= 1 && fy <= 16) spec = {(void*)upfirdn2d_kernel_small, 1,64,8, 1}; + if (fx <= 1 && fy <= 12) spec = {(void*)upfirdn2d_kernel_small, 1,64,8, 1}; + if (fx <= 1 && fy <= 8 ) spec = {(void*)upfirdn2d_kernel_small, 1,64,8, 1}; + } + return spec; +} + +//------------------------------------------------------------------------ +// Template specializations. + +template upfirdn2d_kernel_spec choose_upfirdn2d_kernel (const upfirdn2d_kernel_params& p); +template upfirdn2d_kernel_spec choose_upfirdn2d_kernel (const upfirdn2d_kernel_params& p); +template upfirdn2d_kernel_spec choose_upfirdn2d_kernel(const upfirdn2d_kernel_params& p); + +//------------------------------------------------------------------------ diff --git a/models/stylegan2/torch_utils/ops/upfirdn2d.h b/models/stylegan2/torch_utils/ops/upfirdn2d.h new file mode 100644 index 0000000000000000000000000000000000000000..c9e2032bcac9d2abde7a75eea4d812da348afadd --- /dev/null +++ b/models/stylegan2/torch_utils/ops/upfirdn2d.h @@ -0,0 +1,59 @@ +// Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +// +// NVIDIA CORPORATION and its licensors retain all intellectual property +// and proprietary rights in and to this software, related documentation +// and any modifications thereto. Any use, reproduction, disclosure or +// distribution of this software and related documentation without an express +// license agreement from NVIDIA CORPORATION is strictly prohibited. + +#include + +//------------------------------------------------------------------------ +// CUDA kernel parameters. + +struct upfirdn2d_kernel_params +{ + const void* x; + const float* f; + void* y; + + int2 up; + int2 down; + int2 pad0; + int flip; + float gain; + + int4 inSize; // [width, height, channel, batch] + int4 inStride; + int2 filterSize; // [width, height] + int2 filterStride; + int4 outSize; // [width, height, channel, batch] + int4 outStride; + int sizeMinor; + int sizeMajor; + + int loopMinor; + int loopMajor; + int loopX; + int launchMinor; + int launchMajor; +}; + +//------------------------------------------------------------------------ +// CUDA kernel specialization. + +struct upfirdn2d_kernel_spec +{ + void* kernel; + int tileOutW; + int tileOutH; + int loopMinor; + int loopX; +}; + +//------------------------------------------------------------------------ +// CUDA kernel selection. + +template upfirdn2d_kernel_spec choose_upfirdn2d_kernel(const upfirdn2d_kernel_params& p); + +//------------------------------------------------------------------------ diff --git a/models/stylegan2/torch_utils/ops/upfirdn2d.py b/models/stylegan2/torch_utils/ops/upfirdn2d.py new file mode 100644 index 0000000000000000000000000000000000000000..ceeac2b9834e33b7c601c28bf27f32aa91c69256 --- /dev/null +++ b/models/stylegan2/torch_utils/ops/upfirdn2d.py @@ -0,0 +1,384 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +# +# NVIDIA CORPORATION and its licensors retain all intellectual property +# and proprietary rights in and to this software, related documentation +# and any modifications thereto. Any use, reproduction, disclosure or +# distribution of this software and related documentation without an express +# license agreement from NVIDIA CORPORATION is strictly prohibited. + +"""Custom PyTorch ops for efficient resampling of 2D images.""" + +import os +import warnings +import numpy as np +import torch +import traceback + +from .. import custom_ops +from .. import misc +from . import conv2d_gradfix + +#---------------------------------------------------------------------------- + +_inited = False +_plugin = None + +def _init(): + global _inited, _plugin + if not _inited: + sources = ['upfirdn2d.cpp', 'upfirdn2d.cu'] + sources = [os.path.join(os.path.dirname(__file__), s) for s in sources] + try: + _plugin = custom_ops.get_plugin('upfirdn2d_plugin', sources=sources, extra_cuda_cflags=['--use_fast_math']) + except: + warnings.warn('Failed to build CUDA kernels for upfirdn2d. Falling back to slow reference implementation. Details:\n\n' + traceback.format_exc()) + return _plugin is not None + +def _parse_scaling(scaling): + if isinstance(scaling, int): + scaling = [scaling, scaling] + assert isinstance(scaling, (list, tuple)) + assert all(isinstance(x, int) for x in scaling) + sx, sy = scaling + assert sx >= 1 and sy >= 1 + return sx, sy + +def _parse_padding(padding): + if isinstance(padding, int): + padding = [padding, padding] + assert isinstance(padding, (list, tuple)) + assert all(isinstance(x, int) for x in padding) + if len(padding) == 2: + padx, pady = padding + padding = [padx, padx, pady, pady] + padx0, padx1, pady0, pady1 = padding + return padx0, padx1, pady0, pady1 + +def _get_filter_size(f): + if f is None: + return 1, 1 + assert isinstance(f, torch.Tensor) and f.ndim in [1, 2] + fw = f.shape[-1] + fh = f.shape[0] + with misc.suppress_tracer_warnings(): + fw = int(fw) + fh = int(fh) + misc.assert_shape(f, [fh, fw][:f.ndim]) + assert fw >= 1 and fh >= 1 + return fw, fh + +#---------------------------------------------------------------------------- + +def setup_filter(f, device=torch.device('cpu'), normalize=True, flip_filter=False, gain=1, separable=None): + r"""Convenience function to setup 2D FIR filter for `upfirdn2d()`. + + Args: + f: Torch tensor, numpy array, or python list of the shape + `[filter_height, filter_width]` (non-separable), + `[filter_taps]` (separable), + `[]` (impulse), or + `None` (identity). + device: Result device (default: cpu). + normalize: Normalize the filter so that it retains the magnitude + for constant input signal (DC)? (default: True). + flip_filter: Flip the filter? (default: False). + gain: Overall scaling factor for signal magnitude (default: 1). + separable: Return a separable filter? (default: select automatically). + + Returns: + Float32 tensor of the shape + `[filter_height, filter_width]` (non-separable) or + `[filter_taps]` (separable). + """ + # Validate. + if f is None: + f = 1 + f = torch.as_tensor(f, dtype=torch.float32) + assert f.ndim in [0, 1, 2] + assert f.numel() > 0 + if f.ndim == 0: + f = f[np.newaxis] + + # Separable? + if separable is None: + separable = (f.ndim == 1 and f.numel() >= 8) + if f.ndim == 1 and not separable: + f = f.ger(f) + assert f.ndim == (1 if separable else 2) + + # Apply normalize, flip, gain, and device. + if normalize: + f /= f.sum() + if flip_filter: + f = f.flip(list(range(f.ndim))) + f = f * (gain ** (f.ndim / 2)) + f = f.to(device=device) + return f + +#---------------------------------------------------------------------------- + +def upfirdn2d(x, f, up=1, down=1, padding=0, flip_filter=False, gain=1, impl='cuda'): + r"""Pad, upsample, filter, and downsample a batch of 2D images. + + Performs the following sequence of operations for each channel: + + 1. Upsample the image by inserting N-1 zeros after each pixel (`up`). + + 2. Pad the image with the specified number of zeros on each side (`padding`). + Negative padding corresponds to cropping the image. + + 3. Convolve the image with the specified 2D FIR filter (`f`), shrinking it + so that the footprint of all output pixels lies within the input image. + + 4. Downsample the image by keeping every Nth pixel (`down`). + + This sequence of operations bears close resemblance to scipy.signal.upfirdn(). + The fused op is considerably more efficient than performing the same calculation + using standard PyTorch ops. It supports gradients of arbitrary order. + + Args: + x: Float32/float64/float16 input tensor of the shape + `[batch_size, num_channels, in_height, in_width]`. + f: Float32 FIR filter of the shape + `[filter_height, filter_width]` (non-separable), + `[filter_taps]` (separable), or + `None` (identity). + up: Integer upsampling factor. Can be a single int or a list/tuple + `[x, y]` (default: 1). + down: Integer downsampling factor. Can be a single int or a list/tuple + `[x, y]` (default: 1). + padding: Padding with respect to the upsampled image. Can be a single number + or a list/tuple `[x, y]` or `[x_before, x_after, y_before, y_after]` + (default: 0). + flip_filter: False = convolution, True = correlation (default: False). + gain: Overall scaling factor for signal magnitude (default: 1). + impl: Implementation to use. Can be `'ref'` or `'cuda'` (default: `'cuda'`). + + Returns: + Tensor of the shape `[batch_size, num_channels, out_height, out_width]`. + """ + assert isinstance(x, torch.Tensor) + assert impl in ['ref', 'cuda'] + if impl == 'cuda' and x.device.type == 'cuda' and _init(): + return _upfirdn2d_cuda(up=up, down=down, padding=padding, flip_filter=flip_filter, gain=gain).apply(x, f) + return _upfirdn2d_ref(x, f, up=up, down=down, padding=padding, flip_filter=flip_filter, gain=gain) + +#---------------------------------------------------------------------------- + +@misc.profiled_function +def _upfirdn2d_ref(x, f, up=1, down=1, padding=0, flip_filter=False, gain=1): + """Slow reference implementation of `upfirdn2d()` using standard PyTorch ops. + """ + # Validate arguments. + assert isinstance(x, torch.Tensor) and x.ndim == 4 + if f is None: + f = torch.ones([1, 1], dtype=torch.float32, device=x.device) + assert isinstance(f, torch.Tensor) and f.ndim in [1, 2] + assert f.dtype == torch.float32 and not f.requires_grad + batch_size, num_channels, in_height, in_width = x.shape + upx, upy = _parse_scaling(up) + downx, downy = _parse_scaling(down) + padx0, padx1, pady0, pady1 = _parse_padding(padding) + + # Upsample by inserting zeros. + x = x.reshape([batch_size, num_channels, in_height, 1, in_width, 1]) + x = torch.nn.functional.pad(x, [0, upx - 1, 0, 0, 0, upy - 1]) + x = x.reshape([batch_size, num_channels, in_height * upy, in_width * upx]) + + # Pad or crop. + x = torch.nn.functional.pad(x, [max(padx0, 0), max(padx1, 0), max(pady0, 0), max(pady1, 0)]) + x = x[:, :, max(-pady0, 0) : x.shape[2] - max(-pady1, 0), max(-padx0, 0) : x.shape[3] - max(-padx1, 0)] + + # Setup filter. + f = f * (gain ** (f.ndim / 2)) + f = f.to(x.dtype) + if not flip_filter: + f = f.flip(list(range(f.ndim))) + + # Convolve with the filter. + f = f[np.newaxis, np.newaxis].repeat([num_channels, 1] + [1] * f.ndim) + if f.ndim == 4: + x = conv2d_gradfix.conv2d(input=x, weight=f, groups=num_channels) + else: + x = conv2d_gradfix.conv2d(input=x, weight=f.unsqueeze(2), groups=num_channels) + x = conv2d_gradfix.conv2d(input=x, weight=f.unsqueeze(3), groups=num_channels) + + # Downsample by throwing away pixels. + x = x[:, :, ::downy, ::downx] + return x + +#---------------------------------------------------------------------------- + +_upfirdn2d_cuda_cache = dict() + +def _upfirdn2d_cuda(up=1, down=1, padding=0, flip_filter=False, gain=1): + """Fast CUDA implementation of `upfirdn2d()` using custom ops. + """ + # Parse arguments. + upx, upy = _parse_scaling(up) + downx, downy = _parse_scaling(down) + padx0, padx1, pady0, pady1 = _parse_padding(padding) + + # Lookup from cache. + key = (upx, upy, downx, downy, padx0, padx1, pady0, pady1, flip_filter, gain) + if key in _upfirdn2d_cuda_cache: + return _upfirdn2d_cuda_cache[key] + + # Forward op. + class Upfirdn2dCuda(torch.autograd.Function): + @staticmethod + def forward(ctx, x, f): # pylint: disable=arguments-differ + assert isinstance(x, torch.Tensor) and x.ndim == 4 + if f is None: + f = torch.ones([1, 1], dtype=torch.float32, device=x.device) + assert isinstance(f, torch.Tensor) and f.ndim in [1, 2] + y = x + if f.ndim == 2: + y = _plugin.upfirdn2d(y, f, upx, upy, downx, downy, padx0, padx1, pady0, pady1, flip_filter, gain) + else: + y = _plugin.upfirdn2d(y, f.unsqueeze(0), upx, 1, downx, 1, padx0, padx1, 0, 0, flip_filter, np.sqrt(gain)) + y = _plugin.upfirdn2d(y, f.unsqueeze(1), 1, upy, 1, downy, 0, 0, pady0, pady1, flip_filter, np.sqrt(gain)) + ctx.save_for_backward(f) + ctx.x_shape = x.shape + return y + + @staticmethod + def backward(ctx, dy): # pylint: disable=arguments-differ + f, = ctx.saved_tensors + _, _, ih, iw = ctx.x_shape + _, _, oh, ow = dy.shape + fw, fh = _get_filter_size(f) + p = [ + fw - padx0 - 1, + iw * upx - ow * downx + padx0 - upx + 1, + fh - pady0 - 1, + ih * upy - oh * downy + pady0 - upy + 1, + ] + dx = None + df = None + + if ctx.needs_input_grad[0]: + dx = _upfirdn2d_cuda(up=down, down=up, padding=p, flip_filter=(not flip_filter), gain=gain).apply(dy, f) + + assert not ctx.needs_input_grad[1] + return dx, df + + # Add to cache. + _upfirdn2d_cuda_cache[key] = Upfirdn2dCuda + return Upfirdn2dCuda + +#---------------------------------------------------------------------------- + +def filter2d(x, f, padding=0, flip_filter=False, gain=1, impl='cuda'): + r"""Filter a batch of 2D images using the given 2D FIR filter. + + By default, the result is padded so that its shape matches the input. + User-specified padding is applied on top of that, with negative values + indicating cropping. Pixels outside the image are assumed to be zero. + + Args: + x: Float32/float64/float16 input tensor of the shape + `[batch_size, num_channels, in_height, in_width]`. + f: Float32 FIR filter of the shape + `[filter_height, filter_width]` (non-separable), + `[filter_taps]` (separable), or + `None` (identity). + padding: Padding with respect to the output. Can be a single number or a + list/tuple `[x, y]` or `[x_before, x_after, y_before, y_after]` + (default: 0). + flip_filter: False = convolution, True = correlation (default: False). + gain: Overall scaling factor for signal magnitude (default: 1). + impl: Implementation to use. Can be `'ref'` or `'cuda'` (default: `'cuda'`). + + Returns: + Tensor of the shape `[batch_size, num_channels, out_height, out_width]`. + """ + padx0, padx1, pady0, pady1 = _parse_padding(padding) + fw, fh = _get_filter_size(f) + p = [ + padx0 + fw // 2, + padx1 + (fw - 1) // 2, + pady0 + fh // 2, + pady1 + (fh - 1) // 2, + ] + return upfirdn2d(x, f, padding=p, flip_filter=flip_filter, gain=gain, impl=impl) + +#---------------------------------------------------------------------------- + +def upsample2d(x, f, up=2, padding=0, flip_filter=False, gain=1, impl='cuda'): + r"""Upsample a batch of 2D images using the given 2D FIR filter. + + By default, the result is padded so that its shape is a multiple of the input. + User-specified padding is applied on top of that, with negative values + indicating cropping. Pixels outside the image are assumed to be zero. + + Args: + x: Float32/float64/float16 input tensor of the shape + `[batch_size, num_channels, in_height, in_width]`. + f: Float32 FIR filter of the shape + `[filter_height, filter_width]` (non-separable), + `[filter_taps]` (separable), or + `None` (identity). + up: Integer upsampling factor. Can be a single int or a list/tuple + `[x, y]` (default: 1). + padding: Padding with respect to the output. Can be a single number or a + list/tuple `[x, y]` or `[x_before, x_after, y_before, y_after]` + (default: 0). + flip_filter: False = convolution, True = correlation (default: False). + gain: Overall scaling factor for signal magnitude (default: 1). + impl: Implementation to use. Can be `'ref'` or `'cuda'` (default: `'cuda'`). + + Returns: + Tensor of the shape `[batch_size, num_channels, out_height, out_width]`. + """ + upx, upy = _parse_scaling(up) + padx0, padx1, pady0, pady1 = _parse_padding(padding) + fw, fh = _get_filter_size(f) + p = [ + padx0 + (fw + upx - 1) // 2, + padx1 + (fw - upx) // 2, + pady0 + (fh + upy - 1) // 2, + pady1 + (fh - upy) // 2, + ] + return upfirdn2d(x, f, up=up, padding=p, flip_filter=flip_filter, gain=gain*upx*upy, impl=impl) + +#---------------------------------------------------------------------------- + +def downsample2d(x, f, down=2, padding=0, flip_filter=False, gain=1, impl='cuda'): + r"""Downsample a batch of 2D images using the given 2D FIR filter. + + By default, the result is padded so that its shape is a fraction of the input. + User-specified padding is applied on top of that, with negative values + indicating cropping. Pixels outside the image are assumed to be zero. + + Args: + x: Float32/float64/float16 input tensor of the shape + `[batch_size, num_channels, in_height, in_width]`. + f: Float32 FIR filter of the shape + `[filter_height, filter_width]` (non-separable), + `[filter_taps]` (separable), or + `None` (identity). + down: Integer downsampling factor. Can be a single int or a list/tuple + `[x, y]` (default: 1). + padding: Padding with respect to the input. Can be a single number or a + list/tuple `[x, y]` or `[x_before, x_after, y_before, y_after]` + (default: 0). + flip_filter: False = convolution, True = correlation (default: False). + gain: Overall scaling factor for signal magnitude (default: 1). + impl: Implementation to use. Can be `'ref'` or `'cuda'` (default: `'cuda'`). + + Returns: + Tensor of the shape `[batch_size, num_channels, out_height, out_width]`. + """ + downx, downy = _parse_scaling(down) + padx0, padx1, pady0, pady1 = _parse_padding(padding) + fw, fh = _get_filter_size(f) + p = [ + padx0 + (fw - downx + 1) // 2, + padx1 + (fw - downx) // 2, + pady0 + (fh - downy + 1) // 2, + pady1 + (fh - downy) // 2, + ] + return upfirdn2d(x, f, down=down, padding=p, flip_filter=flip_filter, gain=gain, impl=impl) + +#---------------------------------------------------------------------------- diff --git a/models/stylegan2/torch_utils/persistence.py b/models/stylegan2/torch_utils/persistence.py new file mode 100644 index 0000000000000000000000000000000000000000..566d12b3ae8806e873fea6c8eac53f4ffedab892 --- /dev/null +++ b/models/stylegan2/torch_utils/persistence.py @@ -0,0 +1,252 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +# +# NVIDIA CORPORATION and its licensors retain all intellectual property +# and proprietary rights in and to this software, related documentation +# and any modifications thereto. Any use, reproduction, disclosure or +# distribution of this software and related documentation without an express +# license agreement from NVIDIA CORPORATION is strictly prohibited. + +"""Facilities for pickling Python code alongside other data. + +The pickled code is automatically imported into a separate Python module +during unpickling. This way, any previously exported pickles will remain +usable even if the original code is no longer available, or if the current +version of the code is not consistent with what was originally pickled.""" + +import sys +import pickle +import io +import inspect +import copy +import uuid +import types +import models.stylegan2.dnnlib as dnnlib + +#---------------------------------------------------------------------------- + +_version = 6 # internal version number +_decorators = set() # {decorator_class, ...} +_import_hooks = [] # [hook_function, ...] +_module_to_src_dict = dict() # {module: src, ...} +_src_to_module_dict = dict() # {src: module, ...} + +#---------------------------------------------------------------------------- + +def persistent_class(orig_class): + r"""Class decorator that extends a given class to save its source code + when pickled. + + Example: + + from torch_utils import persistence + + @persistence.persistent_class + class MyNetwork(torch.nn.Module): + def __init__(self, num_inputs, num_outputs): + super().__init__() + self.fc = MyLayer(num_inputs, num_outputs) + ... + + @persistence.persistent_class + class MyLayer(torch.nn.Module): + ... + + When pickled, any instance of `MyNetwork` and `MyLayer` will save its + source code alongside other internal state (e.g., parameters, buffers, + and submodules). This way, any previously exported pickle will remain + usable even if the class definitions have been modified or are no + longer available. + + The decorator saves the source code of the entire Python module + containing the decorated class. It does *not* save the source code of + any imported modules. Thus, the imported modules must be available + during unpickling, also including `torch_utils.persistence` itself. + + It is ok to call functions defined in the same module from the + decorated class. However, if the decorated class depends on other + classes defined in the same module, they must be decorated as well. + This is illustrated in the above example in the case of `MyLayer`. + + It is also possible to employ the decorator just-in-time before + calling the constructor. For example: + + cls = MyLayer + if want_to_make_it_persistent: + cls = persistence.persistent_class(cls) + layer = cls(num_inputs, num_outputs) + + As an additional feature, the decorator also keeps track of the + arguments that were used to construct each instance of the decorated + class. The arguments can be queried via `obj.init_args` and + `obj.init_kwargs`, and they are automatically pickled alongside other + object state. A typical use case is to first unpickle a previous + instance of a persistent class, and then upgrade it to use the latest + version of the source code: + + with open('old_pickle.pkl', 'rb') as f: + old_net = pickle.load(f) + new_net = MyNetwork(*old_obj.init_args, **old_obj.init_kwargs) + misc.copy_params_and_buffers(old_net, new_net, require_all=True) + """ + assert isinstance(orig_class, type) + if is_persistent(orig_class): + return orig_class + + assert orig_class.__module__ in sys.modules + orig_module = sys.modules[orig_class.__module__] + orig_module_src = _module_to_src(orig_module) + + class Decorator(orig_class): + _orig_module_src = orig_module_src + _orig_class_name = orig_class.__name__ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._init_args = copy.deepcopy(args) + self._init_kwargs = copy.deepcopy(kwargs) + assert orig_class.__name__ in orig_module.__dict__ + _check_pickleable(self.__reduce__()) + + @property + def init_args(self): + return copy.deepcopy(self._init_args) + + @property + def init_kwargs(self): + return dnnlib.EasyDict(copy.deepcopy(self._init_kwargs)) + + def __reduce__(self): + fields = list(super().__reduce__()) + fields += [None] * max(3 - len(fields), 0) + if fields[0] is not _reconstruct_persistent_obj: + meta = dict(type='class', version=_version, module_src=self._orig_module_src, class_name=self._orig_class_name, state=fields[2]) + fields[0] = _reconstruct_persistent_obj # reconstruct func + fields[1] = (meta,) # reconstruct args + fields[2] = None # state dict + return tuple(fields) + + Decorator.__name__ = orig_class.__name__ + _decorators.add(Decorator) + return Decorator + +#---------------------------------------------------------------------------- + +def is_persistent(obj): + r"""Test whether the given object or class is persistent, i.e., + whether it will save its source code when pickled. + """ + try: + if obj in _decorators: + return True + except TypeError: + pass + return type(obj) in _decorators # pylint: disable=unidiomatic-typecheck + +#---------------------------------------------------------------------------- + +def import_hook(hook): + r"""Register an import hook that is called whenever a persistent object + is being unpickled. A typical use case is to patch the pickled source + code to avoid errors and inconsistencies when the API of some imported + module has changed. + + The hook should have the following signature: + + hook(meta) -> modified meta + + `meta` is an instance of `dnnlib.EasyDict` with the following fields: + + type: Type of the persistent object, e.g. `'class'`. + version: Internal version number of `torch_utils.persistence`. + module_src Original source code of the Python module. + class_name: Class name in the original Python module. + state: Internal state of the object. + + Example: + + @persistence.import_hook + def wreck_my_network(meta): + if meta.class_name == 'MyNetwork': + print('MyNetwork is being imported. I will wreck it!') + meta.module_src = meta.module_src.replace("True", "False") + return meta + """ + assert callable(hook) + _import_hooks.append(hook) + +#---------------------------------------------------------------------------- + +def _reconstruct_persistent_obj(meta): + r"""Hook that is called internally by the `pickle` module to unpickle + a persistent object. + """ + meta = dnnlib.EasyDict(meta) + meta.state = dnnlib.EasyDict(meta.state) + for hook in _import_hooks: + meta = hook(meta) + assert meta is not None + + assert meta.version == _version + meta.module_src = meta.module_src.replace('torch_utils', 'models.stylegan2.torch_utils') + module = _src_to_module(meta.module_src) + + assert meta.type == 'class' + orig_class = module.__dict__[meta.class_name] + decorator_class = persistent_class(orig_class) + obj = decorator_class.__new__(decorator_class) + + setstate = getattr(obj, '__setstate__', None) + if callable(setstate): + setstate(meta.state) # pylint: disable=not-callable + else: + obj.__dict__.update(meta.state) + return obj + +#---------------------------------------------------------------------------- + +def _module_to_src(module): + r"""Query the source code of a given Python module. + """ + src = _module_to_src_dict.get(module, None) + if src is None: + src = inspect.getsource(module) + _module_to_src_dict[module] = src + _src_to_module_dict[src] = module + return src + +def _src_to_module(src): + r"""Get or create a Python module for the given source code. + """ + module = _src_to_module_dict.get(src, None) + if module is None: + module_name = "_imported_module_" + uuid.uuid4().hex + module = types.ModuleType(module_name) + sys.modules[module_name] = module + _module_to_src_dict[module] = src + _src_to_module_dict[src] = module + exec(src, module.__dict__) # pylint: disable=exec-used + return module + +#---------------------------------------------------------------------------- + +def _check_pickleable(obj): + r"""Check that the given object is pickleable, raising an exception if + it is not. This function is expected to be considerably more efficient + than actually pickling the object. + """ + def recurse(obj): + if isinstance(obj, (list, tuple, set)): + return [recurse(x) for x in obj] + if isinstance(obj, dict): + return [[recurse(x), recurse(y)] for x, y in obj.items()] + if isinstance(obj, (str, int, float, bool, bytes, bytearray)): + return None # Python primitive types are pickleable. + if f'{type(obj).__module__}.{type(obj).__name__}' in ['numpy.ndarray', 'torch.Tensor']: + return None # NumPy arrays and PyTorch tensors are pickleable. + if is_persistent(obj): + return None # Persistent objects are pickleable, by virtue of the constructor check. + return obj + with io.BytesIO() as f: + pickle.dump(recurse(obj), f) + +#---------------------------------------------------------------------------- diff --git a/models/stylegan2/torch_utils/training_stats.py b/models/stylegan2/torch_utils/training_stats.py new file mode 100644 index 0000000000000000000000000000000000000000..132b355853df8f5d1817fe50fcc3021cc6da9e98 --- /dev/null +++ b/models/stylegan2/torch_utils/training_stats.py @@ -0,0 +1,268 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +# +# NVIDIA CORPORATION and its licensors retain all intellectual property +# and proprietary rights in and to this software, related documentation +# and any modifications thereto. Any use, reproduction, disclosure or +# distribution of this software and related documentation without an express +# license agreement from NVIDIA CORPORATION is strictly prohibited. + +"""Facilities for reporting and collecting training statistics across +multiple processes and devices. The interface is designed to minimize +synchronization overhead as well as the amount of boilerplate in user +code.""" + +import re +import numpy as np +import torch +import models.stylegan2.dnnlib as dnnlib + +from . import misc + +#---------------------------------------------------------------------------- + +_num_moments = 3 # [num_scalars, sum_of_scalars, sum_of_squares] +_reduce_dtype = torch.float32 # Data type to use for initial per-tensor reduction. +_counter_dtype = torch.float64 # Data type to use for the internal counters. +_rank = 0 # Rank of the current process. +_sync_device = None # Device to use for multiprocess communication. None = single-process. +_sync_called = False # Has _sync() been called yet? +_counters = dict() # Running counters on each device, updated by report(): name => device => torch.Tensor +_cumulative = dict() # Cumulative counters on the CPU, updated by _sync(): name => torch.Tensor + +#---------------------------------------------------------------------------- + +def init_multiprocessing(rank, sync_device): + r"""Initializes `torch_utils.training_stats` for collecting statistics + across multiple processes. + + This function must be called after + `torch.distributed.init_process_group()` and before `Collector.update()`. + The call is not necessary if multi-process collection is not needed. + + Args: + rank: Rank of the current process. + sync_device: PyTorch device to use for inter-process + communication, or None to disable multi-process + collection. Typically `torch.device('cuda', rank)`. + """ + global _rank, _sync_device + assert not _sync_called + _rank = rank + _sync_device = sync_device + +#---------------------------------------------------------------------------- + +@misc.profiled_function +def report(name, value): + r"""Broadcasts the given set of scalars to all interested instances of + `Collector`, across device and process boundaries. + + This function is expected to be extremely cheap and can be safely + called from anywhere in the training loop, loss function, or inside a + `torch.nn.Module`. + + Warning: The current implementation expects the set of unique names to + be consistent across processes. Please make sure that `report()` is + called at least once for each unique name by each process, and in the + same order. If a given process has no scalars to broadcast, it can do + `report(name, [])` (empty list). + + Args: + name: Arbitrary string specifying the name of the statistic. + Averages are accumulated separately for each unique name. + value: Arbitrary set of scalars. Can be a list, tuple, + NumPy array, PyTorch tensor, or Python scalar. + + Returns: + The same `value` that was passed in. + """ + if name not in _counters: + _counters[name] = dict() + + elems = torch.as_tensor(value) + if elems.numel() == 0: + return value + + elems = elems.detach().flatten().to(_reduce_dtype) + moments = torch.stack([ + torch.ones_like(elems).sum(), + elems.sum(), + elems.square().sum(), + ]) + assert moments.ndim == 1 and moments.shape[0] == _num_moments + moments = moments.to(_counter_dtype) + + device = moments.device + if device not in _counters[name]: + _counters[name][device] = torch.zeros_like(moments) + _counters[name][device].add_(moments) + return value + +#---------------------------------------------------------------------------- + +def report0(name, value): + r"""Broadcasts the given set of scalars by the first process (`rank = 0`), + but ignores any scalars provided by the other processes. + See `report()` for further details. + """ + report(name, value if _rank == 0 else []) + return value + +#---------------------------------------------------------------------------- + +class Collector: + r"""Collects the scalars broadcasted by `report()` and `report0()` and + computes their long-term averages (mean and standard deviation) over + user-defined periods of time. + + The averages are first collected into internal counters that are not + directly visible to the user. They are then copied to the user-visible + state as a result of calling `update()` and can then be queried using + `mean()`, `std()`, `as_dict()`, etc. Calling `update()` also resets the + internal counters for the next round, so that the user-visible state + effectively reflects averages collected between the last two calls to + `update()`. + + Args: + regex: Regular expression defining which statistics to + collect. The default is to collect everything. + keep_previous: Whether to retain the previous averages if no + scalars were collected on a given round + (default: True). + """ + def __init__(self, regex='.*', keep_previous=True): + self._regex = re.compile(regex) + self._keep_previous = keep_previous + self._cumulative = dict() + self._moments = dict() + self.update() + self._moments.clear() + + def names(self): + r"""Returns the names of all statistics broadcasted so far that + match the regular expression specified at construction time. + """ + return [name for name in _counters if self._regex.fullmatch(name)] + + def update(self): + r"""Copies current values of the internal counters to the + user-visible state and resets them for the next round. + + If `keep_previous=True` was specified at construction time, the + operation is skipped for statistics that have received no scalars + since the last update, retaining their previous averages. + + This method performs a number of GPU-to-CPU transfers and one + `torch.distributed.all_reduce()`. It is intended to be called + periodically in the main training loop, typically once every + N training steps. + """ + if not self._keep_previous: + self._moments.clear() + for name, cumulative in _sync(self.names()): + if name not in self._cumulative: + self._cumulative[name] = torch.zeros([_num_moments], dtype=_counter_dtype) + delta = cumulative - self._cumulative[name] + self._cumulative[name].copy_(cumulative) + if float(delta[0]) != 0: + self._moments[name] = delta + + def _get_delta(self, name): + r"""Returns the raw moments that were accumulated for the given + statistic between the last two calls to `update()`, or zero if + no scalars were collected. + """ + assert self._regex.fullmatch(name) + if name not in self._moments: + self._moments[name] = torch.zeros([_num_moments], dtype=_counter_dtype) + return self._moments[name] + + def num(self, name): + r"""Returns the number of scalars that were accumulated for the given + statistic between the last two calls to `update()`, or zero if + no scalars were collected. + """ + delta = self._get_delta(name) + return int(delta[0]) + + def mean(self, name): + r"""Returns the mean of the scalars that were accumulated for the + given statistic between the last two calls to `update()`, or NaN if + no scalars were collected. + """ + delta = self._get_delta(name) + if int(delta[0]) == 0: + return float('nan') + return float(delta[1] / delta[0]) + + def std(self, name): + r"""Returns the standard deviation of the scalars that were + accumulated for the given statistic between the last two calls to + `update()`, or NaN if no scalars were collected. + """ + delta = self._get_delta(name) + if int(delta[0]) == 0 or not np.isfinite(float(delta[1])): + return float('nan') + if int(delta[0]) == 1: + return float(0) + mean = float(delta[1] / delta[0]) + raw_var = float(delta[2] / delta[0]) + return np.sqrt(max(raw_var - np.square(mean), 0)) + + def as_dict(self): + r"""Returns the averages accumulated between the last two calls to + `update()` as an `dnnlib.EasyDict`. The contents are as follows: + + dnnlib.EasyDict( + NAME = dnnlib.EasyDict(num=FLOAT, mean=FLOAT, std=FLOAT), + ... + ) + """ + stats = dnnlib.EasyDict() + for name in self.names(): + stats[name] = dnnlib.EasyDict(num=self.num(name), mean=self.mean(name), std=self.std(name)) + return stats + + def __getitem__(self, name): + r"""Convenience getter. + `collector[name]` is a synonym for `collector.mean(name)`. + """ + return self.mean(name) + +#---------------------------------------------------------------------------- + +def _sync(names): + r"""Synchronize the global cumulative counters across devices and + processes. Called internally by `Collector.update()`. + """ + if len(names) == 0: + return [] + global _sync_called + _sync_called = True + + # Collect deltas within current rank. + deltas = [] + device = _sync_device if _sync_device is not None else torch.device('cpu') + for name in names: + delta = torch.zeros([_num_moments], dtype=_counter_dtype, device=device) + for counter in _counters[name].values(): + delta.add_(counter.to(device)) + counter.copy_(torch.zeros_like(counter)) + deltas.append(delta) + deltas = torch.stack(deltas) + + # Sum deltas across ranks. + if _sync_device is not None: + torch.distributed.all_reduce(deltas) + + # Update cumulative values. + deltas = deltas.cpu() + for idx, name in enumerate(names): + if name not in _cumulative: + _cumulative[name] = torch.zeros([_num_moments], dtype=_counter_dtype) + _cumulative[name].add_(deltas[idx]) + + # Return name-value pairs. + return [(name, _cumulative[name]) for name in names] + +#---------------------------------------------------------------------------- diff --git a/poetry.lock b/poetry.lock new file mode 100644 index 0000000000000000000000000000000000000000..eb35a9caec59cdedd227d3e381c6e00eeaaf3410 --- /dev/null +++ b/poetry.lock @@ -0,0 +1,1819 @@ +# This file is automatically @generated by Poetry 1.8.2 and should not be changed by hand. + +[[package]] +name = "addict" +version = "2.4.0" +description = "Addict is a dictionary whose items can be set using both attribute and item syntax." +optional = false +python-versions = "*" +files = [ + {file = "addict-2.4.0-py3-none-any.whl", hash = "sha256:249bb56bbfd3cdc2a004ea0ff4c2b6ddc84d53bc2194761636eb314d5cfa5dfc"}, + {file = "addict-2.4.0.tar.gz", hash = "sha256:b3b2210e0e067a281f5646c8c5db92e99b7231ea8b0eb5f74dbdf9e259d4e494"}, +] + +[[package]] +name = "appdirs" +version = "1.4.4" +description = "A small Python module for determining appropriate platform-specific dirs, e.g. a \"user data dir\"." +optional = false +python-versions = "*" +files = [ + {file = "appdirs-1.4.4-py2.py3-none-any.whl", hash = "sha256:a841dacd6b99318a741b166adb07e19ee71a274450e68237b4650ca1055ab128"}, + {file = "appdirs-1.4.4.tar.gz", hash = "sha256:7d5d0167b2b1ba821647616af46a749d1c653740dd0d2415100fe26e27afdf41"}, +] + +[[package]] +name = "certifi" +version = "2024.2.2" +description = "Python package for providing Mozilla's CA Bundle." +optional = false +python-versions = ">=3.6" +files = [ + {file = "certifi-2024.2.2-py3-none-any.whl", hash = "sha256:dc383c07b76109f368f6106eee2b593b04a011ea4d55f652c6ca24a754d1cdd1"}, + {file = "certifi-2024.2.2.tar.gz", hash = "sha256:0569859f95fc761b18b45ef421b1290a0f65f147e92a1e5eb3e635f9a5e4e66f"}, +] + +[[package]] +name = "charset-normalizer" +version = "3.3.2" +description = "The Real First Universal Charset Detector. Open, modern and actively maintained alternative to Chardet." +optional = false +python-versions = ">=3.7.0" +files = [ + {file = "charset-normalizer-3.3.2.tar.gz", hash = "sha256:f30c3cb33b24454a82faecaf01b19c18562b1e89558fb6c56de4d9118a032fd5"}, + {file = "charset_normalizer-3.3.2-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:25baf083bf6f6b341f4121c2f3c548875ee6f5339300e08be3f2b2ba1721cdd3"}, + {file = "charset_normalizer-3.3.2-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:06435b539f889b1f6f4ac1758871aae42dc3a8c0e24ac9e60c2384973ad73027"}, + {file = "charset_normalizer-3.3.2-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:9063e24fdb1e498ab71cb7419e24622516c4a04476b17a2dab57e8baa30d6e03"}, + {file = "charset_normalizer-3.3.2-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:6897af51655e3691ff853668779c7bad41579facacf5fd7253b0133308cf000d"}, + {file = "charset_normalizer-3.3.2-cp310-cp310-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:1d3193f4a680c64b4b6a9115943538edb896edc190f0b222e73761716519268e"}, + {file = "charset_normalizer-3.3.2-cp310-cp310-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:cd70574b12bb8a4d2aaa0094515df2463cb429d8536cfb6c7ce983246983e5a6"}, + {file = "charset_normalizer-3.3.2-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:8465322196c8b4d7ab6d1e049e4c5cb460d0394da4a27d23cc242fbf0034b6b5"}, + {file = "charset_normalizer-3.3.2-cp310-cp310-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:a9a8e9031d613fd2009c182b69c7b2c1ef8239a0efb1df3f7c8da66d5dd3d537"}, + {file = "charset_normalizer-3.3.2-cp310-cp310-musllinux_1_1_aarch64.whl", hash = "sha256:beb58fe5cdb101e3a055192ac291b7a21e3b7ef4f67fa1d74e331a7f2124341c"}, + {file = "charset_normalizer-3.3.2-cp310-cp310-musllinux_1_1_i686.whl", hash = "sha256:e06ed3eb3218bc64786f7db41917d4e686cc4856944f53d5bdf83a6884432e12"}, + {file = "charset_normalizer-3.3.2-cp310-cp310-musllinux_1_1_ppc64le.whl", hash = "sha256:2e81c7b9c8979ce92ed306c249d46894776a909505d8f5a4ba55b14206e3222f"}, + {file = "charset_normalizer-3.3.2-cp310-cp310-musllinux_1_1_s390x.whl", hash = "sha256:572c3763a264ba47b3cf708a44ce965d98555f618ca42c926a9c1616d8f34269"}, + {file = "charset_normalizer-3.3.2-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:fd1abc0d89e30cc4e02e4064dc67fcc51bd941eb395c502aac3ec19fab46b519"}, + {file = "charset_normalizer-3.3.2-cp310-cp310-win32.whl", hash = "sha256:3d47fa203a7bd9c5b6cee4736ee84ca03b8ef23193c0d1ca99b5089f72645c73"}, + {file = "charset_normalizer-3.3.2-cp310-cp310-win_amd64.whl", hash = "sha256:10955842570876604d404661fbccbc9c7e684caf432c09c715ec38fbae45ae09"}, + {file = "charset_normalizer-3.3.2-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:802fe99cca7457642125a8a88a084cef28ff0cf9407060f7b93dca5aa25480db"}, + {file = "charset_normalizer-3.3.2-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:573f6eac48f4769d667c4442081b1794f52919e7edada77495aaed9236d13a96"}, + {file = "charset_normalizer-3.3.2-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:549a3a73da901d5bc3ce8d24e0600d1fa85524c10287f6004fbab87672bf3e1e"}, + {file = "charset_normalizer-3.3.2-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:f27273b60488abe721a075bcca6d7f3964f9f6f067c8c4c605743023d7d3944f"}, + {file = "charset_normalizer-3.3.2-cp311-cp311-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:1ceae2f17a9c33cb48e3263960dc5fc8005351ee19db217e9b1bb15d28c02574"}, + {file = "charset_normalizer-3.3.2-cp311-cp311-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:65f6f63034100ead094b8744b3b97965785388f308a64cf8d7c34f2f2e5be0c4"}, + {file = "charset_normalizer-3.3.2-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:753f10e867343b4511128c6ed8c82f7bec3bd026875576dfd88483c5c73b2fd8"}, + {file = "charset_normalizer-3.3.2-cp311-cp311-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:4a78b2b446bd7c934f5dcedc588903fb2f5eec172f3d29e52a9096a43722adfc"}, + {file = "charset_normalizer-3.3.2-cp311-cp311-musllinux_1_1_aarch64.whl", hash = "sha256:e537484df0d8f426ce2afb2d0f8e1c3d0b114b83f8850e5f2fbea0e797bd82ae"}, + {file = "charset_normalizer-3.3.2-cp311-cp311-musllinux_1_1_i686.whl", hash = "sha256:eb6904c354526e758fda7167b33005998fb68c46fbc10e013ca97f21ca5c8887"}, + {file = "charset_normalizer-3.3.2-cp311-cp311-musllinux_1_1_ppc64le.whl", hash = "sha256:deb6be0ac38ece9ba87dea880e438f25ca3eddfac8b002a2ec3d9183a454e8ae"}, + {file = "charset_normalizer-3.3.2-cp311-cp311-musllinux_1_1_s390x.whl", hash = "sha256:4ab2fe47fae9e0f9dee8c04187ce5d09f48eabe611be8259444906793ab7cbce"}, + {file = "charset_normalizer-3.3.2-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:80402cd6ee291dcb72644d6eac93785fe2c8b9cb30893c1af5b8fdd753b9d40f"}, + {file = "charset_normalizer-3.3.2-cp311-cp311-win32.whl", hash = "sha256:7cd13a2e3ddeed6913a65e66e94b51d80a041145a026c27e6bb76c31a853c6ab"}, + {file = "charset_normalizer-3.3.2-cp311-cp311-win_amd64.whl", hash = "sha256:663946639d296df6a2bb2aa51b60a2454ca1cb29835324c640dafb5ff2131a77"}, + {file = "charset_normalizer-3.3.2-cp312-cp312-macosx_10_9_universal2.whl", hash = "sha256:0b2b64d2bb6d3fb9112bafa732def486049e63de9618b5843bcdd081d8144cd8"}, + {file = "charset_normalizer-3.3.2-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:ddbb2551d7e0102e7252db79ba445cdab71b26640817ab1e3e3648dad515003b"}, + {file = "charset_normalizer-3.3.2-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:55086ee1064215781fff39a1af09518bc9255b50d6333f2e4c74ca09fac6a8f6"}, + {file = "charset_normalizer-3.3.2-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:8f4a014bc36d3c57402e2977dada34f9c12300af536839dc38c0beab8878f38a"}, + {file = "charset_normalizer-3.3.2-cp312-cp312-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:a10af20b82360ab00827f916a6058451b723b4e65030c5a18577c8b2de5b3389"}, + {file = "charset_normalizer-3.3.2-cp312-cp312-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:8d756e44e94489e49571086ef83b2bb8ce311e730092d2c34ca8f7d925cb20aa"}, + {file = "charset_normalizer-3.3.2-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:90d558489962fd4918143277a773316e56c72da56ec7aa3dc3dbbe20fdfed15b"}, + {file = "charset_normalizer-3.3.2-cp312-cp312-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:6ac7ffc7ad6d040517be39eb591cac5ff87416c2537df6ba3cba3bae290c0fed"}, + {file = "charset_normalizer-3.3.2-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:7ed9e526742851e8d5cc9e6cf41427dfc6068d4f5a3bb03659444b4cabf6bc26"}, + {file = "charset_normalizer-3.3.2-cp312-cp312-musllinux_1_1_i686.whl", hash = "sha256:8bdb58ff7ba23002a4c5808d608e4e6c687175724f54a5dade5fa8c67b604e4d"}, + {file = "charset_normalizer-3.3.2-cp312-cp312-musllinux_1_1_ppc64le.whl", hash = "sha256:6b3251890fff30ee142c44144871185dbe13b11bab478a88887a639655be1068"}, + {file = "charset_normalizer-3.3.2-cp312-cp312-musllinux_1_1_s390x.whl", hash = "sha256:b4a23f61ce87adf89be746c8a8974fe1c823c891d8f86eb218bb957c924bb143"}, + {file = "charset_normalizer-3.3.2-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:efcb3f6676480691518c177e3b465bcddf57cea040302f9f4e6e191af91174d4"}, + {file = "charset_normalizer-3.3.2-cp312-cp312-win32.whl", hash = "sha256:d965bba47ddeec8cd560687584e88cf699fd28f192ceb452d1d7ee807c5597b7"}, + {file = "charset_normalizer-3.3.2-cp312-cp312-win_amd64.whl", hash = "sha256:96b02a3dc4381e5494fad39be677abcb5e6634bf7b4fa83a6dd3112607547001"}, + {file = "charset_normalizer-3.3.2-cp37-cp37m-macosx_10_9_x86_64.whl", hash = "sha256:95f2a5796329323b8f0512e09dbb7a1860c46a39da62ecb2324f116fa8fdc85c"}, + {file = "charset_normalizer-3.3.2-cp37-cp37m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:c002b4ffc0be611f0d9da932eb0f704fe2602a9a949d1f738e4c34c75b0863d5"}, + {file = "charset_normalizer-3.3.2-cp37-cp37m-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:a981a536974bbc7a512cf44ed14938cf01030a99e9b3a06dd59578882f06f985"}, + {file = "charset_normalizer-3.3.2-cp37-cp37m-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:3287761bc4ee9e33561a7e058c72ac0938c4f57fe49a09eae428fd88aafe7bb6"}, + {file = "charset_normalizer-3.3.2-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:42cb296636fcc8b0644486d15c12376cb9fa75443e00fb25de0b8602e64c1714"}, + {file = "charset_normalizer-3.3.2-cp37-cp37m-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:0a55554a2fa0d408816b3b5cedf0045f4b8e1a6065aec45849de2d6f3f8e9786"}, + {file = "charset_normalizer-3.3.2-cp37-cp37m-musllinux_1_1_aarch64.whl", hash = "sha256:c083af607d2515612056a31f0a8d9e0fcb5876b7bfc0abad3ecd275bc4ebc2d5"}, + {file = "charset_normalizer-3.3.2-cp37-cp37m-musllinux_1_1_i686.whl", hash = "sha256:87d1351268731db79e0f8e745d92493ee2841c974128ef629dc518b937d9194c"}, + {file = "charset_normalizer-3.3.2-cp37-cp37m-musllinux_1_1_ppc64le.whl", hash = "sha256:bd8f7df7d12c2db9fab40bdd87a7c09b1530128315d047a086fa3ae3435cb3a8"}, + {file = "charset_normalizer-3.3.2-cp37-cp37m-musllinux_1_1_s390x.whl", hash = "sha256:c180f51afb394e165eafe4ac2936a14bee3eb10debc9d9e4db8958fe36afe711"}, + {file = "charset_normalizer-3.3.2-cp37-cp37m-musllinux_1_1_x86_64.whl", hash = "sha256:8c622a5fe39a48f78944a87d4fb8a53ee07344641b0562c540d840748571b811"}, + {file = "charset_normalizer-3.3.2-cp37-cp37m-win32.whl", hash = "sha256:db364eca23f876da6f9e16c9da0df51aa4f104a972735574842618b8c6d999d4"}, + {file = "charset_normalizer-3.3.2-cp37-cp37m-win_amd64.whl", hash = "sha256:86216b5cee4b06df986d214f664305142d9c76df9b6512be2738aa72a2048f99"}, + {file = "charset_normalizer-3.3.2-cp38-cp38-macosx_10_9_universal2.whl", hash = "sha256:6463effa3186ea09411d50efc7d85360b38d5f09b870c48e4600f63af490e56a"}, + {file = "charset_normalizer-3.3.2-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:6c4caeef8fa63d06bd437cd4bdcf3ffefe6738fb1b25951440d80dc7df8c03ac"}, + {file = "charset_normalizer-3.3.2-cp38-cp38-macosx_11_0_arm64.whl", hash = "sha256:37e55c8e51c236f95b033f6fb391d7d7970ba5fe7ff453dad675e88cf303377a"}, + {file = "charset_normalizer-3.3.2-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:fb69256e180cb6c8a894fee62b3afebae785babc1ee98b81cdf68bbca1987f33"}, + {file = "charset_normalizer-3.3.2-cp38-cp38-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:ae5f4161f18c61806f411a13b0310bea87f987c7d2ecdbdaad0e94eb2e404238"}, + {file = "charset_normalizer-3.3.2-cp38-cp38-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:b2b0a0c0517616b6869869f8c581d4eb2dd83a4d79e0ebcb7d373ef9956aeb0a"}, + {file = "charset_normalizer-3.3.2-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:45485e01ff4d3630ec0d9617310448a8702f70e9c01906b0d0118bdf9d124cf2"}, + {file = "charset_normalizer-3.3.2-cp38-cp38-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:eb00ed941194665c332bf8e078baf037d6c35d7c4f3102ea2d4f16ca94a26dc8"}, + {file = "charset_normalizer-3.3.2-cp38-cp38-musllinux_1_1_aarch64.whl", hash = "sha256:2127566c664442652f024c837091890cb1942c30937add288223dc895793f898"}, + {file = "charset_normalizer-3.3.2-cp38-cp38-musllinux_1_1_i686.whl", hash = "sha256:a50aebfa173e157099939b17f18600f72f84eed3049e743b68ad15bd69b6bf99"}, + {file = "charset_normalizer-3.3.2-cp38-cp38-musllinux_1_1_ppc64le.whl", hash = "sha256:4d0d1650369165a14e14e1e47b372cfcb31d6ab44e6e33cb2d4e57265290044d"}, + {file = "charset_normalizer-3.3.2-cp38-cp38-musllinux_1_1_s390x.whl", hash = "sha256:923c0c831b7cfcb071580d3f46c4baf50f174be571576556269530f4bbd79d04"}, + {file = "charset_normalizer-3.3.2-cp38-cp38-musllinux_1_1_x86_64.whl", hash = "sha256:06a81e93cd441c56a9b65d8e1d043daeb97a3d0856d177d5c90ba85acb3db087"}, + {file = "charset_normalizer-3.3.2-cp38-cp38-win32.whl", hash = "sha256:6ef1d82a3af9d3eecdba2321dc1b3c238245d890843e040e41e470ffa64c3e25"}, + {file = "charset_normalizer-3.3.2-cp38-cp38-win_amd64.whl", hash = "sha256:eb8821e09e916165e160797a6c17edda0679379a4be5c716c260e836e122f54b"}, + {file = "charset_normalizer-3.3.2-cp39-cp39-macosx_10_9_universal2.whl", hash = "sha256:c235ebd9baae02f1b77bcea61bce332cb4331dc3617d254df3323aa01ab47bd4"}, + {file = "charset_normalizer-3.3.2-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:5b4c145409bef602a690e7cfad0a15a55c13320ff7a3ad7ca59c13bb8ba4d45d"}, + {file = "charset_normalizer-3.3.2-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:68d1f8a9e9e37c1223b656399be5d6b448dea850bed7d0f87a8311f1ff3dabb0"}, + {file = "charset_normalizer-3.3.2-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:22afcb9f253dac0696b5a4be4a1c0f8762f8239e21b99680099abd9b2b1b2269"}, + {file = "charset_normalizer-3.3.2-cp39-cp39-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:e27ad930a842b4c5eb8ac0016b0a54f5aebbe679340c26101df33424142c143c"}, + {file = "charset_normalizer-3.3.2-cp39-cp39-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:1f79682fbe303db92bc2b1136016a38a42e835d932bab5b3b1bfcfbf0640e519"}, + {file = "charset_normalizer-3.3.2-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:b261ccdec7821281dade748d088bb6e9b69e6d15b30652b74cbbac25e280b796"}, + {file = "charset_normalizer-3.3.2-cp39-cp39-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:122c7fa62b130ed55f8f285bfd56d5f4b4a5b503609d181f9ad85e55c89f4185"}, + {file = "charset_normalizer-3.3.2-cp39-cp39-musllinux_1_1_aarch64.whl", hash = "sha256:d0eccceffcb53201b5bfebb52600a5fb483a20b61da9dbc885f8b103cbe7598c"}, + {file = "charset_normalizer-3.3.2-cp39-cp39-musllinux_1_1_i686.whl", hash = "sha256:9f96df6923e21816da7e0ad3fd47dd8f94b2a5ce594e00677c0013018b813458"}, + {file = "charset_normalizer-3.3.2-cp39-cp39-musllinux_1_1_ppc64le.whl", hash = "sha256:7f04c839ed0b6b98b1a7501a002144b76c18fb1c1850c8b98d458ac269e26ed2"}, + {file = "charset_normalizer-3.3.2-cp39-cp39-musllinux_1_1_s390x.whl", hash = "sha256:34d1c8da1e78d2e001f363791c98a272bb734000fcef47a491c1e3b0505657a8"}, + {file = "charset_normalizer-3.3.2-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:ff8fa367d09b717b2a17a052544193ad76cd49979c805768879cb63d9ca50561"}, + {file = "charset_normalizer-3.3.2-cp39-cp39-win32.whl", hash = "sha256:aed38f6e4fb3f5d6bf81bfa990a07806be9d83cf7bacef998ab1a9bd660a581f"}, + {file = "charset_normalizer-3.3.2-cp39-cp39-win_amd64.whl", hash = "sha256:b01b88d45a6fcb69667cd6d2f7a9aeb4bf53760d7fc536bf679ec94fe9f3ff3d"}, + {file = "charset_normalizer-3.3.2-py3-none-any.whl", hash = "sha256:3e4d1f6587322d2788836a99c69062fbb091331ec940e02d12d179c1d53e25fc"}, +] + +[[package]] +name = "click" +version = "8.1.7" +description = "Composable command line interface toolkit" +optional = false +python-versions = ">=3.7" +files = [ + {file = "click-8.1.7-py3-none-any.whl", hash = "sha256:ae74fb96c20a0277a1d615f1e4d73c8414f5a98db8b799a7931d1582f3390c28"}, + {file = "click-8.1.7.tar.gz", hash = "sha256:ca9853ad459e787e2192211578cc907e7594e294c7ccc834310722b41b9ca6de"}, +] + +[package.dependencies] +colorama = {version = "*", markers = "platform_system == \"Windows\""} + +[[package]] +name = "clip" +version = "1.0" +description = "" +optional = false +python-versions = "*" +files = [] +develop = false + +[package.dependencies] +ftfy = "*" +regex = "*" +torch = "*" +torchvision = "*" +tqdm = "*" + +[package.extras] +dev = ["pytest"] + +[package.source] +type = "git" +url = "https://github.com/openai/CLIP.git" +reference = "HEAD" +resolved_reference = "a1d071733d7111c9c014f024669f959182114e33" + +[[package]] +name = "colorama" +version = "0.4.6" +description = "Cross-platform colored terminal text." +optional = false +python-versions = "!=3.0.*,!=3.1.*,!=3.2.*,!=3.3.*,!=3.4.*,!=3.5.*,!=3.6.*,>=2.7" +files = [ + {file = "colorama-0.4.6-py2.py3-none-any.whl", hash = "sha256:4f1d9991f5acc0ca119f9d443620b77f9d6b33703e51011c16baf57afb285fc6"}, + {file = "colorama-0.4.6.tar.gz", hash = "sha256:08695f5cb7ed6e0531a20572697297273c47b8cae5a63ffc6d6ed5c201be6e44"}, +] + +[[package]] +name = "contourpy" +version = "1.2.0" +description = "Python library for calculating contours of 2D quadrilateral grids" +optional = false +python-versions = ">=3.9" +files = [ + {file = "contourpy-1.2.0-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:0274c1cb63625972c0c007ab14dd9ba9e199c36ae1a231ce45d725cbcbfd10a8"}, + {file = "contourpy-1.2.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:ab459a1cbbf18e8698399c595a01f6dcc5c138220ca3ea9e7e6126232d102bb4"}, + {file = "contourpy-1.2.0-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:6fdd887f17c2f4572ce548461e4f96396681212d858cae7bd52ba3310bc6f00f"}, + {file = "contourpy-1.2.0-cp310-cp310-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:5d16edfc3fc09968e09ddffada434b3bf989bf4911535e04eada58469873e28e"}, + {file = "contourpy-1.2.0-cp310-cp310-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:1c203f617abc0dde5792beb586f827021069fb6d403d7f4d5c2b543d87edceb9"}, + {file = "contourpy-1.2.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:b69303ceb2e4d4f146bf82fda78891ef7bcd80c41bf16bfca3d0d7eb545448aa"}, + {file = "contourpy-1.2.0-cp310-cp310-musllinux_1_1_aarch64.whl", hash = "sha256:884c3f9d42d7218304bc74a8a7693d172685c84bd7ab2bab1ee567b769696df9"}, + {file = "contourpy-1.2.0-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:4a1b1208102be6e851f20066bf0e7a96b7d48a07c9b0cfe6d0d4545c2f6cadab"}, + {file = "contourpy-1.2.0-cp310-cp310-win32.whl", hash = "sha256:34b9071c040d6fe45d9826cbbe3727d20d83f1b6110d219b83eb0e2a01d79488"}, + {file = "contourpy-1.2.0-cp310-cp310-win_amd64.whl", hash = "sha256:bd2f1ae63998da104f16a8b788f685e55d65760cd1929518fd94cd682bf03e41"}, + {file = "contourpy-1.2.0-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:dd10c26b4eadae44783c45ad6655220426f971c61d9b239e6f7b16d5cdaaa727"}, + {file = "contourpy-1.2.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:5c6b28956b7b232ae801406e529ad7b350d3f09a4fde958dfdf3c0520cdde0dd"}, + {file = "contourpy-1.2.0-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:ebeac59e9e1eb4b84940d076d9f9a6cec0064e241818bcb6e32124cc5c3e377a"}, + {file = "contourpy-1.2.0-cp311-cp311-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:139d8d2e1c1dd52d78682f505e980f592ba53c9f73bd6be102233e358b401063"}, + {file = "contourpy-1.2.0-cp311-cp311-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:1e9dc350fb4c58adc64df3e0703ab076f60aac06e67d48b3848c23647ae4310e"}, + {file = "contourpy-1.2.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:18fc2b4ed8e4a8fe849d18dce4bd3c7ea637758c6343a1f2bae1e9bd4c9f4686"}, + {file = "contourpy-1.2.0-cp311-cp311-musllinux_1_1_aarch64.whl", hash = "sha256:16a7380e943a6d52472096cb7ad5264ecee36ed60888e2a3d3814991a0107286"}, + {file = "contourpy-1.2.0-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:8d8faf05be5ec8e02a4d86f616fc2a0322ff4a4ce26c0f09d9f7fb5330a35c95"}, + {file = "contourpy-1.2.0-cp311-cp311-win32.whl", hash = "sha256:67b7f17679fa62ec82b7e3e611c43a016b887bd64fb933b3ae8638583006c6d6"}, + {file = "contourpy-1.2.0-cp311-cp311-win_amd64.whl", hash = "sha256:99ad97258985328b4f207a5e777c1b44a83bfe7cf1f87b99f9c11d4ee477c4de"}, + {file = "contourpy-1.2.0-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:575bcaf957a25d1194903a10bc9f316c136c19f24e0985a2b9b5608bdf5dbfe0"}, + {file = "contourpy-1.2.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:9e6c93b5b2dbcedad20a2f18ec22cae47da0d705d454308063421a3b290d9ea4"}, + {file = "contourpy-1.2.0-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:464b423bc2a009088f19bdf1f232299e8b6917963e2b7e1d277da5041f33a779"}, + {file = "contourpy-1.2.0-cp312-cp312-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:68ce4788b7d93e47f84edd3f1f95acdcd142ae60bc0e5493bfd120683d2d4316"}, + {file = "contourpy-1.2.0-cp312-cp312-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:3d7d1f8871998cdff5d2ff6a087e5e1780139abe2838e85b0b46b7ae6cc25399"}, + {file = "contourpy-1.2.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6e739530c662a8d6d42c37c2ed52a6f0932c2d4a3e8c1f90692ad0ce1274abe0"}, + {file = "contourpy-1.2.0-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:247b9d16535acaa766d03037d8e8fb20866d054d3c7fbf6fd1f993f11fc60ca0"}, + {file = "contourpy-1.2.0-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:461e3ae84cd90b30f8d533f07d87c00379644205b1d33a5ea03381edc4b69431"}, + {file = "contourpy-1.2.0-cp312-cp312-win32.whl", hash = "sha256:1c2559d6cffc94890b0529ea7eeecc20d6fadc1539273aa27faf503eb4656d8f"}, + {file = "contourpy-1.2.0-cp312-cp312-win_amd64.whl", hash = "sha256:491b1917afdd8638a05b611a56d46587d5a632cabead889a5440f7c638bc6ed9"}, + {file = "contourpy-1.2.0-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:5fd1810973a375ca0e097dee059c407913ba35723b111df75671a1976efa04bc"}, + {file = "contourpy-1.2.0-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:999c71939aad2780f003979b25ac5b8f2df651dac7b38fb8ce6c46ba5abe6ae9"}, + {file = "contourpy-1.2.0-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:b7caf9b241464c404613512d5594a6e2ff0cc9cb5615c9475cc1d9b514218ae8"}, + {file = "contourpy-1.2.0-cp39-cp39-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:266270c6f6608340f6c9836a0fb9b367be61dde0c9a9a18d5ece97774105ff3e"}, + {file = "contourpy-1.2.0-cp39-cp39-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:dbd50d0a0539ae2e96e537553aff6d02c10ed165ef40c65b0e27e744a0f10af8"}, + {file = "contourpy-1.2.0-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:11f8d2554e52f459918f7b8e6aa20ec2a3bce35ce95c1f0ef4ba36fbda306df5"}, + {file = "contourpy-1.2.0-cp39-cp39-musllinux_1_1_aarch64.whl", hash = "sha256:ce96dd400486e80ac7d195b2d800b03e3e6a787e2a522bfb83755938465a819e"}, + {file = "contourpy-1.2.0-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:6d3364b999c62f539cd403f8123ae426da946e142312a514162adb2addd8d808"}, + {file = "contourpy-1.2.0-cp39-cp39-win32.whl", hash = "sha256:1c88dfb9e0c77612febebb6ac69d44a8d81e3dc60f993215425b62c1161353f4"}, + {file = "contourpy-1.2.0-cp39-cp39-win_amd64.whl", hash = "sha256:78e6ad33cf2e2e80c5dfaaa0beec3d61face0fb650557100ee36db808bfa6843"}, + {file = "contourpy-1.2.0-pp39-pypy39_pp73-macosx_10_9_x86_64.whl", hash = "sha256:be16975d94c320432657ad2402f6760990cb640c161ae6da1363051805fa8108"}, + {file = "contourpy-1.2.0-pp39-pypy39_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:b95a225d4948b26a28c08307a60ac00fb8671b14f2047fc5476613252a129776"}, + {file = "contourpy-1.2.0-pp39-pypy39_pp73-win_amd64.whl", hash = "sha256:0d7e03c0f9a4f90dc18d4e77e9ef4ec7b7bbb437f7f675be8e530d65ae6ef956"}, + {file = "contourpy-1.2.0.tar.gz", hash = "sha256:171f311cb758de7da13fc53af221ae47a5877be5a0843a9fe150818c51ed276a"}, +] + +[package.dependencies] +numpy = ">=1.20,<2.0" + +[package.extras] +bokeh = ["bokeh", "selenium"] +docs = ["furo", "sphinx (>=7.2)", "sphinx-copybutton"] +mypy = ["contourpy[bokeh,docs]", "docutils-stubs", "mypy (==1.6.1)", "types-Pillow"] +test = ["Pillow", "contourpy[test-no-images]", "matplotlib"] +test-no-images = ["pytest", "pytest-cov", "pytest-xdist", "wurlitzer"] + +[[package]] +name = "cycler" +version = "0.12.1" +description = "Composable style cycles" +optional = false +python-versions = ">=3.8" +files = [ + {file = "cycler-0.12.1-py3-none-any.whl", hash = "sha256:85cef7cff222d8644161529808465972e51340599459b8ac3ccbac5a854e0d30"}, + {file = "cycler-0.12.1.tar.gz", hash = "sha256:88bb128f02ba341da8ef447245a9e138fae777f6a23943da4540077d3601eb1c"}, +] + +[package.extras] +docs = ["ipython", "matplotlib", "numpydoc", "sphinx"] +tests = ["pytest", "pytest-cov", "pytest-xdist"] + +[[package]] +name = "dill" +version = "0.3.6" +description = "serialize all of python" +optional = false +python-versions = ">=3.7" +files = [ + {file = "dill-0.3.6-py3-none-any.whl", hash = "sha256:a07ffd2351b8c678dfc4a856a3005f8067aea51d6ba6c700796a4d9e280f39f0"}, + {file = "dill-0.3.6.tar.gz", hash = "sha256:e5db55f3687856d8fbdab002ed78544e1c4559a130302693d839dfe8f93f2373"}, +] + +[package.extras] +graph = ["objgraph (>=1.7.2)"] + +[[package]] +name = "dlib" +version = "19.24.1" +description = "A toolkit for making real world machine learning and data analysis applications" +optional = false +python-versions = "*" +files = [ + {file = "dlib-19.24.1.tar.gz", hash = "sha256:e8d3ca047d525eeb36cc47bce45ed0562dab40b0b8f5fa171600dde2fec509c0"}, +] + +[[package]] +name = "docker-pycreds" +version = "0.4.0" +description = "Python bindings for the docker credentials store API" +optional = false +python-versions = "*" +files = [ + {file = "docker-pycreds-0.4.0.tar.gz", hash = "sha256:6ce3270bcaf404cc4c3e27e4b6c70d3521deae82fb508767870fdbf772d584d4"}, + {file = "docker_pycreds-0.4.0-py2.py3-none-any.whl", hash = "sha256:7266112468627868005106ec19cd0d722702d2b7d5912a28e19b826c3d37af49"}, +] + +[package.dependencies] +six = ">=1.4.0" + +[[package]] +name = "face-alignment" +version = "1.3.4" +description = "Detector 2D or 3D face landmarks from Python" +optional = false +python-versions = ">=3" +files = [ + {file = "face_alignment-1.3.4.tar.gz", hash = "sha256:1056e2fffa139484695f16a2ffd6b1d775ad71eb950032e94bb43ab78c88f53f"}, +] + +[package.dependencies] +numba = "*" +numpy = "*" +opencv-python = "*" +scikit-image = "*" +scipy = ">=0.17" +torch = "*" +tqdm = "*" + +[[package]] +name = "filelock" +version = "3.13.1" +description = "A platform independent file lock." +optional = false +python-versions = ">=3.8" +files = [ + {file = "filelock-3.13.1-py3-none-any.whl", hash = "sha256:57dbda9b35157b05fb3e58ee91448612eb674172fab98ee235ccb0b5bee19a1c"}, + {file = "filelock-3.13.1.tar.gz", hash = "sha256:521f5f56c50f8426f5e03ad3b281b490a87ef15bc6c526f168290f0c7148d44e"}, +] + +[package.extras] +docs = ["furo (>=2023.9.10)", "sphinx (>=7.2.6)", "sphinx-autodoc-typehints (>=1.24)"] +testing = ["covdefaults (>=2.3)", "coverage (>=7.3.2)", "diff-cover (>=8)", "pytest (>=7.4.3)", "pytest-cov (>=4.1)", "pytest-mock (>=3.12)", "pytest-timeout (>=2.2)"] +typing = ["typing-extensions (>=4.8)"] + +[[package]] +name = "fonttools" +version = "4.50.0" +description = "Tools to manipulate font files" +optional = false +python-versions = ">=3.8" +files = [ + {file = "fonttools-4.50.0-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:effd303fb422f8ce06543a36ca69148471144c534cc25f30e5be752bc4f46736"}, + {file = "fonttools-4.50.0-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:7913992ab836f621d06aabac118fc258b9947a775a607e1a737eb3a91c360335"}, + {file = "fonttools-4.50.0-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:8e0a1c5bd2f63da4043b63888534b52c5a1fd7ae187c8ffc64cbb7ae475b9dab"}, + {file = "fonttools-4.50.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:d40fc98540fa5360e7ecf2c56ddf3c6e7dd04929543618fd7b5cc76e66390562"}, + {file = "fonttools-4.50.0-cp310-cp310-musllinux_1_1_aarch64.whl", hash = "sha256:9fff65fbb7afe137bac3113827855e0204482727bddd00a806034ab0d3951d0d"}, + {file = "fonttools-4.50.0-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:b1aeae3dd2ee719074a9372c89ad94f7c581903306d76befdaca2a559f802472"}, + {file = "fonttools-4.50.0-cp310-cp310-win32.whl", hash = "sha256:e9623afa319405da33b43c85cceb0585a6f5d3a1d7c604daf4f7e1dd55c03d1f"}, + {file = "fonttools-4.50.0-cp310-cp310-win_amd64.whl", hash = "sha256:778c5f43e7e654ef7fe0605e80894930bc3a7772e2f496238e57218610140f54"}, + {file = "fonttools-4.50.0-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:3dfb102e7f63b78c832e4539969167ffcc0375b013080e6472350965a5fe8048"}, + {file = "fonttools-4.50.0-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:9e58fe34cb379ba3d01d5d319d67dd3ce7ca9a47ad044ea2b22635cd2d1247fc"}, + {file = "fonttools-4.50.0-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:2c673ab40d15a442a4e6eb09bf007c1dda47c84ac1e2eecbdf359adacb799c24"}, + {file = "fonttools-4.50.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:9b3ac35cdcd1a4c90c23a5200212c1bb74fa05833cc7c14291d7043a52ca2aaa"}, + {file = "fonttools-4.50.0-cp311-cp311-musllinux_1_1_aarch64.whl", hash = "sha256:8844e7a2c5f7ecf977e82eb6b3014f025c8b454e046d941ece05b768be5847ae"}, + {file = "fonttools-4.50.0-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:f849bd3c5c2249b49c98eca5aaebb920d2bfd92b3c69e84ca9bddf133e9f83f0"}, + {file = "fonttools-4.50.0-cp311-cp311-win32.whl", hash = "sha256:39293ff231b36b035575e81c14626dfc14407a20de5262f9596c2cbb199c3625"}, + {file = "fonttools-4.50.0-cp311-cp311-win_amd64.whl", hash = "sha256:c33d5023523b44d3481624f840c8646656a1def7630ca562f222eb3ead16c438"}, + {file = "fonttools-4.50.0-cp312-cp312-macosx_10_9_universal2.whl", hash = "sha256:b4a886a6dbe60100ba1cd24de962f8cd18139bd32808da80de1fa9f9f27bf1dc"}, + {file = "fonttools-4.50.0-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:b2ca1837bfbe5eafa11313dbc7edada79052709a1fffa10cea691210af4aa1fa"}, + {file = "fonttools-4.50.0-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a0493dd97ac8977e48ffc1476b932b37c847cbb87fd68673dee5182004906828"}, + {file = "fonttools-4.50.0-cp312-cp312-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:77844e2f1b0889120b6c222fc49b2b75c3d88b930615e98893b899b9352a27ea"}, + {file = "fonttools-4.50.0-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:3566bfb8c55ed9100afe1ba6f0f12265cd63a1387b9661eb6031a1578a28bad1"}, + {file = "fonttools-4.50.0-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:35e10ddbc129cf61775d58a14f2d44121178d89874d32cae1eac722e687d9019"}, + {file = "fonttools-4.50.0-cp312-cp312-win32.whl", hash = "sha256:cc8140baf9fa8f9b903f2b393a6c413a220fa990264b215bf48484f3d0bf8710"}, + {file = "fonttools-4.50.0-cp312-cp312-win_amd64.whl", hash = "sha256:0ccc85fd96373ab73c59833b824d7a73846670a0cb1f3afbaee2b2c426a8f931"}, + {file = "fonttools-4.50.0-cp38-cp38-macosx_10_9_universal2.whl", hash = "sha256:e270a406219af37581d96c810172001ec536e29e5593aa40d4c01cca3e145aa6"}, + {file = "fonttools-4.50.0-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:ac2463de667233372e9e1c7e9de3d914b708437ef52a3199fdbf5a60184f190c"}, + {file = "fonttools-4.50.0-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:47abd6669195abe87c22750dbcd366dc3a0648f1b7c93c2baa97429c4dc1506e"}, + {file = "fonttools-4.50.0-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:074841375e2e3d559aecc86e1224caf78e8b8417bb391e7d2506412538f21adc"}, + {file = "fonttools-4.50.0-cp38-cp38-musllinux_1_1_aarch64.whl", hash = "sha256:0743fd2191ad7ab43d78cd747215b12033ddee24fa1e088605a3efe80d6984de"}, + {file = "fonttools-4.50.0-cp38-cp38-musllinux_1_1_x86_64.whl", hash = "sha256:3d7080cce7be5ed65bee3496f09f79a82865a514863197ff4d4d177389e981b0"}, + {file = "fonttools-4.50.0-cp38-cp38-win32.whl", hash = "sha256:a467ba4e2eadc1d5cc1a11d355abb945f680473fbe30d15617e104c81f483045"}, + {file = "fonttools-4.50.0-cp38-cp38-win_amd64.whl", hash = "sha256:f77e048f805e00870659d6318fd89ef28ca4ee16a22b4c5e1905b735495fc422"}, + {file = "fonttools-4.50.0-cp39-cp39-macosx_10_9_universal2.whl", hash = "sha256:b6245eafd553c4e9a0708e93be51392bd2288c773523892fbd616d33fd2fda59"}, + {file = "fonttools-4.50.0-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:a4062cc7e8de26f1603323ef3ae2171c9d29c8a9f5e067d555a2813cd5c7a7e0"}, + {file = "fonttools-4.50.0-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:34692850dfd64ba06af61e5791a441f664cb7d21e7b544e8f385718430e8f8e4"}, + {file = "fonttools-4.50.0-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:678dd95f26a67e02c50dcb5bf250f95231d455642afbc65a3b0bcdacd4e4dd38"}, + {file = "fonttools-4.50.0-cp39-cp39-musllinux_1_1_aarch64.whl", hash = "sha256:4f2ce7b0b295fe64ac0a85aef46a0f2614995774bd7bc643b85679c0283287f9"}, + {file = "fonttools-4.50.0-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:d346f4dc2221bfb7ab652d1e37d327578434ce559baf7113b0f55768437fe6a0"}, + {file = "fonttools-4.50.0-cp39-cp39-win32.whl", hash = "sha256:a51eeaf52ba3afd70bf489be20e52fdfafe6c03d652b02477c6ce23c995222f4"}, + {file = "fonttools-4.50.0-cp39-cp39-win_amd64.whl", hash = "sha256:8639be40d583e5d9da67795aa3eeeda0488fb577a1d42ae11a5036f18fb16d93"}, + {file = "fonttools-4.50.0-py3-none-any.whl", hash = "sha256:48fa36da06247aa8282766cfd63efff1bb24e55f020f29a335939ed3844d20d3"}, + {file = "fonttools-4.50.0.tar.gz", hash = "sha256:fa5cf61058c7dbb104c2ac4e782bf1b2016a8cf2f69de6e4dd6a865d2c969bb5"}, +] + +[package.extras] +all = ["brotli (>=1.0.1)", "brotlicffi (>=0.8.0)", "fs (>=2.2.0,<3)", "lxml (>=4.0)", "lz4 (>=1.7.4.2)", "matplotlib", "munkres", "pycairo", "scipy", "skia-pathops (>=0.5.0)", "sympy", "uharfbuzz (>=0.23.0)", "unicodedata2 (>=15.1.0)", "xattr", "zopfli (>=0.1.4)"] +graphite = ["lz4 (>=1.7.4.2)"] +interpolatable = ["munkres", "pycairo", "scipy"] +lxml = ["lxml (>=4.0)"] +pathops = ["skia-pathops (>=0.5.0)"] +plot = ["matplotlib"] +repacker = ["uharfbuzz (>=0.23.0)"] +symfont = ["sympy"] +type1 = ["xattr"] +ufo = ["fs (>=2.2.0,<3)"] +unicode = ["unicodedata2 (>=15.1.0)"] +woff = ["brotli (>=1.0.1)", "brotlicffi (>=0.8.0)", "zopfli (>=0.1.4)"] + +[[package]] +name = "ftfy" +version = "6.2.0" +description = "Fixes mojibake and other problems with Unicode, after the fact" +optional = false +python-versions = ">=3.8,<4" +files = [ + {file = "ftfy-6.2.0-py3-none-any.whl", hash = "sha256:f94a2c34b76e07475720e3096f5ca80911d152406fbde66fdb45c4d0c9150026"}, + {file = "ftfy-6.2.0.tar.gz", hash = "sha256:5e42143c7025ef97944ca2619d6b61b0619fc6654f98771d39e862c1424c75c0"}, +] + +[package.dependencies] +wcwidth = ">=0.2.12,<0.3.0" + +[[package]] +name = "gdown" +version = "3.12.2" +description = "Google Drive direct download of big files." +optional = false +python-versions = "*" +files = [ + {file = "gdown-3.12.2.tar.gz", hash = "sha256:4b3a1301e57bfd8dce939bf25ef8fbb4b23967fd0f878eede328bdcc41386bac"}, +] + +[package.dependencies] +filelock = "*" +requests = {version = "*", extras = ["socks"]} +six = "*" +tqdm = "*" + +[[package]] +name = "gitdb" +version = "4.0.11" +description = "Git Object Database" +optional = false +python-versions = ">=3.7" +files = [ + {file = "gitdb-4.0.11-py3-none-any.whl", hash = "sha256:81a3407ddd2ee8df444cbacea00e2d038e40150acfa3001696fe0dcf1d3adfa4"}, + {file = "gitdb-4.0.11.tar.gz", hash = "sha256:bf5421126136d6d0af55bc1e7c1af1c397a34f5b7bd79e776cd3e89785c2b04b"}, +] + +[package.dependencies] +smmap = ">=3.0.1,<6" + +[[package]] +name = "gitpython" +version = "3.1.42" +description = "GitPython is a Python library used to interact with Git repositories" +optional = false +python-versions = ">=3.7" +files = [ + {file = "GitPython-3.1.42-py3-none-any.whl", hash = "sha256:1bf9cd7c9e7255f77778ea54359e54ac22a72a5b51288c457c881057b7bb9ecd"}, + {file = "GitPython-3.1.42.tar.gz", hash = "sha256:2d99869e0fef71a73cbd242528105af1d6c1b108c60dfabd994bf292f76c3ceb"}, +] + +[package.dependencies] +gitdb = ">=4.0.1,<5" + +[package.extras] +test = ["black", "coverage[toml]", "ddt (>=1.1.1,!=1.4.3)", "mock", "mypy", "pre-commit", "pytest (>=7.3.1)", "pytest-cov", "pytest-instafail", "pytest-mock", "pytest-sugar"] + +[[package]] +name = "idna" +version = "3.6" +description = "Internationalized Domain Names in Applications (IDNA)" +optional = false +python-versions = ">=3.5" +files = [ + {file = "idna-3.6-py3-none-any.whl", hash = "sha256:c05567e9c24a6b9faaa835c4821bad0590fbb9d5779e7caa6e1cc4978e7eb24f"}, + {file = "idna-3.6.tar.gz", hash = "sha256:9ecdbbd083b06798ae1e86adcbfe8ab1479cf864e4ee30fe4e46a003d12491ca"}, +] + +[[package]] +name = "imageio" +version = "2.34.0" +description = "Library for reading and writing a wide range of image, video, scientific, and volumetric data formats." +optional = false +python-versions = ">=3.8" +files = [ + {file = "imageio-2.34.0-py3-none-any.whl", hash = "sha256:08082bf47ccb54843d9c73fe9fc8f3a88c72452ab676b58aca74f36167e8ccba"}, + {file = "imageio-2.34.0.tar.gz", hash = "sha256:ae9732e10acf807a22c389aef193f42215718e16bd06eed0c5bb57e1034a4d53"}, +] + +[package.dependencies] +numpy = "*" +pillow = ">=8.3.2" + +[package.extras] +all-plugins = ["astropy", "av", "imageio-ffmpeg", "pillow-heif", "psutil", "tifffile"] +all-plugins-pypy = ["av", "imageio-ffmpeg", "pillow-heif", "psutil", "tifffile"] +build = ["wheel"] +dev = ["black", "flake8", "fsspec[github]", "pytest", "pytest-cov"] +docs = ["numpydoc", "pydata-sphinx-theme", "sphinx (<6)"] +ffmpeg = ["imageio-ffmpeg", "psutil"] +fits = ["astropy"] +full = ["astropy", "av", "black", "flake8", "fsspec[github]", "gdal", "imageio-ffmpeg", "itk", "numpydoc", "pillow-heif", "psutil", "pydata-sphinx-theme", "pytest", "pytest-cov", "sphinx (<6)", "tifffile", "wheel"] +gdal = ["gdal"] +itk = ["itk"] +linting = ["black", "flake8"] +pillow-heif = ["pillow-heif"] +pyav = ["av"] +test = ["fsspec[github]", "pytest", "pytest-cov"] +tifffile = ["tifffile"] + +[[package]] +name = "joblib" +version = "1.2.0" +description = "Lightweight pipelining with Python functions" +optional = false +python-versions = ">=3.7" +files = [ + {file = "joblib-1.2.0-py3-none-any.whl", hash = "sha256:091138ed78f800342968c523bdde947e7a305b8594b910a0fea2ab83c3c6d385"}, + {file = "joblib-1.2.0.tar.gz", hash = "sha256:e1cee4a79e4af22881164f218d4311f60074197fb707e082e803b61f6d137018"}, +] + +[[package]] +name = "kiwisolver" +version = "1.4.5" +description = "A fast implementation of the Cassowary constraint solver" +optional = false +python-versions = ">=3.7" +files = [ + {file = "kiwisolver-1.4.5-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:05703cf211d585109fcd72207a31bb170a0f22144d68298dc5e61b3c946518af"}, + {file = "kiwisolver-1.4.5-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:146d14bebb7f1dc4d5fbf74f8a6cb15ac42baadee8912eb84ac0b3b2a3dc6ac3"}, + {file = "kiwisolver-1.4.5-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:6ef7afcd2d281494c0a9101d5c571970708ad911d028137cd558f02b851c08b4"}, + {file = "kiwisolver-1.4.5-cp310-cp310-manylinux_2_12_i686.manylinux2010_i686.whl", hash = "sha256:9eaa8b117dc8337728e834b9c6e2611f10c79e38f65157c4c38e9400286f5cb1"}, + {file = "kiwisolver-1.4.5-cp310-cp310-manylinux_2_12_x86_64.manylinux2010_x86_64.whl", hash = "sha256:ec20916e7b4cbfb1f12380e46486ec4bcbaa91a9c448b97023fde0d5bbf9e4ff"}, + {file = "kiwisolver-1.4.5-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:39b42c68602539407884cf70d6a480a469b93b81b7701378ba5e2328660c847a"}, + {file = "kiwisolver-1.4.5-cp310-cp310-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:aa12042de0171fad672b6c59df69106d20d5596e4f87b5e8f76df757a7c399aa"}, + {file = "kiwisolver-1.4.5-cp310-cp310-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:2a40773c71d7ccdd3798f6489aaac9eee213d566850a9533f8d26332d626b82c"}, + {file = "kiwisolver-1.4.5-cp310-cp310-musllinux_1_1_aarch64.whl", hash = "sha256:19df6e621f6d8b4b9c4d45f40a66839294ff2bb235e64d2178f7522d9170ac5b"}, + {file = "kiwisolver-1.4.5-cp310-cp310-musllinux_1_1_i686.whl", hash = "sha256:83d78376d0d4fd884e2c114d0621624b73d2aba4e2788182d286309ebdeed770"}, + {file = "kiwisolver-1.4.5-cp310-cp310-musllinux_1_1_ppc64le.whl", hash = "sha256:e391b1f0a8a5a10ab3b9bb6afcfd74f2175f24f8975fb87ecae700d1503cdee0"}, + {file = "kiwisolver-1.4.5-cp310-cp310-musllinux_1_1_s390x.whl", hash = "sha256:852542f9481f4a62dbb5dd99e8ab7aedfeb8fb6342349a181d4036877410f525"}, + {file = "kiwisolver-1.4.5-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:59edc41b24031bc25108e210c0def6f6c2191210492a972d585a06ff246bb79b"}, + {file = "kiwisolver-1.4.5-cp310-cp310-win32.whl", hash = "sha256:a6aa6315319a052b4ee378aa171959c898a6183f15c1e541821c5c59beaa0238"}, + {file = "kiwisolver-1.4.5-cp310-cp310-win_amd64.whl", hash = "sha256:d0ef46024e6a3d79c01ff13801cb19d0cad7fd859b15037aec74315540acc276"}, + {file = "kiwisolver-1.4.5-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:11863aa14a51fd6ec28688d76f1735f8f69ab1fabf388851a595d0721af042f5"}, + {file = "kiwisolver-1.4.5-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:8ab3919a9997ab7ef2fbbed0cc99bb28d3c13e6d4b1ad36e97e482558a91be90"}, + {file = "kiwisolver-1.4.5-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:fcc700eadbbccbf6bc1bcb9dbe0786b4b1cb91ca0dcda336eef5c2beed37b797"}, + {file = "kiwisolver-1.4.5-cp311-cp311-manylinux_2_12_i686.manylinux2010_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:dfdd7c0b105af050eb3d64997809dc21da247cf44e63dc73ff0fd20b96be55a9"}, + {file = "kiwisolver-1.4.5-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:76c6a5964640638cdeaa0c359382e5703e9293030fe730018ca06bc2010c4437"}, + {file = "kiwisolver-1.4.5-cp311-cp311-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:bbea0db94288e29afcc4c28afbf3a7ccaf2d7e027489c449cf7e8f83c6346eb9"}, + {file = "kiwisolver-1.4.5-cp311-cp311-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:ceec1a6bc6cab1d6ff5d06592a91a692f90ec7505d6463a88a52cc0eb58545da"}, + {file = "kiwisolver-1.4.5-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:040c1aebeda72197ef477a906782b5ab0d387642e93bda547336b8957c61022e"}, + {file = "kiwisolver-1.4.5-cp311-cp311-musllinux_1_1_aarch64.whl", hash = "sha256:f91de7223d4c7b793867797bacd1ee53bfe7359bd70d27b7b58a04efbb9436c8"}, + {file = "kiwisolver-1.4.5-cp311-cp311-musllinux_1_1_i686.whl", hash = "sha256:faae4860798c31530dd184046a900e652c95513796ef51a12bc086710c2eec4d"}, + {file = "kiwisolver-1.4.5-cp311-cp311-musllinux_1_1_ppc64le.whl", hash = "sha256:b0157420efcb803e71d1b28e2c287518b8808b7cf1ab8af36718fd0a2c453eb0"}, + {file = "kiwisolver-1.4.5-cp311-cp311-musllinux_1_1_s390x.whl", hash = "sha256:06f54715b7737c2fecdbf140d1afb11a33d59508a47bf11bb38ecf21dc9ab79f"}, + {file = "kiwisolver-1.4.5-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:fdb7adb641a0d13bdcd4ef48e062363d8a9ad4a182ac7647ec88f695e719ae9f"}, + {file = "kiwisolver-1.4.5-cp311-cp311-win32.whl", hash = "sha256:bb86433b1cfe686da83ce32a9d3a8dd308e85c76b60896d58f082136f10bffac"}, + {file = "kiwisolver-1.4.5-cp311-cp311-win_amd64.whl", hash = "sha256:6c08e1312a9cf1074d17b17728d3dfce2a5125b2d791527f33ffbe805200a355"}, + {file = "kiwisolver-1.4.5-cp312-cp312-macosx_10_9_universal2.whl", hash = "sha256:32d5cf40c4f7c7b3ca500f8985eb3fb3a7dfc023215e876f207956b5ea26632a"}, + {file = "kiwisolver-1.4.5-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:f846c260f483d1fd217fe5ed7c173fb109efa6b1fc8381c8b7552c5781756192"}, + {file = "kiwisolver-1.4.5-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:5ff5cf3571589b6d13bfbfd6bcd7a3f659e42f96b5fd1c4830c4cf21d4f5ef45"}, + {file = "kiwisolver-1.4.5-cp312-cp312-manylinux_2_12_i686.manylinux2010_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:7269d9e5f1084a653d575c7ec012ff57f0c042258bf5db0954bf551c158466e7"}, + {file = "kiwisolver-1.4.5-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:da802a19d6e15dffe4b0c24b38b3af68e6c1a68e6e1d8f30148c83864f3881db"}, + {file = "kiwisolver-1.4.5-cp312-cp312-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:3aba7311af82e335dd1e36ffff68aaca609ca6290c2cb6d821a39aa075d8e3ff"}, + {file = "kiwisolver-1.4.5-cp312-cp312-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:763773d53f07244148ccac5b084da5adb90bfaee39c197554f01b286cf869228"}, + {file = "kiwisolver-1.4.5-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:2270953c0d8cdab5d422bee7d2007f043473f9d2999631c86a223c9db56cbd16"}, + {file = "kiwisolver-1.4.5-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:d099e745a512f7e3bbe7249ca835f4d357c586d78d79ae8f1dcd4d8adeb9bda9"}, + {file = "kiwisolver-1.4.5-cp312-cp312-musllinux_1_1_i686.whl", hash = "sha256:74db36e14a7d1ce0986fa104f7d5637aea5c82ca6326ed0ec5694280942d1162"}, + {file = "kiwisolver-1.4.5-cp312-cp312-musllinux_1_1_ppc64le.whl", hash = "sha256:7e5bab140c309cb3a6ce373a9e71eb7e4873c70c2dda01df6820474f9889d6d4"}, + {file = "kiwisolver-1.4.5-cp312-cp312-musllinux_1_1_s390x.whl", hash = "sha256:0f114aa76dc1b8f636d077979c0ac22e7cd8f3493abbab152f20eb8d3cda71f3"}, + {file = "kiwisolver-1.4.5-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:88a2df29d4724b9237fc0c6eaf2a1adae0cdc0b3e9f4d8e7dc54b16812d2d81a"}, + {file = "kiwisolver-1.4.5-cp312-cp312-win32.whl", hash = "sha256:72d40b33e834371fd330fb1472ca19d9b8327acb79a5821d4008391db8e29f20"}, + {file = "kiwisolver-1.4.5-cp312-cp312-win_amd64.whl", hash = "sha256:2c5674c4e74d939b9d91dda0fae10597ac7521768fec9e399c70a1f27e2ea2d9"}, + {file = "kiwisolver-1.4.5-cp37-cp37m-macosx_10_9_x86_64.whl", hash = "sha256:3a2b053a0ab7a3960c98725cfb0bf5b48ba82f64ec95fe06f1d06c99b552e130"}, + {file = "kiwisolver-1.4.5-cp37-cp37m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:3cd32d6c13807e5c66a7cbb79f90b553642f296ae4518a60d8d76243b0ad2898"}, + {file = "kiwisolver-1.4.5-cp37-cp37m-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:59ec7b7c7e1a61061850d53aaf8e93db63dce0c936db1fda2658b70e4a1be709"}, + {file = "kiwisolver-1.4.5-cp37-cp37m-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:da4cfb373035def307905d05041c1d06d8936452fe89d464743ae7fb8371078b"}, + {file = "kiwisolver-1.4.5-cp37-cp37m-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:2400873bccc260b6ae184b2b8a4fec0e4082d30648eadb7c3d9a13405d861e89"}, + {file = "kiwisolver-1.4.5-cp37-cp37m-manylinux_2_5_x86_64.manylinux1_x86_64.whl", hash = "sha256:1b04139c4236a0f3aff534479b58f6f849a8b351e1314826c2d230849ed48985"}, + {file = "kiwisolver-1.4.5-cp37-cp37m-musllinux_1_1_aarch64.whl", hash = "sha256:4e66e81a5779b65ac21764c295087de82235597a2293d18d943f8e9e32746265"}, + {file = "kiwisolver-1.4.5-cp37-cp37m-musllinux_1_1_i686.whl", hash = "sha256:7931d8f1f67c4be9ba1dd9c451fb0eeca1a25b89e4d3f89e828fe12a519b782a"}, + {file = "kiwisolver-1.4.5-cp37-cp37m-musllinux_1_1_ppc64le.whl", hash = "sha256:b3f7e75f3015df442238cca659f8baa5f42ce2a8582727981cbfa15fee0ee205"}, + {file = "kiwisolver-1.4.5-cp37-cp37m-musllinux_1_1_s390x.whl", hash = "sha256:bbf1d63eef84b2e8c89011b7f2235b1e0bf7dacc11cac9431fc6468e99ac77fb"}, + {file = "kiwisolver-1.4.5-cp37-cp37m-musllinux_1_1_x86_64.whl", hash = "sha256:4c380469bd3f970ef677bf2bcba2b6b0b4d5c75e7a020fb863ef75084efad66f"}, + {file = "kiwisolver-1.4.5-cp37-cp37m-win32.whl", hash = "sha256:9408acf3270c4b6baad483865191e3e582b638b1654a007c62e3efe96f09a9a3"}, + {file = "kiwisolver-1.4.5-cp37-cp37m-win_amd64.whl", hash = "sha256:5b94529f9b2591b7af5f3e0e730a4e0a41ea174af35a4fd067775f9bdfeee01a"}, + {file = "kiwisolver-1.4.5-cp38-cp38-macosx_10_9_universal2.whl", hash = "sha256:11c7de8f692fc99816e8ac50d1d1aef4f75126eefc33ac79aac02c099fd3db71"}, + {file = "kiwisolver-1.4.5-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:53abb58632235cd154176ced1ae8f0d29a6657aa1aa9decf50b899b755bc2b93"}, + {file = "kiwisolver-1.4.5-cp38-cp38-macosx_11_0_arm64.whl", hash = "sha256:88b9f257ca61b838b6f8094a62418421f87ac2a1069f7e896c36a7d86b5d4c29"}, + {file = "kiwisolver-1.4.5-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:3195782b26fc03aa9c6913d5bad5aeb864bdc372924c093b0f1cebad603dd712"}, + {file = "kiwisolver-1.4.5-cp38-cp38-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:fc579bf0f502e54926519451b920e875f433aceb4624a3646b3252b5caa9e0b6"}, + {file = "kiwisolver-1.4.5-cp38-cp38-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:5a580c91d686376f0f7c295357595c5a026e6cbc3d77b7c36e290201e7c11ecb"}, + {file = "kiwisolver-1.4.5-cp38-cp38-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:cfe6ab8da05c01ba6fbea630377b5da2cd9bcbc6338510116b01c1bc939a2c18"}, + {file = "kiwisolver-1.4.5-cp38-cp38-manylinux_2_5_x86_64.manylinux1_x86_64.whl", hash = "sha256:d2e5a98f0ec99beb3c10e13b387f8db39106d53993f498b295f0c914328b1333"}, + {file = "kiwisolver-1.4.5-cp38-cp38-musllinux_1_1_aarch64.whl", hash = "sha256:a51a263952b1429e429ff236d2f5a21c5125437861baeed77f5e1cc2d2c7c6da"}, + {file = "kiwisolver-1.4.5-cp38-cp38-musllinux_1_1_i686.whl", hash = "sha256:3edd2fa14e68c9be82c5b16689e8d63d89fe927e56debd6e1dbce7a26a17f81b"}, + {file = "kiwisolver-1.4.5-cp38-cp38-musllinux_1_1_ppc64le.whl", hash = "sha256:74d1b44c6cfc897df648cc9fdaa09bc3e7679926e6f96df05775d4fb3946571c"}, + {file = "kiwisolver-1.4.5-cp38-cp38-musllinux_1_1_s390x.whl", hash = "sha256:76d9289ed3f7501012e05abb8358bbb129149dbd173f1f57a1bf1c22d19ab7cc"}, + {file = "kiwisolver-1.4.5-cp38-cp38-musllinux_1_1_x86_64.whl", hash = "sha256:92dea1ffe3714fa8eb6a314d2b3c773208d865a0e0d35e713ec54eea08a66250"}, + {file = "kiwisolver-1.4.5-cp38-cp38-win32.whl", hash = "sha256:5c90ae8c8d32e472be041e76f9d2f2dbff4d0b0be8bd4041770eddb18cf49a4e"}, + {file = "kiwisolver-1.4.5-cp38-cp38-win_amd64.whl", hash = "sha256:c7940c1dc63eb37a67721b10d703247552416f719c4188c54e04334321351ced"}, + {file = "kiwisolver-1.4.5-cp39-cp39-macosx_10_9_universal2.whl", hash = "sha256:9407b6a5f0d675e8a827ad8742e1d6b49d9c1a1da5d952a67d50ef5f4170b18d"}, + {file = "kiwisolver-1.4.5-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:15568384086b6df3c65353820a4473575dbad192e35010f622c6ce3eebd57af9"}, + {file = "kiwisolver-1.4.5-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:0dc9db8e79f0036e8173c466d21ef18e1befc02de8bf8aa8dc0813a6dc8a7046"}, + {file = "kiwisolver-1.4.5-cp39-cp39-manylinux_2_12_i686.manylinux2010_i686.whl", hash = "sha256:cdc8a402aaee9a798b50d8b827d7ecf75edc5fb35ea0f91f213ff927c15f4ff0"}, + {file = "kiwisolver-1.4.5-cp39-cp39-manylinux_2_12_x86_64.manylinux2010_x86_64.whl", hash = "sha256:6c3bd3cde54cafb87d74d8db50b909705c62b17c2099b8f2e25b461882e544ff"}, + {file = "kiwisolver-1.4.5-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:955e8513d07a283056b1396e9a57ceddbd272d9252c14f154d450d227606eb54"}, + {file = "kiwisolver-1.4.5-cp39-cp39-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:346f5343b9e3f00b8db8ba359350eb124b98c99efd0b408728ac6ebf38173958"}, + {file = "kiwisolver-1.4.5-cp39-cp39-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:b9098e0049e88c6a24ff64545cdfc50807818ba6c1b739cae221bbbcbc58aad3"}, + {file = "kiwisolver-1.4.5-cp39-cp39-musllinux_1_1_aarch64.whl", hash = "sha256:00bd361b903dc4bbf4eb165f24d1acbee754fce22ded24c3d56eec268658a5cf"}, + {file = "kiwisolver-1.4.5-cp39-cp39-musllinux_1_1_i686.whl", hash = "sha256:7b8b454bac16428b22560d0a1cf0a09875339cab69df61d7805bf48919415901"}, + {file = "kiwisolver-1.4.5-cp39-cp39-musllinux_1_1_ppc64le.whl", hash = "sha256:f1d072c2eb0ad60d4c183f3fb44ac6f73fb7a8f16a2694a91f988275cbf352f9"}, + {file = "kiwisolver-1.4.5-cp39-cp39-musllinux_1_1_s390x.whl", hash = "sha256:31a82d498054cac9f6d0b53d02bb85811185bcb477d4b60144f915f3b3126342"}, + {file = "kiwisolver-1.4.5-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:6512cb89e334e4700febbffaaa52761b65b4f5a3cf33f960213d5656cea36a77"}, + {file = "kiwisolver-1.4.5-cp39-cp39-win32.whl", hash = "sha256:9db8ea4c388fdb0f780fe91346fd438657ea602d58348753d9fb265ce1bca67f"}, + {file = "kiwisolver-1.4.5-cp39-cp39-win_amd64.whl", hash = "sha256:59415f46a37f7f2efeec758353dd2eae1b07640d8ca0f0c42548ec4125492635"}, + {file = "kiwisolver-1.4.5-pp37-pypy37_pp73-macosx_10_9_x86_64.whl", hash = "sha256:5c7b3b3a728dc6faf3fc372ef24f21d1e3cee2ac3e9596691d746e5a536de920"}, + {file = "kiwisolver-1.4.5-pp37-pypy37_pp73-manylinux_2_12_i686.manylinux2010_i686.whl", hash = "sha256:620ced262a86244e2be10a676b646f29c34537d0d9cc8eb26c08f53d98013390"}, + {file = "kiwisolver-1.4.5-pp37-pypy37_pp73-manylinux_2_12_x86_64.manylinux2010_x86_64.whl", hash = "sha256:378a214a1e3bbf5ac4a8708304318b4f890da88c9e6a07699c4ae7174c09a68d"}, + {file = "kiwisolver-1.4.5-pp37-pypy37_pp73-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:aaf7be1207676ac608a50cd08f102f6742dbfc70e8d60c4db1c6897f62f71523"}, + {file = "kiwisolver-1.4.5-pp37-pypy37_pp73-win_amd64.whl", hash = "sha256:ba55dce0a9b8ff59495ddd050a0225d58bd0983d09f87cfe2b6aec4f2c1234e4"}, + {file = "kiwisolver-1.4.5-pp38-pypy38_pp73-macosx_10_9_x86_64.whl", hash = "sha256:fd32ea360bcbb92d28933fc05ed09bffcb1704ba3fc7942e81db0fd4f81a7892"}, + {file = "kiwisolver-1.4.5-pp38-pypy38_pp73-manylinux_2_12_i686.manylinux2010_i686.whl", hash = "sha256:5e7139af55d1688f8b960ee9ad5adafc4ac17c1c473fe07133ac092310d76544"}, + {file = "kiwisolver-1.4.5-pp38-pypy38_pp73-manylinux_2_12_x86_64.manylinux2010_x86_64.whl", hash = "sha256:dced8146011d2bc2e883f9bd68618b8247387f4bbec46d7392b3c3b032640126"}, + {file = "kiwisolver-1.4.5-pp38-pypy38_pp73-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:c9bf3325c47b11b2e51bca0824ea217c7cd84491d8ac4eefd1e409705ef092bd"}, + {file = "kiwisolver-1.4.5-pp38-pypy38_pp73-win_amd64.whl", hash = "sha256:5794cf59533bc3f1b1c821f7206a3617999db9fbefc345360aafe2e067514929"}, + {file = "kiwisolver-1.4.5-pp39-pypy39_pp73-macosx_10_9_x86_64.whl", hash = "sha256:e368f200bbc2e4f905b8e71eb38b3c04333bddaa6a2464a6355487b02bb7fb09"}, + {file = "kiwisolver-1.4.5-pp39-pypy39_pp73-manylinux_2_12_i686.manylinux2010_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:e5d706eba36b4c4d5bc6c6377bb6568098765e990cfc21ee16d13963fab7b3e7"}, + {file = "kiwisolver-1.4.5-pp39-pypy39_pp73-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:85267bd1aa8880a9c88a8cb71e18d3d64d2751a790e6ca6c27b8ccc724bcd5ad"}, + {file = "kiwisolver-1.4.5-pp39-pypy39_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:210ef2c3a1f03272649aff1ef992df2e724748918c4bc2d5a90352849eb40bea"}, + {file = "kiwisolver-1.4.5-pp39-pypy39_pp73-win_amd64.whl", hash = "sha256:11d011a7574eb3b82bcc9c1a1d35c1d7075677fdd15de527d91b46bd35e935ee"}, + {file = "kiwisolver-1.4.5.tar.gz", hash = "sha256:e57e563a57fb22a142da34f38acc2fc1a5c864bc29ca1517a88abc963e60d6ec"}, +] + +[[package]] +name = "lazy-loader" +version = "0.3" +description = "lazy_loader" +optional = false +python-versions = ">=3.7" +files = [ + {file = "lazy_loader-0.3-py3-none-any.whl", hash = "sha256:1e9e76ee8631e264c62ce10006718e80b2cfc74340d17d1031e0f84af7478554"}, + {file = "lazy_loader-0.3.tar.gz", hash = "sha256:3b68898e34f5b2a29daaaac172c6555512d0f32074f147e2254e4a6d9d838f37"}, +] + +[package.extras] +lint = ["pre-commit (>=3.3)"] +test = ["pytest (>=7.4)", "pytest-cov (>=4.1)"] + +[[package]] +name = "llvmlite" +version = "0.42.0" +description = "lightweight wrapper around basic LLVM functionality" +optional = false +python-versions = ">=3.9" +files = [ + {file = "llvmlite-0.42.0-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:3366938e1bf63d26c34fbfb4c8e8d2ded57d11e0567d5bb243d89aab1eb56098"}, + {file = "llvmlite-0.42.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:c35da49666a21185d21b551fc3caf46a935d54d66969d32d72af109b5e7d2b6f"}, + {file = "llvmlite-0.42.0-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:70f44ccc3c6220bd23e0ba698a63ec2a7d3205da0d848804807f37fc243e3f77"}, + {file = "llvmlite-0.42.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:763f8d8717a9073b9e0246998de89929071d15b47f254c10eef2310b9aac033d"}, + {file = "llvmlite-0.42.0-cp310-cp310-win_amd64.whl", hash = "sha256:8d90edf400b4ceb3a0e776b6c6e4656d05c7187c439587e06f86afceb66d2be5"}, + {file = "llvmlite-0.42.0-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:ae511caed28beaf1252dbaf5f40e663f533b79ceb408c874c01754cafabb9cbf"}, + {file = "llvmlite-0.42.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:81e674c2fe85576e6c4474e8c7e7aba7901ac0196e864fe7985492b737dbab65"}, + {file = "llvmlite-0.42.0-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:bb3975787f13eb97629052edb5017f6c170eebc1c14a0433e8089e5db43bcce6"}, + {file = "llvmlite-0.42.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:c5bece0cdf77f22379f19b1959ccd7aee518afa4afbd3656c6365865f84903f9"}, + {file = "llvmlite-0.42.0-cp311-cp311-win_amd64.whl", hash = "sha256:7e0c4c11c8c2aa9b0701f91b799cb9134a6a6de51444eff5a9087fc7c1384275"}, + {file = "llvmlite-0.42.0-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:08fa9ab02b0d0179c688a4216b8939138266519aaa0aa94f1195a8542faedb56"}, + {file = "llvmlite-0.42.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:b2fce7d355068494d1e42202c7aff25d50c462584233013eb4470c33b995e3ee"}, + {file = "llvmlite-0.42.0-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:ebe66a86dc44634b59a3bc860c7b20d26d9aaffcd30364ebe8ba79161a9121f4"}, + {file = "llvmlite-0.42.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:d47494552559e00d81bfb836cf1c4d5a5062e54102cc5767d5aa1e77ccd2505c"}, + {file = "llvmlite-0.42.0-cp312-cp312-win_amd64.whl", hash = "sha256:05cb7e9b6ce69165ce4d1b994fbdedca0c62492e537b0cc86141b6e2c78d5888"}, + {file = "llvmlite-0.42.0-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:bdd3888544538a94d7ec99e7c62a0cdd8833609c85f0c23fcb6c5c591aec60ad"}, + {file = "llvmlite-0.42.0-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:d0936c2067a67fb8816c908d5457d63eba3e2b17e515c5fe00e5ee2bace06040"}, + {file = "llvmlite-0.42.0-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a78ab89f1924fc11482209f6799a7a3fc74ddc80425a7a3e0e8174af0e9e2301"}, + {file = "llvmlite-0.42.0-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:d7599b65c7af7abbc978dbf345712c60fd596aa5670496561cc10e8a71cebfb2"}, + {file = "llvmlite-0.42.0-cp39-cp39-win_amd64.whl", hash = "sha256:43d65cc4e206c2e902c1004dd5418417c4efa6c1d04df05c6c5675a27e8ca90e"}, + {file = "llvmlite-0.42.0.tar.gz", hash = "sha256:f92b09243c0cc3f457da8b983f67bd8e1295d0f5b3746c7a1861d7a99403854a"}, +] + +[[package]] +name = "matplotlib" +version = "3.7.1" +description = "Python plotting package" +optional = false +python-versions = ">=3.8" +files = [ + {file = "matplotlib-3.7.1-cp310-cp310-macosx_10_12_universal2.whl", hash = "sha256:95cbc13c1fc6844ab8812a525bbc237fa1470863ff3dace7352e910519e194b1"}, + {file = "matplotlib-3.7.1-cp310-cp310-macosx_10_12_x86_64.whl", hash = "sha256:08308bae9e91aca1ec6fd6dda66237eef9f6294ddb17f0d0b3c863169bf82353"}, + {file = "matplotlib-3.7.1-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:544764ba51900da4639c0f983b323d288f94f65f4024dc40ecb1542d74dc0500"}, + {file = "matplotlib-3.7.1-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:56d94989191de3fcc4e002f93f7f1be5da476385dde410ddafbb70686acf00ea"}, + {file = "matplotlib-3.7.1-cp310-cp310-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:e99bc9e65901bb9a7ce5e7bb24af03675cbd7c70b30ac670aa263240635999a4"}, + {file = "matplotlib-3.7.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:eb7d248c34a341cd4c31a06fd34d64306624c8cd8d0def7abb08792a5abfd556"}, + {file = "matplotlib-3.7.1-cp310-cp310-win32.whl", hash = "sha256:ce463ce590f3825b52e9fe5c19a3c6a69fd7675a39d589e8b5fbe772272b3a24"}, + {file = "matplotlib-3.7.1-cp310-cp310-win_amd64.whl", hash = "sha256:3d7bc90727351fb841e4d8ae620d2d86d8ed92b50473cd2b42ce9186104ecbba"}, + {file = "matplotlib-3.7.1-cp311-cp311-macosx_10_12_universal2.whl", hash = "sha256:770a205966d641627fd5cf9d3cb4b6280a716522cd36b8b284a8eb1581310f61"}, + {file = "matplotlib-3.7.1-cp311-cp311-macosx_10_12_x86_64.whl", hash = "sha256:f67bfdb83a8232cb7a92b869f9355d677bce24485c460b19d01970b64b2ed476"}, + {file = "matplotlib-3.7.1-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:2bf092f9210e105f414a043b92af583c98f50050559616930d884387d0772aba"}, + {file = "matplotlib-3.7.1-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:89768d84187f31717349c6bfadc0e0d8c321e8eb34522acec8a67b1236a66332"}, + {file = "matplotlib-3.7.1-cp311-cp311-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:83111e6388dec67822e2534e13b243cc644c7494a4bb60584edbff91585a83c6"}, + {file = "matplotlib-3.7.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:a867bf73a7eb808ef2afbca03bcdb785dae09595fbe550e1bab0cd023eba3de0"}, + {file = "matplotlib-3.7.1-cp311-cp311-win32.whl", hash = "sha256:fbdeeb58c0cf0595efe89c05c224e0a502d1aa6a8696e68a73c3efc6bc354304"}, + {file = "matplotlib-3.7.1-cp311-cp311-win_amd64.whl", hash = "sha256:c0bd19c72ae53e6ab979f0ac6a3fafceb02d2ecafa023c5cca47acd934d10be7"}, + {file = "matplotlib-3.7.1-cp38-cp38-macosx_10_12_universal2.whl", hash = "sha256:6eb88d87cb2c49af00d3bbc33a003f89fd9f78d318848da029383bfc08ecfbfb"}, + {file = "matplotlib-3.7.1-cp38-cp38-macosx_10_12_x86_64.whl", hash = "sha256:cf0e4f727534b7b1457898c4f4ae838af1ef87c359b76dcd5330fa31893a3ac7"}, + {file = "matplotlib-3.7.1-cp38-cp38-macosx_11_0_arm64.whl", hash = "sha256:46a561d23b91f30bccfd25429c3c706afe7d73a5cc64ef2dfaf2b2ac47c1a5dc"}, + {file = "matplotlib-3.7.1-cp38-cp38-manylinux_2_12_i686.manylinux2010_i686.whl", hash = "sha256:8704726d33e9aa8a6d5215044b8d00804561971163563e6e6591f9dcf64340cc"}, + {file = "matplotlib-3.7.1-cp38-cp38-manylinux_2_12_x86_64.manylinux2010_x86_64.whl", hash = "sha256:4cf327e98ecf08fcbb82685acaf1939d3338548620ab8dfa02828706402c34de"}, + {file = "matplotlib-3.7.1-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:617f14ae9d53292ece33f45cba8503494ee199a75b44de7717964f70637a36aa"}, + {file = "matplotlib-3.7.1-cp38-cp38-win32.whl", hash = "sha256:7c9a4b2da6fac77bcc41b1ea95fadb314e92508bf5493ceff058e727e7ecf5b0"}, + {file = "matplotlib-3.7.1-cp38-cp38-win_amd64.whl", hash = "sha256:14645aad967684e92fc349493fa10c08a6da514b3d03a5931a1bac26e6792bd1"}, + {file = "matplotlib-3.7.1-cp39-cp39-macosx_10_12_universal2.whl", hash = "sha256:81a6b377ea444336538638d31fdb39af6be1a043ca5e343fe18d0f17e098770b"}, + {file = "matplotlib-3.7.1-cp39-cp39-macosx_10_12_x86_64.whl", hash = "sha256:28506a03bd7f3fe59cd3cd4ceb2a8d8a2b1db41afede01f66c42561b9be7b4b7"}, + {file = "matplotlib-3.7.1-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:8c587963b85ce41e0a8af53b9b2de8dddbf5ece4c34553f7bd9d066148dc719c"}, + {file = "matplotlib-3.7.1-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:8bf26ade3ff0f27668989d98c8435ce9327d24cffb7f07d24ef609e33d582439"}, + {file = "matplotlib-3.7.1-cp39-cp39-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:def58098f96a05f90af7e92fd127d21a287068202aa43b2a93476170ebd99e87"}, + {file = "matplotlib-3.7.1-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:f883a22a56a84dba3b588696a2b8a1ab0d2c3d41be53264115c71b0a942d8fdb"}, + {file = "matplotlib-3.7.1-cp39-cp39-win32.whl", hash = "sha256:4f99e1b234c30c1e9714610eb0c6d2f11809c9c78c984a613ae539ea2ad2eb4b"}, + {file = "matplotlib-3.7.1-cp39-cp39-win_amd64.whl", hash = "sha256:3ba2af245e36990facf67fde840a760128ddd71210b2ab6406e640188d69d136"}, + {file = "matplotlib-3.7.1-pp38-pypy38_pp73-macosx_10_12_x86_64.whl", hash = "sha256:3032884084f541163f295db8a6536e0abb0db464008fadca6c98aaf84ccf4717"}, + {file = "matplotlib-3.7.1-pp38-pypy38_pp73-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:3a2cb34336110e0ed8bb4f650e817eed61fa064acbefeb3591f1b33e3a84fd96"}, + {file = "matplotlib-3.7.1-pp38-pypy38_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:b867e2f952ed592237a1828f027d332d8ee219ad722345b79a001f49df0936eb"}, + {file = "matplotlib-3.7.1-pp38-pypy38_pp73-win_amd64.whl", hash = "sha256:57bfb8c8ea253be947ccb2bc2d1bb3862c2bccc662ad1b4626e1f5e004557042"}, + {file = "matplotlib-3.7.1-pp39-pypy39_pp73-macosx_10_12_x86_64.whl", hash = "sha256:438196cdf5dc8d39b50a45cb6e3f6274edbcf2254f85fa9b895bf85851c3a613"}, + {file = "matplotlib-3.7.1-pp39-pypy39_pp73-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:21e9cff1a58d42e74d01153360de92b326708fb205250150018a52c70f43c290"}, + {file = "matplotlib-3.7.1-pp39-pypy39_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:75d4725d70b7c03e082bbb8a34639ede17f333d7247f56caceb3801cb6ff703d"}, + {file = "matplotlib-3.7.1-pp39-pypy39_pp73-win_amd64.whl", hash = "sha256:97cc368a7268141afb5690760921765ed34867ffb9655dd325ed207af85c7529"}, + {file = "matplotlib-3.7.1.tar.gz", hash = "sha256:7b73305f25eab4541bd7ee0b96d87e53ae9c9f1823be5659b806cd85786fe882"}, +] + +[package.dependencies] +contourpy = ">=1.0.1" +cycler = ">=0.10" +fonttools = ">=4.22.0" +kiwisolver = ">=1.0.1" +numpy = ">=1.20" +packaging = ">=20.0" +pillow = ">=6.2.0" +pyparsing = ">=2.3.1" +python-dateutil = ">=2.7" + +[[package]] +name = "networkx" +version = "3.2.1" +description = "Python package for creating and manipulating graphs and networks" +optional = false +python-versions = ">=3.9" +files = [ + {file = "networkx-3.2.1-py3-none-any.whl", hash = "sha256:f18c69adc97877c42332c170849c96cefa91881c99a7cb3e95b7c659ebdc1ec2"}, + {file = "networkx-3.2.1.tar.gz", hash = "sha256:9f1bb5cf3409bf324e0a722c20bdb4c20ee39bf1c30ce8ae499c8502b0b5e0c6"}, +] + +[package.extras] +default = ["matplotlib (>=3.5)", "numpy (>=1.22)", "pandas (>=1.4)", "scipy (>=1.9,!=1.11.0,!=1.11.1)"] +developer = ["changelist (==0.4)", "mypy (>=1.1)", "pre-commit (>=3.2)", "rtoml"] +doc = ["nb2plots (>=0.7)", "nbconvert (<7.9)", "numpydoc (>=1.6)", "pillow (>=9.4)", "pydata-sphinx-theme (>=0.14)", "sphinx (>=7)", "sphinx-gallery (>=0.14)", "texext (>=0.6.7)"] +extra = ["lxml (>=4.6)", "pydot (>=1.4.2)", "pygraphviz (>=1.11)", "sympy (>=1.10)"] +test = ["pytest (>=7.2)", "pytest-cov (>=4.0)"] + +[[package]] +name = "numba" +version = "0.59.1" +description = "compiling Python code using LLVM" +optional = false +python-versions = ">=3.9" +files = [ + {file = "numba-0.59.1-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:97385a7f12212c4f4bc28f648720a92514bee79d7063e40ef66c2d30600fd18e"}, + {file = "numba-0.59.1-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:0b77aecf52040de2a1eb1d7e314497b9e56fba17466c80b457b971a25bb1576d"}, + {file = "numba-0.59.1-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:3476a4f641bfd58f35ead42f4dcaf5f132569c4647c6f1360ccf18ee4cda3990"}, + {file = "numba-0.59.1-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:525ef3f820931bdae95ee5379c670d5c97289c6520726bc6937a4a7d4230ba24"}, + {file = "numba-0.59.1-cp310-cp310-win_amd64.whl", hash = "sha256:990e395e44d192a12105eca3083b61307db7da10e093972ca285c85bef0963d6"}, + {file = "numba-0.59.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:43727e7ad20b3ec23ee4fc642f5b61845c71f75dd2825b3c234390c6d8d64051"}, + {file = "numba-0.59.1-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:411df625372c77959570050e861981e9d196cc1da9aa62c3d6a836b5cc338966"}, + {file = "numba-0.59.1-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:2801003caa263d1e8497fb84829a7ecfb61738a95f62bc05693fcf1733e978e4"}, + {file = "numba-0.59.1-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:dd2842fac03be4e5324ebbbd4d2d0c8c0fc6e0df75c09477dd45b288a0777389"}, + {file = "numba-0.59.1-cp311-cp311-win_amd64.whl", hash = "sha256:0594b3dfb369fada1f8bb2e3045cd6c61a564c62e50cf1f86b4666bc721b3450"}, + {file = "numba-0.59.1-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:1cce206a3b92836cdf26ef39d3a3242fec25e07f020cc4feec4c4a865e340569"}, + {file = "numba-0.59.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:8c8b4477763cb1fbd86a3be7050500229417bf60867c93e131fd2626edb02238"}, + {file = "numba-0.59.1-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:7d80bce4ef7e65bf895c29e3889ca75a29ee01da80266a01d34815918e365835"}, + {file = "numba-0.59.1-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:f7ad1d217773e89a9845886401eaaab0a156a90aa2f179fdc125261fd1105096"}, + {file = "numba-0.59.1-cp312-cp312-win_amd64.whl", hash = "sha256:5bf68f4d69dd3a9f26a9b23548fa23e3bcb9042e2935257b471d2a8d3c424b7f"}, + {file = "numba-0.59.1-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:4e0318ae729de6e5dbe64c75ead1a95eb01fabfe0e2ebed81ebf0344d32db0ae"}, + {file = "numba-0.59.1-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:0f68589740a8c38bb7dc1b938b55d1145244c8353078eea23895d4f82c8b9ec1"}, + {file = "numba-0.59.1-cp39-cp39-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:649913a3758891c77c32e2d2a3bcbedf4a69f5fea276d11f9119677c45a422e8"}, + {file = "numba-0.59.1-cp39-cp39-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:9712808e4545270291d76b9a264839ac878c5eb7d8b6e02c970dc0ac29bc8187"}, + {file = "numba-0.59.1-cp39-cp39-win_amd64.whl", hash = "sha256:8d51ccd7008a83105ad6a0082b6a2b70f1142dc7cfd76deb8c5a862367eb8c86"}, + {file = "numba-0.59.1.tar.gz", hash = "sha256:76f69132b96028d2774ed20415e8c528a34e3299a40581bae178f0994a2f370b"}, +] + +[package.dependencies] +llvmlite = "==0.42.*" +numpy = ">=1.22,<1.27" + +[[package]] +name = "numpy" +version = "1.23.5" +description = "NumPy is the fundamental package for array computing with Python." +optional = false +python-versions = ">=3.8" +files = [ + {file = "numpy-1.23.5-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:9c88793f78fca17da0145455f0d7826bcb9f37da4764af27ac945488116efe63"}, + {file = "numpy-1.23.5-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:e9f4c4e51567b616be64e05d517c79a8a22f3606499941d97bb76f2ca59f982d"}, + {file = "numpy-1.23.5-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:7903ba8ab592b82014713c491f6c5d3a1cde5b4a3bf116404e08f5b52f6daf43"}, + {file = "numpy-1.23.5-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:5e05b1c973a9f858c74367553e236f287e749465f773328c8ef31abe18f691e1"}, + {file = "numpy-1.23.5-cp310-cp310-win32.whl", hash = "sha256:522e26bbf6377e4d76403826ed689c295b0b238f46c28a7251ab94716da0b280"}, + {file = "numpy-1.23.5-cp310-cp310-win_amd64.whl", hash = "sha256:dbee87b469018961d1ad79b1a5d50c0ae850000b639bcb1b694e9981083243b6"}, + {file = "numpy-1.23.5-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:ce571367b6dfe60af04e04a1834ca2dc5f46004ac1cc756fb95319f64c095a96"}, + {file = "numpy-1.23.5-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:56e454c7833e94ec9769fa0f86e6ff8e42ee38ce0ce1fa4cbb747ea7e06d56aa"}, + {file = "numpy-1.23.5-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:5039f55555e1eab31124a5768898c9e22c25a65c1e0037f4d7c495a45778c9f2"}, + {file = "numpy-1.23.5-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:58f545efd1108e647604a1b5aa809591ccd2540f468a880bedb97247e72db387"}, + {file = "numpy-1.23.5-cp311-cp311-win32.whl", hash = "sha256:b2a9ab7c279c91974f756c84c365a669a887efa287365a8e2c418f8b3ba73fb0"}, + {file = "numpy-1.23.5-cp311-cp311-win_amd64.whl", hash = "sha256:0cbe9848fad08baf71de1a39e12d1b6310f1d5b2d0ea4de051058e6e1076852d"}, + {file = "numpy-1.23.5-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:f063b69b090c9d918f9df0a12116029e274daf0181df392839661c4c7ec9018a"}, + {file = "numpy-1.23.5-cp38-cp38-macosx_11_0_arm64.whl", hash = "sha256:0aaee12d8883552fadfc41e96b4c82ee7d794949e2a7c3b3a7201e968c7ecab9"}, + {file = "numpy-1.23.5-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:92c8c1e89a1f5028a4c6d9e3ccbe311b6ba53694811269b992c0b224269e2398"}, + {file = "numpy-1.23.5-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:d208a0f8729f3fb790ed18a003f3a57895b989b40ea4dce4717e9cf4af62c6bb"}, + {file = "numpy-1.23.5-cp38-cp38-win32.whl", hash = "sha256:06005a2ef6014e9956c09ba07654f9837d9e26696a0470e42beedadb78c11b07"}, + {file = "numpy-1.23.5-cp38-cp38-win_amd64.whl", hash = "sha256:ca51fcfcc5f9354c45f400059e88bc09215fb71a48d3768fb80e357f3b457e1e"}, + {file = "numpy-1.23.5-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:8969bfd28e85c81f3f94eb4a66bc2cf1dbdc5c18efc320af34bffc54d6b1e38f"}, + {file = "numpy-1.23.5-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:a7ac231a08bb37f852849bbb387a20a57574a97cfc7b6cabb488a4fc8be176de"}, + {file = "numpy-1.23.5-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:bf837dc63ba5c06dc8797c398db1e223a466c7ece27a1f7b5232ba3466aafe3d"}, + {file = "numpy-1.23.5-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:33161613d2269025873025b33e879825ec7b1d831317e68f4f2f0f84ed14c719"}, + {file = "numpy-1.23.5-cp39-cp39-win32.whl", hash = "sha256:af1da88f6bc3d2338ebbf0e22fe487821ea4d8e89053e25fa59d1d79786e7481"}, + {file = "numpy-1.23.5-cp39-cp39-win_amd64.whl", hash = "sha256:09b7847f7e83ca37c6e627682f145856de331049013853f344f37b0c9690e3df"}, + {file = "numpy-1.23.5-pp38-pypy38_pp73-macosx_10_9_x86_64.whl", hash = "sha256:abdde9f795cf292fb9651ed48185503a2ff29be87770c3b8e2a14b0cd7aa16f8"}, + {file = "numpy-1.23.5-pp38-pypy38_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:f9a909a8bae284d46bbfdefbdd4a262ba19d3bc9921b1e76126b1d21c3c34135"}, + {file = "numpy-1.23.5-pp38-pypy38_pp73-win_amd64.whl", hash = "sha256:01dd17cbb340bf0fc23981e52e1d18a9d4050792e8fb8363cecbf066a84b827d"}, + {file = "numpy-1.23.5.tar.gz", hash = "sha256:1b1766d6f397c18153d40015ddfc79ddb715cabadc04d2d228d4e5a8bc4ded1a"}, +] + +[[package]] +name = "opencv-python" +version = "4.7.0.72" +description = "Wrapper package for OpenCV python bindings." +optional = false +python-versions = ">=3.6" +files = [ + {file = "opencv-python-4.7.0.72.tar.gz", hash = "sha256:3424794a711f33284581f3c1e4b071cfc827d02b99d6fd9a35391f517c453306"}, + {file = "opencv_python-4.7.0.72-cp37-abi3-macosx_10_16_x86_64.whl", hash = "sha256:d4f8880440c433a0025d78804dda6901d1e8e541a561dda66892d90290aef881"}, + {file = "opencv_python-4.7.0.72-cp37-abi3-macosx_11_0_arm64.whl", hash = "sha256:7a297e7651e22eb17c265ddbbc80e2ba2a8ff4f4a1696a67c45e5f5798245842"}, + {file = "opencv_python-4.7.0.72-cp37-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:cd08343654c6b88c5a8c25bf425f8025aed2e3189b4d7306b5861d32affaf737"}, + {file = "opencv_python-4.7.0.72-cp37-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ebfc0a3a2f57716e709028b992e4de7fd8752105d7a768531c4f434043c6f9ff"}, + {file = "opencv_python-4.7.0.72-cp37-abi3-win32.whl", hash = "sha256:eda115797b114fc16ca6f182b91c5d984f0015c19bec3145e55d33d708e9bae1"}, + {file = "opencv_python-4.7.0.72-cp37-abi3-win_amd64.whl", hash = "sha256:812af57553ec1c6709060c63f6b7e9ad07ddc0f592f3ccc6d00c71e0fe0e6376"}, +] + +[package.dependencies] +numpy = [ + {version = ">=1.21.4", markers = "python_version >= \"3.10\" and platform_system == \"Darwin\""}, + {version = ">=1.21.2", markers = "platform_system != \"Darwin\" and python_version >= \"3.10\""}, +] + +[[package]] +name = "packaging" +version = "24.0" +description = "Core utilities for Python packages" +optional = false +python-versions = ">=3.7" +files = [ + {file = "packaging-24.0-py3-none-any.whl", hash = "sha256:2ddfb553fdf02fb784c234c7ba6ccc288296ceabec964ad2eae3777778130bc5"}, + {file = "packaging-24.0.tar.gz", hash = "sha256:eb82c5e3e56209074766e6885bb04b8c38a0c015d0a30036ebe7ece34c9989e9"}, +] + +[[package]] +name = "pandas" +version = "2.0.1" +description = "Powerful data structures for data analysis, time series, and statistics" +optional = false +python-versions = ">=3.8" +files = [ + {file = "pandas-2.0.1-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:70a996a1d2432dadedbb638fe7d921c88b0cc4dd90374eab51bb33dc6c0c2a12"}, + {file = "pandas-2.0.1-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:909a72b52175590debbf1d0c9e3e6bce2f1833c80c76d80bd1aa09188be768e5"}, + {file = "pandas-2.0.1-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:fe7914d8ddb2d54b900cec264c090b88d141a1eed605c9539a187dbc2547f022"}, + {file = "pandas-2.0.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:0a514ae436b23a92366fbad8365807fc0eed15ca219690b3445dcfa33597a5cc"}, + {file = "pandas-2.0.1-cp310-cp310-win32.whl", hash = "sha256:12bd6618e3cc737c5200ecabbbb5eaba8ab645a4b0db508ceeb4004bb10b060e"}, + {file = "pandas-2.0.1-cp310-cp310-win_amd64.whl", hash = "sha256:2b6fe5f7ce1cba0e74188c8473c9091ead9b293ef0a6794939f8cc7947057abd"}, + {file = "pandas-2.0.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:00959a04a1d7bbc63d75a768540fb20ecc9e65fd80744c930e23768345a362a7"}, + {file = "pandas-2.0.1-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:af2449e9e984dfad39276b885271ba31c5e0204ffd9f21f287a245980b0e4091"}, + {file = "pandas-2.0.1-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:910df06feaf9935d05247db6de452f6d59820e432c18a2919a92ffcd98f8f79b"}, + {file = "pandas-2.0.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6fa0067f2419f933101bdc6001bcea1d50812afbd367b30943417d67fbb99678"}, + {file = "pandas-2.0.1-cp311-cp311-win32.whl", hash = "sha256:7b8395d335b08bc8b050590da264f94a439b4770ff16bb51798527f1dd840388"}, + {file = "pandas-2.0.1-cp311-cp311-win_amd64.whl", hash = "sha256:8db5a644d184a38e6ed40feeb12d410d7fcc36648443defe4707022da127fc35"}, + {file = "pandas-2.0.1-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:7bbf173d364130334e0159a9a034f573e8b44a05320995127cf676b85fd8ce86"}, + {file = "pandas-2.0.1-cp38-cp38-macosx_11_0_arm64.whl", hash = "sha256:6c0853d487b6c868bf107a4b270a823746175b1932093b537b9b76c639fc6f7e"}, + {file = "pandas-2.0.1-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:f25e23a03f7ad7211ffa30cb181c3e5f6d96a8e4cb22898af462a7333f8a74eb"}, + {file = "pandas-2.0.1-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:e09a53a4fe8d6ae2149959a2d02e1ef2f4d2ceb285ac48f74b79798507e468b4"}, + {file = "pandas-2.0.1-cp38-cp38-win32.whl", hash = "sha256:a2564629b3a47b6aa303e024e3d84e850d36746f7e804347f64229f8c87416ea"}, + {file = "pandas-2.0.1-cp38-cp38-win_amd64.whl", hash = "sha256:03e677c6bc9cfb7f93a8b617d44f6091613a5671ef2944818469be7b42114a00"}, + {file = "pandas-2.0.1-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:3d099ecaa5b9e977b55cd43cf842ec13b14afa1cfa51b7e1179d90b38c53ce6a"}, + {file = "pandas-2.0.1-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:a37ee35a3eb6ce523b2c064af6286c45ea1c7ff882d46e10d0945dbda7572753"}, + {file = "pandas-2.0.1-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:320b180d125c3842c5da5889183b9a43da4ebba375ab2ef938f57bf267a3c684"}, + {file = "pandas-2.0.1-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:18d22cb9043b6c6804529810f492ab09d638ddf625c5dea8529239607295cb59"}, + {file = "pandas-2.0.1-cp39-cp39-win32.whl", hash = "sha256:90d1d365d77d287063c5e339f49b27bd99ef06d10a8843cf00b1a49326d492c1"}, + {file = "pandas-2.0.1-cp39-cp39-win_amd64.whl", hash = "sha256:99f7192d8b0e6daf8e0d0fd93baa40056684e4b4aaaef9ea78dff34168e1f2f0"}, + {file = "pandas-2.0.1.tar.gz", hash = "sha256:19b8e5270da32b41ebf12f0e7165efa7024492e9513fb46fb631c5022ae5709d"}, +] + +[package.dependencies] +numpy = {version = ">=1.21.0", markers = "python_version >= \"3.10\""} +python-dateutil = ">=2.8.2" +pytz = ">=2020.1" +tzdata = ">=2022.1" + +[package.extras] +all = ["PyQt5 (>=5.15.1)", "SQLAlchemy (>=1.4.16)", "beautifulsoup4 (>=4.9.3)", "bottleneck (>=1.3.2)", "brotlipy (>=0.7.0)", "fastparquet (>=0.6.3)", "fsspec (>=2021.07.0)", "gcsfs (>=2021.07.0)", "html5lib (>=1.1)", "hypothesis (>=6.34.2)", "jinja2 (>=3.0.0)", "lxml (>=4.6.3)", "matplotlib (>=3.6.1)", "numba (>=0.53.1)", "numexpr (>=2.7.3)", "odfpy (>=1.4.1)", "openpyxl (>=3.0.7)", "pandas-gbq (>=0.15.0)", "psycopg2 (>=2.8.6)", "pyarrow (>=7.0.0)", "pymysql (>=1.0.2)", "pyreadstat (>=1.1.2)", "pytest (>=7.0.0)", "pytest-asyncio (>=0.17.0)", "pytest-xdist (>=2.2.0)", "python-snappy (>=0.6.0)", "pyxlsb (>=1.0.8)", "qtpy (>=2.2.0)", "s3fs (>=2021.08.0)", "scipy (>=1.7.1)", "tables (>=3.6.1)", "tabulate (>=0.8.9)", "xarray (>=0.21.0)", "xlrd (>=2.0.1)", "xlsxwriter (>=1.4.3)", "zstandard (>=0.15.2)"] +aws = ["s3fs (>=2021.08.0)"] +clipboard = ["PyQt5 (>=5.15.1)", "qtpy (>=2.2.0)"] +compression = ["brotlipy (>=0.7.0)", "python-snappy (>=0.6.0)", "zstandard (>=0.15.2)"] +computation = ["scipy (>=1.7.1)", "xarray (>=0.21.0)"] +excel = ["odfpy (>=1.4.1)", "openpyxl (>=3.0.7)", "pyxlsb (>=1.0.8)", "xlrd (>=2.0.1)", "xlsxwriter (>=1.4.3)"] +feather = ["pyarrow (>=7.0.0)"] +fss = ["fsspec (>=2021.07.0)"] +gcp = ["gcsfs (>=2021.07.0)", "pandas-gbq (>=0.15.0)"] +hdf5 = ["tables (>=3.6.1)"] +html = ["beautifulsoup4 (>=4.9.3)", "html5lib (>=1.1)", "lxml (>=4.6.3)"] +mysql = ["SQLAlchemy (>=1.4.16)", "pymysql (>=1.0.2)"] +output-formatting = ["jinja2 (>=3.0.0)", "tabulate (>=0.8.9)"] +parquet = ["pyarrow (>=7.0.0)"] +performance = ["bottleneck (>=1.3.2)", "numba (>=0.53.1)", "numexpr (>=2.7.1)"] +plot = ["matplotlib (>=3.6.1)"] +postgresql = ["SQLAlchemy (>=1.4.16)", "psycopg2 (>=2.8.6)"] +spss = ["pyreadstat (>=1.1.2)"] +sql-other = ["SQLAlchemy (>=1.4.16)"] +test = ["hypothesis (>=6.34.2)", "pytest (>=7.0.0)", "pytest-asyncio (>=0.17.0)", "pytest-xdist (>=2.2.0)"] +xml = ["lxml (>=4.6.3)"] + +[[package]] +name = "pathtools" +version = "0.1.2" +description = "File system general utilities" +optional = false +python-versions = "*" +files = [ + {file = "pathtools-0.1.2.tar.gz", hash = "sha256:7c35c5421a39bb82e58018febd90e3b6e5db34c5443aaaf742b3f33d4655f1c0"}, +] + +[[package]] +name = "pillow" +version = "10.0.0" +description = "Python Imaging Library (Fork)" +optional = false +python-versions = ">=3.8" +files = [ + {file = "Pillow-10.0.0-cp310-cp310-macosx_10_10_x86_64.whl", hash = "sha256:1f62406a884ae75fb2f818694469519fb685cc7eaff05d3451a9ebe55c646891"}, + {file = "Pillow-10.0.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:d5db32e2a6ccbb3d34d87c87b432959e0db29755727afb37290e10f6e8e62614"}, + {file = "Pillow-10.0.0-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:edf4392b77bdc81f36e92d3a07a5cd072f90253197f4a52a55a8cec48a12483b"}, + {file = "Pillow-10.0.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:520f2a520dc040512699f20fa1c363eed506e94248d71f85412b625026f6142c"}, + {file = "Pillow-10.0.0-cp310-cp310-manylinux_2_28_aarch64.whl", hash = "sha256:8c11160913e3dd06c8ffdb5f233a4f254cb449f4dfc0f8f4549eda9e542c93d1"}, + {file = "Pillow-10.0.0-cp310-cp310-manylinux_2_28_x86_64.whl", hash = "sha256:a74ba0c356aaa3bb8e3eb79606a87669e7ec6444be352870623025d75a14a2bf"}, + {file = "Pillow-10.0.0-cp310-cp310-musllinux_1_1_aarch64.whl", hash = "sha256:d5d0dae4cfd56969d23d94dc8e89fb6a217be461c69090768227beb8ed28c0a3"}, + {file = "Pillow-10.0.0-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:22c10cc517668d44b211717fd9775799ccec4124b9a7f7b3635fc5386e584992"}, + {file = "Pillow-10.0.0-cp310-cp310-win_amd64.whl", hash = "sha256:dffe31a7f47b603318c609f378ebcd57f1554a3a6a8effbc59c3c69f804296de"}, + {file = "Pillow-10.0.0-cp311-cp311-macosx_10_10_x86_64.whl", hash = "sha256:9fb218c8a12e51d7ead2a7c9e101a04982237d4855716af2e9499306728fb485"}, + {file = "Pillow-10.0.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:d35e3c8d9b1268cbf5d3670285feb3528f6680420eafe35cccc686b73c1e330f"}, + {file = "Pillow-10.0.0-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:3ed64f9ca2f0a95411e88a4efbd7a29e5ce2cea36072c53dd9d26d9c76f753b3"}, + {file = "Pillow-10.0.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:0b6eb5502f45a60a3f411c63187db83a3d3107887ad0d036c13ce836f8a36f1d"}, + {file = "Pillow-10.0.0-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:c1fbe7621c167ecaa38ad29643d77a9ce7311583761abf7836e1510c580bf3dd"}, + {file = "Pillow-10.0.0-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:cd25d2a9d2b36fcb318882481367956d2cf91329f6892fe5d385c346c0649629"}, + {file = "Pillow-10.0.0-cp311-cp311-musllinux_1_1_aarch64.whl", hash = "sha256:3b08d4cc24f471b2c8ca24ec060abf4bebc6b144cb89cba638c720546b1cf538"}, + {file = "Pillow-10.0.0-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:d737a602fbd82afd892ca746392401b634e278cb65d55c4b7a8f48e9ef8d008d"}, + {file = "Pillow-10.0.0-cp311-cp311-win_amd64.whl", hash = "sha256:3a82c40d706d9aa9734289740ce26460a11aeec2d9c79b7af87bb35f0073c12f"}, + {file = "Pillow-10.0.0-cp311-cp311-win_arm64.whl", hash = "sha256:bc2ec7c7b5d66b8ec9ce9f720dbb5fa4bace0f545acd34870eff4a369b44bf37"}, + {file = "Pillow-10.0.0-cp312-cp312-macosx_10_10_x86_64.whl", hash = "sha256:d80cf684b541685fccdd84c485b31ce73fc5c9b5d7523bf1394ce134a60c6883"}, + {file = "Pillow-10.0.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:76de421f9c326da8f43d690110f0e79fe3ad1e54be811545d7d91898b4c8493e"}, + {file = "Pillow-10.0.0-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:81ff539a12457809666fef6624684c008e00ff6bf455b4b89fd00a140eecd640"}, + {file = "Pillow-10.0.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ce543ed15570eedbb85df19b0a1a7314a9c8141a36ce089c0a894adbfccb4568"}, + {file = "Pillow-10.0.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:685ac03cc4ed5ebc15ad5c23bc555d68a87777586d970c2c3e216619a5476223"}, + {file = "Pillow-10.0.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:d72e2ecc68a942e8cf9739619b7f408cc7b272b279b56b2c83c6123fcfa5cdff"}, + {file = "Pillow-10.0.0-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:d50b6aec14bc737742ca96e85d6d0a5f9bfbded018264b3b70ff9d8c33485551"}, + {file = "Pillow-10.0.0-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:00e65f5e822decd501e374b0650146063fbb30a7264b4d2744bdd7b913e0cab5"}, + {file = "Pillow-10.0.0-cp312-cp312-win_amd64.whl", hash = "sha256:f31f9fdbfecb042d046f9d91270a0ba28368a723302786c0009ee9b9f1f60199"}, + {file = "Pillow-10.0.0-cp312-cp312-win_arm64.whl", hash = "sha256:1ce91b6ec08d866b14413d3f0bbdea7e24dfdc8e59f562bb77bc3fe60b6144ca"}, + {file = "Pillow-10.0.0-cp38-cp38-macosx_10_10_x86_64.whl", hash = "sha256:349930d6e9c685c089284b013478d6f76e3a534e36ddfa912cde493f235372f3"}, + {file = "Pillow-10.0.0-cp38-cp38-macosx_11_0_arm64.whl", hash = "sha256:3a684105f7c32488f7153905a4e3015a3b6c7182e106fe3c37fbb5ef3e6994c3"}, + {file = "Pillow-10.0.0-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:b4f69b3700201b80bb82c3a97d5e9254084f6dd5fb5b16fc1a7b974260f89f43"}, + {file = "Pillow-10.0.0-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:3f07ea8d2f827d7d2a49ecf1639ec02d75ffd1b88dcc5b3a61bbb37a8759ad8d"}, + {file = "Pillow-10.0.0-cp38-cp38-manylinux_2_28_aarch64.whl", hash = "sha256:040586f7d37b34547153fa383f7f9aed68b738992380ac911447bb78f2abe530"}, + {file = "Pillow-10.0.0-cp38-cp38-manylinux_2_28_x86_64.whl", hash = "sha256:f88a0b92277de8e3ca715a0d79d68dc82807457dae3ab8699c758f07c20b3c51"}, + {file = "Pillow-10.0.0-cp38-cp38-musllinux_1_1_aarch64.whl", hash = "sha256:c7cf14a27b0d6adfaebb3ae4153f1e516df54e47e42dcc073d7b3d76111a8d86"}, + {file = "Pillow-10.0.0-cp38-cp38-musllinux_1_1_x86_64.whl", hash = "sha256:3400aae60685b06bb96f99a21e1ada7bc7a413d5f49bce739828ecd9391bb8f7"}, + {file = "Pillow-10.0.0-cp38-cp38-win_amd64.whl", hash = "sha256:dbc02381779d412145331789b40cc7b11fdf449e5d94f6bc0b080db0a56ea3f0"}, + {file = "Pillow-10.0.0-cp39-cp39-macosx_10_10_x86_64.whl", hash = "sha256:9211e7ad69d7c9401cfc0e23d49b69ca65ddd898976d660a2fa5904e3d7a9baa"}, + {file = "Pillow-10.0.0-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:faaf07ea35355b01a35cb442dd950d8f1bb5b040a7787791a535de13db15ed90"}, + {file = "Pillow-10.0.0-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:c9f72a021fbb792ce98306ffb0c348b3c9cb967dce0f12a49aa4c3d3fdefa967"}, + {file = "Pillow-10.0.0-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:9f7c16705f44e0504a3a2a14197c1f0b32a95731d251777dcb060aa83022cb2d"}, + {file = "Pillow-10.0.0-cp39-cp39-manylinux_2_28_aarch64.whl", hash = "sha256:76edb0a1fa2b4745fb0c99fb9fb98f8b180a1bbceb8be49b087e0b21867e77d3"}, + {file = "Pillow-10.0.0-cp39-cp39-manylinux_2_28_x86_64.whl", hash = "sha256:368ab3dfb5f49e312231b6f27b8820c823652b7cd29cfbd34090565a015e99ba"}, + {file = "Pillow-10.0.0-cp39-cp39-musllinux_1_1_aarch64.whl", hash = "sha256:608bfdee0d57cf297d32bcbb3c728dc1da0907519d1784962c5f0c68bb93e5a3"}, + {file = "Pillow-10.0.0-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:5c6e3df6bdd396749bafd45314871b3d0af81ff935b2d188385e970052091017"}, + {file = "Pillow-10.0.0-cp39-cp39-win_amd64.whl", hash = "sha256:7be600823e4c8631b74e4a0d38384c73f680e6105a7d3c6824fcf226c178c7e6"}, + {file = "Pillow-10.0.0-pp310-pypy310_pp73-macosx_10_10_x86_64.whl", hash = "sha256:92be919bbc9f7d09f7ae343c38f5bb21c973d2576c1d45600fce4b74bafa7ac0"}, + {file = "Pillow-10.0.0-pp310-pypy310_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:8f8182b523b2289f7c415f589118228d30ac8c355baa2f3194ced084dac2dbba"}, + {file = "Pillow-10.0.0-pp310-pypy310_pp73-manylinux_2_28_x86_64.whl", hash = "sha256:38250a349b6b390ee6047a62c086d3817ac69022c127f8a5dc058c31ccef17f3"}, + {file = "Pillow-10.0.0-pp310-pypy310_pp73-win_amd64.whl", hash = "sha256:88af2003543cc40c80f6fca01411892ec52b11021b3dc22ec3bc9d5afd1c5334"}, + {file = "Pillow-10.0.0-pp39-pypy39_pp73-macosx_10_10_x86_64.whl", hash = "sha256:c189af0545965fa8d3b9613cfdb0cd37f9d71349e0f7750e1fd704648d475ed2"}, + {file = "Pillow-10.0.0-pp39-pypy39_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ce7b031a6fc11365970e6a5686d7ba8c63e4c1cf1ea143811acbb524295eabed"}, + {file = "Pillow-10.0.0-pp39-pypy39_pp73-manylinux_2_28_x86_64.whl", hash = "sha256:db24668940f82321e746773a4bc617bfac06ec831e5c88b643f91f122a785684"}, + {file = "Pillow-10.0.0-pp39-pypy39_pp73-win_amd64.whl", hash = "sha256:efe8c0681042536e0d06c11f48cebe759707c9e9abf880ee213541c5b46c5bf3"}, + {file = "Pillow-10.0.0.tar.gz", hash = "sha256:9c82b5b3e043c7af0d95792d0d20ccf68f61a1fec6b3530e718b688422727396"}, +] + +[package.extras] +docs = ["furo", "olefile", "sphinx (>=2.4)", "sphinx-copybutton", "sphinx-inline-tabs", "sphinx-removed-in", "sphinxext-opengraph"] +tests = ["check-manifest", "coverage", "defusedxml", "markdown2", "olefile", "packaging", "pyroma", "pytest", "pytest-cov", "pytest-timeout"] + +[[package]] +name = "protobuf" +version = "4.25.3" +description = "" +optional = false +python-versions = ">=3.8" +files = [ + {file = "protobuf-4.25.3-cp310-abi3-win32.whl", hash = "sha256:d4198877797a83cbfe9bffa3803602bbe1625dc30d8a097365dbc762e5790faa"}, + {file = "protobuf-4.25.3-cp310-abi3-win_amd64.whl", hash = "sha256:209ba4cc916bab46f64e56b85b090607a676f66b473e6b762e6f1d9d591eb2e8"}, + {file = "protobuf-4.25.3-cp37-abi3-macosx_10_9_universal2.whl", hash = "sha256:f1279ab38ecbfae7e456a108c5c0681e4956d5b1090027c1de0f934dfdb4b35c"}, + {file = "protobuf-4.25.3-cp37-abi3-manylinux2014_aarch64.whl", hash = "sha256:e7cb0ae90dd83727f0c0718634ed56837bfeeee29a5f82a7514c03ee1364c019"}, + {file = "protobuf-4.25.3-cp37-abi3-manylinux2014_x86_64.whl", hash = "sha256:7c8daa26095f82482307bc717364e7c13f4f1c99659be82890dcfc215194554d"}, + {file = "protobuf-4.25.3-cp38-cp38-win32.whl", hash = "sha256:f4f118245c4a087776e0a8408be33cf09f6c547442c00395fbfb116fac2f8ac2"}, + {file = "protobuf-4.25.3-cp38-cp38-win_amd64.whl", hash = "sha256:c053062984e61144385022e53678fbded7aea14ebb3e0305ae3592fb219ccfa4"}, + {file = "protobuf-4.25.3-cp39-cp39-win32.whl", hash = "sha256:19b270aeaa0099f16d3ca02628546b8baefe2955bbe23224aaf856134eccf1e4"}, + {file = "protobuf-4.25.3-cp39-cp39-win_amd64.whl", hash = "sha256:e3c97a1555fd6388f857770ff8b9703083de6bf1f9274a002a332d65fbb56c8c"}, + {file = "protobuf-4.25.3-py3-none-any.whl", hash = "sha256:f0700d54bcf45424477e46a9f0944155b46fb0639d69728739c0e47bab83f2b9"}, + {file = "protobuf-4.25.3.tar.gz", hash = "sha256:25b5d0b42fd000320bd7830b349e3b696435f3b329810427a6bcce6a5492cc5c"}, +] + +[[package]] +name = "psutil" +version = "5.9.8" +description = "Cross-platform lib for process and system monitoring in Python." +optional = false +python-versions = ">=2.7, !=3.0.*, !=3.1.*, !=3.2.*, !=3.3.*, !=3.4.*, !=3.5.*" +files = [ + {file = "psutil-5.9.8-cp27-cp27m-macosx_10_9_x86_64.whl", hash = "sha256:26bd09967ae00920df88e0352a91cff1a78f8d69b3ecabbfe733610c0af486c8"}, + {file = "psutil-5.9.8-cp27-cp27m-manylinux2010_i686.whl", hash = "sha256:05806de88103b25903dff19bb6692bd2e714ccf9e668d050d144012055cbca73"}, + {file = "psutil-5.9.8-cp27-cp27m-manylinux2010_x86_64.whl", hash = "sha256:611052c4bc70432ec770d5d54f64206aa7203a101ec273a0cd82418c86503bb7"}, + {file = "psutil-5.9.8-cp27-cp27mu-manylinux2010_i686.whl", hash = "sha256:50187900d73c1381ba1454cf40308c2bf6f34268518b3f36a9b663ca87e65e36"}, + {file = "psutil-5.9.8-cp27-cp27mu-manylinux2010_x86_64.whl", hash = "sha256:02615ed8c5ea222323408ceba16c60e99c3f91639b07da6373fb7e6539abc56d"}, + {file = "psutil-5.9.8-cp27-none-win32.whl", hash = "sha256:36f435891adb138ed3c9e58c6af3e2e6ca9ac2f365efe1f9cfef2794e6c93b4e"}, + {file = "psutil-5.9.8-cp27-none-win_amd64.whl", hash = "sha256:bd1184ceb3f87651a67b2708d4c3338e9b10c5df903f2e3776b62303b26cb631"}, + {file = "psutil-5.9.8-cp36-abi3-macosx_10_9_x86_64.whl", hash = "sha256:aee678c8720623dc456fa20659af736241f575d79429a0e5e9cf88ae0605cc81"}, + {file = "psutil-5.9.8-cp36-abi3-manylinux_2_12_i686.manylinux2010_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:8cb6403ce6d8e047495a701dc7c5bd788add903f8986d523e3e20b98b733e421"}, + {file = "psutil-5.9.8-cp36-abi3-manylinux_2_12_x86_64.manylinux2010_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:d06016f7f8625a1825ba3732081d77c94589dca78b7a3fc072194851e88461a4"}, + {file = "psutil-5.9.8-cp36-cp36m-win32.whl", hash = "sha256:7d79560ad97af658a0f6adfef8b834b53f64746d45b403f225b85c5c2c140eee"}, + {file = "psutil-5.9.8-cp36-cp36m-win_amd64.whl", hash = "sha256:27cc40c3493bb10de1be4b3f07cae4c010ce715290a5be22b98493509c6299e2"}, + {file = "psutil-5.9.8-cp37-abi3-win32.whl", hash = "sha256:bc56c2a1b0d15aa3eaa5a60c9f3f8e3e565303b465dbf57a1b730e7a2b9844e0"}, + {file = "psutil-5.9.8-cp37-abi3-win_amd64.whl", hash = "sha256:8db4c1b57507eef143a15a6884ca10f7c73876cdf5d51e713151c1236a0e68cf"}, + {file = "psutil-5.9.8-cp38-abi3-macosx_11_0_arm64.whl", hash = "sha256:d16bbddf0693323b8c6123dd804100241da461e41d6e332fb0ba6058f630f8c8"}, + {file = "psutil-5.9.8.tar.gz", hash = "sha256:6be126e3225486dff286a8fb9a06246a5253f4c7c53b475ea5f5ac934e64194c"}, +] + +[package.extras] +test = ["enum34", "ipaddress", "mock", "pywin32", "wmi"] + +[[package]] +name = "pyparsing" +version = "3.1.2" +description = "pyparsing module - Classes and methods to define and execute parsing grammars" +optional = false +python-versions = ">=3.6.8" +files = [ + {file = "pyparsing-3.1.2-py3-none-any.whl", hash = "sha256:f9db75911801ed778fe61bb643079ff86601aca99fcae6345aa67292038fb742"}, + {file = "pyparsing-3.1.2.tar.gz", hash = "sha256:a1bac0ce561155ecc3ed78ca94d3c9378656ad4c94c1270de543f621420f94ad"}, +] + +[package.extras] +diagrams = ["jinja2", "railroad-diagrams"] + +[[package]] +name = "pysocks" +version = "1.7.1" +description = "A Python SOCKS client module. See https://github.com/Anorov/PySocks for more information." +optional = false +python-versions = ">=2.7, !=3.0.*, !=3.1.*, !=3.2.*, !=3.3.*" +files = [ + {file = "PySocks-1.7.1-py27-none-any.whl", hash = "sha256:08e69f092cc6dbe92a0fdd16eeb9b9ffbc13cadfe5ca4c7bd92ffb078b293299"}, + {file = "PySocks-1.7.1-py3-none-any.whl", hash = "sha256:2725bd0a9925919b9b51739eea5f9e2bae91e83288108a9ad338b2e3a4435ee5"}, + {file = "PySocks-1.7.1.tar.gz", hash = "sha256:3f8804571ebe159c380ac6de37643bb4685970655d3bba243530d6558b799aa0"}, +] + +[[package]] +name = "python-dateutil" +version = "2.9.0.post0" +description = "Extensions to the standard Python datetime module" +optional = false +python-versions = "!=3.0.*,!=3.1.*,!=3.2.*,>=2.7" +files = [ + {file = "python-dateutil-2.9.0.post0.tar.gz", hash = "sha256:37dd54208da7e1cd875388217d5e00ebd4179249f90fb72437e91a35459a0ad3"}, + {file = "python_dateutil-2.9.0.post0-py2.py3-none-any.whl", hash = "sha256:a8b2bc7bffae282281c8140a97d3aa9c14da0b136dfe83f850eea9a5f7470427"}, +] + +[package.dependencies] +six = ">=1.5" + +[[package]] +name = "pytz" +version = "2024.1" +description = "World timezone definitions, modern and historical" +optional = false +python-versions = "*" +files = [ + {file = "pytz-2024.1-py2.py3-none-any.whl", hash = "sha256:328171f4e3623139da4983451950b28e95ac706e13f3f2630a879749e7a8b319"}, + {file = "pytz-2024.1.tar.gz", hash = "sha256:2a29735ea9c18baf14b448846bde5a48030ed267578472d8955cd0e7443a9812"}, +] + +[[package]] +name = "pyyaml" +version = "6.0.1" +description = "YAML parser and emitter for Python" +optional = false +python-versions = ">=3.6" +files = [ + {file = "PyYAML-6.0.1-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:d858aa552c999bc8a8d57426ed01e40bef403cd8ccdd0fc5f6f04a00414cac2a"}, + {file = "PyYAML-6.0.1-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:fd66fc5d0da6d9815ba2cebeb4205f95818ff4b79c3ebe268e75d961704af52f"}, + {file = "PyYAML-6.0.1-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:69b023b2b4daa7548bcfbd4aa3da05b3a74b772db9e23b982788168117739938"}, + {file = "PyYAML-6.0.1-cp310-cp310-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:81e0b275a9ecc9c0c0c07b4b90ba548307583c125f54d5b6946cfee6360c733d"}, + {file = "PyYAML-6.0.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ba336e390cd8e4d1739f42dfe9bb83a3cc2e80f567d8805e11b46f4a943f5515"}, + {file = "PyYAML-6.0.1-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:326c013efe8048858a6d312ddd31d56e468118ad4cdeda36c719bf5bb6192290"}, + {file = "PyYAML-6.0.1-cp310-cp310-win32.whl", hash = "sha256:bd4af7373a854424dabd882decdc5579653d7868b8fb26dc7d0e99f823aa5924"}, + {file = "PyYAML-6.0.1-cp310-cp310-win_amd64.whl", hash = "sha256:fd1592b3fdf65fff2ad0004b5e363300ef59ced41c2e6b3a99d4089fa8c5435d"}, + {file = "PyYAML-6.0.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:6965a7bc3cf88e5a1c3bd2e0b5c22f8d677dc88a455344035f03399034eb3007"}, + {file = "PyYAML-6.0.1-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:f003ed9ad21d6a4713f0a9b5a7a0a79e08dd0f221aff4525a2be4c346ee60aab"}, + {file = "PyYAML-6.0.1-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:42f8152b8dbc4fe7d96729ec2b99c7097d656dc1213a3229ca5383f973a5ed6d"}, + {file = "PyYAML-6.0.1-cp311-cp311-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:062582fca9fabdd2c8b54a3ef1c978d786e0f6b3a1510e0ac93ef59e0ddae2bc"}, + {file = "PyYAML-6.0.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:d2b04aac4d386b172d5b9692e2d2da8de7bfb6c387fa4f801fbf6fb2e6ba4673"}, + {file = "PyYAML-6.0.1-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:e7d73685e87afe9f3b36c799222440d6cf362062f78be1013661b00c5c6f678b"}, + {file = "PyYAML-6.0.1-cp311-cp311-win32.whl", hash = "sha256:1635fd110e8d85d55237ab316b5b011de701ea0f29d07611174a1b42f1444741"}, + {file = "PyYAML-6.0.1-cp311-cp311-win_amd64.whl", hash = "sha256:bf07ee2fef7014951eeb99f56f39c9bb4af143d8aa3c21b1677805985307da34"}, + {file = "PyYAML-6.0.1-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:855fb52b0dc35af121542a76b9a84f8d1cd886ea97c84703eaa6d88e37a2ad28"}, + {file = "PyYAML-6.0.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:40df9b996c2b73138957fe23a16a4f0ba614f4c0efce1e9406a184b6d07fa3a9"}, + {file = "PyYAML-6.0.1-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a08c6f0fe150303c1c6b71ebcd7213c2858041a7e01975da3a99aed1e7a378ef"}, + {file = "PyYAML-6.0.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6c22bec3fbe2524cde73d7ada88f6566758a8f7227bfbf93a408a9d86bcc12a0"}, + {file = "PyYAML-6.0.1-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:8d4e9c88387b0f5c7d5f281e55304de64cf7f9c0021a3525bd3b1c542da3b0e4"}, + {file = "PyYAML-6.0.1-cp312-cp312-win32.whl", hash = "sha256:d483d2cdf104e7c9fa60c544d92981f12ad66a457afae824d146093b8c294c54"}, + {file = "PyYAML-6.0.1-cp312-cp312-win_amd64.whl", hash = "sha256:0d3304d8c0adc42be59c5f8a4d9e3d7379e6955ad754aa9d6ab7a398b59dd1df"}, + {file = "PyYAML-6.0.1-cp36-cp36m-macosx_10_9_x86_64.whl", hash = "sha256:50550eb667afee136e9a77d6dc71ae76a44df8b3e51e41b77f6de2932bfe0f47"}, + {file = "PyYAML-6.0.1-cp36-cp36m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:1fe35611261b29bd1de0070f0b2f47cb6ff71fa6595c077e42bd0c419fa27b98"}, + {file = "PyYAML-6.0.1-cp36-cp36m-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:704219a11b772aea0d8ecd7058d0082713c3562b4e271b849ad7dc4a5c90c13c"}, + {file = "PyYAML-6.0.1-cp36-cp36m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:afd7e57eddb1a54f0f1a974bc4391af8bcce0b444685d936840f125cf046d5bd"}, + {file = "PyYAML-6.0.1-cp36-cp36m-win32.whl", hash = "sha256:fca0e3a251908a499833aa292323f32437106001d436eca0e6e7833256674585"}, + {file = "PyYAML-6.0.1-cp36-cp36m-win_amd64.whl", hash = "sha256:f22ac1c3cac4dbc50079e965eba2c1058622631e526bd9afd45fedd49ba781fa"}, + {file = "PyYAML-6.0.1-cp37-cp37m-macosx_10_9_x86_64.whl", hash = "sha256:b1275ad35a5d18c62a7220633c913e1b42d44b46ee12554e5fd39c70a243d6a3"}, + {file = "PyYAML-6.0.1-cp37-cp37m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:18aeb1bf9a78867dc38b259769503436b7c72f7a1f1f4c93ff9a17de54319b27"}, + {file = "PyYAML-6.0.1-cp37-cp37m-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:596106435fa6ad000c2991a98fa58eeb8656ef2325d7e158344fb33864ed87e3"}, + {file = "PyYAML-6.0.1-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:baa90d3f661d43131ca170712d903e6295d1f7a0f595074f151c0aed377c9b9c"}, + {file = "PyYAML-6.0.1-cp37-cp37m-win32.whl", hash = "sha256:9046c58c4395dff28dd494285c82ba00b546adfc7ef001486fbf0324bc174fba"}, + {file = "PyYAML-6.0.1-cp37-cp37m-win_amd64.whl", hash = "sha256:4fb147e7a67ef577a588a0e2c17b6db51dda102c71de36f8549b6816a96e1867"}, + {file = "PyYAML-6.0.1-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:1d4c7e777c441b20e32f52bd377e0c409713e8bb1386e1099c2415f26e479595"}, + {file = "PyYAML-6.0.1-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a0cd17c15d3bb3fa06978b4e8958dcdc6e0174ccea823003a106c7d4d7899ac5"}, + {file = "PyYAML-6.0.1-cp38-cp38-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:28c119d996beec18c05208a8bd78cbe4007878c6dd15091efb73a30e90539696"}, + {file = "PyYAML-6.0.1-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:7e07cbde391ba96ab58e532ff4803f79c4129397514e1413a7dc761ccd755735"}, + {file = "PyYAML-6.0.1-cp38-cp38-musllinux_1_1_x86_64.whl", hash = "sha256:49a183be227561de579b4a36efbb21b3eab9651dd81b1858589f796549873dd6"}, + {file = "PyYAML-6.0.1-cp38-cp38-win32.whl", hash = "sha256:184c5108a2aca3c5b3d3bf9395d50893a7ab82a38004c8f61c258d4428e80206"}, + {file = "PyYAML-6.0.1-cp38-cp38-win_amd64.whl", hash = "sha256:1e2722cc9fbb45d9b87631ac70924c11d3a401b2d7f410cc0e3bbf249f2dca62"}, + {file = "PyYAML-6.0.1-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:9eb6caa9a297fc2c2fb8862bc5370d0303ddba53ba97e71f08023b6cd73d16a8"}, + {file = "PyYAML-6.0.1-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:c8098ddcc2a85b61647b2590f825f3db38891662cfc2fc776415143f599bb859"}, + {file = "PyYAML-6.0.1-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:5773183b6446b2c99bb77e77595dd486303b4faab2b086e7b17bc6bef28865f6"}, + {file = "PyYAML-6.0.1-cp39-cp39-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:b786eecbdf8499b9ca1d697215862083bd6d2a99965554781d0d8d1ad31e13a0"}, + {file = "PyYAML-6.0.1-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:bc1bf2925a1ecd43da378f4db9e4f799775d6367bdb94671027b73b393a7c42c"}, + {file = "PyYAML-6.0.1-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:04ac92ad1925b2cff1db0cfebffb6ffc43457495c9b3c39d3fcae417d7125dc5"}, + {file = "PyYAML-6.0.1-cp39-cp39-win32.whl", hash = "sha256:faca3bdcf85b2fc05d06ff3fbc1f83e1391b3e724afa3feba7d13eeab355484c"}, + {file = "PyYAML-6.0.1-cp39-cp39-win_amd64.whl", hash = "sha256:510c9deebc5c0225e8c96813043e62b680ba2f9c50a08d3724c7f28a747d1486"}, + {file = "PyYAML-6.0.1.tar.gz", hash = "sha256:bfdf460b1736c775f2ba9f6a92bca30bc2095067b8a9d77876d1fad6cc3b4a43"}, +] + +[[package]] +name = "regex" +version = "2023.12.25" +description = "Alternative regular expression module, to replace re." +optional = false +python-versions = ">=3.7" +files = [ + {file = "regex-2023.12.25-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:0694219a1d54336fd0445ea382d49d36882415c0134ee1e8332afd1529f0baa5"}, + {file = "regex-2023.12.25-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:b014333bd0217ad3d54c143de9d4b9a3ca1c5a29a6d0d554952ea071cff0f1f8"}, + {file = "regex-2023.12.25-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:d865984b3f71f6d0af64d0d88f5733521698f6c16f445bb09ce746c92c97c586"}, + {file = "regex-2023.12.25-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:1e0eabac536b4cc7f57a5f3d095bfa557860ab912f25965e08fe1545e2ed8b4c"}, + {file = "regex-2023.12.25-cp310-cp310-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:c25a8ad70e716f96e13a637802813f65d8a6760ef48672aa3502f4c24ea8b400"}, + {file = "regex-2023.12.25-cp310-cp310-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:a9b6d73353f777630626f403b0652055ebfe8ff142a44ec2cf18ae470395766e"}, + {file = "regex-2023.12.25-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:a9cc99d6946d750eb75827cb53c4371b8b0fe89c733a94b1573c9dd16ea6c9e4"}, + {file = "regex-2023.12.25-cp310-cp310-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:88d1f7bef20c721359d8675f7d9f8e414ec5003d8f642fdfd8087777ff7f94b5"}, + {file = "regex-2023.12.25-cp310-cp310-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_12_x86_64.manylinux2010_x86_64.whl", hash = "sha256:cb3fe77aec8f1995611f966d0c656fdce398317f850d0e6e7aebdfe61f40e1cd"}, + {file = "regex-2023.12.25-cp310-cp310-musllinux_1_1_aarch64.whl", hash = "sha256:7aa47c2e9ea33a4a2a05f40fcd3ea36d73853a2aae7b4feab6fc85f8bf2c9704"}, + {file = "regex-2023.12.25-cp310-cp310-musllinux_1_1_i686.whl", hash = "sha256:df26481f0c7a3f8739fecb3e81bc9da3fcfae34d6c094563b9d4670b047312e1"}, + {file = "regex-2023.12.25-cp310-cp310-musllinux_1_1_ppc64le.whl", hash = "sha256:c40281f7d70baf6e0db0c2f7472b31609f5bc2748fe7275ea65a0b4601d9b392"}, + {file = "regex-2023.12.25-cp310-cp310-musllinux_1_1_s390x.whl", hash = "sha256:d94a1db462d5690ebf6ae86d11c5e420042b9898af5dcf278bd97d6bda065423"}, + {file = "regex-2023.12.25-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:ba1b30765a55acf15dce3f364e4928b80858fa8f979ad41f862358939bdd1f2f"}, + {file = "regex-2023.12.25-cp310-cp310-win32.whl", hash = "sha256:150c39f5b964e4d7dba46a7962a088fbc91f06e606f023ce57bb347a3b2d4630"}, + {file = "regex-2023.12.25-cp310-cp310-win_amd64.whl", hash = "sha256:09da66917262d9481c719599116c7dc0c321ffcec4b1f510c4f8a066f8768105"}, + {file = "regex-2023.12.25-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:1b9d811f72210fa9306aeb88385b8f8bcef0dfbf3873410413c00aa94c56c2b6"}, + {file = "regex-2023.12.25-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:d902a43085a308cef32c0d3aea962524b725403fd9373dea18110904003bac97"}, + {file = "regex-2023.12.25-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:d166eafc19f4718df38887b2bbe1467a4f74a9830e8605089ea7a30dd4da8887"}, + {file = "regex-2023.12.25-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:c7ad32824b7f02bb3c9f80306d405a1d9b7bb89362d68b3c5a9be53836caebdb"}, + {file = "regex-2023.12.25-cp311-cp311-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:636ba0a77de609d6510235b7f0e77ec494d2657108f777e8765efc060094c98c"}, + {file = "regex-2023.12.25-cp311-cp311-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:0fda75704357805eb953a3ee15a2b240694a9a514548cd49b3c5124b4e2ad01b"}, + {file = "regex-2023.12.25-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:f72cbae7f6b01591f90814250e636065850c5926751af02bb48da94dfced7baa"}, + {file = "regex-2023.12.25-cp311-cp311-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:db2a0b1857f18b11e3b0e54ddfefc96af46b0896fb678c85f63fb8c37518b3e7"}, + {file = "regex-2023.12.25-cp311-cp311-musllinux_1_1_aarch64.whl", hash = "sha256:7502534e55c7c36c0978c91ba6f61703faf7ce733715ca48f499d3dbbd7657e0"}, + {file = "regex-2023.12.25-cp311-cp311-musllinux_1_1_i686.whl", hash = "sha256:e8c7e08bb566de4faaf11984af13f6bcf6a08f327b13631d41d62592681d24fe"}, + {file = "regex-2023.12.25-cp311-cp311-musllinux_1_1_ppc64le.whl", hash = "sha256:283fc8eed679758de38fe493b7d7d84a198b558942b03f017b1f94dda8efae80"}, + {file = "regex-2023.12.25-cp311-cp311-musllinux_1_1_s390x.whl", hash = "sha256:f44dd4d68697559d007462b0a3a1d9acd61d97072b71f6d1968daef26bc744bd"}, + {file = "regex-2023.12.25-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:67d3ccfc590e5e7197750fcb3a2915b416a53e2de847a728cfa60141054123d4"}, + {file = "regex-2023.12.25-cp311-cp311-win32.whl", hash = "sha256:68191f80a9bad283432385961d9efe09d783bcd36ed35a60fb1ff3f1ec2efe87"}, + {file = "regex-2023.12.25-cp311-cp311-win_amd64.whl", hash = "sha256:7d2af3f6b8419661a0c421584cfe8aaec1c0e435ce7e47ee2a97e344b98f794f"}, + {file = "regex-2023.12.25-cp312-cp312-macosx_10_9_universal2.whl", hash = "sha256:8a0ccf52bb37d1a700375a6b395bff5dd15c50acb745f7db30415bae3c2b0715"}, + {file = "regex-2023.12.25-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:c3c4a78615b7762740531c27cf46e2f388d8d727d0c0c739e72048beb26c8a9d"}, + {file = "regex-2023.12.25-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:ad83e7545b4ab69216cef4cc47e344d19622e28aabec61574b20257c65466d6a"}, + {file = "regex-2023.12.25-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:b7a635871143661feccce3979e1727c4e094f2bdfd3ec4b90dfd4f16f571a87a"}, + {file = "regex-2023.12.25-cp312-cp312-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:d498eea3f581fbe1b34b59c697512a8baef88212f92e4c7830fcc1499f5b45a5"}, + {file = "regex-2023.12.25-cp312-cp312-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:43f7cd5754d02a56ae4ebb91b33461dc67be8e3e0153f593c509e21d219c5060"}, + {file = "regex-2023.12.25-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:51f4b32f793812714fd5307222a7f77e739b9bc566dc94a18126aba3b92b98a3"}, + {file = "regex-2023.12.25-cp312-cp312-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:ba99d8077424501b9616b43a2d208095746fb1284fc5ba490139651f971d39d9"}, + {file = "regex-2023.12.25-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:4bfc2b16e3ba8850e0e262467275dd4d62f0d045e0e9eda2bc65078c0110a11f"}, + {file = "regex-2023.12.25-cp312-cp312-musllinux_1_1_i686.whl", hash = "sha256:8c2c19dae8a3eb0ea45a8448356ed561be843b13cbc34b840922ddf565498c1c"}, + {file = "regex-2023.12.25-cp312-cp312-musllinux_1_1_ppc64le.whl", hash = "sha256:60080bb3d8617d96f0fb7e19796384cc2467447ef1c491694850ebd3670bc457"}, + {file = "regex-2023.12.25-cp312-cp312-musllinux_1_1_s390x.whl", hash = "sha256:b77e27b79448e34c2c51c09836033056a0547aa360c45eeeb67803da7b0eedaf"}, + {file = "regex-2023.12.25-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:518440c991f514331f4850a63560321f833979d145d7d81186dbe2f19e27ae3d"}, + {file = "regex-2023.12.25-cp312-cp312-win32.whl", hash = "sha256:e2610e9406d3b0073636a3a2e80db05a02f0c3169b5632022b4e81c0364bcda5"}, + {file = "regex-2023.12.25-cp312-cp312-win_amd64.whl", hash = "sha256:cc37b9aeebab425f11f27e5e9e6cf580be7206c6582a64467a14dda211abc232"}, + {file = "regex-2023.12.25-cp37-cp37m-macosx_10_9_x86_64.whl", hash = "sha256:da695d75ac97cb1cd725adac136d25ca687da4536154cdc2815f576e4da11c69"}, + {file = "regex-2023.12.25-cp37-cp37m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:d126361607b33c4eb7b36debc173bf25d7805847346dd4d99b5499e1fef52bc7"}, + {file = "regex-2023.12.25-cp37-cp37m-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:4719bb05094d7d8563a450cf8738d2e1061420f79cfcc1fa7f0a44744c4d8f73"}, + {file = "regex-2023.12.25-cp37-cp37m-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:5dd58946bce44b53b06d94aa95560d0b243eb2fe64227cba50017a8d8b3cd3e2"}, + {file = "regex-2023.12.25-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:22a86d9fff2009302c440b9d799ef2fe322416d2d58fc124b926aa89365ec482"}, + {file = "regex-2023.12.25-cp37-cp37m-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:2aae8101919e8aa05ecfe6322b278f41ce2994c4a430303c4cd163fef746e04f"}, + {file = "regex-2023.12.25-cp37-cp37m-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_12_x86_64.manylinux2010_x86_64.whl", hash = "sha256:e692296c4cc2873967771345a876bcfc1c547e8dd695c6b89342488b0ea55cd8"}, + {file = "regex-2023.12.25-cp37-cp37m-musllinux_1_1_aarch64.whl", hash = "sha256:263ef5cc10979837f243950637fffb06e8daed7f1ac1e39d5910fd29929e489a"}, + {file = "regex-2023.12.25-cp37-cp37m-musllinux_1_1_i686.whl", hash = "sha256:d6f7e255e5fa94642a0724e35406e6cb7001c09d476ab5fce002f652b36d0c39"}, + {file = "regex-2023.12.25-cp37-cp37m-musllinux_1_1_ppc64le.whl", hash = "sha256:88ad44e220e22b63b0f8f81f007e8abbb92874d8ced66f32571ef8beb0643b2b"}, + {file = "regex-2023.12.25-cp37-cp37m-musllinux_1_1_s390x.whl", hash = "sha256:3a17d3ede18f9cedcbe23d2daa8a2cd6f59fe2bf082c567e43083bba3fb00347"}, + {file = "regex-2023.12.25-cp37-cp37m-musllinux_1_1_x86_64.whl", hash = "sha256:d15b274f9e15b1a0b7a45d2ac86d1f634d983ca40d6b886721626c47a400bf39"}, + {file = "regex-2023.12.25-cp37-cp37m-win32.whl", hash = "sha256:ed19b3a05ae0c97dd8f75a5d8f21f7723a8c33bbc555da6bbe1f96c470139d3c"}, + {file = "regex-2023.12.25-cp37-cp37m-win_amd64.whl", hash = "sha256:a6d1047952c0b8104a1d371f88f4ab62e6275567d4458c1e26e9627ad489b445"}, + {file = "regex-2023.12.25-cp38-cp38-macosx_10_9_universal2.whl", hash = "sha256:b43523d7bc2abd757119dbfb38af91b5735eea45537ec6ec3a5ec3f9562a1c53"}, + {file = "regex-2023.12.25-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:efb2d82f33b2212898f1659fb1c2e9ac30493ac41e4d53123da374c3b5541e64"}, + {file = "regex-2023.12.25-cp38-cp38-macosx_11_0_arm64.whl", hash = "sha256:b7fca9205b59c1a3d5031f7e64ed627a1074730a51c2a80e97653e3e9fa0d415"}, + {file = "regex-2023.12.25-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:086dd15e9435b393ae06f96ab69ab2d333f5d65cbe65ca5a3ef0ec9564dfe770"}, + {file = "regex-2023.12.25-cp38-cp38-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:e81469f7d01efed9b53740aedd26085f20d49da65f9c1f41e822a33992cb1590"}, + {file = "regex-2023.12.25-cp38-cp38-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:34e4af5b27232f68042aa40a91c3b9bb4da0eeb31b7632e0091afc4310afe6cb"}, + {file = "regex-2023.12.25-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:9852b76ab558e45b20bf1893b59af64a28bd3820b0c2efc80e0a70a4a3ea51c1"}, + {file = "regex-2023.12.25-cp38-cp38-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:ff100b203092af77d1a5a7abe085b3506b7eaaf9abf65b73b7d6905b6cb76988"}, + {file = "regex-2023.12.25-cp38-cp38-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_12_x86_64.manylinux2010_x86_64.whl", hash = "sha256:cc038b2d8b1470364b1888a98fd22d616fba2b6309c5b5f181ad4483e0017861"}, + {file = "regex-2023.12.25-cp38-cp38-musllinux_1_1_aarch64.whl", hash = "sha256:094ba386bb5c01e54e14434d4caabf6583334090865b23ef58e0424a6286d3dc"}, + {file = "regex-2023.12.25-cp38-cp38-musllinux_1_1_i686.whl", hash = "sha256:5cd05d0f57846d8ba4b71d9c00f6f37d6b97d5e5ef8b3c3840426a475c8f70f4"}, + {file = "regex-2023.12.25-cp38-cp38-musllinux_1_1_ppc64le.whl", hash = "sha256:9aa1a67bbf0f957bbe096375887b2505f5d8ae16bf04488e8b0f334c36e31360"}, + {file = "regex-2023.12.25-cp38-cp38-musllinux_1_1_s390x.whl", hash = "sha256:98a2636994f943b871786c9e82bfe7883ecdaba2ef5df54e1450fa9869d1f756"}, + {file = "regex-2023.12.25-cp38-cp38-musllinux_1_1_x86_64.whl", hash = "sha256:37f8e93a81fc5e5bd8db7e10e62dc64261bcd88f8d7e6640aaebe9bc180d9ce2"}, + {file = "regex-2023.12.25-cp38-cp38-win32.whl", hash = "sha256:d78bd484930c1da2b9679290a41cdb25cc127d783768a0369d6b449e72f88beb"}, + {file = "regex-2023.12.25-cp38-cp38-win_amd64.whl", hash = "sha256:b521dcecebc5b978b447f0f69b5b7f3840eac454862270406a39837ffae4e697"}, + {file = "regex-2023.12.25-cp39-cp39-macosx_10_9_universal2.whl", hash = "sha256:f7bc09bc9c29ebead055bcba136a67378f03d66bf359e87d0f7c759d6d4ffa31"}, + {file = "regex-2023.12.25-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:e14b73607d6231f3cc4622809c196b540a6a44e903bcfad940779c80dffa7be7"}, + {file = "regex-2023.12.25-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:9eda5f7a50141291beda3edd00abc2d4a5b16c29c92daf8d5bd76934150f3edc"}, + {file = "regex-2023.12.25-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:cc6bb9aa69aacf0f6032c307da718f61a40cf970849e471254e0e91c56ffca95"}, + {file = "regex-2023.12.25-cp39-cp39-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:298dc6354d414bc921581be85695d18912bea163a8b23cac9a2562bbcd5088b1"}, + {file = "regex-2023.12.25-cp39-cp39-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:2f4e475a80ecbd15896a976aa0b386c5525d0ed34d5c600b6d3ebac0a67c7ddf"}, + {file = "regex-2023.12.25-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:531ac6cf22b53e0696f8e1d56ce2396311254eb806111ddd3922c9d937151dae"}, + {file = "regex-2023.12.25-cp39-cp39-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:22f3470f7524b6da61e2020672df2f3063676aff444db1daa283c2ea4ed259d6"}, + {file = "regex-2023.12.25-cp39-cp39-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_12_x86_64.manylinux2010_x86_64.whl", hash = "sha256:89723d2112697feaa320c9d351e5f5e7b841e83f8b143dba8e2d2b5f04e10923"}, + {file = "regex-2023.12.25-cp39-cp39-musllinux_1_1_aarch64.whl", hash = "sha256:0ecf44ddf9171cd7566ef1768047f6e66975788258b1c6c6ca78098b95cf9a3d"}, + {file = "regex-2023.12.25-cp39-cp39-musllinux_1_1_i686.whl", hash = "sha256:905466ad1702ed4acfd67a902af50b8db1feeb9781436372261808df7a2a7bca"}, + {file = "regex-2023.12.25-cp39-cp39-musllinux_1_1_ppc64le.whl", hash = "sha256:4558410b7a5607a645e9804a3e9dd509af12fb72b9825b13791a37cd417d73a5"}, + {file = "regex-2023.12.25-cp39-cp39-musllinux_1_1_s390x.whl", hash = "sha256:7e316026cc1095f2a3e8cc012822c99f413b702eaa2ca5408a513609488cb62f"}, + {file = "regex-2023.12.25-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:3b1de218d5375cd6ac4b5493e0b9f3df2be331e86520f23382f216c137913d20"}, + {file = "regex-2023.12.25-cp39-cp39-win32.whl", hash = "sha256:11a963f8e25ab5c61348d090bf1b07f1953929c13bd2309a0662e9ff680763c9"}, + {file = "regex-2023.12.25-cp39-cp39-win_amd64.whl", hash = "sha256:e693e233ac92ba83a87024e1d32b5f9ab15ca55ddd916d878146f4e3406b5c91"}, + {file = "regex-2023.12.25.tar.gz", hash = "sha256:29171aa128da69afdf4bde412d5bedc335f2ca8fcfe4489038577d05f16181e5"}, +] + +[[package]] +name = "requests" +version = "2.31.0" +description = "Python HTTP for Humans." +optional = false +python-versions = ">=3.7" +files = [ + {file = "requests-2.31.0-py3-none-any.whl", hash = "sha256:58cd2187c01e70e6e26505bca751777aa9f2ee0b7f4300988b709f44e013003f"}, + {file = "requests-2.31.0.tar.gz", hash = "sha256:942c5a758f98d790eaed1a29cb6eefc7ffb0d1cf7af05c3d2791656dbd6ad1e1"}, +] + +[package.dependencies] +certifi = ">=2017.4.17" +charset-normalizer = ">=2,<4" +idna = ">=2.5,<4" +PySocks = {version = ">=1.5.6,<1.5.7 || >1.5.7", optional = true, markers = "extra == \"socks\""} +urllib3 = ">=1.21.1,<3" + +[package.extras] +socks = ["PySocks (>=1.5.6,!=1.5.7)"] +use-chardet-on-py3 = ["chardet (>=3.0.2,<6)"] + +[[package]] +name = "scikit-image" +version = "0.22.0" +description = "Image processing in Python" +optional = false +python-versions = ">=3.9" +files = [ + {file = "scikit_image-0.22.0-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:74ec5c1d4693506842cc7c9487c89d8fc32aed064e9363def7af08b8f8cbb31d"}, + {file = "scikit_image-0.22.0-cp310-cp310-macosx_12_0_arm64.whl", hash = "sha256:a05ae4fe03d802587ed8974e900b943275548cde6a6807b785039d63e9a7a5ff"}, + {file = "scikit_image-0.22.0-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:6a92dca3d95b1301442af055e196a54b5a5128c6768b79fc0a4098f1d662dee6"}, + {file = "scikit_image-0.22.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:3663d063d8bf2fb9bdfb0ca967b9ee3b6593139c860c7abc2d2351a8a8863938"}, + {file = "scikit_image-0.22.0-cp310-cp310-win_amd64.whl", hash = "sha256:ebdbdc901bae14dab637f8d5c99f6d5cc7aaf4a3b6f4003194e003e9f688a6fc"}, + {file = "scikit_image-0.22.0-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:95d6da2d8a44a36ae04437c76d32deb4e3c993ffc846b394b9949fd8ded73cb2"}, + {file = "scikit_image-0.22.0-cp311-cp311-macosx_12_0_arm64.whl", hash = "sha256:2c6ef454a85f569659b813ac2a93948022b0298516b757c9c6c904132be327e2"}, + {file = "scikit_image-0.22.0-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:e87872f067444ee90a00dd49ca897208308645382e8a24bd3e76f301af2352cd"}, + {file = "scikit_image-0.22.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:c5c378db54e61b491b9edeefff87e49fcf7fdf729bb93c777d7a5f15d36f743e"}, + {file = "scikit_image-0.22.0-cp311-cp311-win_amd64.whl", hash = "sha256:2bcb74adb0634258a67f66c2bb29978c9a3e222463e003b67ba12056c003971b"}, + {file = "scikit_image-0.22.0-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:003ca2274ac0fac252280e7179ff986ff783407001459ddea443fe7916e38cff"}, + {file = "scikit_image-0.22.0-cp312-cp312-macosx_12_0_arm64.whl", hash = "sha256:cf3c0c15b60ae3e557a0c7575fbd352f0c3ce0afca562febfe3ab80efbeec0e9"}, + {file = "scikit_image-0.22.0-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:f5b23908dd4d120e6aecb1ed0277563e8cbc8d6c0565bdc4c4c6475d53608452"}, + {file = "scikit_image-0.22.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:be79d7493f320a964f8fcf603121595ba82f84720de999db0fcca002266a549a"}, + {file = "scikit_image-0.22.0-cp312-cp312-win_amd64.whl", hash = "sha256:722b970aa5da725dca55252c373b18bbea7858c1cdb406e19f9b01a4a73b30b2"}, + {file = "scikit_image-0.22.0-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:22318b35044cfeeb63ee60c56fc62450e5fe516228138f1d06c7a26378248a86"}, + {file = "scikit_image-0.22.0-cp39-cp39-macosx_12_0_arm64.whl", hash = "sha256:9e801c44a814afdadeabf4dffdffc23733e393767958b82319706f5fa3e1eaa9"}, + {file = "scikit_image-0.22.0-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:c472a1fb3665ec5c00423684590631d95f9afcbc97f01407d348b821880b2cb3"}, + {file = "scikit_image-0.22.0-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:3b7a6c89e8d6252332121b58f50e1625c35f7d6a85489c0b6b7ee4f5155d547a"}, + {file = "scikit_image-0.22.0-cp39-cp39-win_amd64.whl", hash = "sha256:5071b8f6341bfb0737ab05c8ab4ac0261f9e25dbcc7b5d31e5ed230fd24a7929"}, + {file = "scikit_image-0.22.0.tar.gz", hash = "sha256:018d734df1d2da2719087d15f679d19285fce97cd37695103deadfaef2873236"}, +] + +[package.dependencies] +imageio = ">=2.27" +lazy_loader = ">=0.3" +networkx = ">=2.8" +numpy = ">=1.22" +packaging = ">=21" +pillow = ">=9.0.1" +scipy = ">=1.8" +tifffile = ">=2022.8.12" + +[package.extras] +build = ["Cython (>=0.29.32)", "build", "meson-python (>=0.14)", "ninja", "numpy (>=1.22)", "packaging (>=21)", "pythran", "setuptools (>=67)", "spin (==0.6)", "wheel"] +data = ["pooch (>=1.6.0)"] +developer = ["pre-commit", "tomli"] +docs = ["PyWavelets (>=1.1.1)", "dask[array] (>=2022.9.2)", "ipykernel", "ipywidgets", "kaleido", "matplotlib (>=3.5)", "myst-parser", "numpydoc (>=1.6)", "pandas (>=1.5)", "plotly (>=5.10)", "pooch (>=1.6)", "pydata-sphinx-theme (>=0.14.1)", "pytest-runner", "scikit-learn (>=1.1)", "seaborn (>=0.11)", "sphinx (>=7.2)", "sphinx-copybutton", "sphinx-gallery (>=0.14)", "sphinx_design (>=0.5)", "tifffile (>=2022.8.12)"] +optional = ["PyWavelets (>=1.1.1)", "SimpleITK", "astropy (>=5.0)", "cloudpickle (>=0.2.1)", "dask[array] (>=2021.1.0)", "matplotlib (>=3.5)", "pooch (>=1.6.0)", "pyamg", "scikit-learn (>=1.1)"] +test = ["asv", "matplotlib (>=3.5)", "numpydoc (>=1.5)", "pooch (>=1.6.0)", "pytest (>=7.0)", "pytest-cov (>=2.11.0)", "pytest-faulthandler", "pytest-localserver"] + +[[package]] +name = "scikit-learn" +version = "1.2.2" +description = "A set of python modules for machine learning and data mining" +optional = false +python-versions = ">=3.8" +files = [ + {file = "scikit-learn-1.2.2.tar.gz", hash = "sha256:8429aea30ec24e7a8c7ed8a3fa6213adf3814a6efbea09e16e0a0c71e1a1a3d7"}, + {file = "scikit_learn-1.2.2-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:99cc01184e347de485bf253d19fcb3b1a3fb0ee4cea5ee3c43ec0cc429b6d29f"}, + {file = "scikit_learn-1.2.2-cp310-cp310-macosx_12_0_arm64.whl", hash = "sha256:e6e574db9914afcb4e11ade84fab084536a895ca60aadea3041e85b8ac963edb"}, + {file = "scikit_learn-1.2.2-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:6fe83b676f407f00afa388dd1fdd49e5c6612e551ed84f3b1b182858f09e987d"}, + {file = "scikit_learn-1.2.2-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:2e2642baa0ad1e8f8188917423dd73994bf25429f8893ddbe115be3ca3183584"}, + {file = "scikit_learn-1.2.2-cp310-cp310-win_amd64.whl", hash = "sha256:ad66c3848c0a1ec13464b2a95d0a484fd5b02ce74268eaa7e0c697b904f31d6c"}, + {file = "scikit_learn-1.2.2-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:dfeaf8be72117eb61a164ea6fc8afb6dfe08c6f90365bde2dc16456e4bc8e45f"}, + {file = "scikit_learn-1.2.2-cp311-cp311-macosx_12_0_arm64.whl", hash = "sha256:fe0aa1a7029ed3e1dcbf4a5bc675aa3b1bc468d9012ecf6c6f081251ca47f590"}, + {file = "scikit_learn-1.2.2-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:065e9673e24e0dc5113e2dd2b4ca30c9d8aa2fa90f4c0597241c93b63130d233"}, + {file = "scikit_learn-1.2.2-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:bf036ea7ef66115e0d49655f16febfa547886deba20149555a41d28f56fd6d3c"}, + {file = "scikit_learn-1.2.2-cp311-cp311-win_amd64.whl", hash = "sha256:8b0670d4224a3c2d596fd572fb4fa673b2a0ccfb07152688ebd2ea0b8c61025c"}, + {file = "scikit_learn-1.2.2-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:9c710ff9f9936ba8a3b74a455ccf0dcf59b230caa1e9ba0223773c490cab1e51"}, + {file = "scikit_learn-1.2.2-cp38-cp38-macosx_12_0_arm64.whl", hash = "sha256:2dd3ffd3950e3d6c0c0ef9033a9b9b32d910c61bd06cb8206303fb4514b88a49"}, + {file = "scikit_learn-1.2.2-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:44b47a305190c28dd8dd73fc9445f802b6ea716669cfc22ab1eb97b335d238b1"}, + {file = "scikit_learn-1.2.2-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:953236889928d104c2ef14027539f5f2609a47ebf716b8cbe4437e85dce42744"}, + {file = "scikit_learn-1.2.2-cp38-cp38-win_amd64.whl", hash = "sha256:7f69313884e8eb311460cc2f28676d5e400bd929841a2c8eb8742ae78ebf7c20"}, + {file = "scikit_learn-1.2.2-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:8156db41e1c39c69aa2d8599ab7577af53e9e5e7a57b0504e116cc73c39138dd"}, + {file = "scikit_learn-1.2.2-cp39-cp39-macosx_12_0_arm64.whl", hash = "sha256:fe175ee1dab589d2e1033657c5b6bec92a8a3b69103e3dd361b58014729975c3"}, + {file = "scikit_learn-1.2.2-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:7d5312d9674bed14f73773d2acf15a3272639b981e60b72c9b190a0cffed5bad"}, + {file = "scikit_learn-1.2.2-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ea061bf0283bf9a9f36ea3c5d3231ba2176221bbd430abd2603b1c3b2ed85c89"}, + {file = "scikit_learn-1.2.2-cp39-cp39-win_amd64.whl", hash = "sha256:6477eed40dbce190f9f9e9d0d37e020815825b300121307942ec2110302b66a3"}, +] + +[package.dependencies] +joblib = ">=1.1.1" +numpy = ">=1.17.3" +scipy = ">=1.3.2" +threadpoolctl = ">=2.0.0" + +[package.extras] +benchmark = ["matplotlib (>=3.1.3)", "memory-profiler (>=0.57.0)", "pandas (>=1.0.5)"] +docs = ["Pillow (>=7.1.2)", "matplotlib (>=3.1.3)", "memory-profiler (>=0.57.0)", "numpydoc (>=1.2.0)", "pandas (>=1.0.5)", "plotly (>=5.10.0)", "pooch (>=1.6.0)", "scikit-image (>=0.16.2)", "seaborn (>=0.9.0)", "sphinx (>=4.0.1)", "sphinx-gallery (>=0.7.0)", "sphinx-prompt (>=1.3.0)", "sphinxext-opengraph (>=0.4.2)"] +examples = ["matplotlib (>=3.1.3)", "pandas (>=1.0.5)", "plotly (>=5.10.0)", "pooch (>=1.6.0)", "scikit-image (>=0.16.2)", "seaborn (>=0.9.0)"] +tests = ["black (>=22.3.0)", "flake8 (>=3.8.2)", "matplotlib (>=3.1.3)", "mypy (>=0.961)", "numpydoc (>=1.2.0)", "pandas (>=1.0.5)", "pooch (>=1.6.0)", "pyamg (>=4.0.0)", "pytest (>=5.3.1)", "pytest-cov (>=2.9.0)", "scikit-image (>=0.16.2)"] + +[[package]] +name = "scipy" +version = "1.10.1" +description = "Fundamental algorithms for scientific computing in Python" +optional = false +python-versions = "<3.12,>=3.8" +files = [ + {file = "scipy-1.10.1-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:e7354fd7527a4b0377ce55f286805b34e8c54b91be865bac273f527e1b839019"}, + {file = "scipy-1.10.1-cp310-cp310-macosx_12_0_arm64.whl", hash = "sha256:4b3f429188c66603a1a5c549fb414e4d3bdc2a24792e061ffbd607d3d75fd84e"}, + {file = "scipy-1.10.1-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:1553b5dcddd64ba9a0d95355e63fe6c3fc303a8fd77c7bc91e77d61363f7433f"}, + {file = "scipy-1.10.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:4c0ff64b06b10e35215abce517252b375e580a6125fd5fdf6421b98efbefb2d2"}, + {file = "scipy-1.10.1-cp310-cp310-win_amd64.whl", hash = "sha256:fae8a7b898c42dffe3f7361c40d5952b6bf32d10c4569098d276b4c547905ee1"}, + {file = "scipy-1.10.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:0f1564ea217e82c1bbe75ddf7285ba0709ecd503f048cb1236ae9995f64217bd"}, + {file = "scipy-1.10.1-cp311-cp311-macosx_12_0_arm64.whl", hash = "sha256:d925fa1c81b772882aa55bcc10bf88324dadb66ff85d548c71515f6689c6dac5"}, + {file = "scipy-1.10.1-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:aaea0a6be54462ec027de54fca511540980d1e9eea68b2d5c1dbfe084797be35"}, + {file = "scipy-1.10.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:15a35c4242ec5f292c3dd364a7c71a61be87a3d4ddcc693372813c0b73c9af1d"}, + {file = "scipy-1.10.1-cp311-cp311-win_amd64.whl", hash = "sha256:43b8e0bcb877faf0abfb613d51026cd5cc78918e9530e375727bf0625c82788f"}, + {file = "scipy-1.10.1-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:5678f88c68ea866ed9ebe3a989091088553ba12c6090244fdae3e467b1139c35"}, + {file = "scipy-1.10.1-cp38-cp38-macosx_12_0_arm64.whl", hash = "sha256:39becb03541f9e58243f4197584286e339029e8908c46f7221abeea4b749fa88"}, + {file = "scipy-1.10.1-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:bce5869c8d68cf383ce240e44c1d9ae7c06078a9396df68ce88a1230f93a30c1"}, + {file = "scipy-1.10.1-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:07c3457ce0b3ad5124f98a86533106b643dd811dd61b548e78cf4c8786652f6f"}, + {file = "scipy-1.10.1-cp38-cp38-win_amd64.whl", hash = "sha256:049a8bbf0ad95277ffba9b3b7d23e5369cc39e66406d60422c8cfef40ccc8415"}, + {file = "scipy-1.10.1-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:cd9f1027ff30d90618914a64ca9b1a77a431159df0e2a195d8a9e8a04c78abf9"}, + {file = "scipy-1.10.1-cp39-cp39-macosx_12_0_arm64.whl", hash = "sha256:79c8e5a6c6ffaf3a2262ef1be1e108a035cf4f05c14df56057b64acc5bebffb6"}, + {file = "scipy-1.10.1-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:51af417a000d2dbe1ec6c372dfe688e041a7084da4fdd350aeb139bd3fb55353"}, + {file = "scipy-1.10.1-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:1b4735d6c28aad3cdcf52117e0e91d6b39acd4272f3f5cd9907c24ee931ad601"}, + {file = "scipy-1.10.1-cp39-cp39-win_amd64.whl", hash = "sha256:7ff7f37b1bf4417baca958d254e8e2875d0cc23aaadbe65b3d5b3077b0eb23ea"}, + {file = "scipy-1.10.1.tar.gz", hash = "sha256:2cf9dfb80a7b4589ba4c40ce7588986d6d5cebc5457cad2c2880f6bc2d42f3a5"}, +] + +[package.dependencies] +numpy = ">=1.19.5,<1.27.0" + +[package.extras] +dev = ["click", "doit (>=0.36.0)", "flake8", "mypy", "pycodestyle", "pydevtool", "rich-click", "typing_extensions"] +doc = ["matplotlib (>2)", "numpydoc", "pydata-sphinx-theme (==0.9.0)", "sphinx (!=4.1.0)", "sphinx-design (>=0.2.0)"] +test = ["asv", "gmpy2", "mpmath", "pooch", "pytest", "pytest-cov", "pytest-timeout", "pytest-xdist", "scikit-umfpack", "threadpoolctl"] + +[[package]] +name = "sentry-sdk" +version = "1.43.0" +description = "Python client for Sentry (https://sentry.io)" +optional = false +python-versions = "*" +files = [ + {file = "sentry-sdk-1.43.0.tar.gz", hash = "sha256:41df73af89d22921d8733714fb0fc5586c3461907e06688e6537d01a27e0e0f6"}, + {file = "sentry_sdk-1.43.0-py2.py3-none-any.whl", hash = "sha256:8d768724839ca18d7b4c7463ef7528c40b7aa2bfbf7fe554d5f9a7c044acfd36"}, +] + +[package.dependencies] +certifi = "*" +urllib3 = {version = ">=1.26.11", markers = "python_version >= \"3.6\""} + +[package.extras] +aiohttp = ["aiohttp (>=3.5)"] +arq = ["arq (>=0.23)"] +asyncpg = ["asyncpg (>=0.23)"] +beam = ["apache-beam (>=2.12)"] +bottle = ["bottle (>=0.12.13)"] +celery = ["celery (>=3)"] +celery-redbeat = ["celery-redbeat (>=2)"] +chalice = ["chalice (>=1.16.0)"] +clickhouse-driver = ["clickhouse-driver (>=0.2.0)"] +django = ["django (>=1.8)"] +falcon = ["falcon (>=1.4)"] +fastapi = ["fastapi (>=0.79.0)"] +flask = ["blinker (>=1.1)", "flask (>=0.11)", "markupsafe"] +grpcio = ["grpcio (>=1.21.1)"] +httpx = ["httpx (>=0.16.0)"] +huey = ["huey (>=2)"] +loguru = ["loguru (>=0.5)"] +openai = ["openai (>=1.0.0)", "tiktoken (>=0.3.0)"] +opentelemetry = ["opentelemetry-distro (>=0.35b0)"] +opentelemetry-experimental = ["opentelemetry-distro (>=0.40b0,<1.0)", "opentelemetry-instrumentation-aiohttp-client (>=0.40b0,<1.0)", "opentelemetry-instrumentation-django (>=0.40b0,<1.0)", "opentelemetry-instrumentation-fastapi (>=0.40b0,<1.0)", "opentelemetry-instrumentation-flask (>=0.40b0,<1.0)", "opentelemetry-instrumentation-requests (>=0.40b0,<1.0)", "opentelemetry-instrumentation-sqlite3 (>=0.40b0,<1.0)", "opentelemetry-instrumentation-urllib (>=0.40b0,<1.0)"] +pure-eval = ["asttokens", "executing", "pure-eval"] +pymongo = ["pymongo (>=3.1)"] +pyspark = ["pyspark (>=2.4.4)"] +quart = ["blinker (>=1.1)", "quart (>=0.16.1)"] +rq = ["rq (>=0.6)"] +sanic = ["sanic (>=0.8)"] +sqlalchemy = ["sqlalchemy (>=1.2)"] +starlette = ["starlette (>=0.19.1)"] +starlite = ["starlite (>=1.48)"] +tornado = ["tornado (>=5)"] + +[[package]] +name = "setproctitle" +version = "1.3.3" +description = "A Python module to customize the process title" +optional = false +python-versions = ">=3.7" +files = [ + {file = "setproctitle-1.3.3-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:897a73208da48db41e687225f355ce993167079eda1260ba5e13c4e53be7f754"}, + {file = "setproctitle-1.3.3-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:8c331e91a14ba4076f88c29c777ad6b58639530ed5b24b5564b5ed2fd7a95452"}, + {file = "setproctitle-1.3.3-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:bbbd6c7de0771c84b4aa30e70b409565eb1fc13627a723ca6be774ed6b9d9fa3"}, + {file = "setproctitle-1.3.3-cp310-cp310-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:c05ac48ef16ee013b8a326c63e4610e2430dbec037ec5c5b58fcced550382b74"}, + {file = "setproctitle-1.3.3-cp310-cp310-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:1342f4fdb37f89d3e3c1c0a59d6ddbedbde838fff5c51178a7982993d238fe4f"}, + {file = "setproctitle-1.3.3-cp310-cp310-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:fc74e84fdfa96821580fb5e9c0b0777c1c4779434ce16d3d62a9c4d8c710df39"}, + {file = "setproctitle-1.3.3-cp310-cp310-musllinux_1_1_aarch64.whl", hash = "sha256:9617b676b95adb412bb69645d5b077d664b6882bb0d37bfdafbbb1b999568d85"}, + {file = "setproctitle-1.3.3-cp310-cp310-musllinux_1_1_i686.whl", hash = "sha256:6a249415f5bb88b5e9e8c4db47f609e0bf0e20a75e8d744ea787f3092ba1f2d0"}, + {file = "setproctitle-1.3.3-cp310-cp310-musllinux_1_1_ppc64le.whl", hash = "sha256:38da436a0aaace9add67b999eb6abe4b84397edf4a78ec28f264e5b4c9d53cd5"}, + {file = "setproctitle-1.3.3-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:da0d57edd4c95bf221b2ebbaa061e65b1788f1544977288bdf95831b6e44e44d"}, + {file = "setproctitle-1.3.3-cp310-cp310-win32.whl", hash = "sha256:a1fcac43918b836ace25f69b1dca8c9395253ad8152b625064415b1d2f9be4fb"}, + {file = "setproctitle-1.3.3-cp310-cp310-win_amd64.whl", hash = "sha256:200620c3b15388d7f3f97e0ae26599c0c378fdf07ae9ac5a13616e933cbd2086"}, + {file = "setproctitle-1.3.3-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:334f7ed39895d692f753a443102dd5fed180c571eb6a48b2a5b7f5b3564908c8"}, + {file = "setproctitle-1.3.3-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:950f6476d56ff7817a8fed4ab207727fc5260af83481b2a4b125f32844df513a"}, + {file = "setproctitle-1.3.3-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:195c961f54a09eb2acabbfc90c413955cf16c6e2f8caa2adbf2237d1019c7dd8"}, + {file = "setproctitle-1.3.3-cp311-cp311-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:f05e66746bf9fe6a3397ec246fe481096664a9c97eb3fea6004735a4daf867fd"}, + {file = "setproctitle-1.3.3-cp311-cp311-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:b5901a31012a40ec913265b64e48c2a4059278d9f4e6be628441482dd13fb8b5"}, + {file = "setproctitle-1.3.3-cp311-cp311-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:64286f8a995f2cd934082b398fc63fca7d5ffe31f0e27e75b3ca6b4efda4e353"}, + {file = "setproctitle-1.3.3-cp311-cp311-musllinux_1_1_aarch64.whl", hash = "sha256:184239903bbc6b813b1a8fc86394dc6ca7d20e2ebe6f69f716bec301e4b0199d"}, + {file = "setproctitle-1.3.3-cp311-cp311-musllinux_1_1_i686.whl", hash = "sha256:664698ae0013f986118064b6676d7dcd28fefd0d7d5a5ae9497cbc10cba48fa5"}, + {file = "setproctitle-1.3.3-cp311-cp311-musllinux_1_1_ppc64le.whl", hash = "sha256:e5119a211c2e98ff18b9908ba62a3bd0e3fabb02a29277a7232a6fb4b2560aa0"}, + {file = "setproctitle-1.3.3-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:417de6b2e214e837827067048f61841f5d7fc27926f2e43954567094051aff18"}, + {file = "setproctitle-1.3.3-cp311-cp311-win32.whl", hash = "sha256:6a143b31d758296dc2f440175f6c8e0b5301ced3b0f477b84ca43cdcf7f2f476"}, + {file = "setproctitle-1.3.3-cp311-cp311-win_amd64.whl", hash = "sha256:a680d62c399fa4b44899094027ec9a1bdaf6f31c650e44183b50d4c4d0ccc085"}, + {file = "setproctitle-1.3.3-cp312-cp312-macosx_10_9_universal2.whl", hash = "sha256:d4460795a8a7a391e3567b902ec5bdf6c60a47d791c3b1d27080fc203d11c9dc"}, + {file = "setproctitle-1.3.3-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:bdfd7254745bb737ca1384dee57e6523651892f0ea2a7344490e9caefcc35e64"}, + {file = "setproctitle-1.3.3-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:477d3da48e216d7fc04bddab67b0dcde633e19f484a146fd2a34bb0e9dbb4a1e"}, + {file = "setproctitle-1.3.3-cp312-cp312-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:ab2900d111e93aff5df9fddc64cf51ca4ef2c9f98702ce26524f1acc5a786ae7"}, + {file = "setproctitle-1.3.3-cp312-cp312-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:088b9efc62d5aa5d6edf6cba1cf0c81f4488b5ce1c0342a8b67ae39d64001120"}, + {file = "setproctitle-1.3.3-cp312-cp312-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:a6d50252377db62d6a0bb82cc898089916457f2db2041e1d03ce7fadd4a07381"}, + {file = "setproctitle-1.3.3-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:87e668f9561fd3a457ba189edfc9e37709261287b52293c115ae3487a24b92f6"}, + {file = "setproctitle-1.3.3-cp312-cp312-musllinux_1_1_i686.whl", hash = "sha256:287490eb90e7a0ddd22e74c89a92cc922389daa95babc833c08cf80c84c4df0a"}, + {file = "setproctitle-1.3.3-cp312-cp312-musllinux_1_1_ppc64le.whl", hash = "sha256:4fe1c49486109f72d502f8be569972e27f385fe632bd8895f4730df3c87d5ac8"}, + {file = "setproctitle-1.3.3-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:4a6ba2494a6449b1f477bd3e67935c2b7b0274f2f6dcd0f7c6aceae10c6c6ba3"}, + {file = "setproctitle-1.3.3-cp312-cp312-win32.whl", hash = "sha256:2df2b67e4b1d7498632e18c56722851ba4db5d6a0c91aaf0fd395111e51cdcf4"}, + {file = "setproctitle-1.3.3-cp312-cp312-win_amd64.whl", hash = "sha256:f38d48abc121263f3b62943f84cbaede05749047e428409c2c199664feb6abc7"}, + {file = "setproctitle-1.3.3-cp37-cp37m-macosx_10_9_x86_64.whl", hash = "sha256:816330675e3504ae4d9a2185c46b573105d2310c20b19ea2b4596a9460a4f674"}, + {file = "setproctitle-1.3.3-cp37-cp37m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:68f960bc22d8d8e4ac886d1e2e21ccbd283adcf3c43136161c1ba0fa509088e0"}, + {file = "setproctitle-1.3.3-cp37-cp37m-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:00e6e7adff74796ef12753ff399491b8827f84f6c77659d71bd0b35870a17d8f"}, + {file = "setproctitle-1.3.3-cp37-cp37m-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:53bc0d2358507596c22b02db079618451f3bd720755d88e3cccd840bafb4c41c"}, + {file = "setproctitle-1.3.3-cp37-cp37m-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ad6d20f9541f5f6ac63df553b6d7a04f313947f550eab6a61aa758b45f0d5657"}, + {file = "setproctitle-1.3.3-cp37-cp37m-musllinux_1_1_aarch64.whl", hash = "sha256:c1c84beab776b0becaa368254801e57692ed749d935469ac10e2b9b825dbdd8e"}, + {file = "setproctitle-1.3.3-cp37-cp37m-musllinux_1_1_i686.whl", hash = "sha256:507e8dc2891021350eaea40a44ddd887c9f006e6b599af8d64a505c0f718f170"}, + {file = "setproctitle-1.3.3-cp37-cp37m-musllinux_1_1_ppc64le.whl", hash = "sha256:b1067647ac7aba0b44b591936118a22847bda3c507b0a42d74272256a7a798e9"}, + {file = "setproctitle-1.3.3-cp37-cp37m-musllinux_1_1_x86_64.whl", hash = "sha256:2e71f6365744bf53714e8bd2522b3c9c1d83f52ffa6324bd7cbb4da707312cd8"}, + {file = "setproctitle-1.3.3-cp37-cp37m-win32.whl", hash = "sha256:7f1d36a1e15a46e8ede4e953abb104fdbc0845a266ec0e99cc0492a4364f8c44"}, + {file = "setproctitle-1.3.3-cp37-cp37m-win_amd64.whl", hash = "sha256:c9a402881ec269d0cc9c354b149fc29f9ec1a1939a777f1c858cdb09c7a261df"}, + {file = "setproctitle-1.3.3-cp38-cp38-macosx_10_9_universal2.whl", hash = "sha256:ff814dea1e5c492a4980e3e7d094286077054e7ea116cbeda138819db194b2cd"}, + {file = "setproctitle-1.3.3-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:accb66d7b3ccb00d5cd11d8c6e07055a4568a24c95cf86109894dcc0c134cc89"}, + {file = "setproctitle-1.3.3-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:554eae5a5b28f02705b83a230e9d163d645c9a08914c0ad921df363a07cf39b1"}, + {file = "setproctitle-1.3.3-cp38-cp38-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:a911b26264dbe9e8066c7531c0591cfab27b464459c74385b276fe487ca91c12"}, + {file = "setproctitle-1.3.3-cp38-cp38-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:2982efe7640c4835f7355fdb4da313ad37fb3b40f5c69069912f8048f77b28c8"}, + {file = "setproctitle-1.3.3-cp38-cp38-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:df3f4274b80709d8bcab2f9a862973d453b308b97a0b423a501bcd93582852e3"}, + {file = "setproctitle-1.3.3-cp38-cp38-musllinux_1_1_aarch64.whl", hash = "sha256:af2c67ae4c795d1674a8d3ac1988676fa306bcfa1e23fddb5e0bd5f5635309ca"}, + {file = "setproctitle-1.3.3-cp38-cp38-musllinux_1_1_i686.whl", hash = "sha256:af4061f67fd7ec01624c5e3c21f6b7af2ef0e6bab7fbb43f209e6506c9ce0092"}, + {file = "setproctitle-1.3.3-cp38-cp38-musllinux_1_1_ppc64le.whl", hash = "sha256:37a62cbe16d4c6294e84670b59cf7adcc73faafe6af07f8cb9adaf1f0e775b19"}, + {file = "setproctitle-1.3.3-cp38-cp38-musllinux_1_1_x86_64.whl", hash = "sha256:a83ca086fbb017f0d87f240a8f9bbcf0809f3b754ee01cec928fff926542c450"}, + {file = "setproctitle-1.3.3-cp38-cp38-win32.whl", hash = "sha256:059f4ce86f8cc92e5860abfc43a1dceb21137b26a02373618d88f6b4b86ba9b2"}, + {file = "setproctitle-1.3.3-cp38-cp38-win_amd64.whl", hash = "sha256:ab92e51cd4a218208efee4c6d37db7368fdf182f6e7ff148fb295ecddf264287"}, + {file = "setproctitle-1.3.3-cp39-cp39-macosx_10_9_universal2.whl", hash = "sha256:c7951820b77abe03d88b114b998867c0f99da03859e5ab2623d94690848d3e45"}, + {file = "setproctitle-1.3.3-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:5bc94cf128676e8fac6503b37763adb378e2b6be1249d207630f83fc325d9b11"}, + {file = "setproctitle-1.3.3-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:1f5d9027eeda64d353cf21a3ceb74bb1760bd534526c9214e19f052424b37e42"}, + {file = "setproctitle-1.3.3-cp39-cp39-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:2e4a8104db15d3462e29d9946f26bed817a5b1d7a47eabca2d9dc2b995991503"}, + {file = "setproctitle-1.3.3-cp39-cp39-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:c32c41ace41f344d317399efff4cffb133e709cec2ef09c99e7a13e9f3b9483c"}, + {file = "setproctitle-1.3.3-cp39-cp39-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:cbf16381c7bf7f963b58fb4daaa65684e10966ee14d26f5cc90f07049bfd8c1e"}, + {file = "setproctitle-1.3.3-cp39-cp39-musllinux_1_1_aarch64.whl", hash = "sha256:e18b7bd0898398cc97ce2dfc83bb192a13a087ef6b2d5a8a36460311cb09e775"}, + {file = "setproctitle-1.3.3-cp39-cp39-musllinux_1_1_i686.whl", hash = "sha256:69d565d20efe527bd8a9b92e7f299ae5e73b6c0470f3719bd66f3cd821e0d5bd"}, + {file = "setproctitle-1.3.3-cp39-cp39-musllinux_1_1_ppc64le.whl", hash = "sha256:ddedd300cd690a3b06e7eac90ed4452348b1348635777ce23d460d913b5b63c3"}, + {file = "setproctitle-1.3.3-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:415bfcfd01d1fbf5cbd75004599ef167a533395955305f42220a585f64036081"}, + {file = "setproctitle-1.3.3-cp39-cp39-win32.whl", hash = "sha256:21112fcd2195d48f25760f0eafa7a76510871bbb3b750219310cf88b04456ae3"}, + {file = "setproctitle-1.3.3-cp39-cp39-win_amd64.whl", hash = "sha256:5a740f05d0968a5a17da3d676ce6afefebeeeb5ce137510901bf6306ba8ee002"}, + {file = "setproctitle-1.3.3-pp310-pypy310_pp73-macosx_10_9_x86_64.whl", hash = "sha256:6b9e62ddb3db4b5205c0321dd69a406d8af9ee1693529d144e86bd43bcb4b6c0"}, + {file = "setproctitle-1.3.3-pp310-pypy310_pp73-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:9e3b99b338598de0bd6b2643bf8c343cf5ff70db3627af3ca427a5e1a1a90dd9"}, + {file = "setproctitle-1.3.3-pp310-pypy310_pp73-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:38ae9a02766dad331deb06855fb7a6ca15daea333b3967e214de12cfae8f0ef5"}, + {file = "setproctitle-1.3.3-pp310-pypy310_pp73-win_amd64.whl", hash = "sha256:200ede6fd11233085ba9b764eb055a2a191fb4ffb950c68675ac53c874c22e20"}, + {file = "setproctitle-1.3.3-pp37-pypy37_pp73-macosx_10_9_x86_64.whl", hash = "sha256:0d3a953c50776751e80fe755a380a64cb14d61e8762bd43041ab3f8cc436092f"}, + {file = "setproctitle-1.3.3-pp37-pypy37_pp73-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:e5e08e232b78ba3ac6bc0d23ce9e2bee8fad2be391b7e2da834fc9a45129eb87"}, + {file = "setproctitle-1.3.3-pp37-pypy37_pp73-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:f1da82c3e11284da4fcbf54957dafbf0655d2389cd3d54e4eaba636faf6d117a"}, + {file = "setproctitle-1.3.3-pp37-pypy37_pp73-win_amd64.whl", hash = "sha256:aeaa71fb9568ebe9b911ddb490c644fbd2006e8c940f21cb9a1e9425bd709574"}, + {file = "setproctitle-1.3.3-pp38-pypy38_pp73-macosx_10_9_x86_64.whl", hash = "sha256:59335d000c6250c35989394661eb6287187854e94ac79ea22315469ee4f4c244"}, + {file = "setproctitle-1.3.3-pp38-pypy38_pp73-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:c3ba57029c9c50ecaf0c92bb127224cc2ea9fda057b5d99d3f348c9ec2855ad3"}, + {file = "setproctitle-1.3.3-pp38-pypy38_pp73-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:d876d355c53d975c2ef9c4f2487c8f83dad6aeaaee1b6571453cb0ee992f55f6"}, + {file = "setproctitle-1.3.3-pp38-pypy38_pp73-win_amd64.whl", hash = "sha256:224602f0939e6fb9d5dd881be1229d485f3257b540f8a900d4271a2c2aa4e5f4"}, + {file = "setproctitle-1.3.3-pp39-pypy39_pp73-macosx_10_9_x86_64.whl", hash = "sha256:d7f27e0268af2d7503386e0e6be87fb9b6657afd96f5726b733837121146750d"}, + {file = "setproctitle-1.3.3-pp39-pypy39_pp73-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:f5e7266498cd31a4572378c61920af9f6b4676a73c299fce8ba93afd694f8ae7"}, + {file = "setproctitle-1.3.3-pp39-pypy39_pp73-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:33c5609ad51cd99d388e55651b19148ea99727516132fb44680e1f28dd0d1de9"}, + {file = "setproctitle-1.3.3-pp39-pypy39_pp73-win_amd64.whl", hash = "sha256:eae8988e78192fd1a3245a6f4f382390b61bce6cfcc93f3809726e4c885fa68d"}, + {file = "setproctitle-1.3.3.tar.gz", hash = "sha256:c913e151e7ea01567837ff037a23ca8740192880198b7fbb90b16d181607caae"}, +] + +[package.extras] +test = ["pytest"] + +[[package]] +name = "setuptools" +version = "69.2.0" +description = "Easily download, build, install, upgrade, and uninstall Python packages" +optional = false +python-versions = ">=3.8" +files = [ + {file = "setuptools-69.2.0-py3-none-any.whl", hash = "sha256:c21c49fb1042386df081cb5d86759792ab89efca84cf114889191cd09aacc80c"}, + {file = "setuptools-69.2.0.tar.gz", hash = "sha256:0ff4183f8f42cd8fa3acea16c45205521a4ef28f73c6391d8a25e92893134f2e"}, +] + +[package.extras] +docs = ["furo", "jaraco.packaging (>=9.3)", "jaraco.tidelift (>=1.4)", "pygments-github-lexers (==0.0.5)", "rst.linker (>=1.9)", "sphinx (<7.2.5)", "sphinx (>=3.5)", "sphinx-favicon", "sphinx-inline-tabs", "sphinx-lint", "sphinx-notfound-page (>=1,<2)", "sphinx-reredirects", "sphinxcontrib-towncrier"] +testing = ["build[virtualenv]", "filelock (>=3.4.0)", "importlib-metadata", "ini2toml[lite] (>=0.9)", "jaraco.develop (>=7.21)", "jaraco.envs (>=2.2)", "jaraco.path (>=3.2.0)", "mypy (==1.9)", "packaging (>=23.2)", "pip (>=19.1)", "pytest (>=6)", "pytest-checkdocs (>=2.4)", "pytest-cov", "pytest-enabler (>=2.2)", "pytest-home (>=0.5)", "pytest-mypy (>=0.9.1)", "pytest-perf", "pytest-ruff (>=0.2.1)", "pytest-timeout", "pytest-xdist (>=3)", "tomli", "tomli-w (>=1.0.0)", "virtualenv (>=13.0.0)", "wheel"] +testing-integration = ["build[virtualenv] (>=1.0.3)", "filelock (>=3.4.0)", "jaraco.envs (>=2.2)", "jaraco.path (>=3.2.0)", "packaging (>=23.2)", "pytest", "pytest-enabler", "pytest-xdist", "tomli", "virtualenv (>=13.0.0)", "wheel"] + +[[package]] +name = "six" +version = "1.16.0" +description = "Python 2 and 3 compatibility utilities" +optional = false +python-versions = ">=2.7, !=3.0.*, !=3.1.*, !=3.2.*" +files = [ + {file = "six-1.16.0-py2.py3-none-any.whl", hash = "sha256:8abb2f1d86890a2dfb989f9a77cfcfd3e47c2a354b01111771326f8aa26e0254"}, + {file = "six-1.16.0.tar.gz", hash = "sha256:1e61c37477a1626458e36f7b1d82aa5c9b094fa4802892072e49de9c60c4c926"}, +] + +[[package]] +name = "smmap" +version = "5.0.1" +description = "A pure Python implementation of a sliding window memory map manager" +optional = false +python-versions = ">=3.7" +files = [ + {file = "smmap-5.0.1-py3-none-any.whl", hash = "sha256:e6d8668fa5f93e706934a62d7b4db19c8d9eb8cf2adbb75ef1b675aa332b69da"}, + {file = "smmap-5.0.1.tar.gz", hash = "sha256:dceeb6c0028fdb6734471eb07c0cd2aae706ccaecab45965ee83f11c8d3b1f62"}, +] + +[[package]] +name = "threadpoolctl" +version = "3.4.0" +description = "threadpoolctl" +optional = false +python-versions = ">=3.8" +files = [ + {file = "threadpoolctl-3.4.0-py3-none-any.whl", hash = "sha256:8f4c689a65b23e5ed825c8436a92b818aac005e0f3715f6a1664d7c7ee29d262"}, + {file = "threadpoolctl-3.4.0.tar.gz", hash = "sha256:f11b491a03661d6dd7ef692dd422ab34185d982466c49c8f98c8f716b5c93196"}, +] + +[[package]] +name = "tifffile" +version = "2024.2.12" +description = "Read and write TIFF files" +optional = false +python-versions = ">=3.9" +files = [ + {file = "tifffile-2024.2.12-py3-none-any.whl", hash = "sha256:870998f82fbc94ff7c3528884c1b0ae54863504ff51dbebea431ac3fa8fb7c21"}, + {file = "tifffile-2024.2.12.tar.gz", hash = "sha256:4920a3ec8e8e003e673d3c6531863c99eedd570d1b8b7e141c072ed78ff8030d"}, +] + +[package.dependencies] +numpy = "*" + +[package.extras] +all = ["defusedxml", "fsspec", "imagecodecs (>=2023.8.12)", "lxml", "matplotlib", "zarr"] + +[[package]] +name = "torch" +version = "1.13.1+cu117" +description = "Tensors and Dynamic neural networks in Python with strong GPU acceleration" +optional = false +python-versions = ">=3.7.0" +files = [ + {file = "torch-1.13.1+cu117-cp310-cp310-linux_x86_64.whl", hash = "sha256:14c5c9db09df8cf1b3942a3479c779da6e293a84a162d8a6ac71e2bde30e30c5"}, + {file = "torch-1.13.1+cu117-cp310-cp310-win_amd64.whl", hash = "sha256:978239684c6ec455ad2157ff33d44fdb9dd8d3a93b9d2f4ac7aa57691e990136"}, + {file = "torch-1.13.1+cu117-cp311-cp311-linux_x86_64.whl", hash = "sha256:d151d82462f125e35995ef6701213a499ffbe0994ce66bdad4349049a51a015e"}, + {file = "torch-1.13.1+cu117-cp37-cp37m-linux_x86_64.whl", hash = "sha256:6d783a1f79871724cada3c374ca509886359a8c4ac78d053df89b142ee904440"}, + {file = "torch-1.13.1+cu117-cp37-cp37m-win_amd64.whl", hash = "sha256:0efea7f09642fe382042bfb5fd7b6b83d47b56d4760278e660a4a74c728f1cb9"}, + {file = "torch-1.13.1+cu117-cp38-cp38-linux_x86_64.whl", hash = "sha256:bbf9546f0d0d8b51263ca479637b426a88335fca0034f42cec63d4d32dee05af"}, + {file = "torch-1.13.1+cu117-cp38-cp38-win_amd64.whl", hash = "sha256:99a6a9449adbe6c35c71f4ecb5d8f57a12fff06cda232dfa386ac1aacb8753f4"}, + {file = "torch-1.13.1+cu117-cp39-cp39-linux_x86_64.whl", hash = "sha256:b3ac139e0d4a0b303cc16f51eb77146ec7d14c0b1e702ad638613628f5086af7"}, + {file = "torch-1.13.1+cu117-cp39-cp39-win_amd64.whl", hash = "sha256:e775fa85f412bd1bf816b8798dadb3b852b71e33e8008e9db29b6190ed94fe27"}, +] + +[package.dependencies] +typing-extensions = "*" + +[package.extras] +opt-einsum = ["opt-einsum (>=3.3)"] + +[package.source] +type = "legacy" +url = "https://download.pytorch.org/whl/cu117" +reference = "torch+cu117" + +[[package]] +name = "torchmetrics" +version = "0.11.4" +description = "PyTorch native Metrics" +optional = false +python-versions = ">=3.7" +files = [ + {file = "torchmetrics-0.11.4-py3-none-any.whl", hash = "sha256:45f892f3534e91f3ad9e2488d1b05a93b7cb76b7d037969435a41a1f24750d9a"}, + {file = "torchmetrics-0.11.4.tar.gz", hash = "sha256:1fe45a14b44dd65d90199017dd5a4b5a128d56a8a311da7916c402c18c671494"}, +] + +[package.dependencies] +numpy = ">=1.17.2" +packaging = "*" +torch = ">=1.8.1" + +[package.extras] +all = ["lpips (<=0.1.4)", "nltk (>=3.6)", "pycocotools (>2.0.0)", "pystoi (<=0.3.3)", "regex (>=2021.9.24)", "scipy (>1.0.0)", "torch-fidelity (<=0.3.0)", "torchvision (>=0.8)", "tqdm (>=4.41.0)", "transformers (>=4.10.0)"] +audio = ["pystoi (<=0.3.3)"] +detection = ["pycocotools (>2.0.0)", "torchvision (>=0.8)"] +image = ["lpips (<=0.1.4)", "scipy (>1.0.0)", "torch-fidelity (<=0.3.0)", "torchvision (>=0.8)"] +multimodal = ["transformers (>=4.10.0)"] +test = ["bert-score (==0.3.13)", "cloudpickle (>1.3)", "coverage (>5.2)", "dython (<=0.7.3)", "fast-bss-eval (>=0.1.0)", "fire (<=0.5.0)", "huggingface-hub (<0.7)", "jiwer (>=2.3.0)", "kornia (>=0.6.7)", "mir-eval (>=0.6)", "mypy (==0.982)", "netcal (>1.0.0)", "pandas (>1.0.0)", "phmdoctest (>=1.1.1)", "psutil (<=5.9.4)", "pypesq (>1.2)", "pytest (>=6.0.0)", "pytest-cov (>2.10)", "pytest-doctestplus (>=0.9.0)", "pytest-rerunfailures (>=10.0)", "pytest-timeout (<=2.1.0)", "pytorch-msssim (==0.2.1)", "requests (<=2.28.2)", "rouge-score (>0.1.0)", "sacrebleu (>=2.0.0)", "scikit-image (>0.17.1)", "scikit-learn (>1.0)", "scipy (>1.0.0)", "torch-complex (<=0.4.3)", "transformers (>4.4.0)", "types-PyYAML", "types-emoji", "types-protobuf", "types-requests", "types-setuptools", "types-six", "types-tabulate"] +text = ["nltk (>=3.6)", "regex (>=2021.9.24)", "tqdm (>=4.41.0)"] + +[[package]] +name = "torchvision" +version = "0.14.1+cu117" +description = "image and video datasets and models for torch deep learning" +optional = false +python-versions = ">=3.7" +files = [ + {file = "torchvision-0.14.1+cu117-cp310-cp310-linux_x86_64.whl", hash = "sha256:83d271a530392814b8c7f69382285cac7c7da79b8f5dcc6bcf56c628547bce53"}, + {file = "torchvision-0.14.1+cu117-cp310-cp310-win_amd64.whl", hash = "sha256:b39fc67e7131053d435804d7901e88528611c0832fd9f1cc26476b5a27cc5d81"}, + {file = "torchvision-0.14.1+cu117-cp37-cp37m-linux_x86_64.whl", hash = "sha256:bce3a15aa1ae72572f8cd28e49d1ec1ad8f0866fbb661c7368d9c1285f69b22e"}, + {file = "torchvision-0.14.1+cu117-cp37-cp37m-win_amd64.whl", hash = "sha256:de04a6a5659d8b800c49602697c2534f9be409139cf37b516080300eca3a2487"}, + {file = "torchvision-0.14.1+cu117-cp38-cp38-linux_x86_64.whl", hash = "sha256:7409385139ce8864ceb2c82ffce338ebe15a55193f4b80c42a6d0f9cfd781629"}, + {file = "torchvision-0.14.1+cu117-cp38-cp38-win_amd64.whl", hash = "sha256:39846a1ed944964fe35d98da806b629c9d4294d6a5b069a097b0954f6c7da493"}, + {file = "torchvision-0.14.1+cu117-cp39-cp39-linux_x86_64.whl", hash = "sha256:8a89a0ee007d7cdba53bd56426d8795060a064b8919b61ac39e00e399ded6be2"}, + {file = "torchvision-0.14.1+cu117-cp39-cp39-win_amd64.whl", hash = "sha256:c9fa805e58971ad0d8f4810b8e56f2b260939297ff564a988cf0194f31bd5d69"}, +] + +[package.dependencies] +numpy = "*" +pillow = ">=5.3.0,<8.3.dev0 || >=8.4.dev0" +requests = "*" +torch = "1.13.1" +typing-extensions = "*" + +[package.extras] +scipy = ["scipy"] + +[package.source] +type = "legacy" +url = "https://download.pytorch.org/whl/cu117" +reference = "torch+cu117" + +[[package]] +name = "tqdm" +version = "4.66.1" +description = "Fast, Extensible Progress Meter" +optional = false +python-versions = ">=3.7" +files = [ + {file = "tqdm-4.66.1-py3-none-any.whl", hash = "sha256:d302b3c5b53d47bce91fea46679d9c3c6508cf6332229aa1e7d8653723793386"}, + {file = "tqdm-4.66.1.tar.gz", hash = "sha256:d88e651f9db8d8551a62556d3cff9e3034274ca5d66e93197cf2490e2dcb69c7"}, +] + +[package.dependencies] +colorama = {version = "*", markers = "platform_system == \"Windows\""} + +[package.extras] +dev = ["pytest (>=6)", "pytest-cov", "pytest-timeout", "pytest-xdist"] +notebook = ["ipywidgets (>=6)"] +slack = ["slack-sdk"] +telegram = ["requests"] + +[[package]] +name = "typing-extensions" +version = "4.10.0" +description = "Backported and Experimental Type Hints for Python 3.8+" +optional = false +python-versions = ">=3.8" +files = [ + {file = "typing_extensions-4.10.0-py3-none-any.whl", hash = "sha256:69b1a937c3a517342112fb4c6df7e72fc39a38e7891a5730ed4985b5214b5475"}, + {file = "typing_extensions-4.10.0.tar.gz", hash = "sha256:b0abd7c89e8fb96f98db18d86106ff1d90ab692004eb746cf6eda2682f91b3cb"}, +] + +[[package]] +name = "tzdata" +version = "2024.1" +description = "Provider of IANA time zone data" +optional = false +python-versions = ">=2" +files = [ + {file = "tzdata-2024.1-py2.py3-none-any.whl", hash = "sha256:9068bc196136463f5245e51efda838afa15aaeca9903f49050dfa2679db4d252"}, + {file = "tzdata-2024.1.tar.gz", hash = "sha256:2674120f8d891909751c38abcdfd386ac0a5a1127954fbc332af6b5ceae07efd"}, +] + +[[package]] +name = "urllib3" +version = "2.2.1" +description = "HTTP library with thread-safe connection pooling, file post, and more." +optional = false +python-versions = ">=3.8" +files = [ + {file = "urllib3-2.2.1-py3-none-any.whl", hash = "sha256:450b20ec296a467077128bff42b73080516e71b56ff59a60a02bef2232c4fa9d"}, + {file = "urllib3-2.2.1.tar.gz", hash = "sha256:d0570876c61ab9e520d776c38acbbb5b05a776d3f9ff98a5c8fd5162a444cf19"}, +] + +[package.extras] +brotli = ["brotli (>=1.0.9)", "brotlicffi (>=0.8.0)"] +h2 = ["h2 (>=4,<5)"] +socks = ["pysocks (>=1.5.6,!=1.5.7,<2.0)"] +zstd = ["zstandard (>=0.18.0)"] + +[[package]] +name = "wandb" +version = "0.15.10" +description = "A CLI and library for interacting with the Weights & Biases API." +optional = false +python-versions = ">=3.6" +files = [ + {file = "wandb-0.15.10-py3-none-any.whl", hash = "sha256:bc810879fecd1b62ccd498658e55ade3702939090a94b99418630e77e1f43d50"}, + {file = "wandb-0.15.10.tar.gz", hash = "sha256:5ce6d2d6be904da726b591a19ac22d2e71aa4b6a9ed450d8cd817d1712969b13"}, +] + +[package.dependencies] +appdirs = ">=1.4.3" +Click = ">=7.1,<8.0.0 || >8.0.0" +docker-pycreds = ">=0.4.0" +GitPython = ">=1.0.0,<3.1.29 || >3.1.29" +pathtools = "*" +protobuf = {version = ">=3.19.0,<4.21.0 || >4.21.0,<5", markers = "python_version > \"3.9\" or sys_platform != \"linux\""} +psutil = ">=5.0.0" +PyYAML = "*" +requests = ">=2.0.0,<3" +sentry-sdk = ">=1.0.0" +setproctitle = "*" +setuptools = "*" + +[package.extras] +async = ["httpx (>=0.22.0)"] +aws = ["boto3"] +azure = ["azure-identity", "azure-storage-blob"] +gcp = ["google-cloud-storage"] +kubeflow = ["google-cloud-storage", "kubernetes", "minio", "sh"] +launch = ["awscli", "azure-containerregistry", "azure-identity", "azure-storage-blob", "boto3", "botocore", "chardet", "google-auth", "google-cloud-artifact-registry", "google-cloud-compute", "google-cloud-storage", "iso8601", "kubernetes", "nbconvert", "nbformat", "optuna", "typing-extensions"] +media = ["bokeh", "moviepy", "numpy", "pillow", "plotly", "rdkit-pypi", "soundfile"] +models = ["cloudpickle"] +perf = ["orjson"] +sweeps = ["sweeps (>=0.2.0)"] + +[[package]] +name = "wcwidth" +version = "0.2.13" +description = "Measures the displayed width of unicode strings in a terminal" +optional = false +python-versions = "*" +files = [ + {file = "wcwidth-0.2.13-py2.py3-none-any.whl", hash = "sha256:3da69048e4540d84af32131829ff948f1e022c1c6bdb8d6102117aac784f6859"}, + {file = "wcwidth-0.2.13.tar.gz", hash = "sha256:72ea0c06399eb286d978fdedb6923a9eb47e1c486ce63e9b4e64fc18303972b5"}, +] + +[metadata] +lock-version = "2.0" +python-versions = "~3.10" +content-hash = "1a33af797cf0f0d0c26c9d159e096055bce9f384d3672ee83c9ab4ff6eea274d" diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000000000000000000000000000000000000..e439d5809264660faf102c74ff7dc3f7b2aeafb3 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,33 @@ +[tool.poetry] +package-mode = false + +[tool.poetry.dependencies] +python = "~3.10" +torch = {version = "1.13.1", source = "torch+cu117"} +torchvision = {version = "0.14.1", source = "torch+cu117"} +clip = {git = "https://github.com/openai/CLIP.git"} +torchmetrics = "0.11.4" +numpy = "1.23.5" +pillow = "10.0.0" +pandas = "2.0.1" +wandb = "0.15.10" +addict = "2.4.0" +tqdm = "4.66.1" +gdown = "3.12.2" +dlib = "19.24.1" +scipy = "1.10.1" +opencv-python = "4.7.0.72" +dill = "0.3.6" +face-alignment = "1.3.4" +matplotlib = "3.7.1" +joblib = "1.2.0" +scikit-learn = "1.2.2" + +[[tool.poetry.source]] +name = "torch+cu117" +url = "https://download.pytorch.org/whl/cu117" +priority = "explicit" + +[build-system] +requires = ["poetry-core"] +build-backend = "poetry.core.masonry.api" diff --git a/render.yaml b/render.yaml new file mode 100644 index 0000000000000000000000000000000000000000..d9e98e124d94620a4839b6f9ad643787cd7aab48 --- /dev/null +++ b/render.yaml @@ -0,0 +1,13 @@ +services: + - type: web + name: hairfastgan-backend + runtime: python + buildCommand: | + curl -sSL https://install.python-poetry.org | python3 - + poetry config virtualenvs.create false + poetry install --no-root + startCommand: poetry run python inference_server.py + envVars: + - key: PYTHONUNBUFFERED + value: "TRUE" + plan: free diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..2b6f18e34a7dd2037fb37cd072ef3b73ec06943c --- /dev/null +++ b/requirements.txt @@ -0,0 +1,60 @@ +addict==2.4.0 +appdirs==1.4.4 +certifi==2024.2.2 +charset-normalizer==3.3.2 +click==8.1.7 +clip @ git+https://github.com/openai/CLIP.git@a1d071733d7111c9c014f024669f959182114e33 +contourpy==1.2.0 +cycler==0.12.1 +dill==0.3.6 +dlib==19.24.1 +docker-pycreds==0.4.0 +face_alignment==1.3.4 +filelock==3.13.1 +fonttools==4.50.0 +ftfy==6.2.0 +gdown==3.12.2 +gitdb==4.0.11 +GitPython==3.1.42 +idna==3.6 +imageio==2.34.0 +joblib==1.2.0 +kiwisolver==1.4.5 +lazy_loader==0.3 +llvmlite==0.42.0 +matplotlib==3.7.1 +networkx==3.2.1 +numba==0.59.1 +numpy==1.23.5 +opencv-python==4.7.0.72 +packaging==24.0 +pandas==2.0.1 +pathtools==0.1.2 +Pillow==10.0.0 +protobuf==4.25.3 +psutil==5.9.8 +pyparsing==3.1.2 +PySocks==1.7.1 +python-dateutil==2.9.0.post0 +pytz==2024.1 +PyYAML==6.0.1 +regex==2023.12.25 +requests==2.31.0 +scikit-image==0.22.0 +scikit-learn==1.2.2 +scipy==1.10.1 +sentry-sdk==1.43.0 +setproctitle==1.3.3 +six==1.16.0 +smmap==5.0.1 +threadpoolctl==3.4.0 +tifffile==2024.2.12 +torch==1.13.1 +torchmetrics==0.11.4 +torchvision==0.14.1 +tqdm==4.66.1 +typing_extensions==4.10.0 +tzdata==2024.1 +urllib3==2.2.1 +wandb==0.15.10 +wcwidth==0.2.13 diff --git a/scripts/align_face.py b/scripts/align_face.py new file mode 100644 index 0000000000000000000000000000000000000000..76bf60caacb3b8f85262433d4694a6a28ce27323 --- /dev/null +++ b/scripts/align_face.py @@ -0,0 +1,43 @@ +import argparse +import os +import sys +from pathlib import Path + +from PIL import Image + +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) +from utils.seed import set_seed +from utils.image_utils import list_image_files +from utils.shape_predictor import align_face + + +def alignment_from_path(input_path, output_path, replace_cropped=False): + output_path.mkdir(parents=True, exist_ok=True) + image_files = list_image_files(input_path) + + for img_path in image_files: + if (output_path / img_path).is_file() and not replace_cropped: + continue + + image = Image.open(input_path / img_path) + + # returns the largest face + crop_image = align_face(image, return_tensors=False) + + if len(crop_image) == 1: + crop_image[0].save(output_path / img_path) + elif len(crop_image) > 1: + raise NotImplemented + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description='Align faces') + + parser.add_argument('-unprocessed_dir', type=Path, default='unprocessed', help='directory with unprocessed images') + parser.add_argument('-output_dir', type=Path, default='input', help='output directory') + parser.add_argument('-replace_cropped', action='store_true') + + args = parser.parse_args() + + set_seed(3407) + alignment_from_path(args.unprocessed_dir, args.output_dir, args.replace_cropped) diff --git a/scripts/blending_gen.py b/scripts/blending_gen.py new file mode 100644 index 0000000000000000000000000000000000000000..62c11476e759e8fe5ea5eeb63a31dddce077c3eb --- /dev/null +++ b/scripts/blending_gen.py @@ -0,0 +1,80 @@ +import argparse +import os +import sys +from pathlib import Path + +import numpy as np +from tqdm.auto import tqdm + +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) +from utils.train import seed_everything +from utils.image_utils import list_image_files +from hair_swap import get_parser, HairFast +from utils.save_utils import save_latents + + +def identity_func(align_shape, align_color, name_to_embed, **kwargs): + return align_shape, align_color, name_to_embed + + +def align_instead_shape(hair_fast): + def shape_module(func): + def wrapper(*args, **kwargs): + if kwargs.get('align_flag', False): + return hair_fast.align.align_images(*args, **kwargs) + else: + return func(*args, **kwargs) + + return wrapper + + def align_module(func): + def wrapper(*args, **kwargs): + if 'align_flag' in kwargs: + kwargs = kwargs.copy() + kwargs.pop('align_flag') + return func(*args, **kwargs) + + return wrapper + + hair_fast.align.shape_module = shape_module(hair_fast.align.shape_module) + hair_fast.align.align_images = align_module(hair_fast.align.align_images) + + +def main(args): + seed_everything(args.seed) + + # init HairFast + model_parser = get_parser() + model_args = model_parser.parse_args([]) + hair_fast = HairFast(model_args) + hair_fast.blend.blend_images = identity_func + align_instead_shape(hair_fast) + + # generate dataset + images = list_image_files(args.FFHQ) + face, shape, color = np.array_split(np.random.choice(images, size=3 * args.size), 3) + + os.makedirs(args.output, exist_ok=True) + with open(args.output / 'dataset.exps', 'w') as f_exps: + for imgs in tqdm(zip(face, shape, color)): + im1, im2, im3 = map(lambda im: im.split('.')[0], imgs) + print(im1, im2, im3, file=f_exps, flush=True) + + pt1, pt2, pt3 = map(lambda im: args.FFHQ / im, imgs) + align_shape, align_color, name_to_embed = hair_fast(pt1, pt2, pt3, align_flag=True) + save_latents(args.output, 'FS', f'{im1}.npz', latent_in=name_to_embed['face']['S']) + save_latents(args.output, 'FS', f'{im2}.npz', latent_in=name_to_embed['shape']['S']) + save_latents(args.output, 'FS', f'{im3}.npz', latent_in=name_to_embed['color']['S']) + save_latents(args.output, 'Align', f'{im1}_{im2}.npz', latent_F=align_shape['latent_F_align']) + save_latents(args.output, 'Align', f'{im1}_{im3}.npz', latent_F=align_color['latent_F_align']) + + +if __name__ == '__main__': + parser = argparse.ArgumentParser(description='Blending dataset') + parser.add_argument('--FFHQ', type=Path) + parser.add_argument('--seed', type=int, default=3407) + parser.add_argument('--size', type=int, default=3_000) + parser.add_argument('--output', type=Path, default='input/blending_dataset') + args = parser.parse_args() + + main(args) diff --git a/scripts/blending_train.py b/scripts/blending_train.py new file mode 100644 index 0000000000000000000000000000000000000000..a18751a4adc4e969f673f1e0eb57fd87e3b2ae34 --- /dev/null +++ b/scripts/blending_train.py @@ -0,0 +1,294 @@ +import argparse +import os +import sys +from argparse import Namespace +from pathlib import Path +from tempfile import TemporaryDirectory + +import numpy as np +import torch +import torch.nn.functional as F +import wandb +from PIL import Image +from joblib import Parallel, delayed +from sklearn.model_selection import train_test_split +from torch.utils.data import Dataset, DataLoader +from torchvision import transforms as T +from tqdm.auto import tqdm + +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) +from models.Encoders import ClipBlendingModel as BlendingModel +from models.Net import Net +from models.face_parsing.model import BiSeNet, seg_mean, seg_std +from utils.bicubic import BicubicDownSample +from utils.image_utils import DilateErosion +from utils.train import toggle_grad, WandbLogger, image_grid, seed_everything, get_fid_calc + + +class Trainer: + def __init__(self, + model=None, + optimizer=None, + scheduler=None, + train_dataloader=None, + test_dataloader=None, + logger=None, + ): + self.model = model + self.optimizer = optimizer + self.scheduler = scheduler + self.train_dataloader = train_dataloader + self.test_dataloader = test_dataloader + self.logger = logger + self.device = 'cuda' if torch.cuda.is_available() else 'cpu' + self.dilate_erosion = DilateErosion(device=self.device) + + if self.model is not None: + self.fid_calc = get_fid_calc('input/fid.pkl', args.fid_dataset) + + self.net = Net(Namespace(size=1024, ckpt='pretrained_models/StyleGAN/ffhq.pt', channel_multiplier=2, latent=512, + n_mlp=8, device=self.device)) + self.seg = BiSeNet(n_classes=16) + self.seg.to(self.device) + self.seg.eval() + + self.seg.load_state_dict(torch.load('pretrained_models/BiSeNet/seg.pth')) + toggle_grad(self.seg, False) + toggle_grad(self.net.generator, False) + + self.downsample_512 = BicubicDownSample(factor=2) + self.downsample_256 = BicubicDownSample(factor=4) + self.downsample_128 = BicubicDownSample(factor=8) + + self.best_loss = float('+inf') + self.cur_iter = 0 + + @torch.no_grad() + def generate_mask(self, I): + IM = (self.downsample_512((I + 1) / 2) - seg_mean) / seg_std + down_seg, _, _ = self.seg(IM) + current_mask = torch.argmax(down_seg, dim=1).long().float() + HM_X = torch.where(current_mask == 10, torch.ones_like(current_mask), torch.zeros_like(current_mask)) + HM_X = F.interpolate(HM_X.unsqueeze(1), size=(256, 256), mode='nearest') + + HM_XD, HM_XE = self.dilate_erosion.mask(HM_X) + return HM_XD, HM_XE + + def save_model(self, name, save_online=True): + with TemporaryDirectory() as tmp_dir: + model_state_dict = self.model.state_dict() + + # delete pretrained clip + for key in list(model_state_dict.keys()): + if key.startswith("clip_model."): + del model_state_dict[key] + + torch.save({'model_state_dict': model_state_dict}, f'{tmp_dir}/{name}.pth') + self.logger.save(f'{tmp_dir}/{name}.pth', save_online) + + def calc_loss(self, I_gen, I_face, I_color, mask_face, mask_hair, gen_hair): + gen_embed = self.model.get_image_embed(I_gen * mask_face) + gt_embed = self.model.get_image_embed(I_face * mask_face) + face_loss = (1 - F.cosine_similarity(gen_embed, gt_embed)).mean() + + gen_embed = self.model.get_image_embed(I_gen * mask_hair) + gt_embed = self.model.get_image_embed(I_color * mask_hair) + hair_loss = (1 - F.cosine_similarity(gen_embed, gt_embed)).mean() + + losses = {'face loss': face_loss, 'hair loss': hair_loss, 'loss': face_loss + hair_loss} + return losses['loss'], losses + + def train_one_epoch(self): + self.model.to(self.device).train() + for batch in tqdm(self.train_dataloader): + color_s, align_s, align_f, color_i, face_i, target_mask, HM_3E, HM_XE = map(lambda x: x.to(self.device), + batch) + bsz = color_s.size(0) + + blend_s = self.model(align_s[:, 6:], color_s[:, 6:], face_i * target_mask, color_i * HM_3E) + latent_in = torch.cat((torch.zeros(bsz, 6, 512, device=self.device), blend_s), axis=1) + I_G, _ = self.net.generator([latent_in], input_is_latent=True, return_latents=False, start_layer=4, + end_layer=8, layer_in=align_f) + + loss, info = self.calc_loss(self.downsample_256(I_G), face_i, color_i, target_mask, HM_3E, HM_XE) + + self.optimizer.zero_grad() + loss.backward() + + total_norm = torch.nn.utils.clip_grad_norm_(self.model.parameters(), 5) + self.optimizer.step() + + self.logger.next_step() + for key, val in info.items(): + self.logger.log(key, val.item()) + self.logger.log('grad', total_norm.item()) + self.cur_iter += 1 + + @torch.no_grad() + def validate(self): + self.model.to(self.device).eval() + + sum_losses = lambda x, y: {key: val + x.get(key, 0) for key, val in y.items()} + files = [] + losses = {} + to_299 = T.Resize((299, 299)) + images_to_fid = [] + + for batch in tqdm(self.test_dataloader): + color_s, align_s, align_f, color_i, face_i, target_mask, HM_3E, HM_XE = map(lambda x: x.to(self.device), + batch) + bsz = color_s.size(0) + + blend_s = self.model(align_s[:, 6:], color_s[:, 6:], face_i * target_mask, color_i * HM_3E) + latent_in = torch.cat((torch.zeros(bsz, 6, 512, device=self.device), blend_s), axis=1) + I_G, _ = self.net.generator([latent_in], input_is_latent=True, return_latents=False, start_layer=4, + end_layer=8, layer_in=align_f) + + _, info = self.calc_loss(self.downsample_256(I_G), face_i, color_i, target_mask, HM_3E, HM_XE) + losses = sum_losses(losses, info) + for k in range(bsz): + files.append([color_i[k].cpu(), face_i[k].cpu(), self.downsample_256(I_G)[k].cpu()]) + + images_to_fid.append(to_299((I_G + 1) / 2).clip(0, 1)) + + losses['FID CLIP'] = self.fid_calc(torch.cat(images_to_fid)) + for key, val in losses.items(): + if key != 'FID CLIP': + val = val.item() / len(self.test_dataloader) + self.logger.log(f'val {key}', val) + + np.random.seed(1927) + idxs = np.random.choice(len(files), size=100, replace=False) + images_to_log = [ + image_grid([T.functional.to_pil_image(((img + 1) / 2).clamp(0, 1)) for img in files[idx]], 1, 3) for idx in + idxs] + self.logger.log('val images', [wandb.Image(image) for image in images_to_log]) + + return losses['loss'] + + def train_loop(self, epochs): + self.validate() + for epoch in range(epochs): + self.train_one_epoch() + loss = self.validate() + + self.save_model('last', save_online=False) + if loss <= self.best_loss: + self.best_loss = loss + self.save_model(f'best', save_online=False) + + +def prepare_item(exp, path): + im1, im2, im3 = exp + + try: + color_path = os.path.join(path, 'FS', f'{im3}.npz') + Color_S = torch.from_numpy(np.load(color_path)['latent_in']).squeeze(0) + + face_path = os.path.join(path, 'FS', f'{im1}.npz') + Align_S = torch.from_numpy(np.load(face_path)['latent_in']).squeeze(0) + + Color_I = T.functional.normalize(T.functional.to_tensor( + Image.open(os.path.join(args.FFHQ, f'{im3}.png')) + ), [0.5], [0.5]) + Face_I = T.functional.normalize(T.functional.to_tensor( + Image.open(os.path.join(args.FFHQ, f'{im1}.png')) + ), [0.5], [0.5]) + + align_path = os.path.join(path, 'Align') + data = np.load( + os.path.join(align_path, f'{im1}_{im3}.npz') + ) + Align_F = torch.from_numpy(data['latent_F']).squeeze(0) + + return (Color_S, Align_S, Align_F, Color_I, Face_I) + except Exception as e: + print(e, file=sys.stderr) + return None + + +class Blending_dataset(Dataset): + def __init__(self, exps, path, net_trainer): + super().__init__() + downsample_256 = BicubicDownSample(factor=4) + data = Parallel(n_jobs=-1)( + delayed(prepare_item)(exp, path) for (p1, p2, p3) in tqdm(exps) for exp in [(p1, p2, p3), (p1, p3, p2)]) + data = [elem for elem in data if elem is not None] + print(f'Load: {len(data)}/{2 * len(exps)}', file=sys.stderr) + + tmp_dataloader = DataLoader(data, batch_size=24, pin_memory=False, shuffle=False) + + self.items = [] + with torch.no_grad(): + for (Color_S, Align_S, Align_F, Color_I, Face_I) in tqdm(tmp_dataloader): + HM_3D, HM_3E = net_trainer.generate_mask(Color_I.to('cuda')) + + HM_1D, _ = net_trainer.generate_mask(Face_I.to('cuda')) + I_X, _ = net_trainer.net.generator([Align_S.to('cuda')], input_is_latent=True, return_latents=False, + start_layer=4, + end_layer=8, layer_in=Align_F.to('cuda')) + HM_XD, HM_XE = net_trainer.generate_mask(I_X) + + target_mask = ((1 - HM_1D) * (1 - HM_3D) * (1 - HM_XD)).cpu() + HM_3E = HM_3E.cpu() + HM_XE = HM_XE + self.items.extend( + [item for item in zip(*list(map(lambda x: [item.squeeze(0) for item in torch.split(x, 1)], + (Color_S, + Align_S, + Align_F, + downsample_256(Color_I.to('cuda')).cpu(), + downsample_256(Face_I.to('cuda')).cpu(), + target_mask, HM_3E, HM_XE))) + ) if item[-2].any() and item[-1].any()] + ) + + print(f'dataset: {len(self.items)}/{len(data)}', file=sys.stderr) + + def __len__(self): + return len(self.items) + + def __getitem__(self, idx): + return self.items[idx] + + +def main(args): + seed_everything() + + exps = [] + with open(os.path.join(args.dataset, 'dataset.exps'), 'r') as file: + for exp in file.readlines(): + exps.append(list(map(lambda x: x.replace('.png', ''), exp.split()))) + + X_train, X_test = train_test_split(exps, test_size=512, random_state=42) + + net_trainer = Trainer() + train_dataset = Blending_dataset(X_train, args.dataset, net_trainer) + test_dataset = Blending_dataset(X_test, args.dataset, net_trainer) + + train_dataloader = DataLoader(train_dataset, batch_size=32, shuffle=True, drop_last=True) + test_dataloader = DataLoader(test_dataset, batch_size=32, shuffle=False) + + logger = WandbLogger(name=args.name_run, project='Barbershop-Blending') + logger.start_logging() + logger.save(__file__) + + model = BlendingModel() + optimizer = torch.optim.Adam(model.parameters(), lr=1e-4, weight_decay=0.000001) + + trainer = Trainer(model, optimizer, None, train_dataloader, test_dataloader, logger) + + trainer.train_loop(1000) + + logger.wandb.finish() + + +if __name__ == '__main__': + parser = argparse.ArgumentParser(description='Blending trainer') + parser.add_argument('--name_run', type=str, default='test') + parser.add_argument('--dataset', type=Path, default='input/blending_dataset') + parser.add_argument('--FFHQ', type=Path) + parser.add_argument('--fid_dataset', type=str, default='input') + args = parser.parse_args() + + main(args) diff --git a/scripts/fid_metric.py b/scripts/fid_metric.py new file mode 100644 index 0000000000000000000000000000000000000000..5d8af41edfb5ae8365d729e7e76363206f6fe879 --- /dev/null +++ b/scripts/fid_metric.py @@ -0,0 +1,82 @@ +import argparse +import os +import sys +from pathlib import Path + +import pandas as pd +import torch +from torch.utils.data import DataLoader, TensorDataset +from torchmetrics.image.fid import FrechetInceptionDistance +from tqdm.auto import tqdm + +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) +from models.Encoders import ClipModel +from utils.seed import set_seed +from utils.train import parallel_load_images +from utils.image_utils import list_image_files + + +def name_path(pair): + name, path = pair.split(',') + return name, Path(path) + + +@torch.inference_mode() +def compute_fid_datasets(datasets, target='celeba', device=torch.device('cuda'), CLIP=False, seed=3407): + set_seed(seed) + result = {} + + if CLIP: + fid = FrechetInceptionDistance(feature=ClipModel(), reset_real_features=False, normalize=False) + else: + fid = FrechetInceptionDistance(reset_real_features=False, normalize=False) + fid.to(device).eval() + + real_dataloader = DataLoader(TensorDataset(datasets[target]), batch_size=128) + for batch in tqdm(real_dataloader): + batch = batch[0].to(device) + fid.update(batch, real=True) + + for key, tensor in datasets.items(): + if key == target: + continue + fid.reset() + + fake_dataloader = DataLoader(TensorDataset(tensor), batch_size=128) + for batch in tqdm(fake_dataloader): + batch = batch[0].to(device) + fid.update(batch, real=False) + result[key] = fid.compute().item() + return result + + +def main(args): + datasets = {} + + source = args.source_dataset.name + datasets[source] = parallel_load_images(args.source_dataset, list_image_files(args.source_dataset)) + + for method, path_dataset in args.methods_dataset: + datasets[method] = parallel_load_images(path_dataset, list_image_files(path_dataset)) + + FIDs = compute_fid_datasets(datasets, target=source, CLIP=False) + df_fid = pd.DataFrame.from_dict(FIDs, orient='index', columns=['FID']) + + FIDs_CLIP = compute_fid_datasets(datasets, target=source, CLIP=True) + df_clip = pd.DataFrame.from_dict(FIDs_CLIP, orient='index', columns=['FID_CLIP']) + + df_result = pd.concat([df_fid, df_clip], axis=1).round(2) + print(df_result) + + os.makedirs(args.output.parent, exist_ok=True) + df_result.to_csv(args.output, index=True) + + +if __name__ == '__main__': + parser = argparse.ArgumentParser(description='Compute metrics') + parser.add_argument('--source_dataset', type=Path, help='Dataset with real faces') + parser.add_argument('--methods_dataset', type=name_path, nargs='+', help='Datasets after applying the method') + parser.add_argument('--output', type=Path, default='logs/metric.csv', help='Folder for saving logs') + args = parser.parse_args() + + main(args) diff --git a/scripts/pp_gen.py b/scripts/pp_gen.py new file mode 100644 index 0000000000000000000000000000000000000000..2001d586c60a53d644efff18b5c6579e82f58525 --- /dev/null +++ b/scripts/pp_gen.py @@ -0,0 +1,129 @@ +import argparse +import os +import sys +import tempfile +from pathlib import Path + +import numpy as np +import torch +import torch.nn as nn +from PIL import Image +from joblib import Parallel, delayed +from torch.utils.data import DataLoader +from torchvision import transforms as T +from torchvision.utils import save_image +from tqdm.auto import tqdm + +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) +from hair_swap import get_parser, HairFast +from scripts.pp_train import Trainer +from utils.train import seed_everything +from utils.image_utils import list_image_files + + +class ImageException(Exception): + def __init__(self, image, message="Return image before PP"): + self.image = image + self.message = message + super().__init__(self.message) + + +def hairfast_wo_pp(hair_fast): + class RaiseDownsample(nn.Module): + def __init__(self): + super().__init__() + + def forward(self, image): + image = ((image[0] + 1) / 2).clip(0, 1) + raise ImageException(image) + + def blend_images(func): + def wrapper(*args, **kwargs): + try: + func(*args, **kwargs) + except ImageException as e: + return e.image + + return wrapper + + hair_fast.blend.downsample_256 = RaiseDownsample() + hair_fast.blend.blend_images = blend_images(hair_fast.blend.blend_images) + + +def load_image(path): + return T.functional.to_tensor(Image.open(path)) + + +@torch.no_grad() +def load_dataset_images(imgs, dataset_path): + net_trainer = Trainer() + source_images = Parallel(n_jobs=-1)(delayed(load_image)(os.path.join(args.FFHQ, img[0])) for img in tqdm(imgs)) + target_images = Parallel(n_jobs=-1)(delayed(load_image)(os.path.join(dataset_path, img[1])) for img in tqdm(imgs)) + tensors_dataloader = DataLoader(list(zip(source_images, target_images)), batch_size=64, pin_memory=False, + shuffle=False, drop_last=False) + source_files = [os.path.join(args.FFHQ, img[0]) for img in imgs] + + data = [] + total = 0 + for batch in tqdm(tensors_dataloader): + source, target = [elem.to('cuda') for elem in batch] + + HS_D, _ = net_trainer.generate_mask(source) + HT_D, HT_E = net_trainer.generate_mask(target) + target_mask = (1 - HS_D) * (1 - HT_D) + + data.extend(list(zip( + [source_files[total + i] for i in range(len(source))], + net_trainer.downsample_256(target).clip(0, 1).cpu(), + target_mask.cpu(), + HT_E.cpu(), + ))) + total += len(source) + + return data + + +def main(args): + seed_everything(args.seed) + + # init HairFast + model_parser = get_parser() + model_args = model_parser.parse_args([]) + hair_fast = HairFast(model_args) + hairfast_wo_pp(hair_fast) + + # generate dataset + os.makedirs(args.output, exist_ok=True) + images = list_image_files(args.FFHQ) + face, shape, color = np.array_split(np.random.choice(images, size=3 * args.size), 3) + + exps = [] + for exp in zip(face, shape, color): + imgs = map(lambda im: im.split('.')[0], exp) + exps.append([exp[0], f"{'_'.join(imgs)}.png", exp]) + + batch = 5_000 + left, right, idx = 0, min(len(exps), batch), 1 + while left < len(exps): + with tempfile.TemporaryDirectory() as temp_dir: + for exp in tqdm(exps[left:right]): + im1, im2, im3 = exp[-1] + image = hair_fast(args.FFHQ / im1, args.FFHQ / im2, args.FFHQ / im3) + save_image(image, os.path.join(temp_dir, exp[1])) + + batch_data = load_dataset_images(exps[left:right], temp_dir) + torch.save(batch_data, os.path.join(args.output, f'pp_part_{idx}.dataset')) + left = right + right = min(len(exps), right + batch) + idx += 1 + + +if __name__ == '__main__': + parser = argparse.ArgumentParser(description='Blending dataset') + parser.add_argument('--FFHQ', type=Path) + parser.add_argument('--seed', type=int, default=3407) + parser.add_argument('--size', type=int, default=10_000) + parser.add_argument('--output', type=Path, default='input/pp_dataset') + args = parser.parse_args() + + main(args) diff --git a/scripts/pp_train.py b/scripts/pp_train.py new file mode 100644 index 0000000000000000000000000000000000000000..6cb86fec9e7bbb10235d8938c697ead19b6cc57e --- /dev/null +++ b/scripts/pp_train.py @@ -0,0 +1,386 @@ +import argparse +import os +import random +import sys +from argparse import Namespace +from pathlib import Path +from tempfile import TemporaryDirectory + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +import wandb +from PIL import Image +from sklearn.model_selection import train_test_split +from torch.utils.data import Dataset, DataLoader +from torchvision import transforms as T +from tqdm.auto import tqdm + +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) +from losses.pp_losses import LossBuilder, LossBuilderMulti +from models.Encoders import ModulationModule, FeatureiResnet, FeatureEncoderMult +from models.Net import Net +from models.face_parsing.model import BiSeNet, seg_mean, seg_std +from models.stylegan2 import dnnlib +from models.stylegan2.model import PixelNorm +from utils.bicubic import BicubicDownSample +from utils.image_utils import DilateErosion +from utils.train import image_grid, WandbLogger, toggle_grad, _LegacyUnpickler, seed_everything, get_fid_calc + + +class Trainer: + def __init__(self, + model=None, + args=None, + optimizer=None, + scheduler=None, + train_dataloader=None, + test_dataloader=None, + logger=None + ): + self.model = model + self.args = args + self.optimizer = optimizer + self.scheduler = scheduler + self.train_dataloader = train_dataloader + self.test_dataloader = test_dataloader + self.logger = logger + self.device = 'cuda' if torch.cuda.is_available() else 'cpu' + self.dilate_erosion = DilateErosion(device=self.device) + self.normalize = T.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]) + + if self.model is not None: + self.fid_calc = get_fid_calc('input/fid.pkl', args.fid_dataset) + + self.net = Net(Namespace(size=1024, ckpt='pretrained_models/StyleGAN/ffhq.pt', channel_multiplier=2, latent=512, + n_mlp=8, device=self.device)) + + with dnnlib.util.open_url("pretrained_models/StyleGAN/ffhq.pkl") as f: + data = _LegacyUnpickler(f).load() + self.discriminator = data['D'].cuda().eval() + self.disc_optim = torch.optim.Adam(self.discriminator.parameters(), lr=3e-4, betas=(0.9, 0.999), amsgrad=False, + weight_decay=0) + + self.seg = BiSeNet(n_classes=16) + self.seg.to(self.device) + self.seg.load_state_dict(torch.load('pretrained_models/BiSeNet/seg.pth')) + self.seg.eval() + + toggle_grad(self.discriminator, False) + toggle_grad(self.net.generator, False) + toggle_grad(self.seg, False) + + self.downsample_512 = BicubicDownSample(factor=2) + self.downsample_256 = BicubicDownSample(factor=4) + self.downsample_128 = BicubicDownSample(factor=8) + + self.best_loss = float('+inf') + if self.args is not None: + if self.args.pretrain: + self.LossBuilder = LossBuilder( + {'lpips_scale': 0.8, 'id': 0.1, 'landmark': 0, 'feat_rec': 0.01, 'adv': self.args.adv_coef}) + else: + self.LossBuilder = LossBuilderMulti( + {'lpips_scale': 0.8, 'id': 0.1, 'landmark': 0.1, 'feat_rec': 0.01, 'adv': self.args.adv_coef, + 'inpaint': self.args.inpaint}) + self.cur_iter = 1 + + @torch.no_grad() + def generate_mask(self, I): + IM = (self.downsample_512(I) - seg_mean) / seg_std + down_seg, _, _ = self.seg(IM) + current_mask = torch.argmax(down_seg, dim=1).long().float() + HM_X = torch.where(current_mask == 10, torch.ones_like(current_mask), torch.zeros_like(current_mask)) + HM_X = F.interpolate(HM_X.unsqueeze(1), size=(256, 256), mode='nearest') + + HM_XD, HM_XE = self.dilate_erosion.mask(HM_X) + return HM_XD, HM_XE + + def save_model(self, name, save_online=True): + with TemporaryDirectory() as tmp_dir: + model_state_dict = self.model.state_dict() + + # delete pretrained clip + for key in list(model_state_dict.keys()): + if key.startswith("clip_model."): + del model_state_dict[key] + + torch.save( + {'model_state_dict': model_state_dict, 'D': self.discriminator.state_dict(), 'cur_iter': self.cur_iter}, + f'{tmp_dir}/{name}.pth') + self.logger.save(f'{tmp_dir}/{name}.pth', save_online) + + def load_model(self, checkpoint_path): + checkpoint = torch.load(checkpoint_path) + if 'D' in checkpoint: + self.discriminator.load_state_dict(checkpoint['D'], strict=False) + if 'model_state_dict' in checkpoint: + self.model.load_state_dict(checkpoint['model_state_dict'], strict=False) + + def train_one_epoch(self): + self.model.to(self.device).train() + for batch in tqdm(self.train_dataloader): + source, target, target_mask, HT_E = map(lambda x: x.to(self.device), batch) + source, source_1024 = self.downsample_256(source).clip(0, 1), self.normalize(source) + + latent_s, latent_f = self.model(self.normalize(source), self.normalize(target), target_mask, HT_E) + + gen_im_W, _ = self.net.generator([latent_s], input_is_latent=True, return_latents=False) + F_w, _ = self.net.generator([latent_s], input_is_latent=True, return_latents=False, + start_layer=0, end_layer=4) + + if self.args.pretrain: + alpha = min(1, self.cur_iter / self.args.iter_before) + latent_f_gen = alpha * latent_f + (1 - alpha) * F_w + else: + latent_f_gen = latent_f + + gen_im_F, _ = self.net.generator([latent_s], input_is_latent=True, return_latents=False, + start_layer=5, end_layer=8, layer_in=latent_f_gen) + + losses = self.LossBuilder(source, target, target_mask, HT_E, gen_im_W, F_w, gen_im_F, latent_f) + + if self.args.use_adv and self.cur_iter >= self.args.iter_before: + losses.update(self.LossBuilder.CalcAdvLoss(self.discriminator, gen_im_F)) + + losses['loss'] = sum(losses.values()) + + self.optimizer.zero_grad() + losses['loss'].backward() + + total_norm = torch.nn.utils.clip_grad_norm_(self.model.parameters(), 0.5) + self.optimizer.step() + + if self.args.use_adv and self.cur_iter >= self.args.iter_before: + if self.cur_iter == self.args.iter_before: + print('Start scripts discr') + + toggle_grad(self.discriminator, True) + self.discriminator.train() + + disc_loss = self.LossBuilder.CalcDisLoss(self.discriminator, source_1024, gen_im_F.detach()) + if self.cur_iter % self.args.d_reg_every: + disc_loss.update(self.LossBuilder.CalcR1Loss(self.discriminator, source_1024)) + + total_loss = sum(disc_loss.values()) + + self.disc_optim.zero_grad() + + total_loss.backward() + total_norm_d = torch.nn.utils.clip_grad_norm_(self.discriminator.parameters(), 0.5) + disc_loss['grad disc'] = total_norm_d + + self.disc_optim.step() + + toggle_grad(self.discriminator, False) + self.discriminator.eval() + losses.update(disc_loss) + + losses['scripts grad'] = total_norm + self.logger.next_step() + self.logger.log_scalars({f'scripts {key}': val for key, val in losses.items()}) + self.cur_iter += 1 + + @torch.no_grad() + def validate(self): + self.model.to(self.device).eval() + + sum_losses = lambda x, y: {key: y.get(key, 0) + x.get(key, 0) for key in set(x.keys()) | set(y.keys())} + files = [] + val_losses = {} + to_299 = T.Resize((299, 299)) + images_to_fid = [] + + for batch in tqdm(self.test_dataloader): + source, target, target_mask, HT_E = map(lambda x: x.to(self.device), batch) + source = self.downsample_256(source).clip(0, 1) + bsz = source.size(0) + + latent_s, latent_f = self.model(self.normalize(source), self.normalize(target), target_mask, HT_E) + + gen_im_W, _ = self.net.generator([latent_s], input_is_latent=True, return_latents=False) + F_w, _ = self.net.generator([latent_s], input_is_latent=True, return_latents=False, + start_layer=0, end_layer=4) + gen_im_F, _ = self.net.generator([latent_s], input_is_latent=True, return_latents=False, + start_layer=5, end_layer=8, layer_in=latent_f) + + losses = self.LossBuilder(source, target, target_mask, HT_E, gen_im_W, F_w, gen_im_F, latent_f) + losses['loss'] = sum(losses.values()) + + gen_w_256 = self.downsample_256((gen_im_W + 1) / 2).clip(0, 1) + gen_f_256 = self.downsample_256((gen_im_F + 1) / 2).clip(0, 1) + + images_to_fid.append(to_299((gen_im_F + 1) / 2).clip(0, 1)) + + val_losses = sum_losses(val_losses, losses) + for k in range(bsz): + files.append([source[k].cpu(), target[k].cpu(), gen_w_256[k].cpu(), gen_f_256[k].cpu()]) + + val_losses['FID CLIP'] = self.fid_calc(torch.cat(images_to_fid)) + for key, val in val_losses.items(): + if key != 'FID CLIP': + val = val.item() / len(self.test_dataloader) + self.logger.log_scalars({f'val {key}': val}) + + np.random.seed(1927) + idxs = np.random.choice(len(files), size=min(len(files), 100), replace=False) + images_to_log = [image_grid(list(map(T.functional.to_pil_image, files[idx])), 1, len(files[idx])) for idx in + idxs] + self.logger.log_scalars({'val images': [wandb.Image(image) for image in images_to_log]}) + + return val_losses['loss'] + + def train_loop(self, epochs): + self.validate() + for epoch in range(epochs): + self.train_one_epoch() + loss = self.validate() + + self.save_model('last', save_online=False) + if loss <= self.best_loss: + self.best_loss = loss + self.save_model(f'best_{epoch}', save_online=False) + + +class PP_dataset(Dataset): + def __init__(self, source, target, target_mask, HT_E, is_test=False): + super().__init__() + self.source = source + self.target = target + self.target_mask = target_mask + self.HT_E = HT_E + self.is_test = is_test + + def __len__(self): + return len(self.source) + + def load_image(self, path): + return T.functional.to_tensor(Image.open(path)) + + def __transform__(self, img1, img2, mask1, mask2): + if self.is_test: + return img1, img2, mask1, mask2 + + if random.random() > 0.5: + img1 = T.functional.hflip(img1) + img2 = T.functional.hflip(img2) + mask1 = T.functional.hflip(mask1) + mask2 = T.functional.hflip(mask2) + + return img1, img2, mask1, mask2 + + def __getitem__(self, idx): + return self.__transform__(self.load_image(self.source[idx]), self.target[idx], self.target_mask[idx], + self.HT_E[idx]) + + +class PostProcessModel(nn.Module): + def __init__(self, args): + super().__init__() + self.args = args + + self.encoder_face = FeatureEncoderMult(fs_layers=[9], opts=argparse.Namespace( + **{'arcface_model_path': "pretrained_models/ArcFace/backbone_ir50.pth"})) + if not self.args.finetune: + toggle_grad(self.encoder_face, False) + + self.latent_avg = torch.load('pretrained_models/PostProcess/latent_avg.pt', map_location=torch.device('cuda')) + self.to_feature = FeatureiResnet([[1024, 2], [768, 2], [512, 2]]) + + if self.args.use_mod: + self.to_latent_1 = nn.ModuleList([ModulationModule(18, i == 4) for i in range(5)]) + self.to_latent_2 = nn.ModuleList([ModulationModule(18, i == 4) for i in range(5)]) + self.pixelnorm = PixelNorm() + else: + self.to_latent = nn.Sequential(nn.Linear(1024, 1024), nn.LayerNorm([1024]), nn.LeakyReLU(), + nn.Linear(1024, 512)) + + def forward(self, source, target, target_mask=None, *args, **kwargs): + s_face, [f_face] = self.encoder_face(source) + if self.args.pretrain: + return self.latent_avg + s_face, f_face + + s_hair, [f_hair] = self.encoder_face(target) + + if self.args.use_mod: + dt_latent_face = self.pixelnorm(s_face) + dt_latent_hair = self.pixelnorm(s_hair) + + for mod_module in self.to_latent_1: + dt_latent_face = mod_module(dt_latent_face, s_hair) + + for mod_module in self.to_latent_2: + dt_latent_hair = mod_module(dt_latent_hair, s_face) + finall_s = self.latent_avg + 0.1 * (dt_latent_face + dt_latent_hair) + else: + cat_s = torch.cat((s_face, s_hair), dim=-1) + finall_s = self.latent_avg + self.to_latent(cat_s) + + if self.args.use_full: + cat_f = torch.cat((f_face, f_hair), dim=1) + else: + t_mask = F.interpolate(target_mask, size=(64, 64), mode='nearest') + cat_f = torch.cat((f_face * t_mask, f_hair * (1 - t_mask)), dim=1) + + finall_f = self.to_feature(cat_f) + return finall_s, finall_f + + +def main(args): + seed_everything() + dataset = [] + + idx = 1 + while os.path.isfile(args.dataset / f'pp_part_{idx}.dataset'): + batch_data = torch.load(args.dataset / f'pp_part_{idx}.dataset') + dataset.extend(batch_data) + idx += 1 + + X_train, X_test = train_test_split(dataset, test_size=1024, random_state=42) + + train_dataset = PP_dataset(*list(zip(*X_train))) + test_dataset = PP_dataset(*list(zip(*X_test)), is_test=True) + + train_dataloader = DataLoader(train_dataset, batch_size=args.batch_size, num_workers=10, pin_memory=True, + shuffle=True, + drop_last=True) + test_dataloader = DataLoader(test_dataset, batch_size=args.batch_size, num_workers=10, pin_memory=True, + shuffle=False) + + logger = WandbLogger(name=args.name_run, project='HairFast-PostProcess') + logger.start_logging() + logger.save(__file__) + + model = PostProcessModel(args) + if args.pretrain: + optimizer = torch.optim.Adam(model.parameters(), lr=2e-4, weight_decay=0) + else: + optimizer = torch.optim.Adam(model.parameters(), lr=1e-4, weight_decay=0) + + trainer = Trainer(model, args, optimizer, None, train_dataloader, test_dataloader, logger) + if not args.pretrain: + trainer.load_model(args.checkpoint) + trainer.train_loop(1000) + + +if __name__ == '__main__': + parser = argparse.ArgumentParser(description='Post Process trainer') + parser.add_argument('--name_run', type=str, default='test') + parser.add_argument('--FFHQ', type=Path) + parser.add_argument('--dataset', type=Path, default='input/pp_dataset') + parser.add_argument('--fid_dataset', type=str, default='input') + parser.add_argument('--batch_size', type=int, default=16) + parser.add_argument('--iter_before', type=int, default=10_000) + parser.add_argument('--d_reg_every', type=int, default=16) + parser.add_argument('--inpaint', type=float, default=0.) + parser.add_argument('--use_adv', action='store_true') + parser.add_argument('--adv_coef', type=float, default=0.05) + parser.add_argument('--checkpoint', type=str) + parser.add_argument('--use_mod', action='store_true') + parser.add_argument('--use_full', action='store_true') + parser.add_argument('--pretrain', action='store_true') + parser.add_argument('--finetune', action='store_true') + args = parser.parse_args() + + main(args) diff --git a/scripts/rotate_gen.py b/scripts/rotate_gen.py new file mode 100644 index 0000000000000000000000000000000000000000..dcc7061f75e6d7b5c9a82bfa7d2be73d0b250a92 --- /dev/null +++ b/scripts/rotate_gen.py @@ -0,0 +1,64 @@ +import argparse +import os +import random +import sys +from pathlib import Path + +import torch +from PIL import Image +from joblib import Parallel, delayed +from torch.utils.data import DataLoader +from torchvision import transforms as T +from tqdm.auto import tqdm + +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) +from scripts.rotate_train import Trainer +from utils.train import seed_everything +from utils.image_utils import list_image_files + +toTensor = T.ToTensor() +net_trainer = Trainer() + + +def load_image(path): + return toTensor(Image.open(path)) + + +@torch.no_grad() +def load_dataset_images(imgs): + tensors_images = Parallel(n_jobs=-1)( + delayed(load_image)(os.path.join(args.FFHQ, str(img))) for img in tqdm(imgs)) + tensors_dataloader = DataLoader(tensors_images, batch_size=32, pin_memory=False, shuffle=False, drop_last=False) + + images, key_points, latents = [], [], [] + for batch in tqdm(tensors_dataloader): + batch = batch.to(net_trainer.device) + + images_256 = net_trainer.downsample_256(batch).clip(0, 1) + images.extend(images_256.cpu()) + latents.extend(net_trainer.generate_latents(images_256 * 2 - 1).cpu()) + key_points.extend(net_trainer.generate_key_points(batch).cpu()) + + return images, key_points, latents + + +def main(args): + seed_everything(args.seed) + + images = list_image_files(args.FFHQ) + random.shuffle(images) + + images, key_points, latents = load_dataset_images(images[:args.size]) + + torch.save({'images': images, 'key_points': key_points, 'latents': latents}, args.output) + + +if __name__ == '__main__': + parser = argparse.ArgumentParser(description='Rotate dataset') + parser.add_argument('--FFHQ', type=Path) + parser.add_argument('--seed', type=int, default=3407) + parser.add_argument('--size', type=int, default=10_000) + parser.add_argument('--output', type=Path, default='input/rotate_dataset.pkl') + args = parser.parse_args() + + main(args) diff --git a/scripts/rotate_train.py b/scripts/rotate_train.py new file mode 100644 index 0000000000000000000000000000000000000000..357e6bdd1aef61c2d58fe53acc2f0b364dd87ef9 --- /dev/null +++ b/scripts/rotate_train.py @@ -0,0 +1,351 @@ +import argparse +import os +import sys +from argparse import Namespace +from tempfile import TemporaryDirectory + +import numpy as np +import torch +import torch.nn.functional as F +import wandb +from sklearn.model_selection import train_test_split +from torch.utils.data import Dataset, DataLoader +from torchvision import transforms as T +from tqdm.auto import tqdm + +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) +from models.STAR.lib import utility +from models.Encoders import RotateModel +from models.Net import Net +from models.Net import iresnet100 +from models.encoder4editing.utils.model_utils import setup_model, get_latents +from utils.bicubic import BicubicDownSample +from utils.train import image_grid, WandbLogger, seed_everything, toggle_grad + + +class MovingAverageLoss: + def __init__(self, weights: dict, alpha=0.02): + self.alpha = alpha + self.weights = weights + self.vals = {} + + def reset(self): + self.vals = {} + + def update(self, cur_vals): + for key, val in cur_vals.items(): + self.vals[key] = self.alpha * val + (1 - self.alpha) * self.vals.get(key, val) + + def calc_loss(self, losses): + loss = 0. + for key, val in losses.items(): + loss += self.weights.get(key, 1) * val / self.vals.get(key, 1) + return loss + + +class Trainer: + def __init__(self, + model=None, + args=None, + optimizer=None, + scheduler=None, + train_dataloader=None, + test_dataloader=None, + logger=None + ): + self.model = model + self.args = args + self.optimizer = optimizer + self.scheduler = scheduler + self.train_dataloader = train_dataloader + self.test_dataloader = test_dataloader + self.logger = logger + self.device = 'cuda' if torch.cuda.is_available() else 'cpu' + + self.net = Net(Namespace(size=1024, ckpt='pretrained_models/StyleGAN/ffhq.pt', channel_multiplier=2, latent=512, + n_mlp=8, device=self.device)) + self.e4e = setup_model('pretrained_models/encoder4editing/e4e_ffhq_encode.pt', 'cuda')[0] + self.arc_face = iresnet100() + self.arc_face.load_state_dict(torch.load("pretrained_models/ArcFace/backbone_r100.pth")) + self.arc_face.eval().cuda() + self.toArcface = T.Compose([ + T.Resize((112, 112)), + T.Normalize(0.5, 0.5) + ]) + + # init landmarks + config = utility.get_config(utility.landmarks_arg) + self.kp_extractor = utility.get_net(config) + model_path = utility.landmarks_arg.pretrained_weight + checkpoint = torch.load(model_path) + self.kp_extractor.load_state_dict(checkpoint["net"]) + self.kp_extractor = self.kp_extractor.float().to('cuda') + self.kp_extractor.eval() + self.toLandmarks = T.Compose([ + T.Resize((256, 256)), + T.Normalize(0.5, 0.5) + ]) + + toggle_grad(self.arc_face, False) + toggle_grad(self.kp_extractor, False) + toggle_grad(self.net.generator, False) + toggle_grad(self.e4e.encoder, False) + + self.downsample_512 = BicubicDownSample(factor=2) + self.downsample_256 = BicubicDownSample(factor=4) + self.downsample_128 = BicubicDownSample(factor=8) + + self.MAL = MovingAverageLoss({'mse points to': 6, 'mse latents': 2}) + self.best_loss = float('+inf') + + def generate_key_points(self, batch): + _, _, landmarks = self.kp_extractor(self.toLandmarks(batch)) + final_marks_2D = (landmarks[:, :76] + 1) / 2 * torch.tensor([256 - 1, 256 - 1]).to('cuda').view(1, 1, 2) + return final_marks_2D + + @torch.no_grad() + def generate_latents(self, batch): + return get_latents(self.e4e, batch) + + def save_model(self, name, save_online=True): + with TemporaryDirectory() as tmp_dir: + model_state_dict = self.model.state_dict() + + # delete pretrained clip + for key in list(model_state_dict.keys()): + if key.startswith("clip_model."): + del model_state_dict[key] + + torch.save({'model_state_dict': model_state_dict}, f'{tmp_dir}/{name}.pth') + self.logger.save(f'{tmp_dir}/{name}.pth', save_online) + + def load_model(self, checkpoint_path): + self.model.load_state_dict(torch.load(checkpoint_path)['model_state_dict'], strict=False) + + def calc_loss(self, + I_to, + I_from, + key_points_to, + latents_from, + latents_to, + ret_images=False, + normalize=True + ): + # rotate + rotate_to = self.model(latents_from[:, :6], latents_to[:, :6]) + latent_in = torch.cat((rotate_to, latents_from[:, 6:]), axis=1) + I_G_to, _ = self.net.generator([latent_in], input_is_latent=True, return_latents=False) + I_G_to_0_1 = ((I_G_to + 1) / 2) + I_gen_to = self.downsample_256(I_G_to_0_1).clip(0, 1) + + # key_point_loss + key_points_gen_to = self.generate_key_points(I_gen_to) + key_point_loss_to = F.mse_loss(key_points_gen_to, key_points_to) + + # arcface loss + gen_embed = self.arc_face(self.toArcface(I_gen_to)) + gt_embed = self.arc_face(self.toArcface(I_from)) + arc_face_loss = 20 * (1 - F.cosine_similarity(gen_embed, gt_embed)).mean() + + losses = { + 'mse points to': key_point_loss_to, + 'arc face': arc_face_loss + } + + if normalize: + losses['loss'] = self.MAL.calc_loss(losses) + else: + losses['loss'] = sum(losses.values()) + + if ret_images: + return losses['loss'], {key: val.item() for key, val in losses.items()}, I_gen_to, latent_in + else: + return losses['loss'], {key: val.item() for key, val in losses.items()} + + def calc_hair_loss(self, + latents_from, + latents_to, + ret_images=False, + normalize=True + ): + # rotate + rotate_to = self.model(latents_from[:, :6], latents_to[:, :6]) + mse_latents = 300 * F.mse_loss(rotate_to, latents_to[:, :6]) + + losses = { + 'mse latents': mse_latents + } + + if normalize: + losses['loss'] = self.MAL.calc_loss(losses) + else: + losses['loss'] = sum(losses.values()) + + if ret_images: + latent_in = torch.cat((rotate_to, latents_from[:, 6:]), axis=1) + I_G_to, _ = self.net.generator([latent_in], input_is_latent=True, return_latents=False) + I_G_to_0_1 = ((I_G_to + 1) / 2) + I_gen_to = self.downsample_256(I_G_to_0_1).clip(0, 1) + + return losses['loss'], {key: val.item() for key, val in losses.items()}, I_gen_to + else: + return losses['loss'], {key: val.item() for key, val in losses.items()} + + def train_one_epoch(self): + self.model.to(self.device).train() + sum_losses = lambda x, y: {key: y.get(key, 0) + x.get(key, 0) for key in set(x.keys()) | set(y.keys())} + + dataloader_to = iter(self.train_dataloader) + for batch in tqdm(self.train_dataloader): + I_from, key_points_from, latents_from = map(lambda x: x.to(self.device), batch) + I_to, key_points_to, latents_to = map(lambda x: x.to(self.device), next(dataloader_to)) + + self.optimizer.zero_grad() + + loss, info, _, gen_latent = self.calc_loss( + I_to, + I_from, + key_points_to, + latents_from, + latents_to, + ret_images=True + ) + + if self.args.use_hair_loss: + hair_loss, info2 = self.calc_hair_loss( + gen_latent, + latents_from + ) + loss += hair_loss + info = sum_losses(info, info2) + loss.backward() + + self.MAL.update(info) + + total_norm = torch.nn.utils.clip_grad_norm_(self.model.parameters(), 5) + self.optimizer.step() + + self.logger.next_step() + for key, val in info.items(): + self.logger.log(key, val) + self.logger.log('grad', total_norm.item()) + + @torch.no_grad() + def validate(self): + self.model.to(self.device).eval() + + sum_losses = lambda x, y: {key: y.get(key, 0) + x.get(key, 0) for key in set(x.keys()) | set(y.keys())} + files = [] + losses = {} + + for batch in tqdm(self.test_dataloader): + I_from, key_points_from, latents_from, \ + I_to, key_points_to, latents_to, = map(lambda x: x.to(self.device), batch) + bsz = I_from.size(0) + + loss, info, I_gen_to, gen_latent = self.calc_loss( + I_to, + I_from, + key_points_to, + latents_from, + latents_to, + ret_images=True, + normalize=False + ) + + if args.use_hair_loss: + loss, info2, I_gen_to_rec = self.calc_hair_loss( + gen_latent, + latents_from, + ret_images=True, + normalize=False + ) + losses = sum_losses(losses, info2) + else: + I_G_from, _ = self.net.generator([latents_from], input_is_latent=True, return_latents=False) + I_G_from_0_1 = ((I_G_from + 1) / 2) + I_gen_to_rec = self.downsample_256(I_G_from_0_1).clip(0, 1) + + losses = sum_losses(losses, info) + for k in range(bsz): + files.append([I_from[k].cpu(), I_gen_to_rec[k].cpu(), I_gen_to[k].cpu(), I_to[k].cpu()]) + + for key, val in losses.items(): + val /= len(self.test_dataloader) + self.logger.log(f'val {key}', val) + + np.random.seed(1927) + idxs = np.random.choice(len(files), size=min(len(files), 100), replace=False) + images_to_log = [image_grid(list(map(T.functional.to_pil_image, files[idx])), 1, 4) for idx in idxs] + self.logger.log('val images', [wandb.Image(image) for image in images_to_log]) + + return losses['loss'] / len(self.test_dataloader) + + def train_loop(self, epochs): + # self.validate() + for epoch in range(epochs): + self.train_one_epoch() + loss = self.validate() + + self.save_model(f'rotate_{epoch}', save_online=False) + self.save_model('last') + if loss <= self.best_loss: + self.best_loss = loss + self.save_model(f'best', save_online=False) + + +class Rotate_dataset(Dataset): + def __init__(self, tensors_images, key_points, latents, is_test=False): + super().__init__() + self.tensors_images = tensors_images + self.key_points = key_points + self.latents = latents + self.is_test = is_test + + def __len__(self): + return len(self.tensors_images) + + def __get_elem__(self, idx): + return self.tensors_images[idx], self.key_points[idx], self.latents[idx] + + def __getitem__(self, idx): + if self.is_test: + return *self.__get_elem__(idx), *self.__get_elem__(-idx) + else: + return self.__get_elem__(idx) + + +def main(args): + seed_everything() + data = list(torch.load(args.dataset).values()) + + X_train, X_test = train_test_split(list(zip(data[0], data[1], data[2])), test_size=512, random_state=42) + + train_dataset = Rotate_dataset(*list(zip(*X_train))) + test_dataset = Rotate_dataset(*list(zip(*X_test)), is_test=True) + + train_dataloader = DataLoader(train_dataset, batch_size=args.batch_size, pin_memory=True, shuffle=True, + drop_last=True, num_workers=4) + test_dataloader = DataLoader(test_dataset, batch_size=args.batch_size, pin_memory=True, shuffle=False, + num_workers=4) + + logger = WandbLogger(name=args.name_run, project='HairFast-Rotate') + logger.start_logging() + logger.save(__file__) + + model = RotateModel() + optimizer = torch.optim.Adam(model.parameters(), lr=1e-4, weight_decay=0.000001) + + trainer = Trainer(model, args, optimizer, None, train_dataloader, test_dataloader, logger) + trainer.train_loop(1000) + + +if __name__ == '__main__': + parser = argparse.ArgumentParser(description='Rotate trainer') + parser.add_argument('--name_run', type=str, default='test') + parser.add_argument('--dataset', type=str, default='input/rotate_dataset.pkl') + parser.add_argument('--use_hair_loss', action='store_false') + parser.add_argument('--batch_size', type=int, default=16) + args = parser.parse_args() + + main(args) diff --git a/utils/PCA_utils.py b/utils/PCA_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..c4ea063dda3a3f5adaacf32142ac1249611e9d82 --- /dev/null +++ b/utils/PCA_utils.py @@ -0,0 +1,29 @@ +from sklearn.decomposition import IncrementalPCA +import numpy as np +class IPCAEstimator(): + def __init__(self, n_components): + self.n_components = n_components + self.whiten = False + self.transformer = IncrementalPCA(n_components, whiten=self.whiten, batch_size=max(100, 5*n_components)) + self.batch_support = True + + def get_param_str(self): + return "ipca_c{}{}".format(self.n_components, '_w' if self.whiten else '') + + def fit(self, X): + self.transformer.fit(X) + + def fit_partial(self, X): + try: + self.transformer.partial_fit(X) + self.transformer.n_samples_seen_ = \ + self.transformer.n_samples_seen_.astype(np.int64) # avoid overflow + return True + except ValueError as e: + print(f'\nIPCA error:', e) + return False + + def get_components(self): + stdev = np.sqrt(self.transformer.explained_variance_) # already sorted + var_ratio = self.transformer.explained_variance_ratio_ + return self.transformer.components_, stdev, var_ratio # PCA outputs are normalized \ No newline at end of file diff --git a/utils/__init__.py b/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/utils/bicubic.py b/utils/bicubic.py new file mode 100644 index 0000000000000000000000000000000000000000..bf1cc312fd3ee984b7ce3e85faa2e0c3d4508bec --- /dev/null +++ b/utils/bicubic.py @@ -0,0 +1,75 @@ +import torch +from torch import nn +from torch.nn import functional as F + + +class BicubicDownSample(nn.Module): + def bicubic_kernel(self, x, a=-0.50): + """ + This equation is exactly copied from the website below: + https://clouard.users.greyc.fr/Pantheon/experiments/rescaling/index-en.html#bicubic + """ + abs_x = torch.abs(x) + if abs_x <= 1.: + return (a + 2.) * torch.pow(abs_x, 3.) - (a + 3.) * torch.pow(abs_x, 2.) + 1 + elif 1. < abs_x < 2.: + return a * torch.pow(abs_x, 3) - 5. * a * torch.pow(abs_x, 2.) + 8. * a * abs_x - 4. * a + else: + return 0.0 + + def __init__(self, factor=4, cuda=True, padding='reflect'): + super().__init__() + self.factor = factor + size = factor * 4 + k = torch.tensor([self.bicubic_kernel((i - torch.floor(torch.tensor(size / 2)) + 0.5) / factor) + for i in range(size)], dtype=torch.float32) + k = k / torch.sum(k) + # k = torch.einsum('i,j->ij', (k, k)) + k1 = torch.reshape(k, shape=(1, 1, size, 1)) + self.k1 = torch.cat([k1, k1, k1], dim=0) + k2 = torch.reshape(k, shape=(1, 1, 1, size)) + self.k2 = torch.cat([k2, k2, k2], dim=0) + self.cuda = '.cuda' if cuda else '' + self.padding = padding + for param in self.parameters(): + param.requires_grad = False + + def forward(self, x, nhwc=False, clip_round=False, byte_output=False): + # x = torch.from_numpy(x).type('torch.FloatTensor') + filter_height = self.factor * 4 + filter_width = self.factor * 4 + stride = self.factor + + pad_along_height = max(filter_height - stride, 0) + pad_along_width = max(filter_width - stride, 0) + filters1 = self.k1.type('torch{}.FloatTensor'.format(self.cuda)) + filters2 = self.k2.type('torch{}.FloatTensor'.format(self.cuda)) + + # compute actual padding values for each side + pad_top = pad_along_height // 2 + pad_bottom = pad_along_height - pad_top + pad_left = pad_along_width // 2 + pad_right = pad_along_width - pad_left + + # apply mirror padding + if nhwc: + x = torch.transpose(torch.transpose( + x, 2, 3), 1, 2) # NHWC to NCHW + + # downscaling performed by 1-d convolution + x = F.pad(x, (0, 0, pad_top, pad_bottom), self.padding) + x = F.conv2d(input=x, weight=filters1, stride=(stride, 1), groups=3) + if clip_round: + x = torch.clamp(torch.round(x), 0.0, 255.) + + x = F.pad(x, (pad_left, pad_right, 0, 0), self.padding) + x = F.conv2d(input=x, weight=filters2, stride=(1, stride), groups=3) + if clip_round: + x = torch.clamp(torch.round(x), 0.0, 255.) + + if nhwc: + x = torch.transpose(torch.transpose(x, 1, 3), 1, 2) + if byte_output: + return x.type('torch.ByteTensor'.format(self.cuda)) + else: + return x diff --git a/utils/drive.py b/utils/drive.py new file mode 100644 index 0000000000000000000000000000000000000000..ba07cbe9a630f22cf2f8a3094796011e5cceb75e --- /dev/null +++ b/utils/drive.py @@ -0,0 +1,110 @@ +# URL helpers, see https://github.com/NVlabs/stylegan +# ------------------------------------------------------------------------------------------ + +import requests +import html +import hashlib +import gdown +import glob +import os +import io +from typing import Any +import re +import uuid + +weight_dic = {'afhqwild.pt': 'https://drive.google.com/file/d/14OnzO4QWaAytKXVqcfWo_o2MzoR4ygnr/view?usp=sharing', + 'afhqdog.pt': 'https://drive.google.com/file/d/16v6jPtKVlvq8rg2Sdi3-R9qZEVDgvvEA/view?usp=sharing', + 'afhqcat.pt': 'https://drive.google.com/file/d/1HXLER5R3EMI8DSYDBZafoqpX4EtyOf2R/view?usp=sharing', + 'ffhq.pt': 'https://drive.google.com/file/d/1AT6bNR2ppK8f2ETL_evT27f3R_oyWNHS/view?usp=sharing', + 'metfaces.pt': 'https://drive.google.com/file/d/16wM2PwVWzaMsRgPExvRGsq6BWw_muKbf/view?usp=sharing', + 'seg.pth': 'https://drive.google.com/file/d/1lIKvQaFKHT5zC7uS4p17O9ZpfwmwlS62/view?usp=sharing'} + + +def download_weight(weight_path): + gdown.download(weight_dic[os.path.basename(weight_path)], + output=weight_path, fuzzy=True) + + +def is_url(obj: Any) -> bool: + """Determine whether the given object is a valid URL string.""" + if not isinstance(obj, str) or not "://" in obj: + return False + try: + res = requests.compat.urlparse(obj) + if not res.scheme or not res.netloc or not "." in res.netloc: + return False + res = requests.compat.urlparse(requests.compat.urljoin(obj, "/")) + if not res.scheme or not res.netloc or not "." in res.netloc: + return False + except: + return False + return True + + +def open_url(url: str, cache_dir: str = None, num_attempts: int = 10, verbose: bool = True, + return_path: bool = False) -> Any: + """Download the given URL and return a binary-mode file object to access the data.""" + assert is_url(url) + assert num_attempts >= 1 + + # Lookup from cache. + url_md5 = hashlib.md5(url.encode("utf-8")).hexdigest() + if cache_dir is not None: + cache_files = glob.glob(os.path.join(cache_dir, url_md5 + "_*")) + if len(cache_files) == 1: + if (return_path): + return cache_files[0] + else: + return open(cache_files[0], "rb") + + # Download. + url_name = None + url_data = None + with requests.Session() as session: + if verbose: + print("Downloading %s ..." % url, end="", flush=True) + for attempts_left in reversed(range(num_attempts)): + try: + with session.get(url) as res: + res.raise_for_status() + if len(res.content) == 0: + raise IOError("No data received") + + if len(res.content) < 8192: + content_str = res.content.decode("utf-8") + if "download_warning" in res.headers.get("Set-Cookie", ""): + links = [html.unescape(link) for link in content_str.split('"') if + "export=download" in link] + if len(links) == 1: + url = requests.compat.urljoin(url, links[0]) + raise IOError("Google Drive virus checker nag") + if "Google Drive - Quota exceeded" in content_str: + raise IOError("Google Drive quota exceeded") + + match = re.search(r'filename="([^"]*)"', res.headers.get("Content-Disposition", "")) + url_name = match[1] if match else url + url_data = res.content + if verbose: + print(" done") + break + except: + if not attempts_left: + if verbose: + print(" failed") + raise + if verbose: + print(".", end="", flush=True) + + # Save to cache. + if cache_dir is not None: + safe_name = re.sub(r"[^0-9a-zA-Z-._]", "_", url_name) + cache_file = os.path.join(cache_dir, url_md5 + "_" + safe_name) + temp_file = os.path.join(cache_dir, "tmp_" + uuid.uuid4().hex + "_" + url_md5 + "_" + safe_name) + os.makedirs(cache_dir, exist_ok=True) + with open(temp_file, "wb") as f: + f.write(url_data) + os.replace(temp_file, cache_file) # atomic + if (return_path): return cache_file + + # Return data as file object. + return io.BytesIO(url_data) diff --git a/utils/image_utils.py b/utils/image_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..f44dc5b0509b03aacbad6471d94b0e5e17317ae7 --- /dev/null +++ b/utils/image_utils.py @@ -0,0 +1,108 @@ +import os +import subprocess +import tempfile +from pathlib import Path + +import torch +import torch.nn.functional as F +from PIL import Image +from torchvision.transforms import transforms +from torchvision.utils import save_image + +from models.Net import get_segmentation + + +def equal_replacer(images: list[torch.Tensor]) -> list[torch.Tensor]: + for i in range(len(images)): + if images[i].dtype is torch.uint8: + images[i] = images[i] / 255 + + for i in range(len(images)): + for j in range(i + 1, len(images)): + if torch.allclose(images[i], images[j]): + images[j] = images[i] + return images + + +class DilateErosion: + def __init__(self, dilate_erosion=5, device='cuda'): + self.dilate_erosion = dilate_erosion + self.weight = torch.Tensor([ + [False, True, False], + [True, True, True], + [False, True, False] + ]).float()[None, None, ...].to(device) + + def hair_from_mask(self, mask): + mask = torch.where(mask == 13, torch.ones_like(mask), torch.zeros_like(mask)) + mask = F.interpolate(mask, size=(256, 256), mode='nearest') + dilate, erosion = self.mask(mask) + return dilate, erosion + + def mask(self, mask): + masks = mask.clone().repeat(*([2] + [1] * (len(mask.shape) - 1))).float() + sum_w = self.weight.sum().item() + n = len(mask) + + for _ in range(self.dilate_erosion): + masks = F.conv2d(masks, self.weight, + bias=None, stride=1, padding='same', dilation=1, groups=1) + masks[:n] = (masks[:n] > 0).float() + masks[n:] = (masks[n:] == sum_w).float() + + hair_mask_dilate, hair_mask_erode = masks[:n], masks[n:] + + return hair_mask_dilate, hair_mask_erode + + +def poisson_image_blending(final_image, face_image, dilate_erosion=30, maxn=115): + dilate_erosion = DilateErosion(dilate_erosion=dilate_erosion) + transform = transforms.ToTensor() + + if isinstance(face_image, str): + face_image = transform(Image.open(face_image)) + elif not isinstance(face_image, torch.Tensor): + face_image = transform(face_image) + + final_mask = get_segmentation(final_image.cuda().unsqueeze(0), resize=False) + face_mask = get_segmentation(face_image.cuda().unsqueeze(0), resize=False) + + hair_target = torch.where(final_mask == 13, torch.ones_like(final_mask), + torch.zeros_like(final_mask)) + hair_face = torch.where(face_mask == 13, torch.ones_like(face_mask), + torch.zeros_like(face_mask)) + + final_mask = F.interpolate(((1 - hair_target) * (1 - hair_face)).float(), size=(1024, 1024), mode='bicubic') + dilation, _ = dilate_erosion.mask(1 - final_mask) + mask_save = 1 - dilation[0] + + with tempfile.TemporaryDirectory() as temp_dir: + final_image_path = os.path.join(temp_dir, 'final_image.png') + face_image_path = os.path.join(temp_dir, 'face_image.png') + mask_path = os.path.join(temp_dir, 'mask_save.png') + save_image(final_image, final_image_path) + save_image(face_image, face_image_path) + save_image(mask_save, mask_path) + + out_image_path = os.path.join(temp_dir, 'out_image_path.png') + result = subprocess.run( + ["fpie", "-s", face_image_path, "-m", mask_path, "-t", final_image_path, "-o", out_image_path, "-n", + str(maxn), "-b", "taichi-gpu", "-g", "max"], + check=True + ) + + return Image.open(out_image_path), Image.open(mask_path) + + +def list_image_files(directory): + image_extensions = ['.jpg', '.jpeg', '.png'] + image_files = [] + + for entry in sorted(os.listdir(directory)): + file_path = os.path.join(directory, entry) + if os.path.isfile(file_path): + file_extension = Path(file_path).suffix.lower() + if file_extension in image_extensions: + image_files.append(entry) + + return image_files diff --git a/utils/save_utils.py b/utils/save_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..999ff79722509534436af130c09b4cf5ba741ba2 --- /dev/null +++ b/utils/save_utils.py @@ -0,0 +1,38 @@ +import os + +import numpy as np +import torchvision.transforms as T +from PIL import Image + +from models.CtrlHair.util.mask_color_util import mask_to_rgb + +toPIL = T.ToPILImage() + + +def save_gen_image(output_dir, path, name, gen_im): + if len(gen_im.shape) == 4: + gen_im = gen_im[0] + save_im = toPIL(((gen_im + 1) / 2).detach().cpu().clamp(0, 1)) + + save_dir = output_dir / path + os.makedirs(save_dir, exist_ok=True) + + image_path = save_dir / name + save_im.save(image_path) + + +def save_vis_mask(output_dir, path, name, mask): + out_dir = output_dir / path + os.makedirs(out_dir, exist_ok=True) + out_mask_path = out_dir / name + + rgb_img = Image.fromarray(mask_to_rgb(mask.detach().cpu().squeeze(), 0)) + rgb_img.save(out_mask_path) + + +def save_latents(output_dir, path, file_name, **latents): + save_dir = output_dir / path + os.makedirs(save_dir, exist_ok=True) + + latent_path = save_dir / file_name + np.savez(latent_path, **{key: latent.detach().cpu().numpy() for key, latent in latents.items()}) diff --git a/utils/seed.py b/utils/seed.py new file mode 100644 index 0000000000000000000000000000000000000000..e45c5263343bced8c73a21e748d677bcdac3d771 --- /dev/null +++ b/utils/seed.py @@ -0,0 +1,31 @@ +import functools +import random + +import numpy as np +import torch + + +def set_seed(seed): + torch.manual_seed(seed) + torch.cuda.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + torch.backends.cudnn.benchmark = False + torch.backends.cudnn.deterministic = True + np.random.seed(seed) + random.seed(seed) + + +def seed_setter(func): + default_seed = 3407 + + @functools.wraps(func) + def wraps(*args, **kwargs): + seed = kwargs.pop('seed', None) + if seed is None: + seed = default_seed + set_seed(seed) + + result = func(*args, **kwargs) + return result + + return wraps diff --git a/utils/shape_predictor.py b/utils/shape_predictor.py new file mode 100644 index 0000000000000000000000000000000000000000..8596ebe3739db9209a07c168a54abb8786f3de42 --- /dev/null +++ b/utils/shape_predictor.py @@ -0,0 +1,194 @@ +import os +from pathlib import Path + +import PIL +import dlib +import numpy as np +import scipy +import scipy.ndimage +import torch +from PIL import Image +from torchvision import transforms as T + +from utils.drive import open_url + +""" +brief: face alignment with FFHQ method (https://github.com/NVlabs/ffhq-dataset) +author: lzhbrian (https://lzhbrian.me) +date: 2020.1.5 +note: code is heavily borrowed from + https://github.com/NVlabs/ffhq-dataset + http://dlib.net/face_landmark_detection.py.html + +requirements: + apt install cmake + conda install Pillow numpy scipy + pip install dlib + # download face landmark model from: + # http://dlib.net/files/shape_predictor_68_face_landmarks.dat.bz2 +""" + + +def get_landmark(filepath, predictor): + """get landmark with dlib + :return: np.array shape=(68, 2) + """ + detector = dlib.get_frontal_face_detector() + + img = dlib.load_rgb_image(filepath) + dets = detector(img, 1) + filepath = Path(filepath) + print(f"{filepath.name}: Number of faces detected: {len(dets)}") + shapes = [predictor(img, d) for k, d in enumerate(dets)] + + lms = [np.array([[tt.x, tt.y] for tt in shape.parts()]) for shape in shapes] + + return lms + + +def get_landmark_from_tensors(tensors: list[torch.Tensor | Image.Image | np.ndarray], predictor): + detector = dlib.get_frontal_face_detector() + transform = T.ToPILImage() + images = [] + lms = [] + + for k, tensor in enumerate(tensors): + if isinstance(tensor, torch.Tensor): + img_pil = transform(tensor) + else: + img_pil = tensor + img = np.array(img_pil) + images.append(img_pil) + + dets = detector(img, 1) + if len(dets) == 0: + raise ValueError(f"No faces detected in the image {k}.") + elif len(dets) == 1: + print(f"Number of faces detected: {len(dets)}") + else: + print(f"Number of faces detected: {len(dets)}, get largest face") + + # Find the largest face + dets = sorted(dets, key=lambda det: det.width() * det.height(), reverse=True) + shape = predictor(img, dets[0]) + lm = np.array([[tt.x, tt.y] for tt in shape.parts()]) + lms.append(lm) + + return images, lms + + +def align_face(data, predictor=None, is_filepath=False, return_tensors=True): + """ + :param data: filepath or list torch Tensors + :return: list of PIL Images + """ + if predictor is None: + predictor_path = 'pretrained_models/ShapeAdaptor/shape_predictor_68_face_landmarks.dat' + + if not os.path.isfile(predictor_path): + print("Downloading Shape Predictor") + data_io = open_url("https://drive.google.com/uc?id=1huhv8PYpNNKbGCLOaYUjOgR1pY5pmbJx") + with open(predictor_path, 'wb') as f: + f.write(data_io.getbuffer()) + + predictor = dlib.shape_predictor(predictor_path) + + if is_filepath: + lms = get_landmark(data, predictor) + else: + if not isinstance(data, list): + data = [data] + images, lms = get_landmark_from_tensors(data, predictor) + + imgs = [] + for num_img, lm in enumerate(lms): + lm_chin = lm[0: 17] # left-right + lm_eyebrow_left = lm[17: 22] # left-right + lm_eyebrow_right = lm[22: 27] # left-right + lm_nose = lm[27: 31] # top-down + lm_nostrils = lm[31: 36] # top-down + lm_eye_left = lm[36: 42] # left-clockwise + lm_eye_right = lm[42: 48] # left-clockwise + lm_mouth_outer = lm[48: 60] # left-clockwise + lm_mouth_inner = lm[60: 68] # left-clockwise + + # Calculate auxiliary vectors. + eye_left = np.mean(lm_eye_left, axis=0) + eye_right = np.mean(lm_eye_right, axis=0) + eye_avg = (eye_left + eye_right) * 0.5 + eye_to_eye = eye_right - eye_left + mouth_left = lm_mouth_outer[0] + mouth_right = lm_mouth_outer[6] + mouth_avg = (mouth_left + mouth_right) * 0.5 + eye_to_mouth = mouth_avg - eye_avg + + # Choose oriented crop rectangle. + x = eye_to_eye - np.flipud(eye_to_mouth) * [-1, 1] + x /= np.hypot(*x) + x *= max(np.hypot(*eye_to_eye) * 2.0, np.hypot(*eye_to_mouth) * 1.8) + y = np.flipud(x) * [-1, 1] + c = eye_avg + eye_to_mouth * 0.1 + quad = np.stack([c - x - y, c - x + y, c + x + y, c + x - y]) + qsize = np.hypot(*x) * 2 + + # read image + if is_filepath: + img = PIL.Image.open(data) + else: + img = images[num_img] + + output_size = 1024 + # output_size = 256 + transform_size = 4096 + enable_padding = True + + # Shrink. + shrink = int(np.floor(qsize / output_size * 0.5)) + if shrink > 1: + rsize = (int(np.rint(float(img.size[0]) / shrink)), int(np.rint(float(img.size[1]) / shrink))) + img = img.resize(rsize, PIL.Image.ANTIALIAS) + quad /= shrink + qsize /= shrink + + # Crop. + border = max(int(np.rint(qsize * 0.1)), 3) + crop = (int(np.floor(min(quad[:, 0]))), int(np.floor(min(quad[:, 1]))), int(np.ceil(max(quad[:, 0]))), + int(np.ceil(max(quad[:, 1])))) + crop = (max(crop[0] - border, 0), max(crop[1] - border, 0), min(crop[2] + border, img.size[0]), + min(crop[3] + border, img.size[1])) + if crop[2] - crop[0] < img.size[0] or crop[3] - crop[1] < img.size[1]: + img = img.crop(crop) + quad -= crop[0:2] + + # Pad. + pad = (int(np.floor(min(quad[:, 0]))), int(np.floor(min(quad[:, 1]))), int(np.ceil(max(quad[:, 0]))), + int(np.ceil(max(quad[:, 1])))) + pad = (max(-pad[0] + border, 0), max(-pad[1] + border, 0), max(pad[2] - img.size[0] + border, 0), + max(pad[3] - img.size[1] + border, 0)) + if enable_padding and max(pad) > border - 4: + pad = np.maximum(pad, int(np.rint(qsize * 0.3))) + img = np.pad(np.float32(img), ((pad[1], pad[3]), (pad[0], pad[2]), (0, 0)), 'reflect') + h, w, _ = img.shape + y, x, _ = np.ogrid[:h, :w, :1] + mask = np.maximum(1.0 - np.minimum(np.float32(x) / pad[0], np.float32(w - 1 - x) / pad[2]), + 1.0 - np.minimum(np.float32(y) / pad[1], np.float32(h - 1 - y) / pad[3])) + blur = qsize * 0.02 + img += (scipy.ndimage.gaussian_filter(img, [blur, blur, 0]) - img) * np.clip(mask * 3.0 + 1.0, 0.0, 1.0) + img += (np.median(img, axis=(0, 1)) - img) * np.clip(mask, 0.0, 1.0) + img = PIL.Image.fromarray(np.uint8(np.clip(np.rint(img), 0, 255)), 'RGB') + quad += pad[:2] + + # Transform. + img = img.transform((transform_size, transform_size), PIL.Image.QUAD, (quad + 0.5).flatten(), + PIL.Image.BILINEAR) + if output_size < transform_size: + img = img.resize((output_size, output_size), PIL.Image.LANCZOS) + + # Save aligned image. + imgs.append(img) + + if return_tensors: + transform = T.ToTensor() + tensors = [transform(img).clamp(0, 1) for img in imgs] + return tensors + return imgs diff --git a/utils/time.py b/utils/time.py new file mode 100644 index 0000000000000000000000000000000000000000..0fbb5c2d5f69a607e4fd923780b4177b01dff80b --- /dev/null +++ b/utils/time.py @@ -0,0 +1,36 @@ +import functools +import sys +import time + +import numpy as np +import torch + + +def get_time(): + torch.cuda.current_stream().synchronize() + return time.time() + + +def bench_session(func): + times = [] + + @functools.wraps(func) + def wraps(*args, **kwargs): + if kwargs.pop('benchmark', False): + nonlocal times + start = get_time() + + result = func(*args, **kwargs) + + eval_time = get_time() - start + times.append(eval_time) + + print(f'\n{len(times)} experiment ended in {eval_time:.3f}(s)', file=sys.stderr) + print(f'min time: {np.min(times):.3f}(s),' + f' median time: {np.median(times):.3f}(s),' + f' std time: {np.std(times):.3f}(s)', file=sys.stderr) + return result + else: + return func(*args, **kwargs) + + return wraps diff --git a/utils/train.py b/utils/train.py new file mode 100644 index 0000000000000000000000000000000000000000..9ae6b132e1a0891912757846713145270d8c0a96 --- /dev/null +++ b/utils/train.py @@ -0,0 +1,161 @@ +import os +import pickle +import random +import shutil +import typing as tp + +import numpy as np +import torch +import torchvision.transforms as T +import wandb +from PIL import Image +from joblib import Parallel, delayed +from torch.utils.data import DataLoader, TensorDataset +from torchmetrics.image.fid import FrechetInceptionDistance +from tqdm.auto import tqdm + +from models.Encoders import ClipModel + + +def image_grid(imgs, rows, cols): + assert len(imgs) == rows * cols + + w, h = imgs[0].size + grid = Image.new('RGB', size=(cols * w, rows * h)) + + for i, img in enumerate(imgs): + grid.paste(img, box=(i % cols * w, i // cols * h)) + return grid + + +class WandbLogger: + def __init__(self, name='base-name', project='HairFast'): + self.name = name + self.project = project + + def start_logging(self): + wandb.login(key=os.environ['WANDB_KEY'].strip(), relogin=True) + wandb.init( + project=self.project, + name=self.name + ) + self.wandb = wandb + self.run_dir = self.wandb.run.dir + self.train_step = 0 + + def log(self, scalar_name: str, scalar: tp.Any): + self.wandb.log({scalar_name: scalar}, step=self.train_step, commit=False) + + def log_scalars(self, scalars: dict): + self.wandb.log(scalars, step=self.train_step, commit=False) + + def next_step(self): + self.train_step += 1 + + def save(self, file_path, save_online=True): + file = os.path.basename(file_path) + new_path = os.path.join(self.run_dir, file) + shutil.copy2(file_path, new_path) + if save_online: + self.wandb.save(new_path) + + def __del__(self): + self.wandb.finish() + + +def toggle_grad(model, flag=True): + for p in model.parameters(): + p.requires_grad = flag + + +class _LegacyUnpickler(pickle.Unpickler): + def find_class(self, module, name): + if module == 'dnnlib.tflib.network' and name == 'Network': + return _TFNetworkStub + module = module.replace('torch_utils', 'models.stylegan2.torch_utils') + module = module.replace('dnnlib', 'models.stylegan2.dnnlib') + return super().find_class(module, name) + + +def seed_everything(seed: int = 1729) -> None: + random.seed(seed) + os.environ["PYTHONHASHSEED"] = str(seed) + np.random.seed(seed) + torch.manual_seed(seed) + torch.cuda.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + torch.backends.cudnn.deterministic = True + + +def load_images_to_torch(paths, imgs=None, use_tqdm=True): + transform = T.PILToTensor() + tensor = [] + for path in paths: + if imgs is None: + pbar = sorted(os.listdir(path)) + else: + pbar = imgs + + if use_tqdm: + pbar = tqdm(pbar) + + for img_name in pbar: + if '.jpg' in img_name or '.png' in img_name: + img_path = os.path.join(path, img_name) + img = Image.open(img_path).resize((299, 299), resample=Image.LANCZOS) + tensor.append(transform(img)) + try: + return torch.stack(tensor) + except: + print(paths, imgs) + return torch.tensor([], dtype=torch.uint8) + + +def parallel_load_images(paths, imgs): + assert imgs is not None + if not isinstance(paths, list): + paths = [paths] + + list_torch_images = Parallel(n_jobs=-1)(delayed(load_images_to_torch)( + paths, [i], use_tqdm=False + ) for i in tqdm(imgs)) + return torch.cat(list_torch_images) + + +def get_fid_calc(instance='fid.pkl', dataset_path='', device=torch.device('cuda')): + if os.path.isfile(instance): + with open(instance, 'rb') as f: + fid = pickle.load(f) + else: + fid = FrechetInceptionDistance(feature=ClipModel(), reset_real_features=False, normalize=True) + fid.to(device).eval() + + imgs_file = [] + for file in os.listdir(dataset_path): + if 'flip' not in file and os.path.splitext(file)[1] in ['.png', '.jpg']: + imgs_file.append(file) + + tensor_images = parallel_load_images([dataset_path], imgs_file).float().div(255) + real_dataloader = DataLoader(TensorDataset(tensor_images), batch_size=128) + with torch.inference_mode(): + for batch in tqdm(real_dataloader): + batch = batch[0].to(device) + fid.update(batch, real=True) + + with open(instance, 'wb') as f: + pickle.dump(fid.cpu(), f) + fid.to(device).eval() + + @torch.inference_mode() + def compute_fid_datasets(images): + nonlocal fid, device + fid.reset() + + fake_dataloader = DataLoader(TensorDataset(images), batch_size=128) + for batch in tqdm(fake_dataloader): + batch = batch[0].to(device) + fid.update(batch, real=False) + + return fid.compute() + + return compute_fid_datasets