Upload 10 files

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

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

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

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()})

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

utils/shape_predictor.py

@@ -19,6 +19,7 @@ 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
@@ -82,7 +83,7 @@ def align_face(data, predictor=None, is_filepath=False, return_tensors=True):
     :return: list of PIL Images
     """
     if predictor is None:
-        predictor_path = 'shape_predictor_68_face_landmarks.dat'
+        predictor_path = 'pretrained_models/ShapeAdaptor/shape_predictor_68_face_landmarks.dat'
 
         if not os.path.isfile(predictor_path):
             print("Downloading Shape Predictor")

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

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