| | import cv2
|
| | import os
|
| | import os.path as osp
|
| | import numpy as np
|
| | from PIL import Image
|
| | import torch
|
| | from torch.hub import download_url_to_file, get_dir
|
| | from urllib.parse import urlparse
|
| |
|
| |
|
| | ROOT_DIR = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
| |
|
| |
|
| | def download_pretrained_models(file_ids, save_path_root):
|
| | import gdown
|
| |
|
| | os.makedirs(save_path_root, exist_ok=True)
|
| |
|
| | for file_name, file_id in file_ids.items():
|
| | file_url = 'https://drive.google.com/uc?id='+file_id
|
| | save_path = osp.abspath(osp.join(save_path_root, file_name))
|
| | if osp.exists(save_path):
|
| | user_response = input(f'{file_name} already exist. Do you want to cover it? Y/N\n')
|
| | if user_response.lower() == 'y':
|
| | print(f'Covering {file_name} to {save_path}')
|
| | gdown.download(file_url, save_path, quiet=False)
|
| |
|
| | elif user_response.lower() == 'n':
|
| | print(f'Skipping {file_name}')
|
| | else:
|
| | raise ValueError('Wrong input. Only accepts Y/N.')
|
| | else:
|
| | print(f'Downloading {file_name} to {save_path}')
|
| | gdown.download(file_url, save_path, quiet=False)
|
| |
|
| |
|
| |
|
| | def imwrite(img, file_path, params=None, auto_mkdir=True):
|
| | """Write image to file.
|
| |
|
| | Args:
|
| | img (ndarray): Image array to be written.
|
| | file_path (str): Image file path.
|
| | params (None or list): Same as opencv's :func:`imwrite` interface.
|
| | auto_mkdir (bool): If the parent folder of `file_path` does not exist,
|
| | whether to create it automatically.
|
| |
|
| | Returns:
|
| | bool: Successful or not.
|
| | """
|
| | if auto_mkdir:
|
| | dir_name = os.path.abspath(os.path.dirname(file_path))
|
| | os.makedirs(dir_name, exist_ok=True)
|
| | return cv2.imwrite(file_path, img, params)
|
| |
|
| |
|
| | def img2tensor(imgs, bgr2rgb=True, float32=True):
|
| | """Numpy array to tensor.
|
| |
|
| | Args:
|
| | imgs (list[ndarray] | ndarray): Input images.
|
| | bgr2rgb (bool): Whether to change bgr to rgb.
|
| | float32 (bool): Whether to change to float32.
|
| |
|
| | Returns:
|
| | list[tensor] | tensor: Tensor images. If returned results only have
|
| | one element, just return tensor.
|
| | """
|
| |
|
| | def _totensor(img, bgr2rgb, float32):
|
| | if img.shape[2] == 3 and bgr2rgb:
|
| | if img.dtype == 'float64':
|
| | img = img.astype('float32')
|
| | img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
|
| | img = torch.from_numpy(img.transpose(2, 0, 1))
|
| | if float32:
|
| | img = img.float()
|
| | return img
|
| |
|
| | if isinstance(imgs, list):
|
| | return [_totensor(img, bgr2rgb, float32) for img in imgs]
|
| | else:
|
| | return _totensor(imgs, bgr2rgb, float32)
|
| |
|
| |
|
| | def load_file_from_url(url, model_dir=None, progress=True, file_name=None):
|
| | """Ref:https://github.com/1adrianb/face-alignment/blob/master/face_alignment/utils.py
|
| | """
|
| | if model_dir is None:
|
| | hub_dir = get_dir()
|
| | model_dir = os.path.join(hub_dir, 'checkpoints')
|
| |
|
| | os.makedirs(os.path.join(ROOT_DIR, model_dir), exist_ok=True)
|
| |
|
| | parts = urlparse(url)
|
| | filename = os.path.basename(parts.path)
|
| | if file_name is not None:
|
| | filename = file_name
|
| | cached_file = os.path.abspath(os.path.join(ROOT_DIR, model_dir, filename))
|
| | if not os.path.exists(cached_file):
|
| | print(f'Downloading: "{url}" to {cached_file}\n')
|
| | download_url_to_file(url, cached_file, hash_prefix=None, progress=progress)
|
| | return cached_file
|
| |
|
| |
|
| | def scandir(dir_path, suffix=None, recursive=False, full_path=False):
|
| | """Scan a directory to find the interested files.
|
| | Args:
|
| | dir_path (str): Path of the directory.
|
| | suffix (str | tuple(str), optional): File suffix that we are
|
| | interested in. Default: None.
|
| | recursive (bool, optional): If set to True, recursively scan the
|
| | directory. Default: False.
|
| | full_path (bool, optional): If set to True, include the dir_path.
|
| | Default: False.
|
| | Returns:
|
| | A generator for all the interested files with relative paths.
|
| | """
|
| |
|
| | if (suffix is not None) and not isinstance(suffix, (str, tuple)):
|
| | raise TypeError('"suffix" must be a string or tuple of strings')
|
| |
|
| | root = dir_path
|
| |
|
| | def _scandir(dir_path, suffix, recursive):
|
| | for entry in os.scandir(dir_path):
|
| | if not entry.name.startswith('.') and entry.is_file():
|
| | if full_path:
|
| | return_path = entry.path
|
| | else:
|
| | return_path = osp.relpath(entry.path, root)
|
| |
|
| | if suffix is None:
|
| | yield return_path
|
| | elif return_path.endswith(suffix):
|
| | yield return_path
|
| | else:
|
| | if recursive:
|
| | yield from _scandir(entry.path, suffix=suffix, recursive=recursive)
|
| | else:
|
| | continue
|
| |
|
| | return _scandir(dir_path, suffix=suffix, recursive=recursive)
|
| |
|
| |
|
| | def is_gray(img, threshold=10):
|
| | img = Image.fromarray(img)
|
| | if len(img.getbands()) == 1:
|
| | return True
|
| | img1 = np.asarray(img.getchannel(channel=0), dtype=np.int16)
|
| | img2 = np.asarray(img.getchannel(channel=1), dtype=np.int16)
|
| | img3 = np.asarray(img.getchannel(channel=2), dtype=np.int16)
|
| | diff1 = (img1 - img2).var()
|
| | diff2 = (img2 - img3).var()
|
| | diff3 = (img3 - img1).var()
|
| | diff_sum = (diff1 + diff2 + diff3) / 3.0
|
| | if diff_sum <= threshold:
|
| | return True
|
| | else:
|
| | return False
|
| |
|
| | def rgb2gray(img, out_channel=3):
|
| | r, g, b = img[:,:,0], img[:,:,1], img[:,:,2]
|
| | gray = 0.2989 * r + 0.5870 * g + 0.1140 * b
|
| | if out_channel == 3:
|
| | gray = gray[:,:,np.newaxis].repeat(3, axis=2)
|
| | return gray
|
| |
|
| | def bgr2gray(img, out_channel=3):
|
| | b, g, r = img[:,:,0], img[:,:,1], img[:,:,2]
|
| | gray = 0.2989 * r + 0.5870 * g + 0.1140 * b
|
| | if out_channel == 3:
|
| | gray = gray[:,:,np.newaxis].repeat(3, axis=2)
|
| | return gray
|
| |
|
| |
|
| | def calc_mean_std(feat, eps=1e-5):
|
| | """
|
| | Args:
|
| | feat (numpy): 3D [w h c]s
|
| | """
|
| | size = feat.shape
|
| | assert len(size) == 3, 'The input feature should be 3D tensor.'
|
| | c = size[2]
|
| | feat_var = feat.reshape(-1, c).var(axis=0) + eps
|
| | feat_std = np.sqrt(feat_var).reshape(1, 1, c)
|
| | feat_mean = feat.reshape(-1, c).mean(axis=0).reshape(1, 1, c)
|
| | return feat_mean, feat_std
|
| |
|
| |
|
| | def adain_npy(content_feat, style_feat):
|
| | """Adaptive instance normalization for numpy.
|
| |
|
| | Args:
|
| | content_feat (numpy): The input feature.
|
| | style_feat (numpy): The reference feature.
|
| | """
|
| | size = content_feat.shape
|
| | style_mean, style_std = calc_mean_std(style_feat)
|
| | content_mean, content_std = calc_mean_std(content_feat)
|
| | normalized_feat = (content_feat - np.broadcast_to(content_mean, size)) / np.broadcast_to(content_std, size)
|
| | return normalized_feat * np.broadcast_to(style_std, size) + np.broadcast_to(style_mean, size) |
