|
1 | 1 | """ |
2 | 2 | Lens blur generator |
3 | 3 | """ |
| 4 | +from typing import Tuple, Dict, List |
4 | 5 | import os |
5 | 6 | import math |
6 | 7 | from functools import reduce |
|
50 | 51 | [2.201904, 19.032909, -0.152784, -0.107988]]] |
51 | 52 |
|
52 | 53 | # Obtain specific parameters and scale for a given component count |
53 | | -def get_parameters(component_count = 2): |
| 54 | +def get_parameters(component_count: int = 2) -> Tuple[List[Dict[str, float]], float]: |
| 55 | + """ |
| 56 | + Obtain specific parameters and scale for a given component count. |
| 57 | + """ |
54 | 58 | parameter_index = max(0, min(component_count - 1, len(kernel_params))) |
55 | | - parameter_dictionaries = [dict(zip(['a','b','A','B'], b)) for b in kernel_params[parameter_index]] |
56 | | - return (parameter_dictionaries, kernel_scales[parameter_index]) |
| 59 | + parameter_dictionaries = [dict(zip(['a', 'b', 'A', 'B'], b)) for b in kernel_params[parameter_index]] |
| 60 | + return parameter_dictionaries, kernel_scales[parameter_index] |
57 | 61 |
|
58 | 62 | # Produces a complex kernel of a given radius and scale (adjusts radius to be more accurate) |
59 | 63 | # a and b are parameters of this complex kernel |
60 | | -def complex_kernel_1d(radius, scale, a, b): |
61 | | - kernel_radius = radius |
| 64 | +def complex_kernel_1d(radius: float, scale: float, a: float, b: float) -> np.ndarray: |
| 65 | + """ |
| 66 | + Produces a complex kernel of a given radius and scale (adjusts radius to be more accurate). |
| 67 | + """ |
| 68 | + kernel_radius = int(math.ceil(radius)) |
62 | 69 | kernel_size = kernel_radius * 2 + 1 |
63 | | - ax = np.arange(-kernel_radius, kernel_radius + 1., dtype=np.float32) |
64 | | - ax = ax * scale * (1 / kernel_radius) |
65 | | - kernel_complex = np.zeros((kernel_size), dtype=np.complex64) |
| 70 | + ax = np.linspace(-radius, radius, kernel_size, dtype=np.float32) |
| 71 | + ax = ax * scale * (1 / radius) |
| 72 | + kernel_complex = np.zeros((kernel_size,), dtype=np.complex64) |
66 | 73 | kernel_complex.real = np.exp(-a * (ax**2)) * np.cos(b * (ax**2)) |
67 | 74 | kernel_complex.imag = np.exp(-a * (ax**2)) * np.sin(b * (ax**2)) |
68 | 75 | return kernel_complex.reshape((1, kernel_size)) |
69 | 76 |
|
70 | | -def normalise_kernels(kernels, params): |
71 | | - # Normalises with respect to A*real+B*imag |
| 77 | + |
| 78 | +def normalise_kernels(kernels: List[np.ndarray], params: List[Dict[str, float]]) -> np.ndarray: |
| 79 | + """ |
| 80 | + Normalises the kernels with respect to A*real + B*imag. |
| 81 | + """ |
72 | 82 | total = 0 |
73 | 83 |
|
74 | | - for k,p in zip(kernels, params): |
75 | | - # 1D kernel - applied in 2D |
| 84 | + for k, p in zip(kernels, params): |
76 | 85 | for i in range(k.shape[1]): |
77 | 86 | for j in range(k.shape[1]): |
78 | | - # Complex multiply and weighted sum |
79 | | - total += p['A'] * (k[0,i].real*k[0,j].real - k[0,i].imag*k[0,j].imag) + p['B'] * (k[0,i].real*k[0,j].imag + k[0,i].imag*k[0,j].real) |
| 87 | + total += p['A'] * (k[0, i].real * k[0, j].real - k[0, i].imag * k[0, j].imag) + \ |
| 88 | + p['B'] * (k[0, i].real * k[0, j].imag + k[0, i].imag * k[0, j].real) |
80 | 89 |
|
81 | 90 | scalar = 1 / math.sqrt(total) |
82 | 91 | kernels = np.asarray(kernels) * scalar |
83 | 92 |
|
84 | 93 | return kernels |
85 | 94 |
|
86 | | -# Combine the real and imaginary parts of an image, weighted by A and B |
87 | | -def weighted_sum(kernel, params): |
| 95 | +def weighted_sum(kernel: np.ndarray, params: Dict[str, float]) -> np.ndarray: |
| 96 | + """ |
| 97 | + Combine the real and imaginary parts of an image, weighted by A and B. |
| 98 | + """ |
88 | 99 | return np.add(kernel.real * params['A'], kernel.imag * params['B']) |
89 | 100 |
|
90 | 101 | # Produce a 2D kernel by self-multiplying a 1d kernel. This would be slower to use |
91 | 102 | # than the separable approach, mostly for visualisation below |
92 | | -def multiply_kernel(kernel): |
| 103 | +def multiply_kernel(kernel: np.ndarray) -> np.ndarray: |
| 104 | + """ |
| 105 | + Produce a 2D kernel by self-multiplying a 1D kernel. |
| 106 | + """ |
93 | 107 | kernel_size = kernel.shape[1] |
94 | 108 | a = np.repeat(kernel, kernel_size, 0) |
95 | 109 | b = np.repeat(kernel.transpose(), kernel_size, 1) |
96 | | - return np.multiply(a,b) |
| 110 | + return np.multiply(a, b) |
97 | 111 |
|
98 | 112 | # ---------------------------------------------------------------- |
99 | 113 |
|
100 | | -def filter_task(idx, channel, img_channel, component, component_params): |
| 114 | +def filter_task(idx: int, channel: int, img_channel: np.ndarray, component: np.ndarray, component_params: Dict[str, float]) -> Tuple[int, int, np.ndarray]: |
101 | 115 | """ |
102 | | - https://github.com/Davide-sd/GIMP-lens-blur/blob/master/GIMP-lens-blur.py#L188 |
| 116 | + Perform convolution with the complex kernel components on the image channel. |
103 | 117 | """ |
104 | 118 | component_real = np.real(component) |
105 | 119 | component_imag = np.imag(component) |
106 | 120 |
|
107 | 121 | component_real_t = component_real.transpose() |
108 | 122 | component_imag_t = component_imag.transpose() |
109 | | - # first convolution |
110 | | - inter_real = cv2.filter2D(img_channel, -1, component_real) |
111 | | - inter_imag = cv2.filter2D(img_channel, -1, component_imag) |
112 | | - # second convolution (see NOTE above, here inter_ is f, component_ is g) |
113 | | - final_1 = cv2.filter2D(inter_real, -1, component_real_t) |
114 | | - final_2 = cv2.filter2D(inter_real, -1, component_imag_t) |
115 | | - final_3 = cv2.filter2D(inter_imag, -1, component_real_t) |
116 | | - final_4 = cv2.filter2D(inter_imag, -1, component_imag_t) |
| 123 | + |
| 124 | + inter_real = cv2.filter2D(img_channel, -1, component_real, borderType=cv2.BORDER_REPLICATE) |
| 125 | + inter_imag = cv2.filter2D(img_channel, -1, component_imag, borderType=cv2.BORDER_REPLICATE) |
| 126 | + |
| 127 | + final_1 = cv2.filter2D(inter_real, -1, component_real_t, borderType=cv2.BORDER_REPLICATE) |
| 128 | + final_2 = cv2.filter2D(inter_real, -1, component_imag_t, borderType=cv2.BORDER_REPLICATE) |
| 129 | + final_3 = cv2.filter2D(inter_imag, -1, component_real_t, borderType=cv2.BORDER_REPLICATE) |
| 130 | + final_4 = cv2.filter2D(inter_imag, -1, component_imag_t, borderType=cv2.BORDER_REPLICATE) |
| 131 | + |
117 | 132 | final = final_1 - final_4 + 1j * (final_2 + final_3) |
118 | 133 | weight_sum = weighted_sum(final, component_params) |
119 | | - # return index, channel No. and sum of weights |
| 134 | + |
120 | 135 | return idx, channel, weight_sum |
121 | 136 |
|
122 | | -def lens_blur(img, radius=3, components=5, exposure_gamma=5): |
| 137 | +def lens_blur(img: np.ndarray, radius: float = 3.0, components: int = 5, exposure_gamma: float = 5.0) -> np.ndarray: |
| 138 | + """ |
| 139 | + Apply lens blur to the input image. |
| 140 | + """ |
123 | 141 |
|
124 | 142 | img = img/255. |
125 | 143 |
|
|
0 commit comments