sczhou/CodeFormer
2.4k
1import math2import numpy as np3import torch4 5 6def cubic(x):7 """cubic function used for calculate_weights_indices."""8 absx = torch.abs(x)9 absx2 = absx**210 absx3 = absx**311 return (1.5 * absx3 - 2.5 * absx2 + 1) * (12 (absx <= 1).type_as(absx)) + (-0.5 * absx3 + 2.5 * absx2 - 4 * absx + 2) * (((absx > 1) *13 (absx <= 2)).type_as(absx))14 15 16def calculate_weights_indices(in_length, out_length, scale, kernel, kernel_width, antialiasing):17 """Calculate weights and indices, used for imresize function.18 19 Args:20 in_length (int): Input length.21 out_length (int): Output length.22 scale (float): Scale factor.23 kernel_width (int): Kernel width.24 antialisaing (bool): Whether to apply anti-aliasing when downsampling.25 """26 27 if (scale < 1) and antialiasing:28 # Use a modified kernel (larger kernel width) to simultaneously29 # interpolate and antialias30 kernel_width = kernel_width / scale31 32 # Output-space coordinates33 x = torch.linspace(1, out_length, out_length)34 35 # Input-space coordinates. Calculate the inverse mapping such that 0.536 # in output space maps to 0.5 in input space, and 0.5 + scale in output37 # space maps to 1.5 in input space.38 u = x / scale + 0.5 * (1 - 1 / scale)39 40 # What is the left-most pixel that can be involved in the computation?41 left = torch.floor(u - kernel_width / 2)42 43 # What is the maximum number of pixels that can be involved in the44 # computation? Note: it's OK to use an extra pixel here; if the45 # corresponding weights are all zero, it will be eliminated at the end46 # of this function.47 p = math.ceil(kernel_width) + 248 49 # The indices of the input pixels involved in computing the k-th output50 # pixel are in row k of the indices matrix.51 indices = left.view(out_length, 1).expand(out_length, p) + torch.linspace(0, p - 1, p).view(1, p).expand(52 out_length, p)53 54 # The weights used to compute the k-th output pixel are in row k of the55 # weights matrix.56 distance_to_center = u.view(out_length, 1).expand(out_length, p) - indices57 58 # apply cubic kernel59 if (scale < 1) and antialiasing:60 weights = scale * cubic(distance_to_center * scale)61 else:62 weights = cubic(distance_to_center)63 64 # Normalize the weights matrix so that each row sums to 1.65 weights_sum = torch.sum(weights, 1).view(out_length, 1)66 weights = weights / weights_sum.expand(out_length, p)67 68 # If a column in weights is all zero, get rid of it. only consider the69 # first and last column.70 weights_zero_tmp = torch.sum((weights == 0), 0)71 if not math.isclose(weights_zero_tmp[0], 0, rel_tol=1e-6):72 indices = indices.narrow(1, 1, p - 2)73 weights = weights.narrow(1, 1, p - 2)74 if not math.isclose(weights_zero_tmp[-1], 0, rel_tol=1e-6):75 indices = indices.narrow(1, 0, p - 2)76 weights = weights.narrow(1, 0, p - 2)77 weights = weights.contiguous()78 indices = indices.contiguous()79 sym_len_s = -indices.min() + 180 sym_len_e = indices.max() - in_length81 indices = indices + sym_len_s - 182 return weights, indices, int(sym_len_s), int(sym_len_e)83 84 85@torch.no_grad()86def imresize(img, scale, antialiasing=True):87 """imresize function same as MATLAB.88 89 It now only supports bicubic.90 The same scale applies for both height and width.91 92 Args:93 img (Tensor | Numpy array):94 Tensor: Input image with shape (c, h, w), [0, 1] range.95 Numpy: Input image with shape (h, w, c), [0, 1] range.96 scale (float): Scale factor. The same scale applies for both height97 and width.98 antialisaing (bool): Whether to apply anti-aliasing when downsampling.99 Default: True.100 101 Returns:102 Tensor: Output image with shape (c, h, w), [0, 1] range, w/o round.103 """104 if type(img).__module__ == np.__name__: # numpy type105 numpy_type = True106 img = torch.from_numpy(img.transpose(2, 0, 1)).float()107 else:108 numpy_type = False109 110 in_c, in_h, in_w = img.size()111 out_h, out_w = math.ceil(in_h * scale), math.ceil(in_w * scale)112 kernel_width = 4113 kernel = 'cubic'114 115 # get weights and indices116 weights_h, indices_h, sym_len_hs, sym_len_he = calculate_weights_indices(in_h, out_h, scale, kernel, kernel_width,117 antialiasing)118 weights_w, indices_w, sym_len_ws, sym_len_we = calculate_weights_indices(in_w, out_w, scale, kernel, kernel_width,119 antialiasing)120 # process H dimension121 # symmetric copying122 img_aug = torch.FloatTensor(in_c, in_h + sym_len_hs + sym_len_he, in_w)123 img_aug.narrow(1, sym_len_hs, in_h).copy_(img)124 125 sym_patch = img[:, :sym_len_hs, :]126 inv_idx = torch.arange(sym_patch.size(1) - 1, -1, -1).long()127 sym_patch_inv = sym_patch.index_select(1, inv_idx)128 img_aug.narrow(1, 0, sym_len_hs).copy_(sym_patch_inv)129 130 sym_patch = img[:, -sym_len_he:, :]131 inv_idx = torch.arange(sym_patch.size(1) - 1, -1, -1).long()132 sym_patch_inv = sym_patch.index_select(1, inv_idx)133 img_aug.narrow(1, sym_len_hs + in_h, sym_len_he).copy_(sym_patch_inv)134 135 out_1 = torch.FloatTensor(in_c, out_h, in_w)136 kernel_width = weights_h.size(1)137 for i in range(out_h):138 idx = int(indices_h[i][0])139 for j in range(in_c):140 out_1[j, i, :] = img_aug[j, idx:idx + kernel_width, :].transpose(0, 1).mv(weights_h[i])141 142 # process W dimension143 # symmetric copying144 out_1_aug = torch.FloatTensor(in_c, out_h, in_w + sym_len_ws + sym_len_we)145 out_1_aug.narrow(2, sym_len_ws, in_w).copy_(out_1)146 147 sym_patch = out_1[:, :, :sym_len_ws]148 inv_idx = torch.arange(sym_patch.size(2) - 1, -1, -1).long()149 sym_patch_inv = sym_patch.index_select(2, inv_idx)150 out_1_aug.narrow(2, 0, sym_len_ws).copy_(sym_patch_inv)151 152 sym_patch = out_1[:, :, -sym_len_we:]153 inv_idx = torch.arange(sym_patch.size(2) - 1, -1, -1).long()154 sym_patch_inv = sym_patch.index_select(2, inv_idx)155 out_1_aug.narrow(2, sym_len_ws + in_w, sym_len_we).copy_(sym_patch_inv)156 157 out_2 = torch.FloatTensor(in_c, out_h, out_w)158 kernel_width = weights_w.size(1)159 for i in range(out_w):160 idx = int(indices_w[i][0])161 for j in range(in_c):162 out_2[j, :, i] = out_1_aug[j, :, idx:idx + kernel_width].mv(weights_w[i])163 164 if numpy_type:165 out_2 = out_2.numpy().transpose(1, 2, 0)166 return out_2167 168 169def rgb2ycbcr(img, y_only=False):170 """Convert a RGB image to YCbCr image.171 172 This function produces the same results as Matlab's `rgb2ycbcr` function.173 It implements the ITU-R BT.601 conversion for standard-definition174 television. See more details in175 https://en.wikipedia.org/wiki/YCbCr#ITU-R_BT.601_conversion.176 177 It differs from a similar function in cv2.cvtColor: `RGB <-> YCrCb`.178 In OpenCV, it implements a JPEG conversion. See more details in179 https://en.wikipedia.org/wiki/YCbCr#JPEG_conversion.180 181 Args:182 img (ndarray): The input image. It accepts:183 1. np.uint8 type with range [0, 255];184 2. np.float32 type with range [0, 1].185 y_only (bool): Whether to only return Y channel. Default: False.186 187 Returns:188 ndarray: The converted YCbCr image. The output image has the same type189 and range as input image.190 """191 img_type = img.dtype192 img = _convert_input_type_range(img)193 if y_only:194 out_img = np.dot(img, [65.481, 128.553, 24.966]) + 16.0195 else:196 out_img = np.matmul(197 img, [[65.481, -37.797, 112.0], [128.553, -74.203, -93.786], [24.966, 112.0, -18.214]]) + [16, 128, 128]198 out_img = _convert_output_type_range(out_img, img_type)199 return out_img200 201 202def bgr2ycbcr(img, y_only=False):203 """Convert a BGR image to YCbCr image.204 205 The bgr version of rgb2ycbcr.206 It implements the ITU-R BT.601 conversion for standard-definition207 television. See more details in208 https://en.wikipedia.org/wiki/YCbCr#ITU-R_BT.601_conversion.209 210 It differs from a similar function in cv2.cvtColor: `BGR <-> YCrCb`.211 In OpenCV, it implements a JPEG conversion. See more details in212 https://en.wikipedia.org/wiki/YCbCr#JPEG_conversion.213 214 Args:215 img (ndarray): The input image. It accepts:216 1. np.uint8 type with range [0, 255];217 2. np.float32 type with range [0, 1].218 y_only (bool): Whether to only return Y channel. Default: False.219 220 Returns:221 ndarray: The converted YCbCr image. The output image has the same type222 and range as input image.223 """224 img_type = img.dtype225 img = _convert_input_type_range(img)226 if y_only:227 out_img = np.dot(img, [24.966, 128.553, 65.481]) + 16.0228 else:229 out_img = np.matmul(230 img, [[24.966, 112.0, -18.214], [128.553, -74.203, -93.786], [65.481, -37.797, 112.0]]) + [16, 128, 128]231 out_img = _convert_output_type_range(out_img, img_type)232 return out_img233 234 235def ycbcr2rgb(img):236 """Convert a YCbCr image to RGB image.237 238 This function produces the same results as Matlab's ycbcr2rgb function.239 It implements the ITU-R BT.601 conversion for standard-definition240 television. See more details in241 https://en.wikipedia.org/wiki/YCbCr#ITU-R_BT.601_conversion.242 243 It differs from a similar function in cv2.cvtColor: `YCrCb <-> RGB`.244 In OpenCV, it implements a JPEG conversion. See more details in245 https://en.wikipedia.org/wiki/YCbCr#JPEG_conversion.246 247 Args:248 img (ndarray): The input image. It accepts:249 1. np.uint8 type with range [0, 255];250 2. np.float32 type with range [0, 1].251 252 Returns:253 ndarray: The converted RGB image. The output image has the same type254 and range as input image.255 """256 img_type = img.dtype257 img = _convert_input_type_range(img) * 255258 out_img = np.matmul(img, [[0.00456621, 0.00456621, 0.00456621], [0, -0.00153632, 0.00791071],259 [0.00625893, -0.00318811, 0]]) * 255.0 + [-222.921, 135.576, -276.836] # noqa: E126260 out_img = _convert_output_type_range(out_img, img_type)261 return out_img262 263 264def ycbcr2bgr(img):265 """Convert a YCbCr image to BGR image.266 267 The bgr version of ycbcr2rgb.268 It implements the ITU-R BT.601 conversion for standard-definition269 television. See more details in270 https://en.wikipedia.org/wiki/YCbCr#ITU-R_BT.601_conversion.271 272 It differs from a similar function in cv2.cvtColor: `YCrCb <-> BGR`.273 In OpenCV, it implements a JPEG conversion. See more details in274 https://en.wikipedia.org/wiki/YCbCr#JPEG_conversion.275 276 Args:277 img (ndarray): The input image. It accepts:278 1. np.uint8 type with range [0, 255];279 2. np.float32 type with range [0, 1].280 281 Returns:282 ndarray: The converted BGR image. The output image has the same type283 and range as input image.284 """285 img_type = img.dtype286 img = _convert_input_type_range(img) * 255287 out_img = np.matmul(img, [[0.00456621, 0.00456621, 0.00456621], [0.00791071, -0.00153632, 0],288 [0, -0.00318811, 0.00625893]]) * 255.0 + [-276.836, 135.576, -222.921] # noqa: E126289 out_img = _convert_output_type_range(out_img, img_type)290 return out_img291 292 293def _convert_input_type_range(img):294 """Convert the type and range of the input image.295 296 It converts the input image to np.float32 type and range of [0, 1].297 It is mainly used for pre-processing the input image in colorspace298 convertion functions such as rgb2ycbcr and ycbcr2rgb.299 300 Args:301 img (ndarray): The input image. It accepts:302 1. np.uint8 type with range [0, 255];303 2. np.float32 type with range [0, 1].304 305 Returns:306 (ndarray): The converted image with type of np.float32 and range of307 [0, 1].308 """309 img_type = img.dtype310 img = img.astype(np.float32)311 if img_type == np.float32:312 pass313 elif img_type == np.uint8:314 img /= 255.315 else:316 raise TypeError('The img type should be np.float32 or np.uint8, ' f'but got {img_type}')317 return img318 319 320def _convert_output_type_range(img, dst_type):321 """Convert the type and range of the image according to dst_type.322 323 It converts the image to desired type and range. If `dst_type` is np.uint8,324 images will be converted to np.uint8 type with range [0, 255]. If325 `dst_type` is np.float32, it converts the image to np.float32 type with326 range [0, 1].327 It is mainly used for post-processing images in colorspace convertion328 functions such as rgb2ycbcr and ycbcr2rgb.329 330 Args:331 img (ndarray): The image to be converted with np.float32 type and332 range [0, 255].333 dst_type (np.uint8 | np.float32): If dst_type is np.uint8, it334 converts the image to np.uint8 type with range [0, 255]. If335 dst_type is np.float32, it converts the image to np.float32 type336 with range [0, 1].337 338 Returns:339 (ndarray): The converted image with desired type and range.340 """341 if dst_type not in (np.uint8, np.float32):342 raise TypeError('The dst_type should be np.float32 or np.uint8, ' f'but got {dst_type}')343 if dst_type == np.uint8:344 img = img.round()345 else:346 img /= 255.347 return img.astype(dst_type)348