Skip to content

Commit 3aacf84

Browse files
alexlibalexlib
committed
more numba
1 parent ed9b0e9 commit 3aacf84

2 files changed

Lines changed: 82 additions & 48 deletions

File tree

‎openptv_python/trafo.py‎

Lines changed: 39 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,9 @@
11
"""Module for coordinate transformations."""
2-
from math import cos, sin, sqrt
32
from typing import Tuple
43

54
import numpy as np
6-
from numba import njit
5+
from numba import float64, int32, njit
6+
from numpy import cos, sin, sqrt
77

88
from .calibration import Calibration, ap_52
99
from .parameters import ControlPar
@@ -38,22 +38,29 @@ def fast_pixel_to_metric(x_pixel, y_pixel, imx, imy, pix_x, pix_y) -> Tuple[floa
3838

3939
return (x_metric, y_metric)
4040

41-
42-
def arr_pixel_to_metric(pixel: np.ndarray, parameters: ControlPar) -> np.ndarray:
41+
@njit(float64[:,:](int32[:,:],int32,int32,float64,float64))
42+
def arr_pixel_to_metric(pixel: np.ndarray,
43+
imx: int,
44+
imy: int,
45+
pix_x: float,
46+
pix_y: float) -> np.ndarray:
4347
"""Convert pixel coordinates to metric coordinates.
4448
4549
Arguments:
4650
---------
51+
imx (float): image width in pixels.
52+
imy (float): image height in pixels.
53+
pix_x (float): pixel size in x-direction.
54+
pix_y (float): pixel size in y-direction.
55+
56+
Returns
57+
-------
4758
metric (np.ndarray): output metric coordinates.
48-
pixel (np.ndarray): input pixel coordinates.
49-
parameters (ControlPar): control structure holding image and pixel sizes.
59+
5060
"""
51-
pixel = np.atleast_2d(np.array(pixel))
52-
metric = np.empty_like(pixel)
53-
metric[:, 0] = (pixel[:, 0] - float(parameters.imx) /
54-
2.0) * parameters.pix_x
55-
metric[:, 1] = (float(parameters.imy) / 2.0 -
56-
pixel[:, 1]) * parameters.pix_y
61+
metric = np.empty_like(pixel, dtype=np.float64)
62+
metric[:, 0] = (pixel[:, 0] - imx / 2.0) * pix_x
63+
metric[:, 1] = (imy / 2.0 - pixel[:, 1]) * pix_y
5764

5865
return metric
5966

@@ -81,6 +88,7 @@ def metric_to_pixel(
8188
parameters.pix_y
8289
)
8390

91+
8492
@njit
8593
def fast_metric_to_pixel(
8694
x_metric,
@@ -89,15 +97,16 @@ def fast_metric_to_pixel(
8997
imy,
9098
pix_x,
9199
pix_y
92-
) -> Tuple[float, float]:
100+
) -> Tuple[float, float]:
93101
"""Convert metric coordinates to pixel coordinates."""
94102
x_pixel = (x_metric / pix_x) + (float(imx) / 2.0)
95103
y_pixel = (float(imy) / 2.0) - (y_metric / pix_y)
96104

97105
return x_pixel, y_pixel
98106

99107

100-
def arr_metric_to_pixel(metric: np.ndarray, parameters: ControlPar) -> np.ndarray:
108+
def arr_metric_to_pixel(metric: np.ndarray,
109+
parameters: ControlPar) -> np.ndarray:
101110
"""Convert an array of metric coordinates to pixel coordinates.
102111
103112
Arguments:
@@ -109,7 +118,7 @@ def arr_metric_to_pixel(metric: np.ndarray, parameters: ControlPar) -> np.ndarra
109118
-------
110119
pixel (np.ndarray): output array of pixel coordinates.
111120
"""
112-
metric = np.atleast_2d(np.array(metric))
121+
metric = np.atleast_2d(metric)
113122

114123
return fast_arr_metric_to_pixel(
115124
metric,
@@ -119,18 +128,19 @@ def arr_metric_to_pixel(metric: np.ndarray, parameters: ControlPar) -> np.ndarra
119128
parameters.pix_y
120129
)
121130

122-
@njit
131+
132+
@njit(float64[:,:](float64[:,:],int32,int32,float64,float64))
123133
def fast_arr_metric_to_pixel(
124-
metric,
125-
imx,
126-
imy,
127-
pix_x,
128-
pix_y
129-
) -> np.ndarray:
134+
metric: np.ndarray,
135+
imx: int,
136+
imy: int,
137+
pix_x: float,
138+
pix_y: float
139+
) -> np.ndarray:
130140
"""Convert an array of metric coordinates to pixel coordinates."""
131141
pixel = np.zeros_like(metric)
132-
pixel[:, 0] = (metric[:, 0] / pix_x) + (float(imx) / 2.0)
133-
pixel[:, 1] = (float(imy) / 2.0) - (metric[:, 1] / pix_y)
142+
pixel[:, 0] = (metric[:, 0] / pix_x) + (imx / 2.0)
143+
pixel[:, 1] = (imy / 2.0) - (metric[:, 1] / pix_y)
134144

