129	
130	    Returns:
131	        Tensor: Warped image or feature map.
132	    """
133	    assert x.size()[-2:] == flow.size()[1:3]
134	    _, _, h, w = x.size()
135	    # create mesh grid
136	    grid_y, grid_x = torch.meshgrid(torch.arange(0, h).type_as(x), torch.arange(0, w).type_as(x))

198	        Tensor: the pixel unshuffled feature.
199	    """
200	    b, c, hh, hw = x.size()
201	    out_channel = c * (scale**2)
202	    assert hh % scale == 0 and hw % scale == 0
203	    h = hh // scale
204	    w = hw // scale
205	    x_view = x.view(b, c, h, scale, w, scale)

17	        eps (float): A small value added to the variance to avoid
18	            divide-by-zero. Default: 1e-5.
19	    """
20	    size = feat.size()
21	    assert len(size) == 4, 'The input feature should be 4D tensor.'
22	    b, c = size[:2]
23	    feat_var = feat.view(b, c, -1).var(dim=2) + eps
24	    feat_std = feat_var.sqrt().view(b, c, 1, 1)

36	        indices = [v % dataset_size for v in indices]
37	
38	        # subsample
39	        indices = indices[self.rank:self.total_size:self.num_replicas]
40	        assert len(indices) == self.num_samples
41	
42	        return iter(indices)
43	

51	
52	    Returns:
53	        list[int]: A list of indices.
54	    """
55	    assert num_frames % 2 == 1, 'num_frames should be an odd number.'
56	    assert padding in ('replicate', 'reflection', 'reflection_circle', 'circle'), f'Wrong padding mode: {padding}.'
57	
58	    max_frame_num = max_frame_num - 1  # start from 0

52	    Returns:
53	        list[int]: A list of indices.
54	    """
55	    assert num_frames % 2 == 1, 'num_frames should be an odd number.'
56	    assert padding in ('replicate', 'reflection', 'reflection_circle', 'circle'), f'Wrong padding mode: {padding}.'
57	
58	    max_frame_num = max_frame_num - 1  # start from 0
59	    num_pad = num_frames // 2

118	
119	    Returns:
120	        list[str]: Returned path list.
121	    """
122	    assert len(folders) == 2, ('The len of folders should be 2 with [input_folder, gt_folder]. '
123	                               f'But got {len(folders)}')
124	    assert len(keys) == 2, ('The len of keys should be 2 with [input_key, gt_key]. ' f'But got {len(keys)}')
125	    input_folder, gt_folder = folders
126	    input_key, gt_key = keys

120	        list[str]: Returned path list.
121	    """
122	    assert len(folders) == 2, ('The len of folders should be 2 with [input_folder, gt_folder]. '
123	                               f'But got {len(folders)}')
124	    assert len(keys) == 2, ('The len of keys should be 2 with [input_key, gt_key]. ' f'But got {len(keys)}')
125	    input_folder, gt_folder = folders
126	    input_key, gt_key = keys
127	

167	
168	    Returns:
169	        list[str]: Returned path list.
170	    """
171	    assert len(folders) == 2, ('The len of folders should be 2 with [input_folder, gt_folder]. '
172	                               f'But got {len(folders)}')
173	    assert len(keys) == 2, ('The len of keys should be 2 with [input_key, gt_key]. ' f'But got {len(keys)}')
174	    input_folder, gt_folder = folders
175	    input_key, gt_key = keys

169	        list[str]: Returned path list.
170	    """
171	    assert len(folders) == 2, ('The len of folders should be 2 with [input_folder, gt_folder]. '
172	                               f'But got {len(folders)}')
173	    assert len(keys) == 2, ('The len of keys should be 2 with [input_key, gt_key]. ' f'But got {len(keys)}')
174	    input_folder, gt_folder = folders
175	    input_key, gt_key = keys
176	

201	
202	    Returns:
203	        list[str]: Returned path list.
204	    """
205	    assert len(folders) == 2, ('The len of folders should be 2 with [input_folder, gt_folder]. '
206	                               f'But got {len(folders)}')
207	    assert len(keys) == 2, ('The len of keys should be 2 with [input_key, gt_key]. ' f'But got {len(keys)}')
208	    input_folder, gt_folder = folders
209	    input_key, gt_key = keys

