generate_upscaling_coefficients.py (8070B)
1 #!/usr/bin/env python3 2 3 # Copyright (c) the JPEG XL Project Authors. All rights reserved. 4 # 5 # Use of this source code is governed by a BSD-style 6 # license that can be found in the LICENSE file. 7 8 """Generates coefficients used in upscaling. 9 10 Given an upscaling factor which can be 2, 4 or 8, we generate coefficients and 11 indices for lib/jxl/image_metadata.cc in the format needed there. 12 """ 13 14 import argparse 15 import itertools 16 import numpy as np 17 18 19 def compute_kernel(sigma): 20 """Gaussian-like kernel with standard deviation sigma.""" 21 # This controls the length of the kernel. 22 m = 2.5 23 diff = int(max(1, m * abs(sigma))) 24 kernel = np.exp(-np.arange(-diff, diff + 1)**2 /(2 * sigma * sigma)) 25 return kernel 26 27 28 def convolution(pixels, kernel): 29 """Computes a horizontal convolution and transposes the result.""" 30 y, x = pixels.shape 31 kernel_len = len(kernel) 32 offset = kernel_len // 2 33 scale = 1 / sum(kernel) 34 out_pixels = np.zeros(shape=(x, y), dtype=pixels.dtype) 35 for i, j in itertools.product(range(x), range(y)): 36 if kernel_len < i < x - kernel_len: 37 out_pixels[i, j] = scale * sum( 38 pixels[j, i - offset + k] * kernel[k] for k in range(kernel_len)) 39 else: 40 out_pixels[i, j] = pixels[j, i] 41 return out_pixels 42 43 44 def _super_sample(pixels, n): 45 return np.repeat(np.repeat(pixels, n, axis=0), n, axis=1) 46 47 48 def _sub_sample(pixels, n): 49 x, y = pixels.shape 50 assert x%n == 0 and y%n == 0 51 return 1 / (n * n) * pixels.reshape(x // n, n, y // n, n).transpose( 52 [0, 2, 1, 3]).sum(axis=(2, 3)) 53 54 55 def smooth_4x4_corners(pixels): 56 """Generates a 4x4 upscaled image, to be smoothed afterwards.""" 57 overshoot = 3.5 58 m = 1.0 / (4.0 - overshoot) 59 y_size, x_size = pixels.shape 60 for y, x in itertools.product(range(3, y_size - 3, 4), 61 range(3, x_size - 3, 4)): 62 ave = ( 63 pixels[y, x] + pixels[y, x + 1] + pixels[y + 1, x] + 64 pixels[y + 1, x + 1]) 65 off = 2 66 other = (ave - overshoot * pixels[y, x]) * m 67 pixels[y - off, x - off] -= (other - pixels[y, x]) 68 pixels[y, x] = other 69 70 other = (ave - overshoot * pixels[y, x + 1]) * m 71 pixels[y - off, x + off + 1] -= (other - pixels[y, x + 1]) 72 pixels[y, x + 1] = other 73 74 other = (ave - overshoot * pixels[y + 1, x]) * m 75 pixels[y + off + 1, x - off] -= (other - pixels[y + 1, x]) 76 pixels[y + 1, x] = other 77 78 other = (ave - overshoot * pixels[y + 1, x + 1]) * m 79 pixels[y + off + 1][x + off + 1] -= (other - pixels[y + 1, x + 1]) 80 pixels[y + 1, x + 1] = other 81 82 return pixels 83 84 85 def smoothing(pixels): 86 new_pixels = smooth_4x4_corners(_super_sample(pixels, 4)) 87 my_kernel = compute_kernel(2.5) 88 smooth_image = convolution(convolution(new_pixels, my_kernel), my_kernel) 89 return smooth_image 90 91 92 upscaling = { 93 2: lambda pixels: _sub_sample(smoothing(pixels), 2), 94 4: smoothing, 95 8: lambda pixels: _sub_sample(smoothing(smoothing(pixels)), 2) 96 } 97 98 99 def get_coeffs(upscaling_factor, kernel_size=5, normalized=True, dtype="float"): 100 """Returns 4-tensor of coefficients. 101 102 Args: 103 upscaling_factor: 2, 4, or 8 104 kernel_size: must be odd 105 normalized: if True, the kernel matrix adds to 1 106 dtype: type of numpy array to return 107 108 Returns: 109 A (upscaling_factor x upscaling_factor) matrix of 110 (kernel_size x kernel_size) matrices, describing the kernel for all pixels. 111 """ 112 113 upscaling_method = upscaling[upscaling_factor] 114 patch_size = 2 * kernel_size + 1 115 matrix_bases = np.eye( 116 patch_size * patch_size, dtype=dtype).reshape(patch_size, patch_size, 117 patch_size, patch_size) 118 119 # takes some time... 120 smoothed_bases = np.array( 121 [[upscaling_method(matrix_bases[a, b]) 122 for a in range(patch_size)] 123 for b in range(patch_size)]) 124 125 middle = patch_size // 2 126 lower = middle - kernel_size // 2 127 upper = middle + kernel_size // 2 + 1 128 assert len(range(lower, upper)) == kernel_size 129 assert sum(range(lower, upper)) == kernel_size * middle 130 131 coefficients = np.array([[[[ 132 smoothed_bases[i, j, upscaling_factor * middle + b, 133 upscaling_factor * middle + a] 134 for i in range(lower, upper) 135 ] 136 for j in range(lower, upper)] 137 for a in range(upscaling_factor)] 138 for b in range(upscaling_factor)]) 139 140 if normalized: 141 return coefficients / coefficients.sum(axis=(2, 3))[..., np.newaxis, 142 np.newaxis] 143 else: 144 return coefficients 145 146 147 def indices_matrix(upscaling_factor, kernel_size=5): 148 """Matrix containing indices with all symmetries.""" 149 matrix = np.zeros( 150 shape=[upscaling_factor * kernel_size] * 2, dtype="int16") 151 # define a fundamental domain 152 counter = 1 153 for i in range((kernel_size * upscaling_factor) // 2): 154 for j in range(i, (kernel_size * upscaling_factor) // 2): 155 matrix[i, j] = counter 156 counter += 1 157 158 matrix_with_transpose = matrix + (matrix.transpose()) * ( 159 matrix != matrix.transpose()) 160 matrix_vertical = matrix_with_transpose + ( 161 np.flip(matrix_with_transpose, axis=0) * 162 (matrix_with_transpose != np.flip(matrix_with_transpose, axis=0))) 163 matrix_horizontal = matrix_vertical + ( 164 np.flip(matrix_vertical, axis=1) * 165 (matrix_vertical != np.flip(matrix_vertical, axis=1))) - 1 166 return matrix_horizontal 167 168 169 def format_indices_matrix(upscaling_factor, kernel_size=5): 170 """Returns string of commented out numbers-only matrices.""" 171 indices = indices_matrix(upscaling_factor) 172 output_str = [] 173 for i in range(upscaling_factor // 2): 174 for j in range(kernel_size): 175 output_str.append("//") 176 for a in range(upscaling_factor // 2): 177 for b in range(kernel_size): 178 output_str.append( 179 f"{'{:x}'.format(int(indices[kernel_size*i + j][kernel_size*a + b])).rjust(2)} " 180 ) 181 output_str.append(" ") 182 output_str.append("\n") 183 output_str.append("\n") 184 return "".join(output_str) 185 186 187 def weights_arrays(upscaling_factor, kernel_size=5): 188 """Returns string describing array of depth 4.""" 189 indices = indices_matrix(upscaling_factor) 190 return ( 191 f"kernel[{upscaling_factor}][{upscaling_factor}][{kernel_size}][{kernel_size}]" 192 f" = {{" + ", \n".join("{\n" + ", \n\n".join( 193 ("{" + ", \n".join("{" + ", ".join( 194 f"weights[{str(indices[kernel_size*i + j][kernel_size*a + b])}]" 195 for b in range(kernel_size)) + "}" 196 for j in range(kernel_size)) + "}" 197 for a in range(upscaling_factor // 2))) + "\n}" 198 for i in range(upscaling_factor // 2)) + "}\n") 199 200 201 def coefficients_list(upscaling_factor, kernel_size=5): 202 """Returns string describing coefficients.""" 203 coeff_tensor = get_coeffs(upscaling_factor, 204 kernel_size).transpose([0, 2, 1, 3]).reshape( 205 kernel_size * upscaling_factor, 206 kernel_size * upscaling_factor) 207 my_weights = [ 208 f'{"{:.8f}".format(coeff_tensor[i][j])}f' 209 for i in range((kernel_size * upscaling_factor) // 2) 210 for j in range(i, (kernel_size * upscaling_factor) // 2) 211 ] 212 return f"kWeights{upscaling_factor} = {{" + ", ".join(my_weights) + "};" 213 214 215 def print_all_output(upscaling_factor): 216 print(format_indices_matrix(upscaling_factor)) 217 print(coefficients_list(upscaling_factor), end="\n\n") 218 print(weights_arrays(upscaling_factor)) 219 220 221 def main(): 222 parser = argparse.ArgumentParser( 223 description="Generates coefficients used in upscaling.") 224 parser.add_argument( 225 "upscaling_factor", 226 type=int, 227 help="upscaling factor, must be 2, 4 or 8.", 228 nargs="?", 229 default=None) 230 231 args = parser.parse_args() 232 upscaling_factor = args.upscaling_factor 233 if upscaling_factor: 234 print_all_output(upscaling_factor) 235 else: 236 for factor in [2, 4, 8]: 237 print(f"upscaling factor = {factor}") 238 print_all_output(factor) 239 240 241 if __name__ == "__main__": 242 main()