135145
return pixel
136146

@@ -143,13 +153,14 @@ def distort_brown_affine(x: float,
143153
if x == 0 and y == 0:
144154
return 0, 0
145155

146-
return fast_distort_brown_affine(x, y, ap.k1, ap.k2, ap.k3,
156+
tmp = fast_distort_brown_affine(x, y, ap.k1, ap.k2, ap.k3,
147157
ap.p1, ap.p2, ap.she, ap.scx)
158+
return tmp[0], tmp[1]
148159

149160
# print(f"x {x}, y {y}")
150161

151162

152-
@njit
163+
@njit(float64[:](float64,float64,float64,float64,float64,float64,float64,float64,float64))
153164
def fast_distort_brown_affine(
154165
x: float,
155166
y: float,
@@ -160,7 +171,7 @@ def fast_distort_brown_affine(
160171
p2: float,
161172
she: float,
162173
scx: float,
163-
) -> Tuple[float, float]:
174+
) -> np.ndarray:
164175
"""Distort a point using the Brown affine model."""
165176
r = sqrt(x**2 + y**2)
166177

@@ -183,7 +194,7 @@ def fast_distort_brown_affine(
183194

184195
# print(f"x1 {x1}, y1 {y1}")
185196

186-
return x1, y1
197+
return np.array([x1, y1])
187198

188199

189200
def correct_brown_affine(

‎tests/test_trafo.py‎

Lines changed: 43 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -224,7 +224,13 @@ def test_pixel_to_metric(self):
224224
y_metric_expected = (float(parameters.imy) / 2.0 - y_pixel) * parameters.pix_y
225225

226226
# call the function to get actual output metric coordinates
227-
metric = arr_pixel_to_metric(np.array([x_pixel, y_pixel]), parameters)
227+
metric = arr_pixel_to_metric(
228+
np.atleast_2d(np.array([x_pixel, y_pixel], dtype=np.int32)),
229+
parameters.imx,
230+
parameters.imy,
231+
parameters.pix_x,
232+
parameters.pix_y
233+
)
228234

229235
# check if the actual output matches the expected output
230236
assert metric[:, 0] == x_metric_expected
@@ -251,23 +257,33 @@ def setUp(self):
251257

252258
def test_transforms_regress(self):
253259
"""Transformed values are as before."""
254-
input_pos = np.full((3, 2), 100.0)
255-
256-
correct_output_pixel_to_metric = [
260+
correct_output_pixel_to_metric = np.array(
261+
[
257262
[-8181.0, 6657.92],
258263
[-8181.0, 6657.92],
259264
[-8181.0, 6657.92],
260-
]
261-
correct_output_metric_to_pixel = [
265+
])
266+
correct_output_metric_to_pixel = np.array([
262267
[646.60066007, 505.81188119],
263268
[646.60066007, 505.81188119],
264269
[646.60066007, 505.81188119],
265-
]
270+
])
271+
272+
input_pos = np.full((3, 2), 100, dtype=np.int32)
273+
266274

267275
# Test when passing a list
268-
output = arr_pixel_to_metric(input_pos, self.control)
276+
output = arr_pixel_to_metric(
277+
input_pos,
278+
self.control.imx,
279+
self.control.imy,
280+
self.control.pix_x,
281+
self.control.pix_y
282+
)
269283
np.testing.assert_array_almost_equal(output, correct_output_pixel_to_metric)
270284

285+
286+
input_pos = np.full((3, 2), 100, dtype=np.float64)
271287
output = arr_metric_to_pixel(input_pos, self.control)
272288
np.testing.assert_array_almost_equal(output, correct_output_metric_to_pixel)
273289

@@ -278,21 +294,28 @@ def test_transforms(self):
278294
cpar.set_pixel_size((0.1, 0.1))
279295

280296
metric_pos = np.array([[1.0, 1.0], [-10.0, 15.0], [20.0, -30.0]])
281-
pixel_pos = np.array([[650.0, 490.0], [540.0, 350.0], [840.0, 800.0]])
297+
pixel_pos = np.array([[650, 490], [540, 350], [840, 800]],dtype=np.int32)
282298

283299
np.testing.assert_array_almost_equal(
284-
pixel_pos, arr_metric_to_pixel(metric_pos, cpar)
300+
pixel_pos,
301+
arr_metric_to_pixel(
302+
metric_pos,
303+
cpar)
285304
)
286305
np.testing.assert_array_almost_equal(
287-
metric_pos, arr_pixel_to_metric(pixel_pos, cpar)
306+
metric_pos, arr_pixel_to_metric(pixel_pos,
307+
cpar.imx,
308+
cpar.imy,
309+
cpar.pix_x,
310+
cpar.pix_y)
288311
)
289312

290313
def test_brown_affine_regress(self):
291314
"""Test that the brown affine transform gives the same results as before."""
292315
input_vec = np.full((3, 2), 100.0)
293316
output = np.zeros((3, 2))
294-
correct_output_corr = [[100.0, 100.0], [100.0, 100.0], [100.0, 100.0]]
295-
correct_output_dist = [[100.0, 100.0], [100.0, 100.0], [100.0, 100.0]]
317+
correct_output_corr = np.array([[100.0, 100.0], [100.0, 100.0], [100.0, 100.0]])
318+
correct_output_dist = np.array([[100.0, 100.0], [100.0, 100.0], [100.0, 100.0]])
296319

297320
# Test when passing an array for output
298321
for i in range(3):
@@ -322,11 +345,11 @@ def test_brown_affine(self):
322345
cal = Calibration()
323346
cal.set_pos([0.0, 0.0, 40.0])
324347
cal.set_angles([0.0, 0.0, 0.0])
325-
cal.set_primary_point([0.0, 0.0, 10.0])
348+
cal.set_primary_point(np.array([0.0, 0.0, 10.0]))
326349
cal.set_glass_vec(np.r_[0.0, 0.0, 20.0])
327-
cal.set_radial_distortion([0,0,0])
328-
cal.set_decentering([0,0])
329-
cal.set_affine_trans([1, 0])
350+
cal.set_radial_distortion(np.zeros(3,))
351+
cal.set_decentering(np.zeros(2,))
352+
cal.set_affine_trans(np.array([1,0]))
330353

331354
# reference metric positions:
332355
ref_pos = np.array([[0.1, 0.1], [1.0, -1.0], [-10.0, 10.0]])
@@ -357,9 +380,9 @@ def test_full_correction(self):
357380
cal.set_angles(np.r_[0.0, 0.0, 0.0])
358381
cal.set_primary_point(np.r_[0.0, 0.0, 10.0])
359382
cal.set_glass_vec(np.r_[0.0, 0.0, 20.0])
360-
cal.set_radial_distortion([0,0,0])
361-
cal.set_decentering([0,0])
362-
cal.set_affine_trans([1, 0])
383+
cal.set_radial_distortion(np.zeros(3,))
384+
cal.set_decentering(np.zeros(2,))
385+
cal.set_affine_trans(np.array([1, 0]))
363386

364387
# reference metric positions:
365388
# Note the last value is different than in test_brown_affine() because

0 commit comments

Comments
 (0)