203	        list[str]: Returned path list.
204	    """
205	    assert len(folders) == 2, ('The len of folders should be 2 with [input_folder, gt_folder]. '
206	                               f'But got {len(folders)}')
207	    assert len(keys) == 2, ('The len of keys should be 2 with [input_key, gt_key]. ' f'But got {len(keys)}')
208	    input_folder, gt_folder = folders
209	    input_key, gt_key = keys
210	

209	    input_key, gt_key = keys
210	
211	    input_paths = list(scandir(input_folder))
212	    gt_paths = list(scandir(gt_folder))
213	    assert len(input_paths) == len(gt_paths), (f'{input_key} and {gt_key} datasets have different number of images: '
214	                                               f'{len(input_paths)}, {len(gt_paths)}.')
215	    paths = []
216	    for gt_path in gt_paths:
217	        basename, ext = osp.splitext(osp.basename(gt_path))

216	    for gt_path in gt_paths:
217	        basename, ext = osp.splitext(osp.basename(gt_path))
218	        input_name = f'{filename_tmpl.format(basename)}{ext}'
219	        input_path = osp.join(input_folder, input_name)
220	        assert input_name in input_paths, (f'{input_name} is not in ' f'{input_key}_paths.')
221	        gt_path = osp.join(gt_folder, gt_path)
222	        paths.append(dict([(f'{input_key}_path', input_path), (f'{gt_key}_path', gt_path)]))
223	    return paths

283	
284	    Returns:
285	        Tensor: DUF downsampled frames.
286	    """
287	    assert scale in (2, 3, 4), f'Only support scale (2, 3, 4), but got {scale}.'
288	
289	    squeeze_flag = False
290	    if x.ndim == 4:

60	                         f'({lq_patch_size}, {lq_patch_size}). '
61	                         f'Please remove {gt_path}.')
62	
63	    # randomly choose top and left coordinates for lq patch
64	    top = random.randint(0, h_lq - lq_patch_size)
65	    left = random.randint(0, w_lq - lq_patch_size)
66	
67	    # crop lq patch

61	                         f'Please remove {gt_path}.')
62	
63	    # randomly choose top and left coordinates for lq patch
64	    top = random.randint(0, h_lq - lq_patch_size)
65	    left = random.randint(0, w_lq - lq_patch_size)
66	
67	    # crop lq patch
68	    img_lqs = [v[top:top + lq_patch_size, left:left + lq_patch_size, ...] for v in img_lqs]

98	        list[ndarray] | ndarray: Augmented images and flows. If returned
99	            results only have one element, just return ndarray.
100	
101	    """
102	    hflip = hflip and random.random() < 0.5
103	    vflip = rotation and random.random() < 0.5
104	    rot90 = rotation and random.random() < 0.5
105	

99	            results only have one element, just return ndarray.
100	
101	    """
102	    hflip = hflip and random.random() < 0.5
103	    vflip = rotation and random.random() < 0.5
104	    rot90 = rotation and random.random() < 0.5
105	
106	    def _augment(img):

100	
101	    """
102	    hflip = hflip and random.random() < 0.5
103	    vflip = rotation and random.random() < 0.5
104	    rot90 = rotation and random.random() < 0.5
105	
106	    def _augment(img):
107	        if hflip:  # horizontal

35	        Tensor: Loss values.
36	    """
37	    # if weight is specified, apply element-wise weight
38	    if weight is not None:
39	        assert weight.dim() == loss.dim()
40	        assert weight.size(1) == 1 or weight.size(1) == loss.size(1)
41	        loss = loss * weight
42	

36	    """
37	    # if weight is specified, apply element-wise weight
38	    if weight is not None:
39	        assert weight.dim() == loss.dim()
40	        assert weight.size(1) == 1 or weight.size(1) == loss.size(1)
41	        loss = loss * weight
42	
43	    # if weight is not specified or reduction is sum, just reduce the loss

23	    Returns:
24	        float: psnr result.
25	    """
26	
27	    assert img1.shape == img2.shape, (f'Image shapes are differnet: {img1.shape}, {img2.shape}.')
28	    if input_order not in ['HWC', 'CHW']:
29	        raise ValueError(f'Wrong input_order {input_order}. Supported input_orders are ' '"HWC" and "CHW"')
30	    img1 = reorder_image(img1, input_order=input_order)

