Skip to content

Commit 86d202d

Browse files
alexlibalexlib
committed
more numba
1 parent b822300 commit 86d202d

1 file changed

Lines changed: 123 additions & 85 deletions

File tree

‎openptv_python/trafo.py‎

Lines changed: 123 additions & 85 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
from typing import Tuple
44

55
import numpy as np
6+
from numba import njit
67

78
from .calibration import Calibration, ap_52
89
from .parameters import ControlPar
@@ -19,8 +20,21 @@ def pixel_to_metric(
1920
x_pixel, y_pixel (float): input pixel coordinates.
2021
parameters (ControlPar): control structure holding image and pixel sizes.
2122
"""
22-
x_metric = (x_pixel - float(parameters.imx) / 2.0) * parameters.pix_x
23-
y_metric = (float(parameters.imy) / 2.0 - y_pixel) * parameters.pix_y
23+
return fast_pixel_to_metric(
24+
x_pixel,
25+
y_pixel,
26+
parameters.imx,
27+
parameters.imy,
28+
parameters.pix_x,
29+
parameters.pix_y
30+
)
31+
32+
33+
@njit
34+
def fast_pixel_to_metric(x_pixel, y_pixel, imx, imy, pix_x, pix_y) -> Tuple[float, float]:
35+
"""Convert pixel coordinates to metric coordinates."""
36+
x_metric = (x_pixel - float(imx) / 2.0) * pix_x
37+
y_metric = (float(imy) / 2.0 - y_pixel) * pix_y
2438

2539
return (x_metric, y_metric)
2640

@@ -36,8 +50,10 @@ def arr_pixel_to_metric(pixel: np.ndarray, parameters: ControlPar) -> np.ndarray
3650
"""
3751
pixel = np.atleast_2d(np.array(pixel))
3852
metric = np.empty_like(pixel)
39-
metric[:, 0] = (pixel[:, 0] - float(parameters.imx) / 2.0) * parameters.pix_x
40-
metric[:, 1] = (float(parameters.imy) / 2.0 - pixel[:, 1]) * parameters.pix_y
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
4157

4258
return metric
4359

@@ -56,8 +72,27 @@ def metric_to_pixel(
5672
-------
5773
x_pixel, y_pixel (float): output pixel coordinates.
5874
"""
59-
x_pixel = (x_metric / parameters.pix_x) + (float(parameters.imx) / 2.0)
60-
y_pixel = (float(parameters.imy) / 2.0) - (y_metric / parameters.pix_y)
75+
return fast_metric_to_pixel(
76+
x_metric,
77+
y_metric,
78+
parameters.imx,
79+
parameters.imy,
80+
parameters.pix_x,
81+
parameters.pix_y
82+
)
83+
84+
@njit
85+
def fast_metric_to_pixel(
86+
x_metric,
87+
y_metric,
88+
imx,
89+
imy,
90+
pix_x,
91+
pix_y
92+
) -> Tuple[float, float]:
93+
"""Convert metric coordinates to pixel coordinates."""
94+
x_pixel = (x_metric / pix_x) + (float(imx) / 2.0)
95+
y_pixel = (float(imy) / 2.0) - (y_metric / pix_y)
6196

6297
return x_pixel, y_pixel
6398

@@ -74,38 +109,77 @@ def arr_metric_to_pixel(metric: np.ndarray, parameters: ControlPar) -> np.ndarra
74109
-------
75110
pixel (np.ndarray): output array of pixel coordinates.
76111
"""
112+
metric = np.atleast_2d(np.array(metric))
113+
114+
return fast_arr_metric_to_pixel(
115+
metric,
116+
parameters.imx,
117+
parameters.imy,
118+
parameters.pix_x,
119+
parameters.pix_y
120+
)
121+
122+
@njit
123+
def fast_arr_metric_to_pixel(
124+
metric,
125+
imx,
126+
imy,
127+
pix_x,
128+
pix_y
129+
) -> np.ndarray:
130+
"""Convert an array of metric coordinates to pixel coordinates."""
77131
pixel = np.zeros_like(metric)
78-
pixel[:, 0] = (metric[:, 0] / parameters.pix_x) + (float(parameters.imx) / 2.0)
79-
pixel[:, 1] = (float(parameters.imy) / 2.0) - (metric[:, 1] / parameters.pix_y)
132+
pixel[:, 0] = (metric[:, 0] / pix_x) + (float(imx) / 2.0)
133+
pixel[:, 1] = (float(imy) / 2.0) - (metric[:, 1] / pix_y)
80134

81135
return pixel
82136

83137

84-
def distort_brown_affine(x: float, y: float, ap: ap_52) -> Tuple[float, float]:
138+
def distort_brown_affine(x: float,
139+
y: float,
140+
ap: ap_52
141+
) -> Tuple[float, float]:
85142
"""Distort a point using the Brown affine model."""
86143
if x == 0 and y == 0:
87144
return 0, 0
88145

146+
return fast_distort_brown_affine(x, y, ap.k1, ap.k2, ap.k3,
147+
ap.p1, ap.p2, ap.she, ap.scx)
148+
89149
# print(f"x {x}, y {y}")
90150

151+
152+
@njit
153+
def fast_distort_brown_affine(
154+
x: float,
155+
y: float,
156+
k1: float,
157+
k2: float,
158+
k3: float,
159+
p1: float,
160+
p2: float,
161+
she: float,
162+
scx: float,
163+
) -> Tuple[float, float]:
164+
"""Distort a point using the Brown affine model."""
91165
r = sqrt(x**2 + y**2)
92166

93167
x += (
94-
x * (ap.k1 * r**2 + ap.k2 * r**4 + ap.k3 * r**6)
95-
+ ap.p1 * (r**2 + 2 * x**2)
96-
+ 2 * ap.p2 * x * y
168+
x * (k1 * r**2 + k2 * r**4 + k3 * r**6)
169+
+ p1 * (r**2 + 2 * x**2)
170+
+ 2 * p2 * x * y
97171
)
98172
y += (
99-
y * (ap.k1 * r**2 + ap.k2 * r**4 + ap.k3 * r**6)
100-
+ ap.p2 * (r**2 + 2 * y**2)
101-
+ 2 * ap.p1 * x * y
173+
y * (k1 * r**2 + k2 * r**4 + k3 * r**6)
174+
+ p2 * (r**2 + 2 * y**2)
175+
+ 2 * p1 * x * y
102176
)
103177

104178
# print(f"x {x}, y {y}")
105179
# print(f"ap.she {ap.she} ap.scx {ap.scx}")
106180

107-
x1 = ap.scx * x - sin(ap.she) * y
108-
y1 = cos(ap.she) * y
181+
x1 = scx * x - sin(she) * y
182+
y1 = cos(she) * y
109183

110184
# print(f"x1 {x1}, y1 {y1}")
111185

@@ -116,6 +190,22 @@ def correct_brown_affine(
116190
x: float, y: float, ap: ap_52, tol: float = 1e-5
117191
) -> Tuple[float, float]:
118192
"""Correct a distorted point using the Brown affine model."""
193+
return fast_correct_brown_affine(x, y, ap.k1, ap.k2, ap.k3, ap.p1, ap.p2, ap.she, ap.scx, tol)
194+
195+
196+
@njit
197+
def fast_correct_brown_affine(
198+
x: float,
199+
y: float,
200+
k1: float,
201+
k2: float,
202+
k3: float,
203+
p1: float,
204+
p2: float,
205+
she: float,
206+
scx: float,
207+
tol: float = 1e-5) -> Tuple[float, float]:
208+
"""Correct a distorted point using the Brown affine model."""
119209
r, rq, xq, yq = 0.0, 0.0, x, y
120210
itnum = 0
121211

@@ -127,17 +217,17 @@ def correct_brown_affine(
127217
while True:
128218
r = rq
129219
xq = (
130-
(x + yq * np.sin(ap.she)) / ap.scx
131-
- xq * (ap.k1 * r**2 + ap.k2 * r**4 + ap.k3 * r**6)
132-
- ap.p1 * (r**2 + 2 * xq**2)
133-
- 2 * ap.p2 * xq * yq
220+
(x + yq * np.sin(she)) / scx
221+
- xq * (k1 * r**2 + k2 * r**4 + k3 * r**6)
222+
- p1 * (r**2 + 2 * xq**2)
223+
- 2 * p2 * xq * yq
134224
)
135225

136226
yq = (
137-
y / np.cos(ap.she)
138-
- yq * (ap.k1 * r**2 + ap.k2 * r**4 + ap.k3 * r**6)
139-
- ap.p2 * (r**2 + 2 * yq**2)
140-
- 2 * ap.p1 * xq * yq
227+
y / np.cos(she)
228+
- yq * (k1 * r**2 + k2 * r**4 + k3 * r**6)
229+
- p2 * (r**2 + 2 * yq**2)
230+
- 2 * p1 * xq * yq
141231
)
142232

143233
rq = np.sqrt(xq**2 + yq**2)
@@ -152,74 +242,22 @@ def correct_brown_affine(
152242

153243
r = rq
154244
x1 = (
155-
(x + yq * np.sin(ap.she)) / ap.scx
156-
- xq * (ap.k1 * r**2 + ap.k2 * r**4 + ap.k3 * r**6)
157-
- ap.p1 * (r**2 + 2 * xq**2)
158-
- 2 * ap.p2 * xq * yq
245+
(x + yq * np.sin(she)) / scx
246+
- xq * (k1 * r**2 + k2 * r**4 + k3 * r**6)
247+
- p1 * (r**2 + 2 * xq**2)
248+
- 2 * p2 * xq * yq
159249
)
160250

161251
y1 = (
162-
y / np.cos(ap.she)
163-
- yq * (ap.k1 * r**2 + ap.k2 * r**4 + ap.k3 * r**6)
164-
- ap.p2 * (r**2 + 2 * yq**2)
165-
- 2 * ap.p1 * xq * yq
252+
y / np.cos(she)
253+
- yq * (k1 * r**2 + k2 * r**4 + k3 * r**6)
254+
- p2 * (r**2 + 2 * yq**2)
255+
- 2 * p1 * xq * yq
166256
)
167257

168258
return x1, y1
169259

170260

171-
# def correct_brown_affine(x, y, ap, tol=1e-5):
172-
# """Correct a distorted point using the Brown affine model."""
173-
# xq, yq = x, y
174-
# rq = np.sqrt(x**2 + y**2)
175-
176-
# def f(r, xq, yq):
177-
# xq = (
178-
# (x + yq * np.sin(ap.she)) / ap.scx
179-
# - xq * (ap.k1 * r**2 + ap.k2 * r**4 + ap.k3 * r**6)
180-
# - ap.p1 * (r**2 + 2 * xq**2)
181-
# - 2 * ap.p2 * xq * yq
182-
# )
183-
# yq = (
184-
# y / np.cos(ap.she)
185-
# - yq * (ap.k1 * r**2 + ap.k2 * r**4 + ap.k3 * r**6)
186-
# - ap.p2 * (r**2 + 2 * yq**2)
187-
# - 2 * ap.p1 * xq * yq
188-
# )
189-
# rq = np.sqrt(xq**2 + yq**2)
190-
# return rq - r
191-
192-
# r = root_scalar(f, args=(xq, yq), bracket=[0, 2 * rq], maxiter=200, xtol=tol).root
193-
194-
# xq = (
195-
# (x + yq * np.sin(ap.she)) / ap.scx
196-
# - xq * (ap.k1 * r**2 + ap.k2 * r**4 + ap.k3 * r**6)
197-
# - ap.p1 * (r**2 + 2 * xq**2)
198-
# - 2 * ap.p2 * xq * yq
199-
# )
200-
# yq = (
201-
# y / np.cos(ap.she)
202-
# - yq * (ap.k1 * r**2 + ap.k2 * r**4 + ap.k3 * r**6)
203-
# - ap.p2 * (r**2 + 2 * yq**2)
204-
# - 2 * ap.p1 * xq * yq
205-
# )
206-
207-
# x1 = (
208-
# (x + yq * np.sin(ap.she)) / ap.scx
209-
# - xq * (ap.k1 * r**2 + ap.k2 * r**4 + ap.k3 * r**6)
210-
# - ap.p1 * (r**2 + 2 * xq**2)
211-
# - 2 * ap.p2 * xq * yq
212-
# )
213-
# y1 = (
214-
# y / np.cos(ap.she)
215-
# - yq * (ap.k1 * r**2 + ap.k2 * r**4 + ap.k3 * r**6)
216-
# - ap.p2 * (r**2 + 2 * yq**2)
217-
# - 2 * ap.p1 * xq * yq
218-
# )
219-
220-
# return x1, y1
221-
222-
223261
def flat_to_dist(flat_x: float, flat_y: float, cal: Calibration) -> Tuple[float, float]:
224262
"""Convert flat-image coordinates to real-image coordinates.
225263

0 commit comments

Comments
 (0)