105	    Returns:
106	        float: ssim result.
107	    """
108	
109	    assert img1.shape == img2.shape, (f'Image shapes are differnet: {img1.shape}, {img2.shape}.')
110	    if input_order not in ['HWC', 'CHW']:
111	        raise ValueError(f'Wrong input_order {input_order}. Supported input_orders are ' '"HWC" and "CHW"')
112	    img1 = reorder_image(img1, input_order=input_order)

55	        if not input.is_cuda:
56	            raise NotImplementedError
57	        else:
58	            cur_im2col_step = min(ctx.im2col_step, input.shape[0])
59	            assert (input.shape[0] % cur_im2col_step) == 0, 'im2col step must divide batchsize'
60	            deform_conv_ext.deform_conv_forward(input, weight,
61	                                                offset, output, ctx.bufs_[0], ctx.bufs_[1], weight.size(3),
62	                                                weight.size(2), ctx.stride[1], ctx.stride[0], ctx.padding[1],

74	        if not grad_output.is_cuda:
75	            raise NotImplementedError
76	        else:
77	            cur_im2col_step = min(ctx.im2col_step, input.shape[0])
78	            assert (input.shape[0] % cur_im2col_step) == 0, 'im2col step must divide batchsize'
79	
80	            if ctx.needs_input_grad[0] or ctx.needs_input_grad[1]:
81	                grad_input = torch.zeros_like(input)

196	                 deformable_groups=1,
197	                 bias=False):
198	        super(DeformConv, self).__init__()
199	
200	        assert not bias
201	        assert in_channels % groups == 0, \
202	            f'in_channels {in_channels} is not divisible by groups {groups}'
203	        assert out_channels % groups == 0, \

197	                 bias=False):
198	        super(DeformConv, self).__init__()
199	
200	        assert not bias
201	        assert in_channels % groups == 0, \
202	            f'in_channels {in_channels} is not divisible by groups {groups}'
203	        assert out_channels % groups == 0, \
204	            f'out_channels {out_channels} is not divisible ' \
205	            f'by groups {groups}'

199	
200	        assert not bias
201	        assert in_channels % groups == 0, \
202	            f'in_channels {in_channels} is not divisible by groups {groups}'
203	        assert out_channels % groups == 0, \
204	            f'out_channels {out_channels} is not divisible ' \
205	            f'by groups {groups}'
206	
207	        self.in_channels = in_channels
208	        self.out_channels = out_channels

2	
3	from setuptools import find_packages, setup
4	
5	import os
6	import subprocess
7	import sys
8	import time
9	import torch

30	        # LANGUAGE is used on win32
31	        env['LANGUAGE'] = 'C'
32	        env['LANG'] = 'C'
33	        env['LC_ALL'] = 'C'
34	        out = subprocess.Popen(cmd, stdout=subprocess.PIPE, env=env).communicate()[0]
35	        return out
36	
37	    try:

44	
45	    # random seed
46	    seed = opt.get('manual_seed')
47	    if seed is None:
48	        seed = random.randint(1, 10000)
49	        opt['manual_seed'] = seed
50	    set_random_seed(seed + opt['rank'])
51	

59	    logger.info(dict2str(opt))
60	
61	    # initialize wandb logger before tensorboard logger to allow proper sync:
62	    if (opt['logger'].get('wandb') is not None) and (opt['logger']['wandb'].get('project') is not None):
63	        assert opt['logger'].get('use_tb_logger') is True, ('should turn on tensorboard when using wandb')
64	        init_wandb_logger(opt)
65	    tb_logger = None
66	    if opt['logger'].get('use_tb_logger'):

1	# Modified from https://github.com/open-mmlab/mmcv/blob/master/mmcv/runner/dist_utils.py  # noqa: E501
2	import functools
3	import os
4	import subprocess
5	import torch
6	import torch.distributed as dist
7	import torch.multiprocessing as mp
8	

103	        if isinstance(db_paths, list):
104	            self.db_paths = [str(v) for v in db_paths]
105	        elif isinstance(db_paths, str):
106	            self.db_paths = [str(db_paths)]
107	        assert len(client_keys) == len(self.db_paths), ('client_keys and db_paths should have the same length, '
108	                                                        f'but received {len(client_keys)} and {len(self.db_paths)}.')
109	
110	        self._client = {}
111	        for client, path in zip(client_keys, self.db_paths):

118	            filepath (str | obj:`Path`): Here, filepath is the lmdb key.
119	            client_key (str): Used for distinguishing differnet lmdb envs.
120	        """
121	        filepath = str(filepath)
122	        assert client_key in self._client, (f'client_key {client_key} is not ' 'in lmdb clients.')
123	        client = self._client[client_key]
124	        with client.begin(write=False) as txn:
125	            value_buf = txn.get(filepath.encode('ascii'))

57	        map_size (int | None): Map size for lmdb env. If None, use the
58	            estimated size from images. Default: None
59	    """
60	
61	    assert len(img_path_list) == len(keys), ('img_path_list and keys should have the same length, '
62	                                             f'but got {len(img_path_list)} and {len(keys)}')
63	    print(f'Create lmdb for {data_path}, save to {lmdb_path}...')
64	    print(f'Totoal images: {len(img_path_list)}')
65	    if not lmdb_path.endswith('.lmdb'):

35	        self._name = name
36	        self._obj_map = {}
37	
38	    def _do_register(self, name, obj):
39	        assert (name not in self._obj_map), (f"An object named '{name}' was already registered "
40	                                             f"in '{self._name}' registry!")
41	        self._obj_map[name] = obj
42	
43	    def register(self, obj=None):

17	        eps (float): A small value added to the variance to avoid
18	            divide-by-zero. Default: 1e-5.
19	    """
20	    size = feat.size()
21	    assert len(size) == 4, 'The input feature should be 4D tensor.'
22	    b, c = size[:2]
23	    feat_var = feat.view(b, c, -1).var(dim=2) + eps
24	    feat_std = feat_var.sqrt().view(b, c, 1, 1)

36	class SSH(nn.Module):
37	
38	    def __init__(self, in_channel, out_channel):
39	        super(SSH, self).__init__()
40	        assert out_channel % 4 == 0
41	        leaky = 0
42	        if (out_channel <= 64):
43	            leaky = 0.1

188	        m = eval(m) if isinstance(m, str) else m  # eval strings
189	        for j, a in enumerate(args):
190	            try:
191	                args[j] = eval(a) if isinstance(a, str) else a  # eval strings
192	            except:
193	                pass
194	
195	        n = max(round(n * gd), 1) if n > 1 else n  # depth gain
196	        if m in [

29	            self.norm = nn.LayerNorm(normalize_shape)
30	        elif norm_type == 'none':
31	            self.norm = lambda x: x * 1.0
32	        else:
33	            assert 1 == 0, f'Norm type {norm_type} not support.'
34	
35	    def forward(self, x, ref=None):
36	        if self.norm_type == 'spade':

64	            self.func = nn.SELU(True)
65	        elif relu_type == 'none':
66	            self.func = lambda x: x * 1.0
67	        else:
68	            assert 1 == 0, f'Relu type {relu_type} not support.'
69	
70	    def forward(self, x):
71	        return self.func(x)

61	        self.template_3points = template_3points  # improve robustness
62	        self.upscale_factor = upscale_factor
63	        # the cropped face ratio based on the square face
64	        self.crop_ratio = crop_ratio  # (h, w)
65	        assert (self.crop_ratio[0] >= 1 and self.crop_ratio[1] >= 1), 'crop ration only supports >=1'
66	        self.face_size = (int(face_size * self.crop_ratio[1]), int(face_size * self.crop_ratio[0]))
67	
68	        if self.template_3points:

253	    def align_warp_face(self, save_cropped_path=None, border_mode='constant'):
254	        """Align and warp faces with face template.
255	        """
256	        if self.pad_blur:
257	            assert len(self.pad_input_imgs) == len(
258	                self.all_landmarks_5), f'Mismatched samples: {len(self.pad_input_imgs)} and {len(self.all_landmarks_5)}'
259	        for idx, landmark in enumerate(self.all_landmarks_5):
260	            # use 5 landmarks to get affine matrix
261	            # use cv2.LMEDS method for the equivalence to skimage transform

309	            upsample_img = cv2.resize(self.input_img, (w_up, h_up), interpolation=cv2.INTER_LINEAR)
310	        else:
311	            upsample_img = cv2.resize(upsample_img, (w_up, h_up), interpolation=cv2.INTER_LANCZOS4)
312	
313	        assert len(self.restored_faces) == len(
314	            self.inverse_affine_matrices), ('length of restored_faces and affine_matrices are different.')
315	        
316	        inv_mask_borders = []
317	        for restored_face, inverse_affine in zip(self.restored_faces, self.inverse_affine_matrices):